Fix ZeroGPU initialisation: spaces before torch, two-phase pipeline load, guidance_scale

Three correctness fixes informed by the BFL reference starter app:

1. Import `spaces` before `torch` — ZeroGPU patches CUDA init at import time;
loading torch first can cause silent GPU mis-behaviour on HF Spaces.
Also unifies the public API to a single @GPU-decorated function using the
same no-op shim pattern as the reference.

2. Two-phase pipeline loading — `_load_pipeline_cpu()` runs at module scope
(outside the @GPU budget) so weights are in CPU RAM before the first
ZeroGPU call. `get_pipeline()` just does `.to(device)` inside @GPU , which
takes ~1s and removes the cold-start timeout risk on first generation.

3. Fix `guidance_scale=1.0` everywhere — the previous `0.0` on local runs
was a copy-paste from schnell; klein-4B distilled uses 1.0 in all envs.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

image_generator.py

@@ -1,5 +1,30 @@
 import random
 import os
+import sys
+import threading
+import asyncio
+
+# ---------------------------------------------------------------------------
+# ZeroGPU shim — spaces MUST be imported BEFORE torch.
+# On HF ZeroGPU, spaces patches CUDA initialisation; that patch must land
+# before torch is imported or GPU calls can silently mis-behave.
+# Locally spaces isn't installed, so we fall back to a no-op @GPU decorator
+# that makes the same code run unchanged on MPS / CPU.
+# ---------------------------------------------------------------------------
+
+try:
+    import spaces  # type: ignore
+    GPU = spaces.GPU
+    ON_ZEROGPU = True
+except Exception:
+    def GPU(*dargs, **dkwargs):  # noqa: N802 — mirror the spaces.GPU API
+        def wrap(fn):
+            return fn
+        if len(dargs) == 1 and callable(dargs[0]) and not dkwargs:
+            return dargs[0]
+        return wrap
+    ON_ZEROGPU = False
+
 import torch
 
 # Patch for torch < 2.4 which lacks torch.xpu (required by diffusers >= 0.30)
@@ -21,10 +46,7 @@ if not hasattr(torch, "xpu"):
 IS_HF_SPACE = os.environ.get("SPACE_ID") is not None
 
 if IS_HF_SPACE:
-    import spaces
-    import sys, threading, asyncio
-
-    # Suppress Python 3.10 asyncio GC bug (Invalid file descriptor: -1)
+    # Suppress Python 3.13 asyncio GC bug (Invalid file descriptor: -1)
     _orig_unraisable = sys.unraisablehook
     def _unraisable_hook(args):
         if args.exc_type is ValueError and "Invalid file descriptor" in str(args.exc_value):
@@ -197,41 +219,54 @@ def build_prompt(animals: list[str], places: list[str]) -> str:
     return f"{build_subject(animals, places)}, {NUMZOO_STYLE}"
 
 # ---------------------------------------------------------------------------
-# Pipeline loader (cached globally — survives between ZeroGPU calls)
+# Pipeline loader — two-phase, following the reference ZeroGPU pattern:
+#
+#   Phase 1 · _load_pipeline_cpu()  — from_pretrained to CPU RAM.
+#     Called at module scope (outside any @GPU function) so the weights are
+#     already resident when the first ZeroGPU call arrives.  This keeps the
+#     model-download cost out of the 60 s GPU-runtime budget and prevents
+#     cold-start timeouts on the very first generation.
+#
+#   Phase 2 · get_pipeline()  — .to(device).
+#     Must be called INSIDE a @GPU-decorated function (where a GPU slice is
+#     guaranteed). Moving already-loaded CPU tensors to CUDA is fast (~1 s)
+#     and comfortably within the budget.
+#
+# Locally (MPS / CPU) both phases happen inside generate_reward_image because
+# the no-op @GPU decorator doesn't impose any budget constraint.
 # ---------------------------------------------------------------------------
 
 _pipe = None
+_use_mps = (not IS_HF_SPACE) and (not torch.cuda.is_available()) and torch.backends.mps.is_available()
+_dtype   = torch.float16 if _use_mps else torch.bfloat16   # float16 on MPS (bfloat16 unsupported)
 
