add YAML-driven pipeline batch mode (vnasset pipeline)
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
This commit is contained in:
Michele Rossi
170
README.md
170
README.md
@@ -53,6 +53,105 @@ point it at.
|
|||||||
|
|
||||||
## Usage
|
## 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)
|
### Generate (SDXL text-to-image)
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
@@ -137,7 +236,10 @@ Either model can be omitted (`None`) for single-model sessions. Properties:
|
|||||||
- `vna.close()` — manual cleanup (automatic with `with`)
|
- `vna.close()` — manual cleanup (automatic with `with`)
|
||||||
|
|
||||||
The standalone `vnasset generate` and `vnasset edit` CLI commands are thin
|
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
|
### Background Removal
|
||||||
|
|
||||||
@@ -219,48 +321,40 @@ a compilation penalty.
|
|||||||
## Architecture
|
## Architecture
|
||||||
|
|
||||||
```
|
```
|
||||||
┌─────────────────────────────────────────────────────────────┐
|
┌──────────────────────────────────────────────────────────────────┐
|
||||||
│ VnAssetsSession │
|
│ VnAssetsSession │
|
||||||
│ │
|
│ │
|
||||||
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
|
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
|
||||||
│ │ SDXL │ │ Qwen │ │ Qwen VL │ │ Qwen │ │
|
│ │ SDXL │ │ Qwen │ │ Qwen VL │ │ Qwen │ │
|
||||||
│ │ UNet │ │ Transf. │ │ 7B TE │ │ VAE │ │
|
│ │ UNet │ │ Transf. │ │ 7B TE │ │ VAE │ │
|
||||||
│ │ (~3.5GB) │ │ (~20GB) │ │ (~14GB) │ │ (~1GB) │ │
|
│ │ (~3.5GB) │ │ (~20GB) │ │ (~14GB) │ │ (~1GB) │ │
|
||||||
│ └──────────┘ └──────────┘ └──────────┘ └──────────┘ │
|
│ └────┬─────┘ └────┬─────┘ └────┬─────┘ └────┬─────┘ │
|
||||||
|
│ │ │ │ │ │
|
||||||
|
│ ▼ ▼ ▼ ▼ │
|
||||||
|
│ ┌────────────────────────────────────────────────────────────┐ │
|
||||||
|
│ │ Pipeline Runner │ │
|
||||||
|
│ │ │ │
|
||||||
|
│ │ stage 1: generate ──► stage 2: edit ──► stage 3: remove_bg │
|
||||||
│ │ │ │ │ │
|
│ │ │ │ │ │
|
||||||
│ ▼ │ │ │ │
|
|
||||||
│ ┌─────────┐ │ │ │ │
|
|
||||||
│ │ Generate│ │ │ │ │
|
|
||||||
│ │ │──────────┼───────────────┼──────────────┤ │
|
|
||||||
│ │ │ base.png │ │ │ │
|
|
||||||
│ └─────────┘ │ │ │ │
|
|
||||||
│ │ ▼ ▼ ▼ │
|
│ │ ▼ ▼ ▼ │
|
||||||
│ │ ┌──────────────────────────────────────┐ │
|
│ │ characters/ expressions/ nobg/ │
|
||||||
│ └──────────► Edit Phase │ │
|
│ │ heroine.png heroine_smile.png ..._nobg.png │
|
||||||
│ │ base.png + prompts[] → variants[] │ │
|
│ │ heroine.json ... ... │
|
||||||
│ └──────────────────────────────────────┘ │
|
│ │ │ │
|
||||||
│ │ │
|
│ │ ▼ │
|
||||||
│ ▼ │
|
│ │ stage 4: upscale │
|
||||||
│ ┌──────────────────┐ │
|
│ │ │ │
|
||||||
│ │ base.png │ │
|
│ │ ▼ │
|
||||||
│ │ happy.png │ │
|
│ │ final/ │
|
||||||
│ │ sad.png │ │
|
│ │ ..._nobg_x2.png │
|
||||||
│ │ angry.png │ │
|
│ └────────────────────────────────────────────────────────────┘ │
|
||||||
│ └────────┬─────────┘ │
|
│ │
|
||||||
│ │ │
|
│ ┌─────────────┐ ┌──────────────┐ │
|
||||||
│ ▼ │
|
│ │ Upscale │ │ Remove BG │ (lazy-loaded on first use) │
|
||||||
│ ┌──────────────────┐ │
|
│ │ Real-ESRGAN │ │ isnet-anime │ │
|
||||||
│ │ Upscale │ │
|
│ │ (~17 MB) │ │ (~176 MB) │ │
|
||||||
│ │ (Real-ESRGAN) │ │
|
│ └─────────────┘ └──────────────┘ │
|
||||||
│ │ ~1.8s each │ │
|
└──────────────────────────────────────────────────────────────────┘
|
||||||
│ └────────┬─────────┘ │
|
|
||||||
│ │ │
|
|
||||||
│ ▼ │
|
|
||||||
│ ┌──────────────────┐ │
|
|
||||||
│ │ Remove BG │ │
|
|
||||||
│ │ (isnet-anime) │ │
|
|
||||||
│ └──────────────────┘ │
|
|
||||||
└─────────────────────────────────────────────────────────────┘
|
|
||||||
```
|
```
|
||||||
|
|
||||||
The Qwen transformer is loaded FP8 → BF16 at construction using
|
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 |
|
| Feature | Status |
|
||||||
|---------|--------|
|
|---------|--------|
|
||||||
|
| `vnasset pipeline` (batch YAML config) | ✅ Working |
|
||||||
| `vnasset generate` | ✅ Working |
|
| `vnasset generate` | ✅ Working |
|
||||||
| `vnasset edit` | ✅ Working |
|
| `vnasset edit` | ✅ Working |
|
||||||
| `VnAssetsSession` (persistent models) | ✅ 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 |
|
| Lightning LoRA fuse-at-load | ✅ Working |
|
||||||
| Flash attention (experimental) | ✅ Working |
|
| Flash attention (experimental) | ✅ Working |
|
||||||
| `vnasset remove-bg` | ✅ Working |
|
| `vnasset remove-bg` | ✅ Working |
|
||||||
| Session background removal | ✅ Working |
|
|
||||||
| `vnasset upscale` | ✅ Working |
|
| `vnasset upscale` | ✅ Working |
|
||||||
| Session upscaling | ✅ Working |
|
|
||||||
| `vnasset pipeline` (batch YAML config) | 🚧 Planned |
|
|
||||||
| `vnasset serve` (daemon/HTTP API) | 🚧 Planned |
|
| `vnasset serve` (daemon/HTTP API) | 🚧 Planned |
|
||||||
| `torch.compile` on UNet | 🚧 Planned |
|
| `torch.compile` on UNet | 🚧 Planned |
|
||||||
| Batch edit loop (shared VAE encode) | 🚧 Planned |
|
| Batch edit loop (shared VAE encode) | 🚧 Planned |
|
||||||
|
|
||||||
## Future Improvements
|
## 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` on UNet** — the UNet forward is identical each step; ROCm's
|
||||||
`torch.compile` support is maturing and could cut per-step time significantly.
|
`torch.compile` support is maturing and could cut per-step time significantly.
|
||||||
- **Shared encode optimization** — for N edit variants of the same input image,
|
- **Shared encode optimization** — for N edit variants of the same input image,
|
||||||
|
|||||||
49
examples/backgrounds.yaml
Normal file
49
examples/backgrounds.yaml
Normal file
@@ -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
|
||||||
74
examples/portrait.yaml
Normal file
74
examples/portrait.yaml
Normal file
@@ -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
|
||||||
188
test_portrait_pipeline.py
Normal file
188
test_portrait_pipeline.py
Normal file
@@ -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}/")
|
||||||
@@ -7,6 +7,7 @@ from .background import remove_background as _remove_bg
|
|||||||
from .background import remove_backgrounds as _remove_bgs
|
from .background import remove_backgrounds as _remove_bgs
|
||||||
from .edit import edit
|
from .edit import edit
|
||||||
from .generate import generate
|
from .generate import generate
|
||||||
|
from .pipeline import load_config, run_pipeline
|
||||||
from .upscale import upscale as _upscale_fn
|
from .upscale import upscale as _upscale_fn
|
||||||
from .upscale import upscales as _upscales_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]
|
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")
|
@main.command("upscale")
|
||||||
@click.option("--input", "input_paths", multiple=True, required=True,
|
@click.option("--input", "input_paths", multiple=True, required=True,
|
||||||
help="Input image path (repeat for batch).")
|
help="Input image path (repeat for batch).")
|
||||||
|
|||||||
524
vnassets/pipeline.py
Normal file
524
vnassets/pipeline.py
Normal file
@@ -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))
|
||||||
Reference in New Issue
Block a user