image_generator.py

import random
import os
import torch

# Patch for torch < 2.4 which lacks torch.xpu (required by diffusers >= 0.30)
if not hasattr(torch, "xpu"):
    class _MockXPU:
        is_available    = staticmethod(lambda: False)
        device_count    = staticmethod(lambda: 0)
        empty_cache     = staticmethod(lambda: None)
        manual_seed     = staticmethod(lambda seed: None)
        reset_peak_memory_stats = staticmethod(lambda: None)
        max_memory_allocated    = staticmethod(lambda: 0)
        synchronize     = staticmethod(lambda: None)
    torch.xpu = _MockXPU()

# ---------------------------------------------------------------------------
# Detect HuggingFace Spaces
# ---------------------------------------------------------------------------

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)
    _orig_unraisable = sys.unraisablehook
    def _unraisable_hook(args):
        if args.exc_type is ValueError and "Invalid file descriptor" in str(args.exc_value):
            return
        _orig_unraisable(args)
    sys.unraisablehook = _unraisable_hook

# ---------------------------------------------------------------------------
# Persistent storage cache — survives sleep/restart on HF Spaces
# ---------------------------------------------------------------------------
# The /data mount repeatedly gets poisoned: partial downloads leave files where
# huggingface_hub/xet expect directories ([Errno 20] Not a directory), and the
# cache root itself can end up as a file with I/O errors. This block is
# bulletproof: any failure setting up /data falls back to the ephemeral
# container cache (slower cold start, but always works) and NEVER crashes the
# import — a dead import takes the whole app down.
_MODEL_ID  = "black-forest-labs/FLUX.2-klein-4B"
_MODEL_DIR = "models--black-forest-labs--FLUX.2-klein-4B"


def _force_remove(path):
    """Remove path whether it's a file, dir, or broken — never raises."""
    import shutil
    try:
        if os.path.isdir(path) and not os.path.islink(path):
            shutil.rmtree(path)
        elif os.path.exists(path) or os.path.islink(path):
            os.remove(path)
    except Exception as e:
        print(f"⚠️  Could not remove {path}: {e}")


def _setup_persistent_cache():
    """Point HF cache at /data if usable. Returns True if persistence is active."""
    cache_dir = "/data/hf_cache_v4"   # bump path to escape any poisoned older cache
    hub_dir   = os.path.join(cache_dir, "hub")

    # Probe: can we create a clean directory tree at cache_dir? If the path is a
    # poisoned file or the mount errors, bail out to ephemeral caching.
    try:
        if os.path.exists(cache_dir) or os.path.islink(cache_dir):
            # Validate existing cache: a real model snapshot with model_index.json
            snaps = os.path.join(hub_dir, _MODEL_DIR, "snapshots")
            valid = os.path.isdir(snaps) and any(
                os.path.isfile(os.path.join(snaps, s, "model_index.json"))
                for s in os.listdir(snaps)
            ) if os.path.isdir(snaps) else False
            if not valid:
                print(f"Cache at {cache_dir} is incomplete/poisoned — wiping")
                _force_remove(cache_dir)
        os.makedirs(hub_dir, exist_ok=True)
    except Exception as e:
        print(f"⚠️  /data cache unusable ({e}) — using ephemeral container cache")
        return False

    os.environ["HF_HOME"]      = cache_dir
    os.environ["HF_HUB_CACHE"] = hub_dir
    print(f"Persistent cache active → {cache_dir}")
    return True


def _is_model_cached():
    snaps = os.path.join(os.environ.get("HF_HUB_CACHE", ""), _MODEL_DIR, "snapshots")
    if not os.path.isdir(snaps):
        return False
    return any(
        os.path.isfile(os.path.join(snaps, s, "model_index.json"))
        for s in os.listdir(snaps)
    )


def _predownload_model():
    """Blocking pre-download so the first @spaces.GPU call is fast. Never raises."""
    if _is_model_cached():
        print(f"✅ FLUX.2-klein-4B already cached — skipping download")
        return
    print(f"Downloading FLUX.2-klein-4B…")
    try:
        from huggingface_hub import snapshot_download
        hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
        result = {}
        def _download():
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                snapshot_download(
                    _MODEL_ID,
                    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}")


