Fix: load fp16 at module level for ZeroGPU compatibility

app.py

@@ -2,7 +2,7 @@
 MedVision Edge — Offline Chest X-ray Analysis
 Gradio demo for Gemma 4 E4B fine-tuned on 112K+ chest X-rays.
 
-Deploy: HuggingFace Space (GPU T4) or local with `python app.py`
+Deploy: HuggingFace Space (ZeroGPU) or local with `python app.py`
 """
 
 import os
@@ -13,15 +13,15 @@ import gradio as gr
 import spaces
 from PIL import Image
 from pathlib import Path
+from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig
 
 # ── Config ──────────────────────────────────────────────────────
-# HF Space: set MEDVISION_MODEL env var to HF repo ID (e.g., "user/medvision-edge-v4")
-# Local: defaults to local path on groundctl
 MODEL_PATH = os.environ.get(
     "MEDVISION_MODEL",
     os.path.expanduser("~/ml-projects/medvision/model_output/final_model"),
 )
-LOAD_IN_4BIT = os.environ.get("MEDVISION_4BIT", "true").lower() == "true"
+IS_SPACES = os.environ.get("SPACE_ID") is not None
+LOAD_IN_4BIT = os.environ.get("MEDVISION_4BIT", "true").lower() == "true" and not IS_SPACES
 PATHOLOGIES = ["Pneumonia", "Consolidation", "Cardiomegaly", "Effusion", "Edema"]
 
 LANGUAGES = {
@@ -48,41 +48,30 @@ TRANSLATE_PROMPT = (
     "Keep medical terminology accurate. Translate only, do not add commentary.\n\n{text}"
 )
 
-# ── Globals (loaded once) ───────────────────────────────────────
-model = None
-processor = None
+# ── Model loading ─────────────────────────────────────────────
+# ZeroGPU: load on CPU at module level, ZeroGPU moves to GPU automatically
+# Local: load with device_map="auto" (optionally 4-bit)
+print(f"Loading model from {MODEL_PATH} (spaces={IS_SPACES}, 4bit={LOAD_IN_4BIT})...")
 
-
-def load_model():
-    global model, processor
-    if model is not None:
-        return
-
-    from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig
-
-    print(f"Loading merged model from {MODEL_PATH} (4bit={LOAD_IN_4BIT})...")
-
-    if LOAD_IN_4BIT:
-        quantization_config = BitsAndBytesConfig(
+if LOAD_IN_4BIT:
+    model = AutoModelForImageTextToText.from_pretrained(
+        MODEL_PATH,
+        quantization_config=BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_compute_dtype=torch.float16,
             bnb_4bit_quant_type="nf4",
-        )
-        model = AutoModelForImageTextToText.from_pretrained(
-            MODEL_PATH,
-            quantization_config=quantization_config,
-            device_map="auto",
-        )
-    else:
-        model = AutoModelForImageTextToText.from_pretrained(
-            MODEL_PATH,
-            torch_dtype=torch.float16,
-            device_map="auto",
-        )
+        ),
+        device_map="auto",
+    )
+else:
+    model = AutoModelForImageTextToText.from_pretrained(
+        MODEL_PATH,
+        torch_dtype=torch.float16,
+    )
 
-    model.eval()
-    processor = AutoProcessor.from_pretrained(MODEL_PATH)
-    print("Model loaded.")
+model.eval()
+processor = AutoProcessor.from_pretrained(MODEL_PATH)
+print("Model loaded.")
 
 
 def parse_response(response: str) -> dict:
@@ -137,6 +126,8 @@ def run_inference(image: Image.Image, prompt: str) -> str:
         return_tensors="pt",
     ).to(model.device)
 
+    print(f"Inference on device: {model.device}, dtype: {model.dtype}")
+
     with torch.no_grad():
         output = model.generate(
             **inputs,
@@ -147,6 +138,7 @@ def run_inference(image: Image.Image, prompt: str) -> str:
         )
 
     decoded = processor.batch_decode(output, skip_special_tokens=True)[0]
+    print(f"Raw decoded (first 500 chars): {decoded[:500]}")
     # Extract only the assistant's response (handle multiple Gemma 4 marker formats)
     for marker in ["<|turn>model", "<start_of_turn>model", "model\n"]:
         if marker in decoded:
@@ -165,8 +157,6 @@ def analyze_xray(image, language, patient_age, patient_weight):
     if image is None:
         return "Please upload a chest X-ray image.", "{}", "", ""
 
-    load_model()
-
     # Convert to PIL if needed
     if not isinstance(image, Image.Image):
         image = Image.fromarray(image).convert("RGB")