make_samples.py

"""
Generate the sample images shipped with the app as reference-node defaults.

They are produced by the studio's own txt2img pipeline rather than downloaded,
which keeps them licence-clean, on-brand, and — for the PNG Info sample —
genuinely carrying an embedded A1111 parameter block, so that pipeline
demonstrates a real round trip out of the box.

    python apps/05_workflow1111/make_samples.py [--force]
"""

import json
import os
import shutil
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
for _s in (sys.stdout, sys.stderr):
    try:
        _s.reconfigure(encoding="utf-8", errors="replace")
    except (AttributeError, ValueError):
        pass

from PIL import Image  # noqa: E402

import nodes as N  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
SAMPLES = os.path.join(HERE, "samples")
os.makedirs(SAMPLES, exist_ok=True)

# (filename, reference node it seeds, prompt, size, what it has to demonstrate)
SPECS = [
    ("init_image.jpg", "ref_init_image",
     "a quiet wooden lakeside cabin at dawn, still water, pine forest, mist",
     (768, 512), "img2img — something worth re-lighting or re-seasoning"),
    ("interrogate.jpg", "ref_interrogate_image",
     "a rain-soaked neon street market at night, stalls, paper lanterns, "
     "reflections on wet asphalt, crowded",
     (768, 512), "interrogate — visually dense, lots for a VLM to name"),
    ("detect.jpg", "ref_detect_image",
     "a photograph of a person walking a dog on a city pavement beside a "
     "parked bicycle, a car in the background, daylight, documentary photo",
     (768, 512), "detect — must contain real COCO classes"),
    ("extras.jpg", "ref_extras_image",
     "studio portrait of a golden retriever sitting, plain light grey seamless "
     "backdrop, soft key light, sharp focus",
     (640, 640), "extras — a clean subject for upscaling and cutout"),
    ("control.jpg", "ref_control_image",
     "an ornate victorian building facade, wrought iron balconies, tall "
     "windows, intricate stonework, straight-on architectural photograph",
     (768, 512), "annotator — hard structural edges for Canny"),
]

PNGINFO_PROMPT = ("a red fox curled asleep in deep snow, soft winter light, "
                  "shallow depth of field")


def shrink(img, box):
    """Fit inside `box` — samples ride in the repo, so keep them small."""
    img = img.convert("RGB")
    img.thumbnail(box, Image.LANCZOS)
    return img


def save_jpg(img, path, quality=86):
    img.save(path, format="JPEG", quality=quality, optimize=True,
             progressive=True, subsampling=1)
    return os.path.getsize(path)


def generate(prompt, width, height, seed):
    out = N.txt2img(prompt, "blurry, watermark, text, low quality", 4, 1.0,
                    seed, width, height, "black-forest-labs/FLUX.1-schnell")
    return N._load_image(out)


def main():
    force = "--force" in sys.argv
    manifest = {}

    for i, (name, ref_id, prompt, box, why) in enumerate(SPECS):
        path = os.path.join(SAMPLES, name)
        if os.path.exists(path) and not force:
            print(f"  skip   {name} (exists)")
            manifest[ref_id] = name
            continue
        img = shrink(generate(prompt, box[0], box[1], 2000 + i * 37), box)
        kb = save_jpg(img, path) // 1024
        print(f"  made   {name:18} {img.width}×{img.height}  {kb} KB   — {why}")
        manifest[ref_id] = name

    # The PNG Info sample must carry a real embedded parameter block.
    png = os.path.join(SAMPLES, "with_parameters.png")
    if not os.path.exists(png) or force:
        seed, steps, cfg, w, h = 987654321, 4, 1.0, 768, 512
        img = shrink(generate(PNGINFO_PROMPT, w, h, seed), (768, 512))
        info = N.generation_info(PNGINFO_PROMPT, "blurry, watermark, text",
                                 steps, cfg, seed, img.width, img.height,
                                 "black-forest-labs/FLUX.1-schnell")
        stamped = N.postprocess(N._emit(img), 1, "Lanczos", 0, 1, 1, 1, 0, 0, 0, "", info)
        # Quantize to a 256-colour palette before shipping. This sample has to
        # stay PNG (a JPEG cannot carry the `parameters` text chunk), and a
        # full-colour 768×512 photo PNG is ~471 KB — slow enough over the Hub
        # (~4.8s) that the reference node looks broken while it loads. The
        # palette version is ~161 KB and keeps the text chunk and the
        # dimensions, so the embedded `Size:` still matches the image.
        from PIL import PngImagePlugin
        full = Image.open(stamped["path"])
        full.load()
        meta = PngImagePlugin.PngInfo()
        for key, value in (full.info or {}).items():
            if isinstance(value, str):
                meta.add_text(key, value)
        quantized = full.convert("RGB").quantize(
            colors=256, method=Image.MEDIANCUT, dither=Image.FLOYDSTEINBERG)
        quantized.save(png, format="PNG", optimize=True, pnginfo=meta)
        print(f"  made   with_parameters.png  {os.path.getsize(png)//1024} KB "
              "— PNG Info, with a real embedded parameter block")
    manifest["ref_pnginfo_image"] = "with_parameters.png"

    with open(os.path.join(SAMPLES, "manifest.json"), "w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2)

    # --- verify the samples actually do their job -------------------------
    print("\nverifying:")
    det = json.loads(__import__("gradio.workflow", fromlist=["x"]).call_model(
        ["facebook/detr-resnet-50", "object_detection",
         json.dumps({"image": {"path": os.path.join(SAMPLES, "detect.jpg")}}),
         __import__("huggingface_hub").get_token(), "auto"]))
    labels = sorted({d["label"] for d in det[0]} if isinstance(det[0], list) else set())
    print(f"  detect.jpg  → DETR finds: {labels or 'NOTHING (bad sample)'}")

    report, fields_json = N.png_info(png)
    fields = json.loads(fields_json)
    ok = fields.get("prompt", "").startswith("a red fox curled")
    print(f"  with_parameters.png → embedded params readable: {ok} "
          f"(seed={fields.get('seed')}, size={fields.get('size')})")

    total = sum(os.path.getsize(os.path.join(SAMPLES, f))
                for f in os.listdir(SAMPLES))
    print(f"\n  samples/ total: {total // 1024} KB")


if __name__ == "__main__":
    main()