scripts/generate_demo_outputs.py

"""
Pre-generate a comprehensive demo gallery for Indic Heritage Studio v2.

Produces:
  - 5 styles × 3 seeds × T2I = 15 images (1024×1024)
  - 5 styles × 2 input images × style transfer = 10 images
  - 5 styles × 1 input × SVD = 5 videos (25-frame MP4)
  - 1 inpainting demo (mask a region, restyle in 1 style)
  - 1 ControlNet demo (canny + 1 style)
  - 1 batch scaling demo (4 GPU vs 1 GPU comparison)

Total: ~30 outputs in ~30 min on 8 × 80GB.
Run this BEFORE recording the demo video so demo recording burns
minimal AMD credits.
"""
from __future__ import annotations

import logging
from pathlib import Path

from config.settings import settings
from config.styles import list_styles

log = logging.getLogger(__name__)

OUTPUT_BASE = settings.examples_dir

PROMPTS = [
    "a young woman reading under a banyan tree at sunset, peacock beside her",
    "a temple festival at dawn with devotees and musicians",
    "a sage meditating by the river, surrounded by lotus flowers",
]

INPUT_IMAGES = [
    "examples/inputs/portrait.jpg",
    "examples/inputs/landscape.jpg",
]


def gen_t2i_gallery() -> None:
    """5 styles × 3 prompts × 3 seeds = 45 images."""
    from core.text_to_image import TextToImagePipeline
    log.info("=== T2I Gallery ===")
    pipe = TextToImagePipeline().load()
    out_dir = OUTPUT_BASE / "1_text_to_image"
    out_dir.mkdir(parents=True, exist_ok=True)
    for style in list_styles():
        for prompt_idx, prompt in enumerate(PROMPTS):
            for seed in (42, 1337, 2024):
                out_path = out_dir / style.id / f"p{prompt_idx}_s{seed}.png"
                if out_path.exists():
                    continue
                log.info("Generating %s", out_path)
                img = pipe.generate(prompt=prompt, style=style, seed=seed)
                out_path.parent.mkdir(parents=True, exist_ok=True)
                img.save(out_path)


def gen_style_gallery() -> None:
    """5 styles × 2 input images = 10 styled images."""
    from core.style_transfer import StyleTransferPipeline
    from utils.image_utils import load_image
    log.info("=== Style Transfer Gallery ===")
    pipe = StyleTransferPipeline().load()
    out_dir = OUTPUT_BASE / "2_style_transfer"
    out_dir.mkdir(parents=True, exist_ok=True)
    for style in list_styles():
        for inp in INPUT_IMAGES:
            inp_path = Path(inp)
            if not inp_path.exists():
                log.warning("Missing input: %s — skipping", inp_path)
                continue
            out_path = out_dir / style.id / inp_path.name
            if out_path.exists():
                continue
            img = load_image(inp_path)
            out = pipe.transfer(image=img, style=style, seed=42)
            out_path.parent.mkdir(parents=True, exist_ok=True)
            out.save(out_path)


def gen_video_gallery() -> None:
    """5 styles × 1 input = 5 videos."""
    from core.image_to_video import ImageToVideoPipeline
    from utils.image_utils import load_image
    log.info("=== Video Gallery ===")
    pipe = ImageToVideoPipeline().load()
    out_dir = OUTPUT_BASE / "3_image_to_video"
    out_dir.mkdir(parents=True, exist_ok=True)
    inp_path = Path(INPUT_IMAGES[0])
    if not inp_path.exists():
        log.warning("Missing input image — skipping video gallery")
        return
    img = load_image(inp_path)
    for style in list_styles():
        out_path = out_dir / f"{style.id}.mp4"
        if out_path.exists():
            continue
        log.info("Generating video for %s", style.id)
        pipe.generate_to_mp4(
            image=img, out_path=out_path, style=style,
            num_frames=25, fps=8, seed=42,
        )


def gen_inpaint_demo() -> None:
    """One inpainting demo."""
    from core.inpainting import InpaintingPipeline
    from utils.image_utils import load_image
    log.info("=== Inpainting Demo ===")
    pipe = InpaintingPipeline().load()
    out_dir = OUTPUT_BASE / "5_inpainting"
    out_dir.mkdir(parents=True, exist_ok=True)
    inp_path = Path(INPUT_IMAGES[0])
    if not inp_path.exists():
        log.warning("Missing input image — skipping inpaint demo")
        return
    # Create a simple mask (right half white)
    from PIL import Image
    img = load_image(inp_path)
    w, h = img.size
    mask = Image.new("L", (w, h), 0)
    for y in range(h):
        for x in range(w // 2, w):
            mask.putpixel((x, y), 255)
    mask = mask.resize((1024, 1024))
    img = img.resize((1024, 1024))

    out = pipe.inpaint(
        image=img, mask=mask,
        prompt="ornate floral border with peacock motifs",
        style=list_styles()[0],  # madhubani
        seed=42,
    )
    out.save(out_dir / "inpaint_demo.png")


def gen_controlnet_demo() -> None:
    """One ControlNet demo (canny + mughal)."""
    from core.controlnet import ControlNetPipeline
    from utils.image_utils import load_image
    log.info("=== ControlNet Demo ===")
    out_dir = OUTPUT_BASE / "6_controlnet"
    out_dir.mkdir(parents=True, exist_ok=True)
    inp_path = Path(INPUT_IMAGES[0])
    if not inp_path.exists():
        log.warning("Missing input image — skipping controlnet demo")
        return
    pipe = ControlNetPipeline(condition_type="canny").load()
    img = load_image(inp_path)
    cond = pipe.detect(img)
    cond.save(out_dir / "canny_condition.png")
    from config.styles import get_style
    out = pipe.generate(
        conditioning_image=cond,
        prompt="a courtly gathering with musicians and dancers",
        style=get_style("mughal"),
        seed=42,
    )
    out.save(out_dir / "controlnet_mughal.png")


def main():
    logging.basicConfig(level=logging.INFO,
                        format="%(asctime)s | %(levelname)s | %(message)s")
    OUTPUT_BASE.mkdir(parents=True, exist_ok=True)
    gen_t2i_gallery()
    gen_style_gallery()
    gen_video_gallery()
    gen_inpaint_demo()
    gen_controlnet_demo()
    log.info("All demo outputs generated.")


if __name__ == "__main__":
    main()