forge: warm model on CPU at startup; GPU window 40s; move to cuda per call

app.py

@@ -40,7 +40,7 @@ def _load():
                 model = AutoModelForCausalLM.from_pretrained(
                     MODEL_ID, quantization_config=qcfg,
                     torch_dtype=torch.float16, low_cpu_mem_usage=True,
-                    device_map="auto")
+                    device_map="cpu")
                 model.eval()
                 _state["model"], _state["tok"] = model, tok
             except Exception as exc:
@@ -49,7 +49,7 @@ def _load():
         return _state["model"], _state["tok"]
 
 
-@spaces.GPU(duration=150)
+@spaces.GPU(duration=40)
 def forge(system_prompt: str, user_prompt: str, max_tokens: float) -> str:
     """One generation. Exposed as Gradio api_name='forge'."""
     model, tok = _load()
@@ -57,18 +57,34 @@ def forge(system_prompt: str, user_prompt: str, max_tokens: float) -> str:
     if system_prompt:
         messages.append({"role": "system", "content": system_prompt})
     messages.append({"role": "user", "content": user_prompt})
+    moved = False
+    if torch.cuda.is_available():
+        try:
+            model.to("cuda"); moved = True
+        except Exception as exc:
+            print("cuda move failed:", repr(exc), flush=True)
     text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = tok(text, return_tensors="pt").to(model.device)
     with torch.inference_mode():
         out = model.generate(**inputs, max_new_tokens=int(max_tokens),
                              do_sample=False, pad_token_id=tok.eos_token_id)
-    return tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+    res = tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+    if moved:
+        try:
+            model.to("cpu"); torch.cuda.empty_cache()
+        except Exception:
+            pass
+    return res
 
 
 def status() -> str:
     return "model=%s loaded=%s error=%s" % (MODEL_ID, _state["model"] is not None, _state["err"])
 
 
+# Warm at startup: the ~2.5 GB download + 4-bit load happen on CPU, outside the GPU
+# window. A free ZeroGPU window is capped (~60 s), far too small to download a model in.
+threading.Thread(target=_load, daemon=True).start()
+
 with gr.Blocks(title="Forge model server") as demo:
     gr.Markdown("## Forge model server\n"
                 "OpenAI-compatible routes are impossible here (see app.py header). "