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)