Update app.py

app.py

@@ -50,24 +50,37 @@ EXAMPLES = [
     "though with a concomitant increase in minor bleeding events.",
 ]
 
-# --- Load model once, at module scope, on CPU. On ZeroGPU there is NO GPU
-#     available at startup — a GPU is only attached inside @spaces.GPU
-#     functions — so we must not touch CUDA here. The model is moved to the
-#     GPU lazily, on the first rewrite() call. Tokenizer comes from the
-#     official base repo to avoid any adapter-side tokenizer.json mismatch. ---
-print("Loading tokenizer and model on CPU...")
+# --- On ZeroGPU, torch is patched at import time and there is NO GPU at
+#     module scope, so we cannot load model weights here (even a CPU load of
+#     safetensors gets intercepted and tries to reach CUDA). We load only the
+#     tokenizer eagerly, and defer ALL model loading to the first GPU call. ---
+print("Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
+print("Tokenizer ready; model loads on first request.")
 
-base_model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,
-    dtype=torch.float16,
-    trust_remote_code=True,
-)
-model = PeftModel.from_pretrained(base_model, ADAPTER_ID, token=HF_TOKEN)
-model.eval()
-_on_gpu = False  # tracks whether the model has been moved to CUDA yet
-print("Model loaded on CPU; will move to GPU on first request.")
+_model = None  # lazily populated inside the @spaces.GPU function
+
+
+def _load_model():
+    """
+    Load base + LoRA adapter and place on CUDA. Called once, lazily, from
+    inside the @spaces.GPU function where a GPU is actually attached.
+    """
+    global _model
+    if _model is not None:
+        return _model
+
+    base_model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        dtype=torch.float16,
+        trust_remote_code=True,
+    )
+    peft_model = PeftModel.from_pretrained(base_model, ADAPTER_ID, token=HF_TOKEN)
+    peft_model = peft_model.to("cuda")
+    peft_model.eval()
+    _model = peft_model
+    return _model
 
 
 def _build_prompt(text: str) -> str:
@@ -87,7 +100,7 @@ def _decode(output_ids, input_len: int) -> str:
     return tokenizer.decode(gen, skip_special_tokens=True).strip()
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=300)
 def rewrite(text: str):
     """
     Produce base ('before') and fine-tuned ('after') rewrites.
@@ -105,12 +118,9 @@ def rewrite(text: str):
     if not text or not text.strip():
         return "", "Please paste some medical text above first."
 
-    # Move the model onto the ZeroGPU-attached GPU on first use. This runs
-    # inside @spaces.GPU, so CUDA is available here (unlike at module scope).
-    global _on_gpu
-    if not _on_gpu:
-        model.to("cuda")
-        _on_gpu = True
+    # Load the model on first use — we are inside @spaces.GPU here, so a GPU
+    # is attached and CUDA is available (unlike at module scope).
+    model = _load_model()
 
     prompt = _build_prompt(text)
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")