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 ZeroGPU
# ---------------------------------------------------------------------------

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.")

# ---------------------------------------------------------------------------
# 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 FluxPipeline

    hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")

    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",
        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()
    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:
        pipe   = get_pipeline()
        prompt = build_prompt(animals, places)
        print(f"Generating | prompt: {prompt}")

        result = pipe(
            prompt=prompt,
            num_inference_steps=4,
            guidance_scale=0.0,
            height=512,
            width=512,
        )
        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:
    # 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)