Upgrade image model to FLUX.2-klein-4B + align training pipeline

image_generator.py:
- Replace FLUX.1-schnell with FLUX.2-klein-4B (Apache 2.0, better quality)
- Use Flux2KleinPipeline, guidance_scale=1.0, float16 on MPS
- Share NUMZOO_STYLE constant so live prompts match training captions exactly
- Log generation time (✅ Generated in Xs)

scripts/generate_dataset.py:
- Switch to FLUX.1-dev via HF Inference API (fal-ai provider, HF Pro token)
- Import NUMZOO_STYLE from image_generator — single source of truth
- Add --count N flag to generate a subset (e.g. --count 5 for a test run)
- Load GEMINI_API_KEY / HF_TOKEN from .env via python-dotenv

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

image_generator.py

@@ -15,24 +15,24 @@ if not hasattr(torch, "xpu"):
     torch.xpu = _MockXPU()
 
 # ---------------------------------------------------------------------------
-# Detect HuggingFace Spaces ZeroGPU
+# Detect HuggingFace Spaces
 # ---------------------------------------------------------------------------
 
 IS_HF_SPACE = os.environ.get("SPACE_ID") is not None
-if IS_HF_SPACE:
-    import spaces
-    import sys
 
-    # Suppress Python 3.10 asyncio GC bug: BaseEventLoop.__del__ crashes with
-    # "Invalid file descriptor: -1" when a loop is collected by the GC.
-    _orig_unraisable = sys.unraisablehook
-    def _unraisable_hook(args):
-        if args.exc_type is ValueError and "Invalid file descriptor" in str(args.exc_value):
-            return  # silently drop — harmless Python 3.10 GC bug
-        _orig_unraisable(args)
-    sys.unraisablehook = _unraisable_hook
-
-    print("HF Space detected — model will load on first GPU call and stay cached in memory.")
+# ---------------------------------------------------------------------------
+# Persistent storage cache — survives sleep/restart on HF Spaces
+# ---------------------------------------------------------------------------
+# Enable in Space Settings → Storage (mount at /data).
+# HF_HOME env var can also be set manually in Space Settings → Variables,
+# but this block handles it automatically when /data is present.
+if IS_HF_SPACE and os.path.isdir("/data"):
+    _cache_dir = "/data/hf_cache"
+    os.makedirs(_cache_dir, exist_ok=True)
+    os.environ.setdefault("HF_HOME", _cache_dir)
+    print(f"Persistent cache active → {_cache_dir}")
+else:
+    print("HF Space detected — model will load on first call and stay cached in memory." if IS_HF_SPACE else "Local mode.")
 
 # ---------------------------------------------------------------------------
 # Emoji → descriptive text maps (fed into FLUX prompt)
@@ -107,31 +107,26 @@ def get_pipeline():
     if _pipe is not None:
         return _pipe
 
-    from diffusers import FluxPipeline
+    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
 
-    use_mps  = (not IS_HF_SPACE) and (not torch.cuda.is_available()) and torch.backends.mps.is_available()
-    # MPS has incomplete bfloat16 support — float16 is faster and uses half the memory there.
-    dtype = torch.float16 if use_mps else torch.bfloat16
-
-    print(f"Loading FLUX.1-schnell pipeline… (token={'set' if hf_token else 'NOT SET'}, dtype={dtype})")
-    _pipe = FluxPipeline.from_pretrained(
-        "black-forest-labs/FLUX.1-schnell",
+    print(f"Loading FLUX.2-klein-4B pipeline… (token={'set' if hf_token else 'NOT SET'}, dtype={dtype})")
+    _pipe = Flux2KleinPipeline.from_pretrained(
+        "black-forest-labs/FLUX.2-klein-4B",
         torch_dtype=dtype,
         token=hf_token,
     )
 
     if IS_HF_SPACE:
-        # ZeroGPU always has CUDA — don't rely on torch.cuda.is_available()
-        # which can return False outside the @spaces.GPU context
         _pipe = _pipe.to("cuda")
     elif torch.cuda.is_available():
         _pipe = _pipe.to("cuda")
     elif use_mps:
-        # Sequential CPU offload avoids MPS OOM: each transformer block is moved
-        # to MPS one at a time, then immediately returned to CPU.
-        _pipe.enable_sequential_cpu_offload()
+        _pipe = _pipe.to("mps")
     else:
         _pipe = _pipe.to("cpu")
 
@@ -144,17 +139,21 @@ def get_pipeline():
 
 def _generate(animals: list[str], places: list[str]):
     try:
+        import time
         pipe   = get_pipeline()
         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=0.0,
+            guidance_scale=guidance,
             height=512,
             width=512,
         )