-def get_pipeline():
+
+def _load_pipeline_cpu() -> None:
+    """Phase 1: load model weights into CPU RAM. Safe to call at module scope."""
     global _pipe
     if _pipe is not None:
-        return _pipe
-
+        return
     from diffusers import Flux2KleinPipeline
-
     hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
-    # float16 on MPS (bfloat16 not fully supported), bfloat16 everywhere else
-    use_mps = (not IS_HF_SPACE) and (not torch.cuda.is_available()) and torch.backends.mps.is_available()
-    dtype   = torch.float16 if use_mps else torch.bfloat16
-
-    print(f"Loading FLUX.2-klein-4B pipeline… (token={'set' if hf_token else 'NOT SET'}, dtype={dtype})")
+    print(f"Loading FLUX.2-klein-4B on CPU… (dtype={_dtype}, token={'set' if hf_token else 'NOT SET'})")
     _pipe = Flux2KleinPipeline.from_pretrained(
-        "black-forest-labs/FLUX.2-klein-4B",
-        torch_dtype=dtype,
+        _MODEL_ID,
+        torch_dtype=_dtype,
         token=hf_token,
     )
+    print("Pipeline loaded on CPU — ready for device placement.")
 
-    if IS_HF_SPACE:
-        _pipe = _pipe.to("cuda")
-    elif torch.cuda.is_available():
-        _pipe = _pipe.to("cuda")
-    elif use_mps:
-        _pipe = _pipe.to("mps")
-    else:
-        _pipe = _pipe.to("cpu")
 
-    print("Pipeline ready.")
-    return _pipe
+def get_pipeline():
+    """Phase 2: move pipeline to the target device. Call inside @GPU on HF Spaces."""
+    global _pipe
+    if _pipe is None:
+        _load_pipeline_cpu()   # fallback for local / first-call safety
+    if torch.cuda.is_available():
+        return _pipe.to("cuda")
+    if _use_mps:
+        return _pipe.to("mps")
+    return _pipe.to("cpu")
 
 # ---------------------------------------------------------------------------
 # Core generation (always wrapped in try/except)
@@ -244,12 +279,11 @@ def _generate(animals: list[str], places: list[str]):
         prompt = build_prompt(animals, places)
         print(f"Generating | prompt: {prompt}")
 
-        guidance = 1.0 if IS_HF_SPACE else 0.0  # klein=1.0, schnell=0.0
         t0 = time.time()
         result = pipe(
             prompt=prompt,
             num_inference_steps=4,
-            guidance_scale=guidance,
+            guidance_scale=1.0,   # klein-4B distilled: always 1.0 (not schnell's 0.0)
             height=512,
             width=512,
         )
@@ -264,15 +298,21 @@ def _generate(animals: list[str], places: list[str]):
 
 
 # ---------------------------------------------------------------------------
-# Public API — two versions depending on environment
+# Module-scope CPU pre-load (HF Spaces only).
+# Runs after the persistent cache is set up and the snapshot is downloaded,
+# so from_pretrained finds the weights locally and completes quickly.
+# Locally this is skipped — the pipeline loads lazily on first generate call.
 # ---------------------------------------------------------------------------
 
 if IS_HF_SPACE:
-    @spaces.GPU(duration=60)
-    def generate_reward_image(animals: list[str], places: list[str]):
-        """Generate reward image on HF Spaces ZeroGPU."""
-        return _generate(animals, places)
-else:
-    def generate_reward_image(animals: list[str], places: list[str]):
-        """Generate reward image locally (MPS/CPU)."""
-        return _generate(animals, places)
+    _load_pipeline_cpu()
+
+# ---------------------------------------------------------------------------
+# Public API — single function, @GPU decorator is a no-op locally.
+# ---------------------------------------------------------------------------
+
+@GPU(duration=60)
+def generate_reward_image(animals: list[str], places: list[str]):
+    """Generate a reward image.  On HF Spaces runs inside a ZeroGPU slice;
+    locally the @GPU decorator is a no-op and MPS/CPU is used instead."""
+    return _generate(animals, places)