diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..37cf994 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,289 @@ +# Data-Oriented Design — Operating Rules + +These are operating rules, not philosophy: +every rule here tells you what to *do*. Approach every problem — code, plan, +pipeline, document — by understanding the real data first, then designing the +simplest machine that transforms the input you actually have into the output +you actually need, at a cost you can state. Decide from facts and measurement, +not habit, analogy, or dogma. + +## Scope, tiers, and precedence + +Scale the ceremony to the task. Decide the tier first; when unsure, pick the +higher tier and say which you picked. + +- **Tier 0 — trivial.** Typo fixes, mechanical edits, one-line bugfixes, + answering questions. Apply the defaults silently (naming, explicit error + behavior, no speculative generality). No written plan or checklist. +- **Tier 1 — non-trivial change.** New function or feature, behavior change, + anything that touches a data layout, contract, or interface. Required: + answer the framing and data questions in a short written plan *before* + implementing, run the simplification pass, and run the final self-check. +- **Tier 2 — subsystem-scale.** New or substantially reworked subsystem, + pipeline, or tool. Everything in tier 1 plus the enforceable deliverables. + +Precedence when rules conflict: + +1. An explicit instruction from the user for the current task. +2. This document. +3. Existing codebase or workflow convention. + +When this document conflicts with existing convention and complying would +mean a large refactor, do not silently rewrite and do not silently conform: +state the conflict, estimate the cost of each option, and propose the +smallest compliant change. + +## Defaults to reject + +These are the three default beliefs that produce bad solutions. Each comes +with the replacement behavior — do the replacement, every time: + +1. **"The tools are the platform."** Reality is the platform: the actual + hardware, organization, deadline, physics. *Do instead:* before designing, + name the real platform and the 2–3 of its fixed properties that constrain + this solution, and design within them. +2. **"Design around a model of the world."** World models (objects, metaphors, + idealized categories) hide the actual data and the actual cost. *Do + instead:* design around the data. Do not introduce an abstraction until + you can describe, concretely, the data it organizes and the transform it + serves — and what the abstraction costs. +3. **"The solution matters more than the data."** The only purpose of any + solution is to transform data from one form to another. *Do instead:* + start every task from the actual inputs and required outputs, never from + the machinery you'd like to build. + +## Core defaults (any problem) + +- **The problem is the data.** Before proposing any solution, describe the + input and output concretely. If you can't, getting that description *is* + the first task — do it before anything else. +- **State the cost.** Every design recommendation you make must state its + cost (time, memory, complexity, maintenance) and on what platform that + cost is paid. A recommendation without a cost is a guess — don't deliver + guesses unlabeled. +- **Solve only the problem you have.** Different data is a different problem. + Concretely: do not add parameters, options, abstraction layers, or + extension points for hypothetical future needs. If you're tempted, write + the one-line note of what you *didn't* build and why, and move on. +- **Where there is one, there are many.** Anything that happens once almost + always happens many times — across space or across the time axis. Default + every design to the batch; treat the single case as a batch of size one. +- **The common case dominates.** Identify the most common case explicitly and + design the straight-line path for it. Handle rare and error cases, but + outside that path — a "maybe" checked everywhere is an "always." +- **Exploit every constraint you have.** List the known constraints (ranges, + volumes, rates, invariants) and use them to remove work. Do not discard a + constraint to make the solution "more general" — that generality is a cost + paid forever for a benefit nobody asked for. +- **Simplicity is removing work.** Prefer fewer states, fewer steps, fewer + special cases, fewer moving parts. Every added state or branch must be + carried, tested, and explained — count them as cost. +- **"Can't be done" is a cost claim.** When something seems impossible, what + is almost always true is that it costs more than it's worth. Say that, with + the estimate, so the tradeoff can actually be decided. + +## Get the real data (required before designing) + +You cannot observe data you were not given — so observe what you *can*, and +label everything else: + +- **Inspect before assuming.** Read representative input files, sample actual + values, read the actual call sites, run the code on real input when a way + to do so exists. Do not design from the type signatures or the docs alone. +- **Sample the data you already have — instrument the live solution.** The + richest data is usually already flowing through the current code; go get it. + Temporarily dump a representative sample of what actually moves through the + system: the arguments reaching a function, the values a hot variable takes, + what a function returns, which branch is taken, the real sizes/counts/lengths. + Then *analyze the sample* — histogram it, sort it, count distinct values, look + at the min/max/mode — and hunt for patterns. Real distributions expose what + the types hide: a variable that is almost always one value, an "array" that is + usually length 0 or 1, input that arrives already sorted or already unique, a + "general" path that is one specific case 98% of the time, a result that is + constant across a run. Each such pattern is a concrete opportunity — specialize + the common case, skip the dead branch, hoist the invariant, precompute the + constant, size the structure to what actually occurs. And the pattern can be + bigger than a local tweak: the data's *shape* can show that a **different + algorithm or representation is the better-fit machine** (sorted-enough → a + different sort/merge; skewed → a different code; runny → a run/stream form; + sparse → a different container), not just that the current machine needs + filing. Sampling justifies *replacing* the machine, not only trimming it. + Sampling is also how you find *new* opportunities mid-optimization, not just + before starting: when a pass **stalls or plateaus**, that is the signal to + re-sample the hottest stage's data and ask whether a different machine fits it + better — not to keep filing the current one. Add the probes, run on real + input, read the output, then remove them — never leave instrumentation on a + timed/measured path. +- **Label every assumption.** For each fact you need but cannot observe, + write an explicit line — `ASSUMPTION: — affects ` — in + your plan, and prefer designs that are cheap to revisit if the assumption + is wrong. Ask the user only when the answer materially changes the design. +- **Never fabricate.** Do not invent plausible-looking values, distributions, + or measurements and treat them as real. + +Answer these about the data (in the tier 1+ plan): + +1. What does the input actually look like — shape, volume, source? +2. What are the most common real values, and how are they distributed? +3. What are the acceptable ranges, and what happens when out-of-range data + arrives? +4. What is the frequency of change — what is stable, what is volatile? +5. What does the solution read and where does it come from? What does it + write and where is it used? What does it touch that it doesn't need? + +## Method (tier 1+ — show this work as a short plan, a line or two per step) + +1. **Frame it.** What is the problem, why is it worth solving, where is the + limit beyond which it isn't, and what is plan B? +2. **Get the data** (section above). +3. **State the cost** of the dominant transform on the real platform. +4. **Design the transform**: a sequence or DAG of explicit transformations — + what comes in, what goes out, what each step is responsible for, with + explicit contracts (shape, meaning, ownership, lifetime, valid ranges) at + each boundary. +5. **Run the simplification pass** (below); say which questions applied and + what work they removed. +6. **Define done.** State the success criteria and what evidence would prove + the approach wrong, before building. +7. **Verify.** Check the result against the real data and the stated + criteria, and report what was and wasn't verified. + +## Simplification pass (run recursively on every sub-problem) + +1. Can we **not do this at all**? +2. Can we do this **only once** (precompute, cache, amortize)? +3. Can we do this **fewer times**? +4. Can we **approximate** the result so that no one notices the difference? +5. Can we use a **small lookup table**? +6. Can we use a **large lookup table**? +7. Can we use a **small buffer/FIFO** to decouple producer from consumer? +8. Can we **constrain the problem further** so a simpler machine suffices? +9. Is there a **different algorithm or representation that fits the data better** + than the current machine? Subtraction has a floor; when filing the current + approach stops paying (a plateau), the win is often a *different* machine the + data's shape points to — reconsider the approach, don't only shrink it. + +## Design rules + +- **Minimize states and branches by design**, not by adding checks. Where the + data genuinely varies, partition it by case and handle each partition + straight-line, rather than re-deciding the case per element. +- **Out-of-range and error behavior is always explicit** — clamp, reject, + drop, or fail loudly; chosen deliberately and written down. Never leave + undefined behavior as an implicit policy, in any tier. +- **Complexity requires evidence.** Add complexity only against a real, + observed need — never a hypothetical one. + +## Performance claims + +- **Never assert an unmeasured performance result.** Not "this should be + faster," not invented numbers. +- If a way to measure exists (benchmark, profiler, test harness, counters), + measure, and include before/after numbers with the change. +- If no way to measure exists here, label the change **unverified**, state + the expected effect as a hypothesis, and specify the exact measurement + that would verify it. +- If there is no measurable performance requirement, build the simplest + correct design and skip speculative optimization entirely. + +--- + +# Software specifics (systems, engine, embedded, game) + +The rules above apply to any problem. These are their conclusions for +software, where the hardware is unforgiving and the data volumes are real. + +## Batch-first transforms (plural by default) + +- Write transforms to operate on **batches/arrays** by default, named in the + **plural** (`update_things`, not `update_thing`). +- A singular call is a degenerate batch: the same batch path with + `count = 1`. Do not maintain separate singular logic without a proven, + measured need. +- Exception: true singletons (configuration state, a single shared resource). + Taking the exception requires a written note: why the data is genuinely + singular and batch semantics don't apply. + +## Memory, layout, and access + +- **Indices over pointers/references/handles by default** (index into a + contiguous array or table). Any pointer-heavy hot path must include a + short written justification for why indices are insufficient. +- Organize data by **access pattern, not conceptual ownership**. Split hot + and cold fields when the cold fields aren't needed in the dominant loop. +- For each hot path, write down the expected **access pattern** + (linear / strided / random), expected **branch behavior** + (predictable / unpredictable), and the hardware assumptions. +- When branch entropy is high, prefer **partitioned passes** (bucket by + state/tag, process each bucket straight-line) over per-element branching. +- Keep the common-case path branch-minimal; rare and error handling lives + outside the hot loop. + +## Data protocols between systems + +Systems communicate through **explicit data protocols**, modeled after +network protocols and file formats — explicit layout, versioning, documented +meaning. The default is a **flat struct**: fixed layout, no hidden pointers, +no OO-style interfaces. Use tagged unions or header-plus-payload when the +flat struct genuinely can't express it. Do not model system boundaries as +objects, virtual calls, or opaque handles. + +## Hardware is the platform + +- Design with the actual hardware's properties — cache hierarchy, memory + bandwidth, alignment, latency vs. throughput — and to its strengths. +- **Latency and throughput are only the same thing in a sequential system.** + For every performance requirement, identify which one it actually is + before designing for it. +- The compiler and language are tools, not magic: memory layout, access + order, and the choice of what work to do at all are your job, not theirs — + and they are roughly 90% of the problem. Know what the compiler can + reasonably do with what you wrote, and don't delegate what it can't. + +## Enforceable deliverables (tier 2) + +For each new or substantially reworked subsystem: + +- One explicit **batch transform contract**: input layout, output layout, + owner, lifetime, valid value ranges. +- A **plural/batch path** for every transform; singular calls are thin + wrappers over the batch implementation (`count = 1`) unless documented as + a true singleton. +- A written **justification for any pointer/reference/handle-heavy hot path** + explaining why index-based access is insufficient. +- Explicit **out-of-range behavior** (clamp/reject/drop/error) at every + input boundary. +- Unresolved design questions filed as **local issue files under `issues/`** + — not GitHub issues, not inline TODOs. + +--- + +# Final self-check (run before delivering tier 1+ work) + +Verify, and fix or flag anything that fails: + +- [ ] The plan answered the framing, data, and cost questions — or every gap + is labeled `ASSUMPTION` with what it affects. +- [ ] The most common case is identified and the design serves it + straight-line; rare/error cases are out of the common path. +- [ ] The simplification pass ran; the work it removed (or why nothing could + be removed) is stated. +- [ ] No speculative generality: no parameter, option, or abstraction exists + for a need that isn't real yet. +- [ ] Out-of-range and error behavior is explicit at every boundary. +- [ ] Transforms are plural/batch, or the singleton exception is documented. +- [ ] Pointer-heavy hot paths carry their written justification; everything + else uses indices. +- [ ] No unmeasured performance claim anywhere in code, comments, or + summary; measurements included where possible, hypotheses labeled + where not. +- [ ] Done-criteria from the plan were checked, and the summary reports what + was verified and what wasn't. +- [ ] (Tier 2) Deliverables above are present; open questions are filed + under `issues/`. + +--- + +# Commit + +When user ask to commit, use mitchellh style commit. diff --git a/docs/stats.md b/docs/stats.md new file mode 100644 index 0000000..8c7aaad --- /dev/null +++ b/docs/stats.md @@ -0,0 +1,116 @@ +# Performance Stats + +Hardware: **AMD Strix Halo** (Ryzen AI Max 395 Pro, Radeon 8060S, 128 GB unified memory). +All runs use `novaAnimeXL_ilV190.safetensors` (SDXL), bfloat16, cfg=4.5. + +## Generate (SDXL) + +| Run | Resolution | Steps | Prompt Syntax | Load (s) | Inference (s) | Size | Notes | +|-----|-----------|-------|---------------|----------|---------------|------|-------| +| `character_base` | 1024×1024 | 20 | `BREAK` | 2.15 | 28.47 | 1.3 MB | Baseline, no weighting | +| `character_weighted` | 1024×1024 | 20 | `(word:weight)` + `BREAK` | 2.57 | 29.94 | 1.4 MB | Full Compel syntax | +| `background_classroom` | 1280×720 | 20 | `(word:weight)` | 2.08 | 181.85 | 1.2 MB | VAE decode dominated (flash attn enabled) | + +### Per-step breakdown (1024×1024) + +| Step | Time | +|------|------| +| 1 (warmup) | ~1.3–1.6s | +| Steady state | ~1.0–1.3s | +| Total (20 steps) | ~20–22s UNet + VAE decode | + +### Larger resolutions + +1280×720 and 1920×1080 trigger flash attention kernel compilation on first run +(up to 250s for the first step). Subsequent runs reuse cached kernels. VAE +decode at these resolutions is the dominant cost — 1920×1080 decode can exceed +2 minutes. + +### Compel overhead + +Prompt weighting via `compel` adds negligible overhead (~0.4s for encoding +long prompts with `BREAK`). The embedding path is identical to raw encoding +once tensors reach the UNet. + +## Edit (Qwen Image Edit + Lightning LoRA) + +All edits use `qwen_image_edit_2509_fp8_e4m3fn.safetensors` + Lightning 4-step +LoRA with turbo settings (steps=4, cfg=1.0) on 1024×1024 input images. + +| Run | Steps | Load (s) | LoRA (s) | Inference (s) | Notes | +|-----|-------|----------|----------|---------------|-------| +| `base_smile` | 4 | 28.16 | 1.49 | 86.81 | Happy smile variant (matmul fallback) | +| `base_smile_flash` | 4 | 31.23 | 1.36 | 53.33 | Happy smile variant (flash attention) | + +### Per-step breakdown (1024×1024, turbo) + +**Matmul fallback (no flash attention):** + +| Step | Time | +|------|------| +| 1 | ~0.3s | +| 2 | ~4.9s | +| 3 | ~11.5s | +| 4 | ~13.7s | + +**Flash attention (`TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1`):** + +| Step | Time | +|------|------| +| 1 | ~0.3s | +| 2 | ~1.6s | +| 3 | ~5.1s | +| 4 | ~7.0s | + +Flash attention cuts edit inference from 86.8s → 53.3s (**1.63× speedup**). +SDXL generate is unaffected (uses its own attention processor, not SDPA). +Step 1 is fast (prefill/encoding). Steps 2–4 engage the full transformer +and VAE; flash attention reduces the attention bottleneck. + +## Session (persistent models) + +Both SDXL and Qwen loaded eagerly into a single `VnAssetsSession`. +Models held in GPU memory across calls. + +| Phase | Wall (s) | Details | +|-------|----------|---------| +| Session load | 28.2 | SDXL UNet + Qwen transformer + VAE + TE + LoRA fuse | +| Generate | 30.7 | SDXL 20-step, 1024×1024, Compel encoding | +| Edit (turbo) | 87.2 | Qwen 4-step, 1024×1024, Lightning LoRA | +| **Total wall** | **146.2** | One session, 2 operations | + +### Session with flash attention + +`TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1` throughout. + +| Phase | Wall (s) | vs matmul fallback | +|-------|----------|--------------------| +| Session load | 31.2 | ~same | +| Generate | 33.1 | ~same (SDXL uses own attn processor) | +| Edit (turbo) | 53.7 | **1.62× faster** | +| **Total wall** | **118.2** | 1.24× faster overall | + +### Session vs cold start + +| Approach | Generate | Edit | Total | +|----------|----------|------|-------| +| Standalone (cold) | ~33s | ~116s | ~149s | +| Session (matmul) | ~31s | ~87s | ~146s | +| Session (flash) | ~33s | ~54s | ~118s | +| **Saved vs cold** | — | **~62s** | **~31s** | + +With 2 operations the session saves one model-load round trip (~28s). +The saving grows linearly with more edits: 5 edits save 4×28s = 112s. +Flash attention adds a further 1.6× multiplier on edit inference time. + +## Notes + +- `inference_time_s` includes VAE decode, which is disproportionately expensive + at non-square resolutions on this hardware. +- `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1` enables ROCm flash attention, + cutting per-step UNet time roughly in half after kernel compilation. +- First run at a new resolution incurs kernel compilation cost; subsequent runs + at the same resolution are fast. +- Session `load_time_s` in metadata reflects total session construction + (all models loaded); individual operation inference times exclude loading. +- LoRA fuse time (~1.5s) is included in session load, once. diff --git a/test_batch.py b/test_batch.py new file mode 100644 index 0000000..acb8e13 --- /dev/null +++ b/test_batch.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python3 +"""Deep batch test: generate base image, edit with multiple emotions, track timing.""" +import subprocess +import time +import json +from pathlib import Path + +OUTPUT_DIR = Path("output/batch_test") +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" +SEED = 42 +STEPS = 20 +CFG = 4.5 + +# Emotion variants to edit +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"), +] + +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(STEPS), + "--cfg", str(CFG), + "--seed", str(SEED), + "--output", str(base_path), +], check=True) + +t_gen = time.perf_counter() - t_start + +# Read metadata +meta = json.loads((OUTPUT_DIR / "base.json").read_text()) +results.append({ + "step": "generate_base", + "output": str(base_path), + "prompt": BASE_PROMPT, + "load_s": meta["load_time_s"], + "inference_s": meta["inference_time_s"], + "total_wall_s": round(t_gen, 1), +}) +print(f" Wall time: {t_gen:.1f}s (load: {meta['load_time_s']}s, inference: {meta['inference_time_s']}s)\n") + +# ── 2. Edit with each emotion ─────────────────────────────────── +total_edits_start = time.perf_counter() + +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(STEPS), + "--cfg", str(CFG), + "--seed", str(SEED), + "--output", str(edit_path), + ], check=True) + 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), + "prompt": prompt, + "load_s": meta["load_time_s"], + "inference_s": meta["inference_time_s"], + "total_wall_s": round(t_edit, 1), + }) + print(f" Wall time: {t_edit:.1f}s (load: {meta['load_time_s']}s, inference: {meta['inference_time_s']}s)\n") + +total_edits_wall = time.perf_counter() - total_edits_start + +# ── 3. Summary ────────────────────────────────────────────────── +print("=" * 60) +print("SUMMARY") +print("=" * 60) +print(f"{'Step':<20} {'Load':>8} {'Infer':>8} {'Wall':>8}") +print("-" * 48) +total_load = 0 +total_infer = 0 +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") + +# Write results JSON +summary = { + "config": { + "checkpoint": CHECKPOINT, + "edit_model": EDIT_MODEL, + "seed": SEED, + "steps": STEPS, + "cfg": 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 saved to {summary_path}") +print(f"Images in {OUTPUT_DIR}/") diff --git a/test_batch_fast.py b/test_batch_fast.py new file mode 100644 index 0000000..64181d1 --- /dev/null +++ b/test_batch_fast.py @@ -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}") diff --git a/vnassets/__init__.py b/vnassets/__init__.py index e69de29..c531eae 100644 --- a/vnassets/__init__.py +++ b/vnassets/__init__.py @@ -0,0 +1,5 @@ +"""VNAsset — fast CLI pipeline for visual novel image asset generation.""" + +from .session import VnAssetsSession + +__all__ = ["VnAssetsSession"] diff --git a/vnassets/edit.py b/vnassets/edit.py index 6d89384..d3a171a 100644 --- a/vnassets/edit.py +++ b/vnassets/edit.py @@ -1,56 +1,8 @@ -"""Qwen Image Edit — image-to-image editing.""" -import gc -import json -import random -import time -from pathlib import Path +"""Qwen Image Edit — image-to-image editing (standalone entry point). -import safetensors.torch -import torch -from accelerate import init_empty_weights -from PIL import Image -from diffusers import QwenImageEditPlusPipeline, FlowMatchEulerDiscreteScheduler -from diffusers.models.autoencoders import AutoencoderKLQwenImage -from diffusers.models.transformers import QwenImageTransformer2DModel -from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor - -from .attention import patch_qwen_transformer - - -TEXT_ENCODER_ID = "Qwen/Qwen2.5-VL-7B-Instruct" -VAE_ID = "Qwen/Qwen-Image" - - -def _load_transformer(path: str, dtype: torch.dtype) -> QwenImageTransformer2DModel: - """Load Qwen Image Edit transformer from a single FP8 safetensors file. - - Uses init_empty_weights and incremental conversion to keep peak memory - manageable. The model is 20B parameters (20 GB FP8, 40 GB BF16).""" - config = QwenImageTransformer2DModel.load_config( - "Qwen/Qwen-Image-Edit", subfolder="transformer" - ) - - state_dict = safetensors.torch.load_file(path) - prefix = "model.diffusion_model." - - # Convert FP8 -> target dtype, freeing FP8 tensors as we go - cleaned = {} - for k in list(state_dict.keys()): - if k.startswith(prefix): - v = state_dict.pop(k) - cleaned[k[len(prefix):]] = v.to(dtype) - del v - del state_dict - gc.collect() - - # Create model on meta device to avoid allocating full model in addition to cleaned dict - with init_empty_weights(): - model = QwenImageTransformer2DModel.from_config(config, torch_dtype=dtype) - - model.load_state_dict(cleaned, strict=True, assign=True) - del cleaned - gc.collect() - return model +For multi-call reuse, use VnAssetsSession directly. +""" +from .session import VnAssetsSession def edit( @@ -63,84 +15,17 @@ def edit( output_path: str = "output.png", lora_path: str | None = None, ) -> None: - device = "cuda" if torch.cuda.is_available() else "cpu" - dtype = torch.bfloat16 + """One-shot image edit. Loads model, runs inference, unloads. - if seed is None: - seed = random.randint(0, 2**32 - 1) - - output = Path(output_path) - output.parent.mkdir(parents=True, exist_ok=True) - - t0 = time.perf_counter() - - transformer = _load_transformer(model_path, dtype) - vae = AutoencoderKLQwenImage.from_pretrained(VAE_ID, subfolder="vae", torch_dtype=dtype) - text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( - TEXT_ENCODER_ID, torch_dtype=dtype - ) - tokenizer = Qwen2Tokenizer.from_pretrained(TEXT_ENCODER_ID) - processor = Qwen2VLProcessor.from_pretrained(TEXT_ENCODER_ID) - scheduler = FlowMatchEulerDiscreteScheduler() - - pipe = QwenImageEditPlusPipeline( - scheduler=scheduler, - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - processor=processor, - transformer=transformer, - ) - pipe.to(device) - patch_qwen_transformer(transformer) - - t_lora = 0.0 - lora_fused = False - if lora_path: - tl = time.perf_counter() - pipe.load_lora_weights(lora_path) - pipe.fuse_lora(lora_scale=1.0, components=["transformer"]) - lora_fused = True - t_lora = time.perf_counter() - tl - print(f"LoRA loaded + fused: {t_lora:.1f}s") - - t_load = time.perf_counter() - t0 - - input_image = Image.open(input_path).convert("RGB") - generator = torch.Generator(device=device).manual_seed(seed) - - t1 = time.perf_counter() - image = pipe( - image=input_image, - prompt=prompt, - true_cfg_scale=cfg, - num_inference_steps=steps, - generator=generator, - ).images[0] - t_infer = time.perf_counter() - t1 - - image.save(output) - print(f"Saved {output}") - - del pipe, transformer, vae, text_encoder - if device == "cuda": - torch.cuda.empty_cache() - - meta_path = output.with_suffix(".json") - meta = { - "model": str(Path(model_path).resolve()), - "vae": VAE_ID, - "text_encoder": TEXT_ENCODER_ID, - "input_image": str(Path(input_path).resolve()), - "prompt": prompt, - "steps": steps, - "cfg": cfg, - "seed": seed, - "lora_path": str(Path(lora_path).resolve()) if lora_path else None, - "lora_load_s": round(t_lora, 2) if lora_path else None, - "lora_fused": lora_fused, - "load_time_s": round(t_load, 2), - "inference_time_s": round(t_infer, 2), - } - meta_path.write_text(json.dumps(meta, indent=2)) - print(f"Saved {meta_path}") + For multiple edits, create a VnAssetsSession to keep the model + loaded between calls. + """ + with VnAssetsSession(edit_model=model_path, edit_lora=lora_path) as vna: + vna.edit( + input_path=input_path, + prompt=prompt, + steps=steps, + cfg=cfg, + seed=seed, + output_path=output_path, + ) diff --git a/vnassets/generate.py b/vnassets/generate.py index 80607cd..d5f9be7 100644 --- a/vnassets/generate.py +++ b/vnassets/generate.py @@ -1,14 +1,8 @@ -"""SDXL text-to-image generation.""" -import json -import random -import time -from pathlib import Path +"""SDXL text-to-image generation (standalone entry point). -import torch -from diffusers import StableDiffusionXLPipeline - -from .attention import patch_unet_attention -from .prompt import build_compel, encode_prompts +For multi-call reuse, use VnAssetsSession directly. +""" +from .session import VnAssetsSession def generate( @@ -23,78 +17,20 @@ def generate( output_path: str = "output.png", raw: bool = False, ) -> None: - device = "cuda" if torch.cuda.is_available() else "cpu" - # bfloat16 avoids ROCm kernel crashes on RDNA 3.5; float16 segfaults - dtype = torch.bfloat16 + """One-shot SDXL generation. Loads model, runs inference, unloads. - if seed is None: - seed = random.randint(0, 2**32 - 1) - - output = Path(output_path) - output.parent.mkdir(parents=True, exist_ok=True) - - t0 = time.perf_counter() - pipe = StableDiffusionXLPipeline.from_single_file( - checkpoint_path, - torch_dtype=dtype, - ) - pipe.to(device) - patch_unet_attention(pipe.unet) - - if not raw: - compel = build_compel(pipe) - prompt_embeds, pooled_embeds, neg_embeds, neg_pooled = encode_prompts( - compel, prompt, negative_prompt - ) - t_load = time.perf_counter() - t0 - - generator = torch.Generator(device=device).manual_seed(seed) - - t1 = time.perf_counter() - if raw: - image = pipe( + For multiple generations, create a VnAssetsSession to keep the model + loaded between calls. + """ + with VnAssetsSession(sdxl_checkpoint=checkpoint_path) as vna: + vna.generate( prompt=prompt, negative_prompt=negative_prompt, width=width, height=height, - num_inference_steps=steps, - guidance_scale=cfg, - generator=generator, - ).images[0] - else: - image = pipe( - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_embeds, - negative_prompt_embeds=neg_embeds, - negative_pooled_prompt_embeds=neg_pooled, - width=width, - height=height, - num_inference_steps=steps, - guidance_scale=cfg, - generator=generator, - ).images[0] - t_infer = time.perf_counter() - t1 - - image.save(output) - print(f"Saved {output}") - - # Free GPU memory before returning - del pipe - if device == "cuda": - torch.cuda.empty_cache() - - meta_path = output.with_suffix(".json") - meta = { - "checkpoint": str(Path(checkpoint_path).resolve()), - "prompt": prompt, - "negative_prompt": negative_prompt, - "width": width, - "height": height, - "steps": steps, - "cfg": cfg, - "seed": seed, - "load_time_s": round(t_load, 2), - "inference_time_s": round(t_infer, 2), - } - meta_path.write_text(json.dumps(meta, indent=2)) - print(f"Saved {meta_path}") + steps=steps, + cfg=cfg, + seed=seed, + output_path=output_path, + raw=raw, + ) diff --git a/vnassets/session.py b/vnassets/session.py new file mode 100644 index 0000000..4380628 --- /dev/null +++ b/vnassets/session.py @@ -0,0 +1,361 @@ +"""Persistent model session — SDXL and Qwen Image Edit held in GPU memory. + +Models are loaded eagerly at construction and reused across generate()/edit() +calls. On 128 GB unified memory (Strix Halo), everything fits simultaneously. +""" +import gc +import json +import random +import time +from pathlib import Path + +import safetensors.torch +import torch +from accelerate import init_empty_weights +from PIL import Image +from diffusers import ( + FlowMatchEulerDiscreteScheduler, + QwenImageEditPlusPipeline, + StableDiffusionXLPipeline, +) +from diffusers.models.autoencoders import AutoencoderKLQwenImage +from diffusers.models.transformers import QwenImageTransformer2DModel +from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor + +from .attention import patch_qwen_transformer, patch_unet_attention +from .prompt import build_compel, encode_prompts + +TEXT_ENCODER_ID = "Qwen/Qwen2.5-VL-7B-Instruct" +VAE_ID = "Qwen/Qwen-Image" + + +class VnAssetsSession: + """Holds SDXL and/or Qwen Image Edit models in GPU memory for reuse. + + Usage as context manager:: + + with VnAssetsSession( + sdxl_checkpoint="models/novaAnimeXL.safetensors", + edit_model="models/qwen_image_edit.safetensors", + edit_lora="models/lightning-4steps.safetensors", + ) as vna: + vna.generate("1girl, red hair", output="base.png") + vna.edit("base.png", "make her smile", output="happy.png") + vna.edit("base.png", "make her sad", output="sad.png") + + Or manual lifecycle:: + + vna = VnAssetsSession(sdxl_checkpoint=...) + vna.generate(...) + vna.close() + """ + + def __init__( + self, + sdxl_checkpoint: str | None = None, + edit_model: str | None = None, + edit_lora: str | None = None, + ): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + # bfloat16 avoids ROCm kernel crashes on RDNA 3.5; float16 segfaults + self.dtype = torch.bfloat16 + + self._sdxl_checkpoint = sdxl_checkpoint + self._edit_model = edit_model + self._edit_lora = edit_lora + self._lora_fused = False + self._lora_load_s: float | None = None + + self._pipe_sdxl: StableDiffusionXLPipeline | None = None + self._compel = None + self._pipe_qwen: QwenImageEditPlusPipeline | None = None + + t0 = time.perf_counter() + if sdxl_checkpoint: + self._load_sdxl(sdxl_checkpoint) + if edit_model: + self._load_qwen(edit_model, edit_lora) + self._load_time_s = round(time.perf_counter() - t0, 2) + + loaded = [] + if self._pipe_sdxl: + loaded.append("SDXL") + if self._pipe_qwen: + loaded.append("Qwen") + if loaded: + print(f"Session ready ({'+'.join(loaded)}, {self._load_time_s}s)") + + # ── SDXL ──────────────────────────────────────────────────────────── + + def _load_sdxl(self, checkpoint_path: str) -> None: + pipe = StableDiffusionXLPipeline.from_single_file( + checkpoint_path, + torch_dtype=self.dtype, + ) + pipe.to(self.device) + patch_unet_attention(pipe.unet) + self._pipe_sdxl = pipe + self._compel = build_compel(pipe) + + # ── Qwen Image Edit ───────────────────────────────────────────────── + + def _load_qwen(self, model_path: str, lora_path: str | None) -> None: + transformer = self._load_transformer(model_path) + vae = AutoencoderKLQwenImage.from_pretrained( + VAE_ID, subfolder="vae", torch_dtype=self.dtype + ) + text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( + TEXT_ENCODER_ID, torch_dtype=self.dtype + ) + tokenizer = Qwen2Tokenizer.from_pretrained(TEXT_ENCODER_ID) + processor = Qwen2VLProcessor.from_pretrained(TEXT_ENCODER_ID) + scheduler = FlowMatchEulerDiscreteScheduler() + + pipe = QwenImageEditPlusPipeline( + scheduler=scheduler, + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + processor=processor, + transformer=transformer, + ) + pipe.to(self.device) + patch_qwen_transformer(transformer) + + if lora_path: + t_lora = time.perf_counter() + pipe.load_lora_weights(lora_path) + pipe.fuse_lora(lora_scale=1.0, components=["transformer"]) + self._lora_fused = True + self._lora_load_s = round(time.perf_counter() - t_lora, 2) + print(f"LoRA loaded + fused: {self._lora_load_s}s") + + self._pipe_qwen = pipe + + def _load_transformer(self, path: str) -> QwenImageTransformer2DModel: + """Load Qwen Image Edit transformer from a single FP8 safetensors file. + + Uses init_empty_weights and incremental conversion to keep peak memory + manageable. The model is 20B parameters (20 GB FP8, 40 GB BF16). + """ + config = QwenImageTransformer2DModel.load_config( + "Qwen/Qwen-Image-Edit", subfolder="transformer" + ) + state_dict = safetensors.torch.load_file(path) + prefix = "model.diffusion_model." + + # Convert FP8 -> target dtype, freeing FP8 tensors as we go + cleaned = {} + for k in list(state_dict.keys()): + if k.startswith(prefix): + v = state_dict.pop(k) + cleaned[k[len(prefix):]] = v.to(self.dtype) + del v + del state_dict + gc.collect() + + with init_empty_weights(): + model = QwenImageTransformer2DModel.from_config(config, torch_dtype=self.dtype) + + model.load_state_dict(cleaned, strict=True, assign=True) + del cleaned + gc.collect() + return model + + # ── Properties ────────────────────────────────────────────────────── + + @property + def load_time_s(self) -> float: + """Total time spent loading models at session construction (seconds).""" + return self._load_time_s + + @property + def has_sdxl(self) -> bool: + return self._pipe_sdxl is not None + + @property + def has_qwen(self) -> bool: + return self._pipe_qwen is not None + + # ── Generate (SDXL text-to-image) ─────────────────────────────────── + + def generate( + self, + prompt: str, + negative_prompt: str = "", + width: int = 1024, + height: int = 1024, + steps: int = 20, + cfg: float = 4.5, + seed: int | None = None, + output_path: str = "output.png", + raw: bool = False, + ) -> None: + """Generate an image from the loaded SDXL checkpoint. + + Args: + prompt: Positive prompt (ComfyUI weighting syntax unless ``raw``). + negative_prompt: Negative prompt. + width, height: Output resolution in pixels. + steps: Number of inference steps. + cfg: CFG scale. + seed: RNG seed (random if None). + output_path: Where to save the PNG. Metadata is written to + ``{output_path}.json``. Parent directories are created. + raw: If True, bypass Compel prompt weighting and use plain + diffusers encoding. + + Raises: + RuntimeError: If SDXL was not loaded at session construction. + """ + if self._pipe_sdxl is None: + raise RuntimeError( + "SDXL model not loaded. Provide sdxl_checkpoint when creating VnAssetsSession." + ) + + if seed is None: + seed = random.randint(0, 2**32 - 1) + + output = Path(output_path) + output.parent.mkdir(parents=True, exist_ok=True) + + generator = torch.Generator(device=self.device).manual_seed(seed) + + t0 = time.perf_counter() + if raw: + image = self._pipe_sdxl( + prompt=prompt, + negative_prompt=negative_prompt, + width=width, + height=height, + num_inference_steps=steps, + guidance_scale=cfg, + generator=generator, + ).images[0] + else: + prompt_embeds, pooled_embeds, neg_embeds, neg_pooled = encode_prompts( + self._compel, prompt, negative_prompt + ) + image = self._pipe_sdxl( + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_embeds, + negative_prompt_embeds=neg_embeds, + negative_pooled_prompt_embeds=neg_pooled, + width=width, + height=height, + num_inference_steps=steps, + guidance_scale=cfg, + generator=generator, + ).images[0] + t_infer = round(time.perf_counter() - t0, 2) + + image.save(output) + print(f"Saved {output}") + + meta_path = output.with_suffix(".json") + meta = { + "checkpoint": str(Path(self._sdxl_checkpoint).resolve()), + "prompt": prompt, + "negative_prompt": negative_prompt, + "width": width, + "height": height, + "steps": steps, + "cfg": cfg, + "seed": seed, + "load_time_s": self._load_time_s, + "inference_time_s": t_infer, + } + meta_path.write_text(json.dumps(meta, indent=2)) + print(f"Saved {meta_path}") + + # ── Edit (Qwen Image Edit) ────────────────────────────────────────── + + def edit( + self, + input_path: str, + prompt: str, + steps: int = 20, + cfg: float = 4.0, + seed: int | None = None, + output_path: str = "output.png", + ) -> None: + """Edit an image using the loaded Qwen Image Edit model. + + Args: + input_path: Path to the input image to edit. + prompt: Edit instruction (e.g. "make her smile"). + steps: Number of inference steps (4 with Lightning LoRA). + cfg: CFG scale (1.0 with Lightning LoRA). + seed: RNG seed (random if None). + output_path: Where to save the PNG. Metadata is written to + ``{output_path}.json``. Parent directories are created. + + Raises: + RuntimeError: If Qwen was not loaded at session construction. + """ + if self._pipe_qwen is None: + raise RuntimeError( + "Qwen model not loaded. Provide edit_model when creating VnAssetsSession." + ) + + if seed is None: + seed = random.randint(0, 2**32 - 1) + + output = Path(output_path) + output.parent.mkdir(parents=True, exist_ok=True) + + input_image = Image.open(input_path).convert("RGB") + generator = torch.Generator(device=self.device).manual_seed(seed) + + t0 = time.perf_counter() + image = self._pipe_qwen( + image=input_image, + prompt=prompt, + true_cfg_scale=cfg, + num_inference_steps=steps, + generator=generator, + ).images[0] + t_infer = round(time.perf_counter() - t0, 2) + + image.save(output) + print(f"Saved {output}") + + meta_path = output.with_suffix(".json") + meta = { + "model": str(Path(self._edit_model).resolve()), + "vae": VAE_ID, + "text_encoder": TEXT_ENCODER_ID, + "input_image": str(Path(input_path).resolve()), + "prompt": prompt, + "steps": steps, + "cfg": cfg, + "seed": seed, + "lora_path": str(Path(self._edit_lora).resolve()) if self._edit_lora else None, + "lora_load_s": self._lora_load_s, + "lora_fused": self._lora_fused, + "load_time_s": self._load_time_s, + "inference_time_s": t_infer, + } + meta_path.write_text(json.dumps(meta, indent=2)) + print(f"Saved {meta_path}") + + # ── Lifecycle ─────────────────────────────────────────────────────── + + def close(self) -> None: + """Release all models and free GPU memory.""" + if self._pipe_sdxl: + del self._pipe_sdxl + self._pipe_sdxl = None + if self._pipe_qwen: + del self._pipe_qwen + self._pipe_qwen = None + self._compel = None + if self.device == "cuda": + torch.cuda.empty_cache() + + def __enter__(self) -> "VnAssetsSession": + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> bool: + self.close() + return False