Update app.py

app.py

@@ -1,34 +1,28 @@
-import os
-import torch
 import spaces
 import gradio as gr
-from PIL import Image
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel, RobertaTokenizer, AutoImageProcessor
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import torch
 
-print("Loading TrOCR model... this may take a minute.")
-tokenizer = RobertaTokenizer.from_pretrained("microsoft/trocr-base-handwritten")
-image_processor = AutoImageProcessor.from_pretrained("microsoft/trocr-base-handwritten")
-processor = TrOCRProcessor(image_processor=image_processor, tokenizer=tokenizer)
-model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
+# Load models
+device = "cuda" if torch.cuda.is_available() else "cpu"
+processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten')
+model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten').to(device)
 
-# This is 100% standard Gradio code, so the ZeroGPU scanner will perfectly detect it!
 @spaces.GPU
 def predict(image):
-    # Check optional API token (Gradio handles this differently, but we can do a basic check if needed, 
-    # though for this simple version we will just leave it open or you can secure it in space settings)
-    
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-    
+    if image is None:
+        return ""
+    # Process the image directly on the ZeroGPU
     pixel_values = processor(image, return_tensors="pt").pixel_values.to(device)
     generated_ids = model.generate(pixel_values)
-    return processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return generated_text
 
-# Standard Gradio Interface
+# The real interface that perfectly matches v5.0.3!
 demo = gr.Interface(
     fn=predict, 
     inputs=gr.Image(type="pil"), 
-    outputs="text", 
+    outputs="text",
     title="TrOCR Hub Backend"
 )