Add training dataset generator + fix MPS dtype for local inference

- scripts/generate_dataset.py: generates 80 kawaii scene images using
Gemini 3.1 Flash Image (Nano Banana 2) — 8 categories × 10 animals,
square 1:1 output, supports --start N resume and --dry-run preview,
loads GEMINI_API_KEY from .env via python-dotenv
- image_generator.py: use float16 on MPS (bfloat16 unsupported) and
enable_sequential_cpu_offload to avoid OOM on local Apple Silicon

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

image_generator.py

@@ -118,10 +118,15 @@ def get_pipeline():
     from diffusers import FluxPipeline
 
     hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
-    print(f"Loading FLUX.1-schnell pipeline… (token={'set' if hf_token else 'NOT SET'})")
+
+    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=torch.bfloat16,
+        torch_dtype=dtype,
         token=hf_token,
     )
 
@@ -131,8 +136,10 @@ def get_pipeline():
         _pipe = _pipe.to("cuda")
     elif torch.cuda.is_available():
         _pipe = _pipe.to("cuda")
-    elif torch.backends.mps.is_available():
-        _pipe = _pipe.to("mps")
+    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")
 

scripts/generate_dataset.py

@@ -0,0 +1,309 @@
+"""
+NumZoo training dataset generator.
+
+Generates 80 images matching the NumZoo aesthetic using Gemini 3.1 Flash Image (Nano Banana 2):
+- Kawaii chibi animals with big sparkling eyes
+- Rich scene backgrounds (no white background)
+- Warm pastel palette, fairy lights, cozy props
+- Landscape 16:9, children's book illustration style
+
+Output: training/image_001.jpg + training/image_001.txt (caption)
+
+Requirements:
+    pip install google-genai pillow
+
+Usage:
+    export GEMINI_API_KEY=your_key
+    python scripts/generate_dataset.py
+    python scripts/generate_dataset.py --start 41  # resume from image 41
+"""
+
+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
+
+# ---------------------------------------------------------------------------
+# 80 prompts — 8 scene categories × 10 animals
+# Style wrapper applied uniformly so the LoRA captures it as a learnable token
+# ---------------------------------------------------------------------------
+
+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"
+)
+
+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 Gemini 3.1 Flash Image (Nano Banana 2)
+# ---------------------------------------------------------------------------
+
+def generate_image(prompt: str):
+    from google import genai
+    from google.genai import types
+
+    api_key = os.environ.get("GEMINI_API_KEY")
+    client = genai.Client(api_key=api_key)
+
+    response = client.models.generate_content(
+        model="gemini-3.1-flash-image",
+        contents=[prompt],
+        config=types.GenerateContentConfig(
+            response_modalities=["IMAGE"],
+            image_config=types.ImageConfig(
+                aspect_ratio="1:1",
+                image_size="1K",
+            ),
+        ),
+    )
+
+    for part in response.parts:
+        if part.inline_data is not None:
+            return part.as_image()
+
+    raise RuntimeError("No image in response")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--start", type=int, default=1, help="Resume from image N (1-based)")
+    parser.add_argument("--dry-run", action="store_true", help="Print prompts without generating")
+    args = parser.parse_args()
+
+    if not os.environ.get("GEMINI_API_KEY") and not args.dry_run:
+        print("❌  Set GEMINI_API_KEY environment variable before running.")
+        print("    export GEMINI_API_KEY=your_key")
+        sys.exit(1)
+
+    out_dir = Path(__file__).parent.parent / "training"
+    out_dir.mkdir(exist_ok=True)
+
+    total = len(PROMPTS)
+
+    print(f"NumZoo dataset generator — {total} images → {out_dir}")
+    print(f"Starting from image {args.start}/{total}")
+    print()
+
+    for i, (name, prompt) in enumerate(PROMPTS):
+        n = i + 1
+        if n < args.start:
+            continue
+
+        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)
+            image.save(img_path, "JPEG", quality=95)
+            txt_path.write_text(prompt)
+            print(f"         ✅  saved {img_path.name}")
+        except Exception as e:
+            print(f"         ❌  failed: {e}")
+            time.sleep(5)  # brief pause on error before continuing
+
+    generated = len(list(out_dir.glob("*.jpg")))
+    print(f"\nDone. {generated}/{total} images in {out_dir}")
+
+
+if __name__ == "__main__":
+    main()