fix: lazy load (avoid Space init OOM with 7B+LoRA)

app.py

@@ -1,14 +1,13 @@
-"""Surrogate-1 ZeroGPU Space — Qwen2.5-Coder-7B + v1 LoRA.
+"""Surrogate-1 ZeroGPU Space — lazy-load variant.
 
-Rewritten 2026-04-30 to use gr.ChatInterface (simpler signature, avoids
-the gradio_client._json_schema_to_python_type recursion bug that broke
-the previous custom-Blocks app.py).
+Don't load the 7B model at import time (causes Space init OOM/timeout).
+Load lazily inside the @spaces.GPU function — ZeroGPU spins a fresh GPU
+worker per request anyway, so loading there is correct.
 """
 import os
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
 
 BASE_MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"
 LORA_REPO  = os.environ.get("LORA_REPO", "axentx/surrogate-1-coder-7b-lora-v1")
@@ -20,79 +19,82 @@ SYSTEM = (
     "'I don't know' rather than confabulate."
 )
 
-print(f"[boot] tokenizer: {BASE_MODEL}")
-tokenizer = AutoTokenizer.from_pretrained(
-    BASE_MODEL, token=HF_TOKEN or None, trust_remote_code=True)
-if tokenizer.pad_token_id is None:
-    tokenizer.pad_token_id = tokenizer.eos_token_id
+# Module-level cache so repeated calls within same GPU worker reuse
+_model = None
+_tokenizer = None
 
-print(f"[boot] base model on CPU: {BASE_MODEL}")
-model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL, torch_dtype=torch.bfloat16,
-    token=HF_TOKEN or None, trust_remote_code=True,
-    device_map="cpu")
 
-LORA_ACTIVE = False
-try:
-    from peft import PeftModel
-    print(f"[boot] LoRA: {LORA_REPO}")
-    model = PeftModel.from_pretrained(model, LORA_REPO, token=HF_TOKEN or None)
-    LORA_ACTIVE = True
-    print("[boot] LoRA applied")
-except Exception as e:
-    print(f"[boot] LoRA failed (using base only): {e}")
+def _load():
+    """Lazy load on first @spaces.GPU call (running on GPU worker)."""
+    global _model, _tokenizer
+    if _model is not None:
+        return _model, _tokenizer
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    print(f"[lazy-load] tokenizer: {BASE_MODEL}")
+    tok = AutoTokenizer.from_pretrained(
+        BASE_MODEL, token=HF_TOKEN or None, trust_remote_code=True)
+    if tok.pad_token_id is None:
+        tok.pad_token_id = tok.eos_token_id
+    print(f"[lazy-load] base model on cuda")
+    m = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL, torch_dtype=torch.bfloat16,
+        token=HF_TOKEN or None, trust_remote_code=True,
+        device_map="cuda")
+    try:
+        from peft import PeftModel
+        print(f"[lazy-load] LoRA: {LORA_REPO}")
+        m = PeftModel.from_pretrained(m, LORA_REPO, token=HF_TOKEN or None)
+        print("[lazy-load] LoRA applied")
+    except Exception as e:
+        print(f"[lazy-load] LoRA failed (using base only): {e}")
+    _model, _tokenizer = m, tok
+    return _model, _tokenizer
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=180)
 def respond(message, history, max_new_tokens=512, temperature=0.4, top_p=0.9):
+    if not message or not message.strip():
+        return ""
+    model, tok = _load()
     msgs = [{"role": "system", "content": SYSTEM}]
     for u, a in (history or []):
         if u: msgs.append({"role": "user", "content": u})
         if a: msgs.append({"role": "assistant", "content": a})
     msgs.append({"role": "user", "content": message})
 
-    prompt = tokenizer.apply_chat_template(
+    prompt = tok.apply_chat_template(
         msgs, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True,
-                        max_length=24000).to("cuda")
-    model.to("cuda")
-
+    inputs = tok(prompt, return_tensors="pt", truncation=True,
+                 max_length=24000).to("cuda")
     out = model.generate(
         **inputs,
         max_new_tokens=int(max_new_tokens),
         temperature=float(temperature) if temperature > 0 else 1e-5,
         top_p=float(top_p),
         do_sample=temperature > 0,
-        pad_token_id=tokenizer.pad_token_id,
-        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tok.pad_token_id,
+        eos_token_id=tok.eos_token_id,
         use_cache=True,
     )
     new_tokens = out[0][inputs["input_ids"].shape[1]:]
-    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
+    return tok.decode(new_tokens, skip_special_tokens=True).strip()
 
 
 desc = (
-    f"**Base**: `{BASE_MODEL}`    "
-    f"**LoRA**: `{LORA_REPO}` "
-    f"{'✅ active' if LORA_ACTIVE else '⚠️ base only'}<br>"
-    f"**Hardware**: ZeroGPU A10G (PRO subscription, 25K min/mo @ $0)"
+    f"**Base**: `{BASE_MODEL}` &nbsp; **LoRA**: `{LORA_REPO}`<br>"
+    f"**Hardware**: ZeroGPU A10G (PRO, 25K min/mo @ $0). "
+    f"First request takes ~30-60s (cold model load), subsequent ~3-10s."
 )
 
 demo = gr.ChatInterface(
     fn=respond,
-    title="Surrogate-1 — DevSecOps + SRE + Code Agent",
+    title="Surrogate-1 — DevSecOps + Code Agent",
     description=desc,
     additional_inputs=[
         gr.Slider(64, 2048, value=512, step=64, label="max new tokens"),
         gr.Slider(0.0, 1.5, value=0.4, step=0.05, label="temperature"),
         gr.Slider(0.5, 1.0, value=0.9, step=0.05, label="top_p"),
     ],
-    examples=[
-        "Write a Terraform module for an S3 bucket with KMS encryption + versioning.",
-        "Diagnose: AWS Lambda cold start latency 3s. Architecture suggestions?",
-        "Review this IAM policy for least-privilege violations: <paste here>",
-        "Implement rate-limit per-API-key in FastAPI with Redis.",
-    ],
 )
 
 if __name__ == "__main__":