Add VnAssetsSession for persistent model lifecycle
- 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
This commit is contained in:
Michele Rossi
132
test_batch_fast.py
Normal file
132
test_batch_fast.py
Normal file
@@ -0,0 +1,132 @@
|
||||
#!/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}")
|
||||
Reference in New Issue
Block a user