Upload streamlit_app.py

src/streamlit_app.py

@@ -110,7 +110,7 @@ div[data-testid="stFileUploader"] * {
 
 /* ═══════════════════════════════════════════════════════════════ */
 
-/* Stats & Output Box */
+/* Stats & DYNAMIC Output Box */
 .stat-card { background: #000; padding: 15px; border-radius: 4px; text-align: center; border: 1px solid rgba(143, 245, 255, 0.1); margin-bottom: 10px; }
 .stat-val { color: #8ff5ff; font-size: 24px; font-weight: 700; font-family: 'Space Grotesk'; }
 .stat-lbl { font-size: 9px; color: #46484d; text-transform: uppercase; letter-spacing: 2px; }
@@ -149,9 +149,12 @@ def load_vision_engine():
 @st.cache_resource(show_spinner=False)
 def load_trocr_model(model_path):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Hugging Face natively downloads the processor via the repo ID
     proc = TrOCRProcessor.from_pretrained(model_path)
     
     if os.path.exists(model_path):
+        # Local Loading Logic
         config = VisionEncoderDecoderConfig.from_pretrained(model_path)
         model = VisionEncoderDecoderModel(config)
         safe_path = os.path.join(model_path, "model.safetensors")
@@ -163,35 +166,41 @@ def load_trocr_model(model_path):
         else:
             model.load_state_dict(torch.load(bin_path, map_location="cpu", weights_only=True), strict=False)
     else:
+        # Cloud Loading Logic: Natively pulls your model from the Hugging Face Hub
         model = VisionEncoderDecoderModel.from_pretrained(model_path)
         
     # Push standard registered parameters/buffers to device
     model.to(device)
     
-    # ─── AGGRESSIVE ROGUE TENSOR MIGRATION (WITH META SAFEGUARD) ───
-    # Snapshot dict to avoid runtime size change errors while finding unregistered weights
+    # ─── BULLETPROOF TENSOR MIGRATION (WITH EXCEPTIONS CATCHER) ───
     for module in model.modules():
-        # 1. Double check parameters (Ensure it's not a meta tensor)
+        # 1. Double check parameters safely
         for name, param in list(module._parameters.items()):
-            if param is not None and not param.is_meta:
-                module._parameters[name] = torch.nn.Parameter(param.to(device))
-        # 2. Double check buffers
+            if param is not None:
+                try: module._parameters[name] = torch.nn.Parameter(param.to(device))
+                except (NotImplementedError, RuntimeError): pass
+                
+        # 2. Double check buffers safely
         for name, buf in list(module._buffers.items()):
-            if buf is not None and not buf.is_meta:
-                module._buffers[name] = buf.to(device)
-        # 3. Hunt down unregistered raw tensors (Fixes the TrOCR positional weights crash)
+            if buf is not None:
+                try: module._buffers[name] = buf.to(device)
+                except (NotImplementedError, RuntimeError): pass
+                
+        # 3. Hunt down unregistered raw tensors safely
         for name, attr in list(module.__dict__.items()):
-            if isinstance(attr, torch.Tensor) and not attr.is_meta:
-                setattr(module, name, attr.to(device))
+            if isinstance(attr, torch.Tensor):
+                try: setattr(module, name, attr.to(device))
+                except (NotImplementedError, RuntimeError): pass
                 
-    # If on GPU, push the entire model to Half precision
+    # If on GPU, push the entire model to Half precision safely
     if device.type == "cuda": 
         model = model.half()
         # Ensure those unregistered raw tensors are ALSO converted to half precision safely
         for module in model.modules():
             for name, attr in list(module.__dict__.items()):
-                if isinstance(attr, torch.Tensor) and not attr.is_meta and attr.is_floating_point():
-                    setattr(module, name, attr.half())
+                if isinstance(attr, torch.Tensor) and attr.is_floating_point():
+                    try: setattr(module, name, attr.half())
+                    except (NotImplementedError, RuntimeError): pass
                     
     model.eval()
     return proc, model, device