Update Gradio app with multiple files

app.py

@@ -41,89 +41,61 @@ def ocr_process(
     if image_input is None:
         return "Please upload an image first."
 
-    try:
-        # Move model to GPU and set dtype
-        model.cuda().to(torch.bfloat16)
-        
-        # Create temp directory for this session
-        with tempfile.TemporaryDirectory() as temp_dir:
-            # Save image with proper format
-            temp_image_path = os.path.join(temp_dir, "input_image.jpg")
-            # Convert RGBA to RGB if necessary
-            if image_input.mode in ('RGBA', 'LA', 'P'):
-                rgb_image = Image.new('RGB', image_input.size, (255, 255, 255))
-                # Handle different image modes
-                if image_input.mode == 'RGBA':
-                    rgb_image.paste(image_input, mask=image_input.split()[3])
-                else:
-                    rgb_image.paste(image_input)
-                rgb_image.save(temp_image_path, 'JPEG', quality=95)
-            else:
-                image_input.save(temp_image_path, 'JPEG', quality=95)
-            
-            # Verify image was saved
-            if not os.path.exists(temp_image_path):
-                return "Error: Failed to save image for processing."
-            
-            # Set parameters based on preset
-            presets = {
-                "tiny": {"base_size": 512, "image_size": 512, "crop_mode": False},
-                "small": {"base_size": 640, "image_size": 640, "crop_mode": False},
-                "base": {"base_size": 1024, "image_size": 1024, "crop_mode": False},
-                "large": {"base_size": 1280, "image_size": 1280, "crop_mode": False},
-                "gundam": {"base_size": 1024, "image_size": 640, "crop_mode": True},
-            }
-            
-            config = presets[preset]
-
-            # Set prompt based on task type
-            if task_type == "markdown":
-                prompt = "<image>\n<|grounding|>Convert the document to markdown. "
+    # Move model to GPU and set dtype
+    model.cuda().to(torch.bfloat16)
+    
+    # Create temp directory for this session
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Save image with proper format
+        temp_image_path = os.path.join(temp_dir, "input_image.jpg")
+        # Convert RGBA to RGB if necessary
+        if image_input.mode in ('RGBA', 'LA', 'P'):
+            rgb_image = Image.new('RGB', image_input.size, (255, 255, 255))
+            # Handle different image modes
+            if image_input.mode == 'RGBA':
+                rgb_image.paste(image_input, mask=image_input.split()[3])
             else:
-                prompt = "<image>\nFree OCR. "
-
-            # Run inference - the model returns the text directly
-            result = model.infer(
-                tokenizer,
-                prompt=prompt,
-                image_file=temp_image_path,
-                output_path=temp_dir,
-                base_size=config["base_size"],
-                image_size=config["image_size"],
-                crop_mode=config["crop_mode"],
-                save_results=False,
-                test_compress=False,
-            )
-
-        # Move model back to CPU to free GPU memory
-        model.to("cpu")
-        torch.cuda.empty_cache()
-
-        # Process the result
-        if result is None:
-            return "No text could be extracted. The image might be too blurry or contain no readable text."
-        
-        # Handle different result types
-        if isinstance(result, str):
-            output_text = result.strip()
-        elif isinstance(result, (list, tuple)) and len(result) > 0:
-            output_text = str(result[0]).strip()
-        elif isinstance(result, dict):
-            # Try to get text from common keys
-            output_text = result.get('text', result.get('output', result.get('result', str(result))))
+                rgb_image.paste(image_input)
+            rgb_image.save(temp_image_path, 'JPEG', quality=95)
         else:
-            output_text = str(result).strip()
+            image_input.save(temp_image_path, 'JPEG', quality=95)
         
-        if not output_text or output_text == "None":
-            return "No text detected. Try adjusting the preset or uploading a clearer image."
+        # Set parameters based on preset
+        presets = {
+            "tiny": {"base_size": 512, "image_size": 512, "crop_mode": False},
+            "small": {"base_size": 640, "image_size": 640, "crop_mode": False},
+            "base": {"base_size": 1024, "image_size": 1024, "crop_mode": False},
+            "large": {"base_size": 1280, "image_size": 1280, "crop_mode": False},
+            "gundam": {"base_size": 1024, "image_size": 640, "crop_mode": True},
+        }
         
-        return output_text
+        config = presets[preset]
 
-    except Exception as e:
-        # Ensure model is moved back to CPU on error
-        model.to("cpu")
-        torch.cuda.empty_cache()
-        return f"Error processing image: {str(e)}\n\nPlease try a different preset or check if the image is valid."
+        # Set prompt based on task type
+        if task_type == "markdown":
+            prompt = "<image>\n<|grounding|>Convert the document to markdown. "
+        else:
+            prompt = "<image>\nFree OCR. "
+
+        # Run inference - return the result directly
+        result = model.infer(
+            tokenizer,
+            prompt=prompt,
+            image_file=temp_image_path,
+            output_path=temp_dir,
+            base_size=config["base_size"],
+            image_size=config["image_size"],
+            crop_mode=config["crop_mode"],
+            save_results=False,
+            test_compress=False,
+        )
+
+    # Move model back to CPU to free GPU memory
+    model.to("cpu")
+    torch.cuda.empty_cache()
+
+    # Return the result directly - the model returns the extracted text
+    return result
 
 
 # Create Gradio interface