""" NumZoo training dataset generator. Generates 80 images matching the NumZoo aesthetic using FLUX.1-dev via the HuggingFace Inference API (fal-ai provider — best quality, free credits on signup). - Kawaii chibi animals with big sparkling eyes - Rich scene backgrounds (no white background) - Warm pastel palette, fairy lights, cozy props - Square 1024×1024 output to match the app layout Output: training/image_001.jpg + training/image_001.txt (caption) Requirements: pip install huggingface_hub pillow python-dotenv Setup (HF Pro — just your existing token): 1. Get your HF token at https://huggingface.co/settings/tokens (fine-grained, with "Make calls to Inference Providers" permission) 2. Add to .env: HF_TOKEN=hf_... With HF Pro your token already has credits on fal-ai, replicate, together etc. No separate provider account needed. Usage: python scripts/generate_dataset.py # all 80 via fal-ai (FLUX.1-dev) python scripts/generate_dataset.py --count 5 # first 5 only (test run) python scripts/generate_dataset.py --start 41 # resume from image 41 python scripts/generate_dataset.py --provider hf-inference # HF native (schnell) """ import os import sys import argparse import time from pathlib import Path # Load .env from project root (GEMINI_API_KEY etc.) try: from dotenv import load_dotenv load_dotenv(Path(__file__).parent.parent / ".env") except ImportError: pass # dotenv optional — can also export GEMINI_API_KEY manually # Import the canonical style string from image_generator so training captions # and live prompts are always identical. sys.path.insert(0, str(Path(__file__).parent.parent)) from image_generator import NUMZOO_STYLE as STYLE # noqa: E402 # --------------------------------------------------------------------------- # 80 prompts — 8 scene categories × 10 animals # STYLE is appended to every prompt so the LoRA learns it as a trigger # --------------------------------------------------------------------------- PROMPTS: list[tuple[str, str]] = [ # ── 1. Cozy cabin interior ────────────────────────────────────────────── ("cozy_cabin_01", f"a fluffy bunny curled on an armchair by a stone fireplace inside a wooden cabin, " f"pastel star blanket, string lights, hot cocoa on a tree-stump table, warm amber glow, {STYLE}"), ("cozy_cabin_02", f"a sleepy kitten reading a tiny book in a cozy log cabin, " f"fairy lights strung above wooden shelves with colorful jars, fireplace crackling, " f"pastel pink and lavender tones, {STYLE}"), ("cozy_cabin_03", f"a baby bear in a rainbow scarf sitting at a round wooden table with pancakes " f"inside a cabin, soft morning light through an arched window, {STYLE}"), ("cozy_cabin_04", f"a hedgehog wrapped in a cloud-print blanket watching raindrops on a cabin window, " f"string lights reflected in the glass, soft purple tones, {STYLE}"), ("cozy_cabin_05", f"a baby fox and a bunny sharing hot tea by a fireplace in a treehouse cabin, " f"wooden bookshelves, pastel mugs, twinkle lights, {STYLE}"), ("cozy_cabin_06", f"a tiny owl perched on a branch reading a glowing book inside a lantern-lit cabin, " f"warm golden light, mossy windowsill with flowers outside, {STYLE}"), ("cozy_cabin_07", f"a baby deer sitting on a plaid cushion inside a snug wooden cabin, " f"potted cactus and tiny succulents on shelves, fairy lights overhead, {STYLE}"), ("cozy_cabin_08", f"a raccoon stirring a sparkling potion in a cozy cabin kitchen, " f"mason jars of colorful ingredients, star-shaped window, warm light, {STYLE}"), ("cozy_cabin_09", f"a baby penguin in a striped scarf napping on a pastel sofa next to a glowing fireplace, " f"snow visible through the cabin window, {STYLE}"), ("cozy_cabin_10", f"a squirrel decorating a tiny pine tree inside a cabin with star ornaments and ribbon, " f"fairy lights everywhere, soft pink tones, {STYLE}"), # ── 2. Enchanted forest — day ─────────────────────────────────────────── ("forest_day_01", f"a bunny in a daisy field inside an enchanted forest, " f"sunbeams through tall trees, floating sparkles, colorful mushrooms, {STYLE}"), ("forest_day_02", f"a kitten playing with a butterfly on a mossy log in a sunlit magical forest, " f"wildflowers in pastel pink and lilac, {STYLE}"), ("forest_day_03", f"a baby bear discovering a tiny glowing door in a tree trunk in an enchanted forest, " f"soft golden afternoon light, flower garlands, {STYLE}"), ("forest_day_04", f"a hedgehog gathering strawberries in a forest clearing, " f"pastel checkered picnic blanket, wicker basket, soft dappled sunlight, {STYLE}"), ("forest_day_05", f"a baby fox sitting on a giant mushroom in a magical forest, " f"fairy dust floating around, colorful birds in branches above, {STYLE}"), ("forest_day_06", f"an owl in a graduation cap perched on a tree branch in a sunlit forest, " f"hanging star lanterns, flowers blooming below, {STYLE}"), ("forest_day_07", f"a baby deer drinking from a sparkling stream in a lush forest, " f"rainbow trout visible in the water, cherry blossoms drifting, {STYLE}"), ("forest_day_08", f"a raccoon painting a tiny canvas in a forest studio made of logs and branches, " f"colorful paint pots, sunlight streaming in, {STYLE}"), ("forest_day_09", f"a baby penguin waddling through a magical spring forest with tulips taller than it, " f"pastel butterflies, {STYLE}"), ("forest_day_10", f"a squirrel and a bunny having a tea party on a giant leaf in the forest, " f"acorn cups, mushroom chairs, sunlit clearing, {STYLE}"), # ── 3. Magical forest — night / campfire ──────────────────────────────── ("forest_night_01", f"a bunny and a kitten roasting marshmallows over a small campfire in a night forest, " f"crescent moon, fireflies, string lights in the trees, {STYLE}"), ("forest_night_02", f"a baby bear stargazing in a moonlit forest clearing, lying on a cloud-print blanket, " f"glowing constellations above, lantern beside, {STYLE}"), ("forest_night_03", f"a baby fox with a tiny lantern walking a forest path at night, " f"glowing mushrooms lining the path, crescent moon, {STYLE}"), ("forest_night_04", f"a hedgehog sleeping under a luminescent mushroom at night in the forest, " f"fireflies around, purple and blue tones, soft moonlight, {STYLE}"), ("forest_night_05", f"an owl sitting on a branch under the stars playing a tiny ukulele for forest friends " f"gathered below, night forest, fairy lights, {STYLE}"), ("forest_night_06", f"a baby deer watching shooting stars from a mossy hilltop at night, " f"forest silhouette below, soft purple sky, sparkles, {STYLE}"), ("forest_night_07", f"a raccoon cooking s'mores at a campfire in a night forest, " f"camping tent with fairy lights, moon and stars above, {STYLE}"), ("forest_night_08", f"a baby penguin looking at its reflection in a glowing moonlit pond in a night forest, " f"lily pads with tiny lights, {STYLE}"), ("forest_night_09", f"a squirrel telling stories around a campfire to bunny and kitten, night forest, " f"warm firelight, stars through the tree canopy, {STYLE}"), ("forest_night_10", f"a kitten catching fireflies in a jar at the edge of a glowing forest at dusk, " f"silhouette trees, warm golden light, {STYLE}"), # ── 4. Treehouse ──────────────────────────────────────────────────────── ("treehouse_01", f"a bunny inside a cozy treehouse at sunset, rope bridge entrance, " f"flower window box, warm light glowing inside, {STYLE}"), ("treehouse_02", f"a kitten at the door of a treehouse with a welcome sign, " f"string lights along the railing, forest view at golden hour, {STYLE}"), ("treehouse_03", f"a baby bear and a hedgehog playing board games on the treehouse porch, " f"paper lanterns, leafy canopy above, {STYLE}"), ("treehouse_04", f"a baby fox gazing at the stars through a treehouse telescope at night, " f"spiral staircase, crescent moon, {STYLE}"), ("treehouse_05", f"an owl's cozy treehouse library with tiny glowing books and a reading nook, " f"ivy covered walls, daytime, {STYLE}"), ("treehouse_06", f"a baby bear watering flower boxes on a treehouse balcony, " f"bees and butterflies, golden afternoon light, leafy forest below, {STYLE}"), ("treehouse_07", f"a hedgehog and a squirrel stringing fairy lights around a treehouse at dusk, " f"warm glow, birds watching from nearby branches, {STYLE}"), ("treehouse_08", f"a baby fox sliding down a spiral slide from a treehouse into a pile of leaves, " f"autumn forest, rope ladder, golden light, {STYLE}"), ("treehouse_09", f"a baby deer reading a glowing map inside a treehouse at night, " f"star-shaped windows, moonlight, cozy lanterns, {STYLE}"), ("treehouse_10", f"a raccoon and a bunny cooking soup in a tiny treehouse kitchen, " f"steam rising, herbs hanging, sunlight through small windows, {STYLE}"), # ── 5. Garden / meadow — day ──────────────────────────────────────────── ("garden_01", f"a bunny watering oversized tulips in a magical garden, " f"pastel blue watering can, butterflies, soft morning light, {STYLE}"), ("garden_02", f"a kitten napping in a flower hammock in a sunny garden, " f"roses and lavender around, bees buzzing, {STYLE}"), ("garden_03", f"a baby bear picking giant strawberries in a garden, " f"ladybugs on leaves, warm afternoon sun, picket fence, {STYLE}"), ("garden_04", f"a hedgehog rolling through a field of sunflowers, " f"pastel sky, birds above, soft breeze, {STYLE}"), ("garden_05", f"a baby fox blowing dandelion seeds in a meadow, " f"seeds glowing like stars, cherry trees in background, {STYLE}"), ("garden_06", f"a baby deer in a rain boots splashing in a puddle in a garden after rain, " f"rainbow in the sky, wet flowers, {STYLE}"), ("garden_07", f"a raccoon painting a garden fence with rainbow stripes, " f"paint buckets, butterflies landing on fresh paint, {STYLE}"), ("garden_08", f"a baby penguin planting seeds in a tiny garden plot, " f"little labeled plant markers, pastel gloves and apron, {STYLE}"), ("garden_09", f"a squirrel collecting colorful autumn leaves in a basket in a golden meadow, " f"maple trees, soft warm light, {STYLE}"), ("garden_10", f"a bunny and kitten flying a star-shaped kite in an open meadow, " f"puffy clouds, wildflowers below, {STYLE}"), # ── 6. Beach / waterfront ─────────────────────────────────────────────── ("beach_01", f"a bunny building a sand castle on a pastel beach at sunset, " f"sea shells, gentle waves, pink and orange sky, {STYLE}"), ("beach_02", f"a kitten surfing a tiny wave on a pastel board, " f"tropical beach, palm trees, clear turquoise water, {STYLE}"), ("beach_03", f"a baby bear making friends with a crab on the beach, " f"rock pools with anemones, soft summer light, {STYLE}"), ("beach_04", f"a baby fox asleep in a hammock strung between two palm trees on a beach, " f"coconuts, ocean at dusk, {STYLE}"), ("beach_05", f"a baby penguin waddling along a sandy beach collecting glowing starfish at night, " f"bioluminescent waves, crescent moon, {STYLE}"), ("beach_06", f"a bunny and a hedgehog sharing a watermelon slice on a colorful beach towel, " f"sunny beach, parasol, gentle waves, {STYLE}"), ("beach_07", f"a baby fox building a sand boat on the beach, " f"pebble sail, seagulls, golden sunset, {STYLE}"), ("beach_08", f"a baby bear snorkeling in a clear lagoon, colorful fish and coral below, " f"tropical beach, bright midday light, {STYLE}"), ("beach_09", f"a kitten watching the sunset from a dock over calm ocean water, " f"orange and pink sky, lantern beside, soft waves, {STYLE}"), ("beach_10", f"a baby deer and a squirrel collecting shells in tide pools at low tide, " f"pastel sunset, starfish, {STYLE}"), # ── 7. Snowy winter ────────────────────────────────────────────────────── ("snow_01", f"a bunny building a snowman in a snowy forest clearing, " f"carrot nose, scarf, snowflakes falling, soft pink sky, {STYLE}"), ("snow_02", f"a kitten ice skating on a frozen pond in a winter forest, " f"snowflakes, pine trees with snow caps, fairy lights on branches, {STYLE}"), ("snow_03", f"a baby bear sledding down a snowy hill at night, " f"moonlight, sparkling snow, cozy village lights below, {STYLE}"), ("snow_04", f"a hedgehog warming up by a bonfire in the snow, " f"hot chocolate in tiny paws, snowflakes floating, pine forest, {STYLE}"), ("snow_05", f"a baby fox wrapped in a giant scarf making snow angels in a winter field, " f"snowflakes, soft blue and white tones, {STYLE}"), ("snow_06", f"an owl delivering a tiny letter through falling snow to a cabin, " f"warm window glow, snow-covered trees, {STYLE}"), ("snow_07", f"a baby deer with snowflakes on its nose in a quiet snowy forest, " f"bare tree branches with icicles, soft grey-blue light, {STYLE}"), ("snow_08", f"a raccoon and a squirrel having a snowball fight in a park, " f"snow-covered benches, winter lanterns, {STYLE}"), ("snow_09", f"a baby penguin sliding down an icy path to the frozen ocean, " f"auroras in the sky, soft blue and green glow, {STYLE}"), ("snow_10", f"a bunny and kitten sharing a pastel umbrella in gently falling snow, " f"fairy lights on trees, warm glow from a café window, {STYLE}"), # ── 8. Autumn / rainy day ──────────────────────────────────────────────── ("autumn_01", f"a bunny jumping in a pile of colorful autumn leaves in a park, " f"orange and red maples, soft afternoon sun, {STYLE}"), ("autumn_02", f"a kitten at a rainy window with a cup of tea, " f"raindrops on glass, cozy sweater, autumn forest visible outside, {STYLE}"), ("autumn_03", f"a baby bear in a pumpkin patch at golden hour, " f"scarecrow, autumn leaves drifting, {STYLE}"), ("autumn_04", f"a hedgehog foraging mushrooms in an autumn forest, " f"golden light through orange leaves, red berries, {STYLE}"), ("autumn_05", f"a baby fox jumping through autumn leaves in a red-orange forest, " f"bokeh light, warm amber tones, {STYLE}"), ("autumn_06", f"an owl in a cozy raincoat standing under a colorful mushroom in the rain, " f"puddles reflecting fairy lights, {STYLE}"), ("autumn_07", f"a baby deer watching the first autumn rain from under a giant leaf, " f"golden forest background, raindrops sparkle, {STYLE}"), ("autumn_08", f"a raccoon selling hot apple cider from a tiny wooden cart in an autumn market, " f"string lights, orange and red leaves, {STYLE}"), ("autumn_09", f"a squirrel cozy in a hollow tree watching autumn rain, " f"acorn collection visible inside, warm golden glow, {STYLE}"), ("autumn_10", f"a bunny and a baby fox under a giant colorful umbrella in a rainy autumn street, " f"café lights, reflections in puddles, {STYLE}"), ] assert len(PROMPTS) == 80, f"Expected 80 prompts, got {len(PROMPTS)}" # --------------------------------------------------------------------------- # Generator using FLUX.1-dev via HuggingFace Inference API # --------------------------------------------------------------------------- # Provider → model routing: # fal-ai → FLUX.1-dev (best quality, free credits at fal.ai) # hf-inference → FLUX.1-schnell (HF free tier, lower quality) PROVIDERS = { "fal-ai": "black-forest-labs/FLUX.1-dev", "hf-inference": "black-forest-labs/FLUX.1-schnell", } def generate_image(prompt: str, provider: str = "fal-ai") -> "PIL.Image.Image": from huggingface_hub import InferenceClient # HF Pro token covers all providers — no separate provider key needed hf_token = os.environ.get("HF_TOKEN") model = PROVIDERS.get(provider, PROVIDERS["fal-ai"]) client = InferenceClient(provider=provider, api_key=hf_token) image = client.text_to_image( prompt, model=model, width=1024, height=1024, ) return image # InferenceClient already returns a PIL Image def main(): parser = argparse.ArgumentParser() parser.add_argument("--start", type=int, default=1, help="Resume from image N (1-based)") parser.add_argument("--count", type=int, default=None, help="Generate at most N images then stop") parser.add_argument("--provider", type=str, default="fal-ai", help="Inference provider: fal-ai | hf-inference") parser.add_argument("--dry-run", action="store_true", help="Print prompts without generating") args = parser.parse_args() if not args.dry_run: if not os.environ.get("HF_TOKEN"): print("❌ HF_TOKEN not found. Add it to .env") print(" Get yours at https://huggingface.co/settings/tokens") sys.exit(1) out_dir = Path(__file__).parent.parent / "training" out_dir.mkdir(exist_ok=True) total = len(PROMPTS) end_at = (args.start - 1 + args.count) if args.count else total # inclusive upper bound (index) count_label = f"{args.count} images" if args.count else f"all {total} images" print(f"NumZoo dataset generator — {count_label} → {out_dir}") print(f"Range: {args.start}–{min(end_at, total)} of {total}") print() generated_this_run = 0 for i, (name, prompt) in enumerate(PROMPTS): n = i + 1 if n < args.start: continue if n > end_at: break img_path = out_dir / f"image_{n:03d}.jpg" txt_path = out_dir / f"image_{n:03d}.txt" if img_path.exists(): print(f"[{n:02d}/{total}] ⏭ {name} — already exists, skipping") continue print(f"[{n:02d}/{total}] 🎨 {name}") if args.dry_run: print(f" {prompt[:120]}…") continue try: image = generate_image(prompt, provider=args.provider) image.save(img_path, "JPEG", quality=95) txt_path.write_text(prompt) generated_this_run += 1 print(f" ✅ saved {img_path.name}") except Exception as e: print(f" ❌ failed: {e}") time.sleep(5) # brief pause on error before continuing total_on_disk = len(list(out_dir.glob("*.jpg"))) print(f"\nDone. {generated_this_run} generated this run · {total_on_disk}/{total} total in {out_dir}") if __name__ == "__main__": main()