c

app.py

@@ -70,7 +70,12 @@ def process_single_line(image, model_name):
     prompt = MODELS[model_name]["prompt"]
     # Add image token to prompt
     prompt = f"<image>{prompt}"
-    model_inputs = processor(text=prompt, images=image, return_tensors="pt").to(torch.bfloat16).to("cuda")
+    
+    # First prepare inputs without moving to CUDA
+    model_inputs = processor(text=prompt, images=image, return_tensors="pt")
+    
+    # Then move to CUDA in a separate step
+    model_inputs = {k: v.to(torch.bfloat16).to("cuda") for k, v in model_inputs.items()}
     
     outputs = model.generate(
         **model_inputs,
@@ -156,8 +161,30 @@ def process_multi_line(image, model_name, progress=gr.Progress()):
         progress(1.0, desc="Done!")
         return "\n".join(all_text), [bbox_image]  # Return as list for gallery
 
+@spaces.GPU
+def process_image(model_name, image, progress=gr.Progress()):
+    """Process a single image"""
+    if image is None:
+        return "", []
+        
+    # Load model if different model selected
+    if model_name != current_model_name:
+        progress(0, desc="Loading model...")
+        load_model(model_name)
+        
+    return process_multi_line(image, model_name, progress)
+
+@spaces.GPU
 def process_pdf(pdf_path, model_name, progress=gr.Progress()):
     """Process a PDF file"""
+    if pdf_path is None:
+        return "", []
+        
+    # Load model if different model selected
+    if model_name != current_model_name:
+        progress(0, desc="Loading model...")
+        load_model(model_name)
+        
     # Create temporary directory
     with tempfile.TemporaryDirectory() as temp_dir:
         # Initialize detector with temp directory
@@ -236,19 +263,6 @@ def process_pdf(pdf_path, model_name, progress=gr.Progress()):
         progress(1.0, desc="Done!")
         return "\n".join(all_text), bbox_images  # Return list of bbox images
 
-@spaces.GPU
-def process_image(model_name, image, progress=gr.Progress()):
-    """Process a single image"""
-    if image is None:
-        return "", None
-        
-    # Load model if different model selected
-    if model_name != current_model_name:
-        progress(0, desc="Loading model...")
-        load_model(model_name)
-        
-    return process_multi_line(image, model_name, progress)
-
 # Example images with descriptions
 examples = [
     ["type_1_sl.png", "Typed Dhivehi text sample 1"],