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
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}/")
|
||||
Reference in New Issue
Block a user