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
289
AGENTS.md
Normal file
289
AGENTS.md
Normal file
@@ -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: <fact> — affects <decision>` — 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.
|
||||
116
docs/stats.md
Normal file
116
docs/stats.md
Normal file
@@ -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.
|
||||
136
test_batch.py
Normal file
136
test_batch.py
Normal file
@@ -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}/")
|
||||
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}")
|
||||
@@ -0,0 +1,5 @@
|
||||
"""VNAsset — fast CLI pipeline for visual novel image asset generation."""
|
||||
|
||||
from .session import VnAssetsSession
|
||||
|
||||
__all__ = ["VnAssetsSession"]
|
||||
|
||||
147
vnassets/edit.py
147
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,
|
||||
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,
|
||||
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}")
|
||||
steps=steps,
|
||||
cfg=cfg,
|
||||
seed=seed,
|
||||
output_path=output_path,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
361
vnassets/session.py
Normal file
361
vnassets/session.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user