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 # --------------------------------------------------------------------------- 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}") # Clean up any stale placeholder files left by a previous failed download # (huggingface_hub creates placeholder files that block re-download) _model_cache = os.path.join(_cache_dir, "models--black-forest-labs--FLUX.2-klein-4B") for _placeholder in [_model_cache, f"{_model_cache}/blobs", f"{_model_cache}/refs", f"{_model_cache}/snapshots"]: if os.path.isfile(_placeholder): os.remove(_placeholder) print(f"Removed stale placeholder: {_placeholder}") 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: 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( "black-forest-labs/FLUX.2-klein-4B", 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("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) # --------------------------------------------------------------------------- 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)