97ac841518
- Extract model loading from generate()/edit() into VnAssetsSession class - Session eagerly loads SDXL + Qwen Image Edit at construction (28s, once) - Both models held in GPU memory across calls; generate()/edit() reuse them - generate.py and edit.py become thin wrappers (backwards compatible CLI) - Context manager (with VnAssetsSession(...) as vna:) for library use - Metadata backwards-compatible: all fields preserved including lora_load_s - load_time_s now reflects total session construction, amortized across calls - Add performance stats for edit path (Qwen Image Edit + Lightning LoRA) - Benchmark matmul fallback (86.8s) vs flash attention (53.3s, 1.63x speedup) - Session vs cold start comparison: 2 ops save one 28s load, 5 edits save 112s - Flash attention via TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 documented
133 lines
4.5 KiB
Python
133 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Batch test with Lightning LoRA + flash attention for max speed."""
|
|
import subprocess
|
|
import time
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
|
|
OUTPUT_DIR = Path("output/batch_fast")
|
|
OUTPUT_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
|
|
|
|
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"),
|
|
]
|
|
|
|
env = os.environ.copy()
|
|
env["TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL"] = "1"
|
|
|
|
results = []
|
|
|
|
# ── 1. Generate base image ─────────────────────────────────────
|
|
print("=" * 60)
|
|
print("STEP 1: Generate base image")
|
|
print("=" * 60)
|
|
|
|
base_path = OUTPUT_DIR / "base.png"
|
|
t_start = time.perf_counter()
|
|
subprocess.run([
|
|
"vnasset", "generate",
|
|
"--checkpoint", CHECKPOINT,
|
|
"--prompt", BASE_PROMPT,
|
|
"--negative-prompt", NEG_PROMPT,
|
|
"--steps", str(GEN_STEPS),
|
|
"--cfg", "4.5",
|
|
"--seed", str(SEED),
|
|
"--output", str(base_path),
|
|
], check=True, env=env)
|
|
t_gen = time.perf_counter() - t_start
|
|
|
|
meta = json.loads((OUTPUT_DIR / "base.json").read_text())
|
|
results.append({
|
|
"step": "generate_base",
|
|
"output": str(base_path),
|
|
"load_s": meta["load_time_s"],
|
|
"inference_s": meta["inference_time_s"],
|
|
"total_wall_s": round(t_gen, 1),
|
|
})
|
|
print(f" Wall: {t_gen:.1f}s (load: {meta['load_time_s']}s, infer: {meta['inference_time_s']}s)\n")
|
|
|
|
# ── 2. Edit with each emotion ───────────────────────────────────
|
|
for i, (name, prompt) in enumerate(EMOTIONS):
|
|
print("=" * 60)
|
|
print(f"STEP {i+2}: Edit → {name}")
|
|
print("=" * 60)
|
|
|
|
edit_path = OUTPUT_DIR / f"base_{name}.png"
|
|
t_start = time.perf_counter()
|
|
subprocess.run([
|
|
"vnasset", "edit",
|
|
"--model", EDIT_MODEL,
|
|
"--input", str(base_path),
|
|
"--prompt", prompt,
|
|
"--steps", str(EDIT_STEPS),
|
|
"--cfg", str(EDIT_CFG),
|
|
"--seed", str(SEED),
|
|
"--lora", LORA,
|
|
"--output", str(edit_path),
|
|
], check=True, env=env)
|
|
t_edit = time.perf_counter() - t_start
|
|
|
|
meta = json.loads((OUTPUT_DIR / f"base_{name}.json").read_text())
|
|
results.append({
|
|
"step": f"edit_{name}",
|
|
"output": str(edit_path),
|
|
"load_s": meta["load_time_s"],
|
|
"inference_s": meta["inference_time_s"],
|
|
"lora_load_s": meta["lora_load_s"],
|
|
"total_wall_s": round(t_edit, 1),
|
|
})
|
|
print(f" Wall: {t_edit:.1f}s (load: {meta['load_time_s']}s, infer: {meta['inference_time_s']}s)\n")
|
|
|
|
# ── 3. Summary ──────────────────────────────────────────────────
|
|
print("=" * 60)
|
|
print("SUMMARY (Lightning LoRA + Flash Attention)")
|
|
print("=" * 60)
|
|
print(f"{'Step':<20} {'Load':>8} {'Infer':>8} {'Wall':>8}")
|
|
print("-" * 48)
|
|
total_load = total_infer = total_wall = 0
|
|
for r in results:
|
|
print(f"{r['step']:<20} {r['load_s']:>7.1f}s {r['inference_s']:>7.1f}s {r['total_wall_s']:>7.1f}s")
|
|
total_load += r["load_s"]
|
|
total_infer += r["inference_s"]
|
|
total_wall += r["total_wall_s"]
|
|
print("-" * 48)
|
|
print(f"{'TOTAL':<20} {total_load:>7.1f}s {total_infer:>7.1f}s {total_wall:>7.1f}s")
|
|
|
|
summary = {
|
|
"config": {
|
|
"checkpoint": CHECKPOINT,
|
|
"edit_model": EDIT_MODEL,
|
|
"lora": LORA,
|
|
"flash_attention": True,
|
|
"seed": SEED,
|
|
"gen_steps": GEN_STEPS,
|
|
"edit_steps": EDIT_STEPS,
|
|
"edit_cfg": EDIT_CFG,
|
|
"base_prompt": BASE_PROMPT,
|
|
},
|
|
"results": results,
|
|
"totals": {
|
|
"load_s": round(total_load, 1),
|
|
"inference_s": round(total_infer, 1),
|
|
"total_wall_s": round(total_wall, 1),
|
|
},
|
|
}
|
|
summary_path = OUTPUT_DIR / "summary.json"
|
|
summary_path.write_text(json.dumps(summary, indent=2))
|
|
print(f"\nSummary → {summary_path}")
|