Update Gradio app with multiple files

app.py

@@ -5,10 +5,7 @@ from PIL import Image
 import io
 import os
 from typing import Optional
-
-# Set device
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-device = "cuda" if torch.cuda.is_available() else "cpu"
+import spaces
 
 # Load model and tokenizer
 model_name = "deepseek-ai/DeepSeek-OCR"
@@ -19,11 +16,10 @@ model = AutoModel.from_pretrained(
     trust_remote_code=True,
     use_safetensors=True,
 )
-model = model.eval().to(device)
-if device == "cuda":
-    model = model.to(torch.bfloat16)
+model = model.eval()
 
 
+@spaces.GPU
 def ocr_process(
     image_input: Image.Image,
     task_type: str = "ocr",
@@ -48,6 +44,9 @@ def ocr_process(
         return "Please upload an image first."
 
     try:
+        # Move model to GPU and set dtype
+        model.to("cuda")
+        model.to(torch.bfloat16)
         # Save image temporarily
         temp_image_path = "/tmp/temp_ocr_image.jpg"
         image_input.save(temp_image_path)
@@ -75,9 +74,16 @@ def ocr_process(
         if os.path.exists(temp_image_path):
             os.remove(temp_image_path)
 
+        # Move model back to CPU to free GPU memory
+        model.to("cpu")
+        torch.cuda.empty_cache()
+
         return output if output else "No text detected in image."
 
     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)}"
 
 

requirements.txt

@@ -5,3 +5,4 @@ Pillow>=10.0.0
 deepseek-ai
 safetensors>=0.4.0
 flash-attn>=2.5.0
+spaces