From bb08efd08c871236dc207b39e640514288aebbee Mon Sep 17 00:00:00 2001 From: Michele Rossi Date: Wed, 8 Jul 2026 14:52:27 +0200 Subject: [PATCH] add YAML-driven pipeline batch mode (vnasset pipeline) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Multi-stage pipelines declared in YAML: generate, edit (cross-product), remove_bg, and upscale stages executed sequentially in a single GPU session. Every intermediate artifact (image + metadata JSON) is saved to {output_dir}/{stage_id}/ for full traceability. Data routing: - generate: list of prompts → list of images (1:1 per item) - edit: fan-out cross-product (each input × each prompt) - remove_bg / upscale: 1:1 passthrough Resume: outputs that already exist are skipped. Use --force to re-run everything — lets you add items to a stage without regenerating. Examples: examples/portrait.yaml, examples/backgrounds.yaml --- README.md | 186 ++++++++++---- examples/backgrounds.yaml | 49 ++++ examples/portrait.yaml | 74 ++++++ test_portrait_pipeline.py | 188 ++++++++++++++ vnassets/cli.py | 18 ++ vnassets/pipeline.py | 524 ++++++++++++++++++++++++++++++++++++++ 6 files changed, 991 insertions(+), 48 deletions(-) create mode 100644 examples/backgrounds.yaml create mode 100644 examples/portrait.yaml create mode 100644 test_portrait_pipeline.py create mode 100644 vnassets/pipeline.py diff --git a/README.md b/README.md index cbccc30..4bc83c7 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,105 @@ point it at. ## Usage +### Pipeline (batch YAML config) + +Declare multi-stage pipelines in YAML — the primary workflow for bulk asset +generation. All intermediate outputs are saved: every stage gets its own +subdirectory with images and metadata JSON files. + +```bash +# Portrait pipeline: generate → emotion variants → remove bg → upscale +vnasset pipeline --config examples/portrait.yaml + +# Background batch: generate many images from independent prompts +vnasset pipeline --config examples/backgrounds.yaml + +# Force re-run all stages (default: skip items whose output files exist) +vnasset pipeline --config examples/portrait.yaml --force +``` + +Config structure: + +```yaml +session: + sdxl_checkpoint: models/novaAnimeXL.safetensors + edit_model: models/qwen_image_edit.safetensors + edit_lora: models/lightning-4steps.safetensors # optional + +output_dir: output/my_pipeline + +defaults: + generate: + steps: 20 + cfg: 4.5 + negative_prompt: "deformed, ugly, bad quality, lowres" + edit: + steps: 4 + cfg: 1.0 + +stages: + # Independent batch: N prompts → N images + - id: characters + generate: + - id: heroine + prompt: "1girl, red hair, school uniform, portrait" + seed: 42 + + # Fan-out cross-product: each input × each prompt + - id: expressions + edit: + input: characters + prompts: + - id: smile + text: "make her smile happily" + - id: angry + text: "make her look angry" + + # 1:1 passthrough: each input → one output + - id: nobg + remove_bg: + input: expressions + + - id: final + upscale: + input: nobg + scale: 2 +``` + +Output structure (every stage saved): + +``` +output/my_pipeline/ + pipeline.json ← summary (stages, timings, skip/done counts) + characters/ ← stage 1 + heroine.png + heroine.json + expressions/ ← stage 2 (cross-product: {input}_{prompt}) + heroine_smile.png + heroine_angry.png + ... + nobg/ ← stage 3 + heroine_smile_nobg.png + ... + final/ ← stage 4 + heroine_smile_nobg_x2.png + ... +``` + +| Stage type | Input | Routing | +|-----------|-------|---------| +| `generate` | none | list of prompts → list of images (1:1 per item) | +| `edit` | previous stage | cross-product: each input × each prompt | +| `remove_bg` | previous stage | 1:1 passthrough | +| `upscale` | previous stage | 1:1 passthrough | + +Resume: if an output file already exists, that item is skipped. Use `--force` +to re-run everything. This lets you add items to a stage and re-run without +regenerating existing work. + +See `examples/portrait.yaml` and `examples/backgrounds.yaml` for ready-to-use +configs. + ### Generate (SDXL text-to-image) ```bash @@ -137,7 +236,10 @@ Either model can be omitted (`None`) for single-model sessions. Properties: - `vna.close()` — manual cleanup (automatic with `with`) The standalone `vnasset generate` and `vnasset edit` CLI commands are thin -wrappers around a one-shot session — same API, backwards compatible. +wrappers around a one-shot session. + +For bulk workflows, use `vnasset pipeline` instead — it handles the +session, outputs, and resume logic declaratively. ### Background Removal @@ -219,48 +321,40 @@ a compilation penalty. ## Architecture ``` -┌─────────────────────────────────────────────────────────────┐ -│ VnAssetsSession │ -│ │ -│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ │ -│ │ SDXL │ │ Qwen │ │ Qwen VL │ │ Qwen │ │ -│ │ UNet │ │ Transf. │ │ 7B TE │ │ VAE │ │ -│ │ (~3.5GB) │ │ (~20GB) │ │ (~14GB) │ │ (~1GB) │ │ -│ └──────────┘ └──────────┘ └──────────┘ └──────────┘ │ -│ │ │ │ │ │ -│ ▼ │ │ │ │ -│ ┌─────────┐ │ │ │ │ -│ │ Generate│ │ │ │ │ -│ │ │──────────┼───────────────┼──────────────┤ │ -│ │ │ base.png │ │ │ │ -│ └─────────┘ │ │ │ │ -│ │ ▼ ▼ ▼ │ -│ │ ┌──────────────────────────────────────┐ │ -│ └──────────► Edit Phase │ │ -│ │ base.png + prompts[] → variants[] │ │ -│ └──────────────────────────────────────┘ │ -│ │ │ -│ ▼ │ -│ ┌──────────────────┐ │ -│ │ base.png │ │ -│ │ happy.png │ │ -│ │ sad.png │ │ -│ │ angry.png │ │ -│ └────────┬─────────┘ │ -│ │ │ -│ ▼ │ -│ ┌──────────────────┐ │ -│ │ Upscale │ │ -│ │ (Real-ESRGAN) │ │ -│ │ ~1.8s each │ │ -│ └────────┬─────────┘ │ -│ │ │ -│ ▼ │ -│ ┌──────────────────┐ │ -│ │ Remove BG │ │ -│ │ (isnet-anime) │ │ -│ └──────────────────┘ │ -└─────────────────────────────────────────────────────────────┘ +┌──────────────────────────────────────────────────────────────────┐ +│ VnAssetsSession │ +│ │ +│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ │ +│ │ SDXL │ │ Qwen │ │ Qwen VL │ │ Qwen │ │ +│ │ UNet │ │ Transf. │ │ 7B TE │ │ VAE │ │ +│ │ (~3.5GB) │ │ (~20GB) │ │ (~14GB) │ │ (~1GB) │ │ +│ └────┬─────┘ └────┬─────┘ └────┬─────┘ └────┬─────┘ │ +│ │ │ │ │ │ +│ ▼ ▼ ▼ ▼ │ +│ ┌────────────────────────────────────────────────────────────┐ │ +│ │ Pipeline Runner │ │ +│ │ │ │ +│ │ stage 1: generate ──► stage 2: edit ──► stage 3: remove_bg │ +│ │ │ │ │ │ +│ │ ▼ ▼ ▼ │ +│ │ characters/ expressions/ nobg/ │ +│ │ heroine.png heroine_smile.png ..._nobg.png │ +│ │ heroine.json ... ... │ +│ │ │ │ +│ │ ▼ │ +│ │ stage 4: upscale │ +│ │ │ │ +│ │ ▼ │ +│ │ final/ │ +│ │ ..._nobg_x2.png │ +│ └────────────────────────────────────────────────────────────┘ │ +│ │ +│ ┌─────────────┐ ┌──────────────┐ │ +│ │ Upscale │ │ Remove BG │ (lazy-loaded on first use) │ +│ │ Real-ESRGAN │ │ isnet-anime │ │ +│ │ (~17 MB) │ │ (~176 MB) │ │ +│ └─────────────┘ └──────────────┘ │ +└──────────────────────────────────────────────────────────────────┘ ``` The Qwen transformer is loaded FP8 → BF16 at construction using @@ -469,6 +563,7 @@ Use `--raw` to bypass weighting and fall back to plain diffusers encoding. | Feature | Status | |---------|--------| +| `vnasset pipeline` (batch YAML config) | ✅ Working | | `vnasset generate` | ✅ Working | | `vnasset edit` | ✅ Working | | `VnAssetsSession` (persistent models) | ✅ Working | @@ -476,18 +571,13 @@ Use `--raw` to bypass weighting and fall back to plain diffusers encoding. | Lightning LoRA fuse-at-load | ✅ Working | | Flash attention (experimental) | ✅ Working | | `vnasset remove-bg` | ✅ Working | -| Session background removal | ✅ Working | | `vnasset upscale` | ✅ Working | -| Session upscaling | ✅ Working | -| `vnasset pipeline` (batch YAML config) | 🚧 Planned | | `vnasset serve` (daemon/HTTP API) | 🚧 Planned | | `torch.compile` on UNet | 🚧 Planned | | Batch edit loop (shared VAE encode) | 🚧 Planned | ## Future Improvements -- **Pipeline batch mode** — `vnasset pipeline --config pipeline.yaml` for - generate + multiple edits in one session from a YAML config file. - **`torch.compile` on UNet** — the UNet forward is identical each step; ROCm's `torch.compile` support is maturing and could cut per-step time significantly. - **Shared encode optimization** — for N edit variants of the same input image, diff --git a/examples/backgrounds.yaml b/examples/backgrounds.yaml new file mode 100644 index 0000000..fabe144 --- /dev/null +++ b/examples/backgrounds.yaml @@ -0,0 +1,49 @@ +# Background batch: generate multiple background images from independent prompts. +# +# Usage: +# vnasset pipeline --config examples/backgrounds.yaml +# +# Output structure: +# output/backgrounds_pipeline/ +# pipeline.json +# backgrounds/ +# classroom.png +# classroom.json +# hallway.png +# hallway.json +# ... + +session: + sdxl_checkpoint: models/novaAnimeXL_ilV190.safetensors + +output_dir: output/backgrounds_pipeline + +defaults: + generate: + steps: 20 + cfg: 4.5 + negative_prompt: "deformed, ugly, bad quality, lowres" + width: 1024 + height: 576 + +stages: + - id: backgrounds + generate: + - id: classroom + prompt: "anime classroom interior, empty, sunlight through windows, desks and chairs, blackboard, warm atmosphere" + seed: 100 + - id: hallway + prompt: "anime school hallway, lockers, windows on one side, afternoon light, clean floor" + seed: 101 + - id: park + prompt: "anime park, cherry blossoms, green grass, bench, path, spring afternoon, peaceful" + seed: 102 + - id: rooftop + prompt: "anime school rooftop, blue sky, chain link fence, water tower, sunset colors" + seed: 103 + - id: cafe + prompt: "anime cafe interior, cozy, warm lighting, wooden tables, counter with pastries, quiet afternoon" + seed: 104 + - id: street + prompt: "anime city street, shops, pedestrians, evening, neon signs starting to glow" + seed: 105 diff --git a/examples/portrait.yaml b/examples/portrait.yaml new file mode 100644 index 0000000..8a162ab --- /dev/null +++ b/examples/portrait.yaml @@ -0,0 +1,74 @@ +# Portrait pipeline: generate base character → emotion variants → remove bg → upscale. +# +# Usage: +# vnasset pipeline --config examples/portrait.yaml +# vnasset pipeline --config examples/portrait.yaml --force # re-run everything +# +# Output structure: +# output/portrait_pipeline/ +# pipeline.json ← pipeline summary +# base/ ← stage 1: generated portraits +# heroine.png +# heroine.json +# expressions/ ← stage 2: emotion edits (cross-product) +# heroine_smile.png +# heroine_smile.json +# heroine_angry.png +# ... +# nobg/ ← stage 3: background removed (RGBA) +# heroine_smile_nobg.png +# ... +# final/ ← stage 4: upscaled 2× +# heroine_smile_nobg_x2.png +# ... + +session: + sdxl_checkpoint: models/novaAnimeXL_ilV190.safetensors + edit_model: models/qwen_image_edit_2509_fp8_e4m3fn.safetensors + edit_lora: models/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors + +output_dir: output/portrait_pipeline + +defaults: + generate: + steps: 20 + cfg: 4.5 + negative_prompt: "deformed, ugly, bad quality, lowres" + edit: + steps: 4 + cfg: 1.0 + +stages: + # ── Stage 1: Generate base character portrait ── + - id: base + generate: + - id: heroine + prompt: "1girl, solo, red hair, glasses, blue eyes, white crop top, standing, portrait" + seed: 42 + + # ── Stage 2: Edit expression variants (fan-out cross-product) ── + - id: expressions + edit: + input: base + prompts: + - id: smile + text: "make her smile happily with a warm genuine smile" + - id: angry + text: "make her look angry and furious, furrowed brow" + - id: sad + text: "make her look sad and crying, tears in her eyes" + - id: surprised + text: "make her look surprised, wide eyes, mouth slightly open" + - id: blushing + text: "make her blush intensely, embarrassed expression, pink cheeks" + + # ── Stage 3: Remove backgrounds (batch 1:1 passthrough) ── + - id: nobg + remove_bg: + input: expressions + + # ── Stage 4: Upscale to 2× resolution ── + - id: final + upscale: + input: nobg + scale: 2 diff --git a/test_portrait_pipeline.py b/test_portrait_pipeline.py new file mode 100644 index 0000000..355ac37 --- /dev/null +++ b/test_portrait_pipeline.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 +"""End-to-end portrait pipeline: generate → emotion variants → background removal. + +Uses VnAssetsSession to keep models warm across all steps. Lightning LoRA + +flash attention for max edit throughput. Background removal uses isnet-anime. + +Usage: + python test_portrait_pipeline.py + TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 python test_portrait_pipeline.py # with flash attention +""" +import json +import os +import time +from pathlib import Path + +from vnassets import VnAssetsSession + +OUTPUT_DIR = Path("output/portrait_pipeline") +OUTPUT_DIR.mkdir(parents=True, exist_ok=True) +NOBG_DIR = OUTPUT_DIR / "nobg" +NOBG_DIR.mkdir(parents=True, exist_ok=True) + +BASE_PROMPT = "1girl, solo, red hair, glasses, blue eyes, white crop top, standing, portrait" +NEG_PROMPT = "deformed, ugly, bad quality, lowres" +CHECKPOINT = "models/novaAnimeXL_ilV190.safetensors" +EDIT_MODEL = "models/qwen_image_edit_2509_fp8_e4m3fn.safetensors" +LORA = "models/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors" +SEED = 42 +GEN_STEPS = 20 +EDIT_STEPS = 4 +EDIT_CFG = 1.0 +REMOVE_BG_MODEL = "isnet-anime" + +EMOTIONS = [ + ("smile", "make her smile happily with a warm genuine smile"), + ("angry", "make her look angry and furious, furrowed brow"), + ("sad", "make her look sad and crying, tears in her eyes"), + ("surprised", "make her look surprised, wide eyes, mouth slightly open"), + ("blushing", "make her blush intensely, embarrassed expression, pink cheeks"), +] + +flash_attn = os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL") == "1" +mode = "flash attention" if flash_attn else "matmul fallback" +print(f"Pipeline test — {mode}") +print(f"Seed: {SEED}, SDXL steps: {GEN_STEPS}, Edit steps: {EDIT_STEPS}, CFG: {EDIT_CFG}") +print() + +results = [] +t_pipeline = time.perf_counter() + +# ── 1. Session (load all models once) ─────────────────────────── +print("=" * 60) +print("Loading session (SDXL + Qwen Edit + Lightning LoRA)") +print("=" * 60) + +t0 = time.perf_counter() +vna = VnAssetsSession( + sdxl_checkpoint=CHECKPOINT, + edit_model=EDIT_MODEL, + edit_lora=LORA, +) +t_session = round(time.perf_counter() - t0, 1) +print(f"Session ready: {t_session}s\n") + +# ── 2. Generate base portrait ─────────────────────────────────── +print("=" * 60) +print("STEP 1: Generate base portrait") +print("=" * 60) + +base_path = str(OUTPUT_DIR / "base.png") +t0 = time.perf_counter() +vna.generate( + prompt=BASE_PROMPT, + negative_prompt=NEG_PROMPT, + steps=GEN_STEPS, + cfg=4.5, + seed=SEED, + output_path=base_path, +) +t_gen = round(time.perf_counter() - t0, 1) + +meta = json.loads((OUTPUT_DIR / "base.json").read_text()) +results.append({ + "step": "generate_base", + "output": base_path, + "inference_s": meta["inference_time_s"], + "wall_s": t_gen, +}) +print(f" Wall: {t_gen}s (infer: {meta['inference_time_s']}s)\n") + +# ── 3. Edit emotion variants ──────────────────────────────────── +variant_paths = [] +for i, (name, prompt) in enumerate(EMOTIONS): + print("=" * 60) + print(f"STEP {i+2}: Edit → {name}") + print("=" * 60) + + edit_path = str(OUTPUT_DIR / f"base_{name}.png") + t0 = time.perf_counter() + vna.edit( + input_path=base_path, + prompt=prompt, + steps=EDIT_STEPS, + cfg=EDIT_CFG, + seed=SEED, + output_path=edit_path, + ) + t_edit = round(time.perf_counter() - t0, 1) + + meta = json.loads((OUTPUT_DIR / f"base_{name}.json").read_text()) + results.append({ + "step": f"edit_{name}", + "output": edit_path, + "inference_s": meta["inference_time_s"], + "lora_load_s": meta.get("lora_load_s"), + "wall_s": t_edit, + }) + variant_paths.append(edit_path) + print(f" Wall: {t_edit}s (infer: {meta['inference_time_s']}s)\n") + +# ── 4. Remove backgrounds (batch) ─────────────────────────────── +print("=" * 60) +print("STEP 7: Remove backgrounds (batch, all variants + base)") +print("=" * 60) + +all_images = [base_path] + variant_paths +t0 = time.perf_counter() +vna.remove_backgrounds(all_images, str(NOBG_DIR), model=REMOVE_BG_MODEL) +t_rmbg = round(time.perf_counter() - t0, 1) + +results.append({ + "step": "remove_backgrounds", + "files": len(all_images), + "inference_s": t_rmbg, + "wall_s": t_rmbg, +}) +print(f" {len(all_images)} images, {t_rmbg}s total\n") + +# ── Cleanup ───────────────────────────────────────────────────── +vna.close() + +t_total = round(time.perf_counter() - t_pipeline, 1) + +# ── Summary ───────────────────────────────────────────────────── +print("=" * 60) +print("SUMMARY") +print("=" * 60) +print(f"{'Step':<22} {'Infer':>8} {'Wall':>8}") +print("-" * 42) +total_infer = 0 +total_wall = 0 +for r in results: + s = r["step"] + infer = r["inference_s"] + wall = r["wall_s"] + extra = "" + if "files" in r: + extra = f" ({r['files']} files)" + print(f"{s+extra:<22} {infer:>7.1f}s {wall:>7.1f}s") + total_infer += infer + total_wall += wall +print("-" * 42) +print(f"{'TOTAL':<22} {total_infer:>7.1f}s {total_wall:>7.1f}s") +print(f"\nPipeline wall time: {t_total}s (session load: {t_session}s)") + +# Write summary +summary = { + "config": { + "checkpoint": CHECKPOINT, + "edit_model": EDIT_MODEL, + "lora": LORA, + "flash_attention": flash_attn, + "seed": SEED, + "gen_steps": GEN_STEPS, + "edit_steps": EDIT_STEPS, + "edit_cfg": EDIT_CFG, + "remove_bg_model": REMOVE_BG_MODEL, + "base_prompt": BASE_PROMPT, + }, + "results": results, + "session_load_s": t_session, + "pipeline_wall_s": t_total, +} +summary_path = OUTPUT_DIR / "summary.json" +summary_path.write_text(json.dumps(summary, indent=2)) +print(f"\nSummary → {summary_path}") +print(f"Images → {OUTPUT_DIR}/") +print(f"Transparent → {NOBG_DIR}/") diff --git a/vnassets/cli.py b/vnassets/cli.py index c660885..ed43e29 100644 --- a/vnassets/cli.py +++ b/vnassets/cli.py @@ -7,6 +7,7 @@ from .background import remove_background as _remove_bg from .background import remove_backgrounds as _remove_bgs from .edit import edit from .generate import generate +from .pipeline import load_config, run_pipeline from .upscale import upscale as _upscale_fn from .upscale import upscales as _upscales_fn @@ -92,6 +93,23 @@ REMOVE_BG_MODELS = ["isnet-anime", "u2net", "u2netp", "u2net_human_seg", "isnet- UPSALE_SCALES = [2, 3, 4] +@main.command("pipeline") +@click.option("--config", "config_path", required=True, help="Path to pipeline YAML config file.") +@click.option("--force", is_flag=True, help="Re-run all stages even if outputs already exist.") +def pipeline_cmd(config_path, force): + """Run a multi-stage pipeline from a YAML config file. + + Stages are executed sequentially in a single GPU session. Every + intermediate output (image + metadata JSON) is saved to + ``{output_dir}/{stage_id}/``. + + Resume: items whose output file already exists are skipped. + Use --force to re-run everything. + """ + config = load_config(config_path) + run_pipeline(config, force=force) + + @main.command("upscale") @click.option("--input", "input_paths", multiple=True, required=True, help="Input image path (repeat for batch).") diff --git a/vnassets/pipeline.py b/vnassets/pipeline.py new file mode 100644 index 0000000..faafc91 --- /dev/null +++ b/vnassets/pipeline.py @@ -0,0 +1,524 @@ +"""Pipeline batch mode — YAML-driven multi-stage image asset generation. + +A pipeline is an ordered sequence of stages. Each stage consumes outputs +from a previous stage and produces its own outputs in a subdirectory. +Every intermediate artifact (image + metadata JSON) is saved. + +Stage types and their data routing: + + generate — no input; list of prompts → list of images (1:1 per item) + edit — input stage → cross-product: each input × each prompt + remove_bg — input stage → 1:1 passthrough (each input → RGBA PNG) + upscale — input stage → 1:1 passthrough (each input → upscaled PNG) + +Resume: if an output file already exists, that item is skipped (unless --force). +""" +from __future__ import annotations + +import json +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import yaml + +from .session import VnAssetsSession + + +# ── Config data types ──────────────────────────────────────────────────── + +@dataclass +class GenerateItem: + id: str + prompt: str + negative_prompt: str | None = None + seed: int | None = None + steps: int | None = None + cfg: float | None = None + width: int | None = None + height: int | None = None + + +@dataclass +class EditPrompt: + id: str + text: str + + +@dataclass +class StageOutput: + id: str + image_path: Path + meta_path: Path + + +@dataclass +class Stage: + id: str + type: str + input_stage: str | None = None + generate_items: list[GenerateItem] = field(default_factory=list) + edit_prompts: list[EditPrompt] = field(default_factory=list) + remove_bg_model: str = "isnet-anime" + upscale_scale: int = 2 + + +@dataclass +class PipelineConfig: + session: dict[str, str | None] + output_dir: Path + stages: list[Stage] + defaults: dict[str, dict[str, Any]] = field(default_factory=dict) + + +# ── Config loading and validation ──────────────────────────────────────── + +def load_config(path: str) -> PipelineConfig: + """Parse and validate a pipeline YAML config file. + + Raises: + FileNotFoundError: If the config file doesn't exist. + ValueError: If the YAML is structurally invalid. + """ + raw = yaml.safe_load(Path(path).read_text()) + + if not isinstance(raw, dict): + raise ValueError("Pipeline config must be a YAML mapping (top-level dict).") + if "session" not in raw: + raise ValueError("Missing required top-level key: 'session'.") + if "output_dir" not in raw: + raise ValueError("Missing required top-level key: 'output_dir'.") + if "stages" not in raw: + raise ValueError("Missing required top-level key: 'stages'.") + if not isinstance(raw["stages"], list) or len(raw["stages"]) == 0: + raise ValueError("'stages' must be a non-empty list.") + + defaults: dict[str, dict[str, Any]] = raw.get("defaults", {}) + + stages: list[Stage] = [] + stage_ids: set[str] = set() + for i, s_raw in enumerate(raw["stages"]): + if not isinstance(s_raw, dict): + raise ValueError(f"stages[{i}]: must be a mapping, got {type(s_raw).__name__}.") + if "id" not in s_raw: + raise ValueError(f"stages[{i}]: missing required field 'id'.") + + sid = s_raw["id"] + if sid in stage_ids: + raise ValueError(f"stages[{i}]: duplicate stage id '{sid}'.") + stage_ids.add(sid) + + # Determine stage type: exactly one of the known operation keys + op_keys = {"generate", "edit", "remove_bg", "upscale"} + found = [k for k in op_keys if k in s_raw] + if len(found) == 0: + raise ValueError( + f"stages[{i}] ('{sid}'): must contain one of: {', '.join(sorted(op_keys))}." + ) + if len(found) > 1: + raise ValueError( + f"stages[{i}] ('{sid}'): contains multiple stage types: {found}. Use exactly one." + ) + op = found[0] + op_raw = s_raw[op] + + input_stage = op_raw.get("input") if isinstance(op_raw, dict) else None + + stage = Stage(id=sid, type=op, input_stage=input_stage) + + if op == "generate": + stage.generate_items = _parse_generate_items(op_raw, defaults.get("generate", {}), sid) + elif op == "edit": + stage.edit_prompts = _parse_edit_prompts(op_raw, sid) + elif op == "remove_bg": + stage.remove_bg_model = op_raw.get("model", defaults.get("remove_bg", {}).get("model", "isnet-anime")) + if input_stage is None: + raise ValueError(f"stages[{i}] ('{sid}'): remove_bg requires 'input' referencing a previous stage.") + elif op == "upscale": + stage.upscale_scale = op_raw.get("scale", defaults.get("upscale", {}).get("scale", 2)) + if stage.upscale_scale not in (2, 3, 4): + raise ValueError(f"stages[{i}] ('{sid}'): upscale scale must be 2, 3, or 4, got {stage.upscale_scale}.") + if input_stage is None: + raise ValueError(f"stages[{i}] ('{sid}'): upscale requires 'input' referencing a previous stage.") + + stages.append(stage) + + # Validate input_stage references + for stage in stages: + if stage.input_stage is not None and stage.input_stage not in stage_ids: + raise ValueError( + f"stage '{stage.id}': input '{stage.input_stage}' references " + f"a stage that does not exist." + ) + + # Prevent circular / forward references: input_stage must come earlier + for i, stage in enumerate(stages): + if stage.input_stage is not None: + ref_idx = next(j for j, s in enumerate(stages) if s.id == stage.input_stage) + if ref_idx >= i: + raise ValueError( + f"stage '{stage.id}': input '{stage.input_stage}' must appear " + f"before '{stage.id}' in the stages list." + ) + + session = raw["session"] + return PipelineConfig( + session=session, + output_dir=Path(raw["output_dir"]), + stages=stages, + defaults=defaults, + ) + + +def _parse_generate_items( + op_raw: dict, + defaults: dict[str, Any], + stage_id: str, +) -> list[GenerateItem]: + items_raw = op_raw.get("items", op_raw) if isinstance(op_raw, dict) else op_raw + if isinstance(items_raw, dict) and "items" in items_raw: + items_raw = items_raw["items"] + + if not isinstance(items_raw, list) or len(items_raw) == 0: + raise ValueError(f"stage '{stage_id}': generate must have a non-empty list of items.") + + items: list[GenerateItem] = [] + seen_ids: set[str] = set() + for j, item_raw in enumerate(items_raw): + if not isinstance(item_raw, dict): + raise ValueError(f"stage '{stage_id}', item {j}: must be a mapping.") + if "id" not in item_raw: + raise ValueError(f"stage '{stage_id}', item {j}: missing required field 'id'.") + if "prompt" not in item_raw: + raise ValueError(f"stage '{stage_id}', item {j}: missing required field 'prompt'.") + iid = item_raw["id"] + if iid in seen_ids: + raise ValueError(f"stage '{stage_id}': duplicate item id '{iid}'.") + seen_ids.add(iid) + + items.append(GenerateItem( + id=iid, + prompt=item_raw["prompt"], + negative_prompt=item_raw.get("negative_prompt", defaults.get("negative_prompt", "")), + seed=item_raw.get("seed", defaults.get("seed")), + steps=item_raw.get("steps", defaults.get("steps", 20)), + cfg=item_raw.get("cfg", defaults.get("cfg", 4.5)), + width=item_raw.get("width", defaults.get("width", 1024)), + height=item_raw.get("height", defaults.get("height", 1024)), + )) + return items + + +def _parse_edit_prompts(op_raw: dict, stage_id: str) -> list[EditPrompt]: + prompts_raw = op_raw.get("prompts", op_raw) if isinstance(op_raw, dict) else op_raw + if isinstance(prompts_raw, dict) and "prompts" in prompts_raw: + prompts_raw = prompts_raw["prompts"] + + if not isinstance(prompts_raw, list) or len(prompts_raw) == 0: + raise ValueError(f"stage '{stage_id}': edit must have a non-empty list of prompts.") + + prompts: list[EditPrompt] = [] + seen_ids: set[str] = set() + for j, p_raw in enumerate(prompts_raw): + if not isinstance(p_raw, dict): + raise ValueError(f"stage '{stage_id}', prompt {j}: must be a mapping.") + if "id" not in p_raw: + p_raw["id"] = f"edit_{j}" + if "text" not in p_raw: + raise ValueError(f"stage '{stage_id}', prompt {j}: missing required field 'text'.") + pid = p_raw["id"] + if pid in seen_ids: + raise ValueError(f"stage '{stage_id}': duplicate prompt id '{pid}'.") + seen_ids.add(pid) + prompts.append(EditPrompt(id=pid, text=p_raw["text"])) + return prompts + + +# ── Pipeline runner ────────────────────────────────────────────────────── + +def run_pipeline(config: PipelineConfig, force: bool = False) -> None: + """Execute all stages in order using a single VnAssetsSession. + + Args: + config: Parsed pipeline configuration. + force: If True, re-run items even when output files exist. + """ + output_dir = config.output_dir + output_dir.mkdir(parents=True, exist_ok=True) + + session_kwargs = { + "sdxl_checkpoint": config.session.get("sdxl_checkpoint"), + "edit_model": config.session.get("edit_model"), + "edit_lora": config.session.get("edit_lora"), + } + + stage_outputs: dict[str, list[StageOutput]] = {} + stage_times: list[dict[str, Any]] = [] + t_pipeline = time.perf_counter() + + with VnAssetsSession(**session_kwargs) as vna: + for stage in config.stages: + t_stage = time.perf_counter() + + stage_dir = output_dir / stage.id + outputs, skipped = _run_stage(vna, stage, stage_dir, stage_outputs, config.defaults, force) + + elapsed = round(time.perf_counter() - t_stage, 2) + stage_outputs[stage.id] = outputs + done = len(outputs) - skipped + stage_times.append({ + "stage": stage.id, + "type": stage.type, + "items": len(outputs), + "done": done, + "skipped": skipped, + "wall_s": elapsed, + }) + + action = "done" if skipped == 0 else f"{done} done, {skipped} skipped" + print(f"[{stage.id}] {action} ({elapsed}s)") + + t_total = round(time.perf_counter() - t_pipeline, 2) + + # Write pipeline summary + summary = { + "output_dir": str(output_dir.resolve()), + "session": config.session, + "stages": stage_times, + "pipeline_wall_s": t_total, + } + _write_json(output_dir / "pipeline.json", summary) + print(f"\nPipeline complete ({t_total}s) — {output_dir.resolve()}/") + + +def _run_stage( + vna: VnAssetsSession, + stage: Stage, + stage_dir: Path, + previous_outputs: dict[str, list[StageOutput]], + defaults: dict[str, dict[str, Any]], + force: bool, +) -> tuple[list[StageOutput], int]: + """Dispatch to the appropriate stage runner based on stage.type. + + Returns (outputs, skipped_count). + """ + stage_dir.mkdir(parents=True, exist_ok=True) + + if stage.type == "generate": + return _run_generate(vna, stage, stage_dir, defaults, force) + elif stage.type == "edit": + inputs = _resolve_inputs(stage, previous_outputs) + return _run_edit(vna, stage, inputs, stage_dir, defaults, force) + elif stage.type == "remove_bg": + inputs = _resolve_inputs(stage, previous_outputs) + return _run_remove_bg(vna, stage, inputs, stage_dir, force) + elif stage.type == "upscale": + inputs = _resolve_inputs(stage, previous_outputs) + return _run_upscale(vna, stage, inputs, stage_dir, force) + else: + raise ValueError(f"Unknown stage type: {stage.type}") + + +def _resolve_inputs( + stage: Stage, + previous_outputs: dict[str, list[StageOutput]], +) -> list[StageOutput]: + """Get the outputs from the stage referenced by stage.input_stage.""" + if stage.input_stage is None: + raise ValueError(f"stage '{stage.id}': no input stage specified.") + inputs = previous_outputs.get(stage.input_stage) + if inputs is None: + raise ValueError( + f"stage '{stage.id}': input stage '{stage.input_stage}' not found " + f"in previous outputs." + ) + return inputs + + +# ── Per-type stage runners ─────────────────────────────────────────────── + +def _run_generate( + vna: VnAssetsSession, + stage: Stage, + stage_dir: Path, + defaults: dict[str, dict[str, Any]], + force: bool, +) -> tuple[list[StageOutput], int]: + g_defaults = defaults.get("generate", {}) + outputs: list[StageOutput] = [] + skipped = 0 + + for item in stage.generate_items: + output_path = stage_dir / f"{item.id}.png" + meta_path = stage_dir / f"{item.id}.json" + + if _skip(output_path, meta_path, force): + outputs.append(StageOutput(id=item.id, image_path=output_path, meta_path=meta_path)) + skipped += 1 + print(f" [{item.id}] skipped (exists)") + continue + + steps = item.steps if item.steps is not None else g_defaults.get("steps", 20) + cfg = item.cfg if item.cfg is not None else g_defaults.get("cfg", 4.5) + neg = item.negative_prompt if item.negative_prompt is not None else g_defaults.get("negative_prompt", "") + width = item.width if item.width is not None else g_defaults.get("width", 1024) + height = item.height if item.height is not None else g_defaults.get("height", 1024) + seed = item.seed if item.seed is not None else g_defaults.get("seed") + + vna.generate( + prompt=item.prompt, + negative_prompt=neg, + width=width, + height=height, + steps=steps, + cfg=cfg, + seed=seed, + output_path=str(output_path), + ) + + outputs.append(StageOutput(id=item.id, image_path=output_path, meta_path=meta_path)) + print(f" [{item.id}] done") + + return outputs, skipped + + +def _run_edit( + vna: VnAssetsSession, + stage: Stage, + inputs: list[StageOutput], + stage_dir: Path, + defaults: dict[str, dict[str, Any]], + force: bool, +) -> tuple[list[StageOutput], int]: + e_defaults = defaults.get("edit", {}) + steps = e_defaults.get("steps", 20) + cfg = e_defaults.get("cfg", 4.0) + outputs: list[StageOutput] = [] + skipped = 0 + + for inp in inputs: + for prompt in stage.edit_prompts: + item_id = f"{inp.id}_{prompt.id}" + output_path = stage_dir / f"{item_id}.png" + meta_path = stage_dir / f"{item_id}.json" + + if _skip(output_path, meta_path, force): + outputs.append(StageOutput(id=item_id, image_path=output_path, meta_path=meta_path)) + skipped += 1 + print(f" [{item_id}] skipped (exists)") + continue + + vna.edit( + input_path=str(inp.image_path), + prompt=prompt.text, + steps=steps, + cfg=cfg, + seed=e_defaults.get("seed"), + output_path=str(output_path), + ) + + outputs.append(StageOutput(id=item_id, image_path=output_path, meta_path=meta_path)) + print(f" [{item_id}] done") + + return outputs, skipped + + +def _run_remove_bg( + vna: VnAssetsSession, + stage: Stage, + inputs: list[StageOutput], + stage_dir: Path, + force: bool, +) -> tuple[list[StageOutput], int]: + outputs: list[StageOutput] = [] + skipped = 0 + + for inp in inputs: + item_id = f"{inp.id}_nobg" + output_path = stage_dir / f"{item_id}.png" + meta_path = stage_dir / f"{item_id}.json" + + if _skip(output_path, meta_path, force): + outputs.append(StageOutput(id=item_id, image_path=output_path, meta_path=meta_path)) + skipped += 1 + print(f" [{item_id}] skipped (exists)") + continue + + vna.remove_background( + input_path=str(inp.image_path), + output_path=str(output_path), + model=stage.remove_bg_model, + ) + _write_json(meta_path, { + "stage": stage.id, + "type": "remove_bg", + "id": item_id, + "input": str(inp.image_path.resolve()), + "output": str(output_path.resolve()), + "model": stage.remove_bg_model, + }) + + outputs.append(StageOutput(id=item_id, image_path=output_path, meta_path=meta_path)) + print(f" [{item_id}] done") + + return outputs, skipped + + +def _run_upscale( + vna: VnAssetsSession, + stage: Stage, + inputs: list[StageOutput], + stage_dir: Path, + force: bool, +) -> tuple[list[StageOutput], int]: + scale = stage.upscale_scale + outputs: list[StageOutput] = [] + skipped = 0 + + for inp in inputs: + item_id = f"{inp.id}_x{scale}" + output_path = stage_dir / f"{item_id}.png" + meta_path = stage_dir / f"{item_id}.json" + + if _skip(output_path, meta_path, force): + outputs.append(StageOutput(id=item_id, image_path=output_path, meta_path=meta_path)) + skipped += 1 + print(f" [{item_id}] skipped (exists)") + continue + + vna.upscale( + input_path=str(inp.image_path), + output_path=str(output_path), + scale=scale, + ) + _write_json(meta_path, { + "stage": stage.id, + "type": "upscale", + "id": item_id, + "input": str(inp.image_path.resolve()), + "output": str(output_path.resolve()), + "scale": scale, + }) + + outputs.append(StageOutput(id=item_id, image_path=output_path, meta_path=meta_path)) + print(f" [{item_id}] done") + + return outputs, skipped + + +# ── Helpers ────────────────────────────────────────────────────────────── + +def _skip(image_path: Path, meta_path: Path, force: bool) -> bool: + if force: + return False + return _exists(image_path) + + +def _exists(path: Path) -> bool: + return path.exists() and path.stat().st_size > 0 + + +def _write_json(path: Path, data: dict) -> None: + path.write_text(json.dumps(data, indent=2))