Restore startup pre-download for FLUX.2-klein-4B on HF Spaces

Model is downloaded to /data/hf_cache at startup (once, then cached).
Subsequent restarts skip the download — cold start is fast.

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

image_generator.py

@@ -31,8 +31,47 @@ if IS_HF_SPACE and os.path.isdir("/data"):
     os.makedirs(_cache_dir, exist_ok=True)
     os.environ.setdefault("HF_HOME", _cache_dir)
     print(f"Persistent cache active → {_cache_dir}")
+
+    # Pre-download model at startup so the first @spaces.GPU call is fast.
+    # /data survives sleep/restart — download only happens once.
+    import threading, asyncio
+    from huggingface_hub import snapshot_download
+
+    _hf_token    = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+    _model_id    = "black-forest-labs/FLUX.2-klein-4B"
+    _model_cache = os.path.join(_cache_dir, "models--black-forest-labs--FLUX.2-klein-4B")
+
+    if os.path.isdir(_model_cache):
+        print(f"✅ FLUX.2-klein-4B already cached — skipping download")
+    else:
+        print(f"Downloading FLUX.2-klein-4B to persistent cache…")
+        try:
+            _result = {}
+            def _download():
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+                try:
+                    snapshot_download(
+                        _model_id,
+                        cache_dir=_cache_dir,
+                        token=_hf_token,
+                        ignore_patterns=["*.msgpack", "*.h5", "flax_model*"],
+                    )
+                    _result["ok"] = True
+                except Exception as e:
+                    _result["error"] = e
+                finally:
+                    try: loop.close()
+                    except: pass
+            _t = threading.Thread(target=_download, daemon=True)
+            _t.start(); _t.join()
+            if "error" in _result:
+                raise _result["error"]
+            print("✅ FLUX.2-klein-4B cached successfully")
+        except Exception as e:
+            print(f"⚠️  Pre-cache warning (will retry at generation time): {e}")
 else:
-    print("HF Space detected — model will load on first call and stay cached in memory." if IS_HF_SPACE else "Local mode.")
+    print("Local mode." if not IS_HF_SPACE else "HF Space — no /data mount, model loads on first call.")
 
 # ---------------------------------------------------------------------------
 # Emoji → descriptive text maps (fed into FLUX prompt)