Update app.py

app.py

@@ -2,12 +2,11 @@
 import subprocess
 import sys
 
-# CRITICAL: Import spaces FIRST before any CUDA initialization
-import spaces
 
-# Now we can import torch and other packages
+import spaces
 import torch
 
+
 # Install flash-attn for GPU only (after spaces import)
 if torch.cuda.is_available():
     print("CUDA detected - installing flash-attn for optimal GPU performance...")
@@ -99,8 +98,13 @@ def extract_text_from_image(image, temperature=0.2):
         return_tensors="pt"
     )
     
-    # Move inputs to device
-    inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
+    # Move inputs to device AND convert to the correct dtype
+    inputs = {
+        k: v.to(device=device, dtype=dtype) if isinstance(v, torch.Tensor) and v.dtype in [torch.float32, torch.float16, torch.bfloat16]
+        else v.to(device) if isinstance(v, torch.Tensor) 
+        else v 
+        for k, v in inputs.items()
+    }
     
     # Generate text with appropriate settings
     with torch.no_grad():  # Disable gradients for inference