if IS_HF_SPACE:
    if os.path.isdir("/data"):
        _setup_persistent_cache()   # falls back to ephemeral on any failure
    _predownload_model()
else:
    print("Local mode.")

# ---------------------------------------------------------------------------
# Emoji → descriptive text maps (fed into FLUX prompt)
# ---------------------------------------------------------------------------

ANIMAL_MAP: dict[str, str] = {
    "🐶": "puppy",        "🐱": "kitten",       "🐰": "bunny",
    "🦊": "baby fox",     "🐼": "baby panda",   "🐨": "baby koala",
    "🦁": "baby lion",    "🐯": "baby tiger",   "🐸": "baby frog",
    "🐧": "baby penguin", "🦋": "butterfly",    "🦄": "unicorn",
}

PLACE_MAP: dict[str, str] = {
    "🌊": "on a sunny beach with ocean waves",
    "🏔️": "on a snowy mountain top",
    "🌸": "in a cherry blossom garden",
    "🌈": "under a rainbow",
    "🌙": "on a glowing crescent moon",
    "⭐": "surrounded by sparkling stars",
    "🌴": "on a tropical island",
    "🏡": "in a cosy cottage garden",
    "🌺": "in a field of tropical flowers",
    "🍄": "in an enchanted mushroom forest",
}

# Shared style suffix — must stay in sync with STYLE in scripts/generate_dataset.py.
# This exact string appears in every training caption so the LoRA learns it as a trigger.
NUMZOO_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"
)

# ---------------------------------------------------------------------------
# Prompt builder
# ---------------------------------------------------------------------------

def build_prompt(animals: list[str], places: list[str]) -> str:
    # Use ALL selected animals and places (up to 3 each)
    animal_list = animals[:3] if animals else [random.choice(list(ANIMAL_MAP))]
    place_list  = places[:3]  if places  else [random.choice(list(PLACE_MAP))]

    # Build animal text: "puppy", "puppy and bunny", "puppy, bunny and kitten"
    animal_names = [ANIMAL_MAP.get(a, "bunny") for a in animal_list]
    if len(animal_names) == 1:
        animal_text = animal_names[0]
    elif len(animal_names) == 2:
        animal_text = f"{animal_names[0]} and {animal_names[1]}"
    else:
        animal_text = f"{animal_names[0]}, {animal_names[1]} and {animal_names[2]}"

    # Build place text: "in a garden", "in a garden and under a rainbow", ...
    place_texts = [PLACE_MAP.get(p, "in a magical garden") for p in place_list]
    if len(place_texts) == 1:
        place_text = place_texts[0]
    elif len(place_texts) == 2:
        place_text = f"{place_texts[0]} and {place_texts[1]}"
    else:
        place_text = f"{place_texts[0]}, {place_texts[1]} and {place_texts[2]}"

    return f"A cute {animal_text} {place_text}, {NUMZOO_STYLE}"

# ---------------------------------------------------------------------------
# Pipeline loader (cached globally — survives between ZeroGPU calls)
# ---------------------------------------------------------------------------

_pipe = None

def get_pipeline():
    global _pipe
    if _pipe is not None:
        return _pipe

    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})")
    _pipe = Flux2KleinPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-klein-4B",
        torch_dtype=dtype,
        token=hf_token,
    )

    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

# ---------------------------------------------------------------------------
# Core generation (always wrapped in try/except)
# ---------------------------------------------------------------------------

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=guidance,
            height=512,
            width=512,
        )
        print(f"✅ Generated in {time.time() - t0:.1f}s")
        return result.images[0], prompt

    except Exception as e:
        import traceback
        print(f"[image_generator] ❌ generation failed: {e}")
        print(traceback.format_exc())
        return None, str(e)


# ---------------------------------------------------------------------------
# Public API — two versions depending on environment
# ---------------------------------------------------------------------------

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)