+        print(f"✅ Generated in {time.time() - t0:.1f}s")
         return result.images[0], prompt
 
     except Exception as e:
@@ -168,14 +167,6 @@ def _generate(animals: list[str], places: list[str]):
 # Public API — two versions depending on environment
 # ---------------------------------------------------------------------------
 
-if IS_HF_SPACE:
-    # duration=120: covers first-run model download (~30s) + load (~10s) + generate (~5s).
-    # On subsequent calls _pipe is already loaded so only ~5s of GPU time is used.
-    @spaces.GPU(duration=120)
-    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."""
-        return _generate(animals, places)
+def generate_reward_image(animals: list[str], places: list[str]):
+    """Generate reward image — works locally (MPS/CPU) and on HF Spaces (persistent GPU)."""
+    return _generate(animals, places)

scripts/generate_dataset.py

@@ -1,21 +1,33 @@
 """
 NumZoo training dataset generator.
 
-Generates 80 images matching the NumZoo aesthetic using Gemini 3.1 Flash Image (Nano Banana 2):
+Generates 80 images matching the NumZoo aesthetic using FLUX.1-dev via the
+HuggingFace Inference API (fal-ai provider — best quality, free credits on signup).
+
 - Kawaii chibi animals with big sparkling eyes
 - Rich scene backgrounds (no white background)
 - Warm pastel palette, fairy lights, cozy props
-- Landscape 16:9, children's book illustration style
+- Square 1024×1024 output to match the app layout
 
 Output: training/image_001.jpg + training/image_001.txt (caption)
 
 Requirements:
-    pip install google-genai pillow
+    pip install huggingface_hub pillow python-dotenv
+
+Setup (HF Pro — just your existing token):
+    1. Get your HF token at https://huggingface.co/settings/tokens
+       (fine-grained, with "Make calls to Inference Providers" permission)
+    2. Add to .env:
+         HF_TOKEN=hf_...
+
+    With HF Pro your token already has credits on fal-ai, replicate, together etc.
+    No separate provider account needed.
 
 Usage:
-    export GEMINI_API_KEY=your_key
-    python scripts/generate_dataset.py
+    python scripts/generate_dataset.py              # all 80 via fal-ai (FLUX.1-dev)
+    python scripts/generate_dataset.py --count 5   # first 5 only (test run)
     python scripts/generate_dataset.py --start 41  # resume from image 41
