Remove startup model pre-caching — let from_pretrained cache in container

The persistent storage placeholder-file issue was fundamental: HF Spaces
git-backed storage creates FILE placeholders at every intermediate directory
path (blobs/, refs/, snapshots/, .locks/) on each mount. Removing them
destroyed access to the cached content behind them, forcing a full 57.8 GB
re-download on every cold start — the opposite of the intended behaviour.

New approach: remove snapshot_download at startup entirely. from_pretrained
downloads to the container's local ~/.cache/huggingface/hub/ on first use.
The Python process (and _pipe global) survive between ZeroGPU GPU-slot
releases, so the model loads once per container lifetime. Increased
@spaces.GPU duration to 120s to cover first-run download + load + generate.

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

image_generator.py

@@ -2,72 +2,6 @@ import random
 import os
 import torch
 
-# ---------------------------------------------------------------------------
-# Persistent storage cache (HF Spaces Pro — /data survives rebuilds)
-# ---------------------------------------------------------------------------
-_PERSISTENT_CACHE = "/data/hf_cache"
-
-
-def _repair_hf_cache(cache_root: str) -> None:
-    """
-    HF Spaces persistent storage mounts placeholder FILES at every path that
-    should be a directory. This breaks huggingface_hub's cache which needs:
-      {cache_root}/models--{org}--{model}/
-      {cache_root}/models--{org}--{model}/blobs/
-      {cache_root}/models--{org}--{model}/refs/
-      {cache_root}/models--{org}--{model}/snapshots/
-      {cache_root}/models--{org}--{model}/snapshots/{40-char-hash}/
-      {cache_root}/.locks/models--{org}--{model}/   ← lock dir used during download
-
-    Walk all known directory levels and unlink any non-directory obstacles.
-    Real content files (safetensors, configs, symlinks) are left untouched.
-    """
-    def _fix(path: str) -> bool:
-        """Remove path if it's not a real directory. Return True if it's a dir."""
-        if os.path.lexists(path) and not os.path.isdir(path):
-            print(f"  fix: removing non-dir placeholder at {path}")
-            os.unlink(path)
-            return False
-        return os.path.isdir(path)
-
-    _fix(cache_root)
-    if not os.path.isdir(cache_root):
-        return
-
-    # Fix models--* tree (blobs, refs, snapshots)
-    for model_name in os.listdir(cache_root):
-        if not model_name.startswith("models--"):
-            continue
-        model_dir = os.path.join(cache_root, model_name)
-        if not _fix(model_dir):
-            continue
-        for sub in ("blobs", "refs", "snapshots"):
-            sub_path = os.path.join(model_dir, sub)
-            if not _fix(sub_path):
-                continue
-            if sub == "snapshots":
-                for h in os.listdir(sub_path):
-                    h_path = os.path.join(sub_path, h)
-                    if not os.path.islink(h_path):   # symlinks are fine
-                        _fix(h_path)
-
-    # Fix .locks tree — huggingface_hub creates per-blob lock files here
-    locks_root = os.path.join(cache_root, ".locks")
-    if _fix(locks_root) and os.path.isdir(locks_root):
-        for model_name in os.listdir(locks_root):
-            if model_name.startswith("models--"):
-                _fix(os.path.join(locks_root, model_name))
-
-
-if os.path.isdir("/data"):
-    if os.path.lexists(_PERSISTENT_CACHE) and not os.path.isdir(_PERSISTENT_CACHE):
-        os.unlink(_PERSISTENT_CACHE)
-    os.makedirs(_PERSISTENT_CACHE, exist_ok=True)
-    _repair_hf_cache(_PERSISTENT_CACHE)
-    os.environ["HF_HUB_CACHE"]         = _PERSISTENT_CACHE
-    os.environ["HUGGINGFACE_HUB_CACHE"] = _PERSISTENT_CACHE
-    print(f"Using persistent cache: {_PERSISTENT_CACHE}")
-
 # Patch for torch < 2.4 which lacks torch.xpu (required by diffusers >= 0.30)
 if not hasattr(torch, "xpu"):
     class _MockXPU:
@@ -87,12 +21,10 @@ 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
+    import sys
 
-    # ── Suppress Python 3.10 asyncio GC bug ──────────────────────────────
-    # BaseEventLoop.__del__ crashes with "Invalid file descriptor: -1" when
-    # an event loop is collected by the GC. sys.unraisablehook is the correct
-    # Python 3.8+ way to intercept these GC-triggered exceptions.
+    # 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):
@@ -100,51 +32,7 @@ if IS_HF_SPACE:
         _orig_unraisable(args)
     sys.unraisablehook = _unraisable_hook
 
-    # Download model files to persistent storage at startup.
-    # /data/hf_cache survives rebuilds → downloads once, loads instantly after.
-    _hf_token  = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
-    _cache_dir = os.environ.get("HF_HUB_CACHE")
-    _model_cache = os.path.join(_cache_dir, "models--black-forest-labs--FLUX.1-schnell")
-    _refs_main   = os.path.join(_model_cache, "refs", "main")
-    print(f"Pre-caching FLUX.1-schnell → {_cache_dir} (token={'set' if _hf_token else 'NOT SET'})")
-    # snapshot_download is idempotent: it checks blob hashes and only fetches
-    # what's missing, so if the cache is intact it completes in seconds.
-    print(f"  model dir : {'dir' if os.path.isdir(_model_cache) else 'absent/file'}")
-    print(f"  refs/main : {'present' if os.path.isfile(_refs_main) else 'absent'}")
-    try:
-
-        from huggingface_hub import snapshot_download
-
-        _result = {}
-        def _download():
-            # Own an explicit event loop so it can be closed cleanly before the
-            # thread exits — prevents the GC from seeing an open loop.
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-            try:
-                snapshot_download(
-                    "black-forest-labs/FLUX.1-schnell",
-                    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 Exception:
-                    pass
-
-        _t = threading.Thread(target=_download, daemon=True)
-        _t.start()
-        _t.join()
-        if "error" in _result:
-            raise _result["error"]
-        print("✅ Model files ready in persistent cache")
-    except Exception as e:
-        print(f"⚠️  Pre-cache warning (will retry at generation time): {e}")
+    print("HF Space detected — model will load on first GPU call and stay cached in memory.")
 
 # ---------------------------------------------------------------------------
 # Emoji → descriptive text maps (fed into FLUX prompt)
@@ -217,7 +105,7 @@ def build_prompt(streak: int, animals: list[str], places: list[str]) -> str:
     return prompt
 
 # ---------------------------------------------------------------------------
-# Pipeline loader (cached)
+# Pipeline loader (cached globally — survives between ZeroGPU calls)
 # ---------------------------------------------------------------------------
 
 _pipe = None
@@ -282,7 +170,9 @@ def _generate(streak: int, animals: list[str], places: list[str]):
 # ---------------------------------------------------------------------------
 
 if IS_HF_SPACE:
-    @spaces.GPU(duration=60)
+    # 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(streak: int, animals: list[str], places: list[str]):
         """Generate reward image on HF Spaces ZeroGPU."""
         return _generate(streak, animals, places)