Update app.py

app.py

@@ -5,35 +5,31 @@ from PIL import Image
 import numpy as np
 import cv2
 from paddleocr import TextDetection
-from spaces import GPU  # ✅ For ZeroGPU
+from spaces import GPU  # ✅ Required for ZeroGPU
 
 MODEL_HUB_ID = "imperiusrex/Handwritten_model"
 
-# --- Load models on CPU first ---
-device_cpu = torch.device("cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+print("🔄 Loading models...")
 
-print("🔄 Loading models on CPU...")
 processor = TrOCRProcessor.from_pretrained(MODEL_HUB_ID)
-model = VisionEncoderDecoderModel.from_pretrained(MODEL_HUB_ID).to(device_cpu)
+model = VisionEncoderDecoderModel.from_pretrained(MODEL_HUB_ID)
+model.to(device)
 model.eval()
 
 ocr_det_model = TextDetection(model_name="PP-OCRv5_server_det")
-print("✅ Models loaded (CPU mode). GPU will be used only during inference.")
 
-@GPU  # ✅ Runs on GPU when called
+print("✅ Models loaded successfully.")
+
+@GPU  # ✅ This tells Hugging Face this function needs the GPU (H200)
 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")
 
-    # Run detection on original image
-    detection_results = ocr_det_model.predict(np.array(image_pil), batch_size=1)
+    detection_results = ocr_det_model.predict(image_input, batch_size=1)
 
     detected_polys = []
     for res in detection_results:
@@ -87,10 +83,10 @@ def build_interface():
         fn=recognize_handwritten_text,
         inputs=gr.Image(type="numpy", label="Upload Handwritten Image"),
         outputs="text",
-        title="✍️ Handwritten Text Recognition",
+        title="✍ Handwritten Text Recognition",
         description="📷 Upload a handwritten image. Uses PaddleOCR (detection) + TrOCR (recognition).",
     )
 
-if __name__ == "__main__":
+if _name_ == "_main_":
     iface = build_interface()
-    iface.launch(enable_queue=True)
+    iface.launch()