+    python scripts/generate_dataset.py --provider hf-inference  # HF native (schnell)
 """
 
 import os
@@ -31,18 +43,16 @@ try:
 except ImportError:
     pass  # dotenv optional — can also export GEMINI_API_KEY manually
 
+# Import the canonical style string from image_generator so training captions
+# and live prompts are always identical.
+sys.path.insert(0, str(Path(__file__).parent.parent))
+from image_generator import NUMZOO_STYLE as STYLE  # noqa: E402
+
 # ---------------------------------------------------------------------------
 # 80 prompts — 8 scene categories × 10 animals
-# Style wrapper applied uniformly so the LoRA captures it as a learnable token
+# STYLE is appended to every prompt so the LoRA learns it as a trigger
 # ---------------------------------------------------------------------------
 
-STYLE = (
-    "kawaii children's book illustration, pastel anime art style, "
-    "soft painterly lighting, detailed rich background with warm fairy lights, "
-    "cozy magical atmosphere, cute chibi character with big sparkling eyes, "
-    "soft pastel color palette, highly detailed scene, no text"
-)
-
 PROMPTS: list[tuple[str, str]] = [
     # ── 1. Cozy cabin interior ──────────────────────────────────────────────
     ("cozy_cabin_01", f"a fluffy bunny curled on an armchair by a stone fireplace inside a wooden cabin, "
@@ -225,59 +235,67 @@ PROMPTS: list[tuple[str, str]] = [
 assert len(PROMPTS) == 80, f"Expected 80 prompts, got {len(PROMPTS)}"
 
 # ---------------------------------------------------------------------------
-# Generator using Gemini 3.1 Flash Image (Nano Banana 2)
+# Generator using FLUX.1-dev via HuggingFace Inference API
 # ---------------------------------------------------------------------------
 
-def generate_image(prompt: str):
-    from google import genai
-    from google.genai import types
-
-    api_key = os.environ.get("GEMINI_API_KEY")
-    client = genai.Client(api_key=api_key)
-
-    response = client.models.generate_content(
-        model="gemini-3.1-flash-image",
-        contents=[prompt],
-        config=types.GenerateContentConfig(
-            response_modalities=["IMAGE"],
-            image_config=types.ImageConfig(
-                aspect_ratio="1:1",
-                image_size="1K",
-            ),
-        ),
+# Provider → model routing:
+#   fal-ai       → FLUX.1-dev  (best quality, free credits at fal.ai)
+#   hf-inference → FLUX.1-schnell (HF free tier, lower quality)
+PROVIDERS = {
+    "fal-ai":       "black-forest-labs/FLUX.1-dev",
+    "hf-inference": "black-forest-labs/FLUX.1-schnell",
+}
+
+def generate_image(prompt: str, provider: str = "fal-ai") -> "PIL.Image.Image":
+    from huggingface_hub import InferenceClient
+
+    # HF Pro token covers all providers — no separate provider key needed
+    hf_token = os.environ.get("HF_TOKEN")
+    model    = PROVIDERS.get(provider, PROVIDERS["fal-ai"])
+    client   = InferenceClient(provider=provider, api_key=hf_token)
+
+    image = client.text_to_image(
+        prompt,
+        model=model,
+        width=1024,
+        height=1024,
     )
-
-    for part in response.parts:
-        if part.inline_data is not None:
-            return part.as_image()
-
-    raise RuntimeError("No image in response")
+    return image  # InferenceClient already returns a PIL Image
 
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--start", type=int, default=1, help="Resume from image N (1-based)")
-    parser.add_argument("--dry-run", action="store_true", help="Print prompts without generating")
+    parser.add_argument("--start",    type=int,  default=1,         help="Resume from image N (1-based)")
+    parser.add_argument("--count",    type=int,  default=None,      help="Generate at most N images then stop")
+    parser.add_argument("--provider", type=str,  default="fal-ai",  help="Inference provider: fal-ai | hf-inference")
+    parser.add_argument("--dry-run",  action="store_true",          help="Print prompts without generating")
     args = parser.parse_args()
 
-    if not os.environ.get("GEMINI_API_KEY") and not args.dry_run:
-        print("❌  Set GEMINI_API_KEY environment variable before running.")
-        print("    export GEMINI_API_KEY=your_key")
-        sys.exit(1)
+    if not args.dry_run:
+        if not os.environ.get("HF_TOKEN"):
+            print("❌  HF_TOKEN not found. Add it to .env")
+            print("    Get yours at https://huggingface.co/settings/tokens")
+            sys.exit(1)
 
     out_dir = Path(__file__).parent.parent / "training"
     out_dir.mkdir(exist_ok=True)
 
-    total = len(PROMPTS)
+    total   = len(PROMPTS)
+    end_at  = (args.start - 1 + args.count) if args.count else total  # inclusive upper bound (index)
 
-    print(f"NumZoo dataset generator — {total} images → {out_dir}")
-    print(f"Starting from image {args.start}/{total}")
+    count_label = f"{args.count} images" if args.count else f"all {total} images"
+    print(f"NumZoo dataset generator — {count_label} → {out_dir}")
+    print(f"Range: {args.start}–{min(end_at, total)} of {total}")
     print()
 
+    generated_this_run = 0
+
     for i, (name, prompt) in enumerate(PROMPTS):
         n = i + 1
         if n < args.start:
             continue
+        if n > end_at:
+            break
 
         img_path = out_dir / f"image_{n:03d}.jpg"
         txt_path = out_dir / f"image_{n:03d}.txt"
@@ -293,16 +311,17 @@ def main():
             continue
 
         try:
-            image = generate_image(prompt)
+            image = generate_image(prompt, provider=args.provider)
             image.save(img_path, "JPEG", quality=95)
             txt_path.write_text(prompt)
+            generated_this_run += 1
             print(f"         ✅  saved {img_path.name}")
         except Exception as e:
             print(f"         ❌  failed: {e}")
             time.sleep(5)  # brief pause on error before continuing
 
-    generated = len(list(out_dir.glob("*.jpg")))
-    print(f"\nDone. {generated}/{total} images in {out_dir}")
+    total_on_disk = len(list(out_dir.glob("*.jpg")))
+    print(f"\nDone. {generated_this_run} generated this run · {total_on_disk}/{total} total in {out_dir}")
 
 
 if __name__ == "__main__":