Update app.py

app.py

@@ -5,31 +5,35 @@ from PIL import Image
 import numpy as np
 import cv2
 from paddleocr import TextDetection
-from spaces import GPU  # ✅ Required for ZeroGPU
+from spaces import GPU  # ✅ For ZeroGPU
 
 MODEL_HUB_ID = "imperiusrex/Handwritten_model"
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-print("🔄 Loading models...")
+# --- Load models on CPU first ---
+device_cpu = torch.device("cpu")
 
+print("🔄 Loading models on CPU...")
 processor = TrOCRProcessor.from_pretrained(MODEL_HUB_ID)
-model = VisionEncoderDecoderModel.from_pretrained(MODEL_HUB_ID)
-model.to(device)
+model = VisionEncoderDecoderModel.from_pretrained(MODEL_HUB_ID).to(device_cpu)
 model.eval()
 
 ocr_det_model = TextDetection(model_name="PP-OCRv5_server_det")
+print("✅ Models loaded (CPU mode). GPU will be used only during inference.")
 
-print("✅ Models loaded successfully.")
-
-@GPU  # ✅ This tells Hugging Face this function needs the GPU (H200)
+@GPU  # ✅ Runs on GPU when called
 def recognize_handwritten_text(image_input):
     if image_input is None:
         return "Please upload an image."
 
+    # Move model to GPU here (only when needed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    torch.zeros(1).to(device)  # CUDA warmup
+
     image_pil = Image.fromarray(image_input).convert("RGB")
 
-    detection_results = ocr_det_model.predict(image_input, batch_size=1)
+    # Run detection on original image
+    detection_results = ocr_det_model.predict(np.array(image_pil), batch_size=1)
 
     detected_polys = []
     for res in detection_results:
@@ -89,4 +93,4 @@ def build_interface():
 
 if __name__ == "__main__":
     iface = build_interface()
-    iface.launch()
+    iface.launch(enable_queue=True)