Fix: use @spaces.GPU for ZeroGPU - load models on CPU, move to GPU per-request

app.py

@@ -2,43 +2,36 @@
 
 import json
 import tempfile
-from pathlib import Path
 
 import cv2
 import gradio as gr
 import numpy as np
+import spaces
 import torch
 from huggingface_hub import hf_hub_download
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from ultralytics import YOLO
 
-# --- Model loading (cached at startup) ---
+# --- Model loading (on CPU at startup, GPU allocated per-request) ---
 
 YOLO_REPO = "prasanthmj/yolov8-braille"
 BYT5_REPO = "prasanthmj/braille-byt5-v3"
 
-def load_models():
-    """Download and load both models."""
-    # YOLOv8 braille detector
-    weights_path = hf_hub_download(YOLO_REPO, "yolov8_braille.pt")
-    braille_map_path = hf_hub_download(YOLO_REPO, "braille_map.json")
-
-    yolo_model = YOLO(weights_path)
-    with open(braille_map_path) as f:
-        dot_to_unicode = json.load(f)
-
-    # ByT5 Grade 2 interpreter
-    tokenizer = AutoTokenizer.from_pretrained(BYT5_REPO)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    byt5_model = AutoModelForSeq2SeqLM.from_pretrained(BYT5_REPO).to(device)
-    byt5_model.eval()
+print("Loading models...")
 
-    return yolo_model, dot_to_unicode, tokenizer, byt5_model, device
+# YOLOv8 braille detector
+weights_path = hf_hub_download(YOLO_REPO, "yolov8_braille.pt")
+braille_map_path = hf_hub_download(YOLO_REPO, "braille_map.json")
+yolo_model = YOLO(weights_path)
+with open(braille_map_path) as f:
+    dot_to_unicode = json.load(f)
 
+# ByT5 Grade 2 interpreter (load on CPU, moved to GPU per-request)
+tokenizer = AutoTokenizer.from_pretrained(BYT5_REPO)
+byt5_model = AutoModelForSeq2SeqLM.from_pretrained(BYT5_REPO)
+byt5_model.eval()
 
-print("Loading models...")
-yolo_model, dot_to_unicode, tokenizer, byt5_model, device = load_models()
-print(f"Models loaded. Device: {device}")
+print("Models loaded (CPU). GPU allocated per-request via ZeroGPU.")
 
 # --- CLAHE Preprocessing ---
 
@@ -102,30 +95,9 @@ def detect_braille(image_path: str, confidence: float = 0.15) -> list[list[dict]
 
     return lines
 
-# --- Stage 2: ByT5 Interpretation ---
-
-def interpret_braille(braille_lines: list[str]) -> list[str]:
-    """Translate braille Unicode lines to English using ByT5."""
-    results = []
-    for line in braille_lines:
-        if not line.strip():
-            results.append("")
-            continue
-
-        input_text = f"translate Braille to English: {line}"
-        inputs = tokenizer(input_text, return_tensors="pt", max_length=1024, truncation=True)
-        inputs = {k: v.to(device) for k, v in inputs.items()}
-
-        with torch.no_grad():
-            outputs = byt5_model.generate(**inputs, max_length=512)
-
-        decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        results.append(decoded)
-
-    return results
-
-# --- Main pipeline ---
+# --- Main pipeline (GPU allocated here) ---
 
+@spaces.GPU
 def transcribe(image) -> str:
     """Full pipeline: image -> detection -> interpretation -> English text."""
     if image is None:
@@ -155,8 +127,25 @@ def transcribe(image) -> str:
     total_cells = sum(len(line) for line in lines)
     avg_conf = np.mean([cell["confidence"] for line in lines for cell in line])
 
-    # Stage 2: Interpret with ByT5
-    english_lines = interpret_braille(braille_lines)
+    # Stage 2: Interpret each line with ByT5 on GPU
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    byt5_model.to(device)
+
+    english_lines = []
+    for line in braille_lines:
+        if not line.strip():
+            english_lines.append("")
+            continue
+
+        input_text = f"translate Braille to English: {line}"
+        inputs = tokenizer(input_text, return_tensors="pt", max_length=1024, truncation=True)
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = byt5_model.generate(**inputs, max_length=512)
+
+        decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        english_lines.append(decoded)
 
     # Format output
     braille_text = "\n".join(braille_lines)