Merge TECH_SPEC.md into README.md, fix all inconsistencies
- Consolidate TECH_SPEC into README as single source of truth - Remove TECH_SPEC.md (375-line README covers everything) - Fix architecture diagram: Qwen transformer ~5GB → ~20GB (20B params) - Remove unimplemented CLI flags from docs: --sampler, --scheduler, --clip, --vae, --turbo (these don't exist in the actual code) - Replace 'FP8 native compute' claim with actual FP8→BF16 conversion - Replace 'peft' with actual diffusers native LoRA loading - Move aspirational optimizations (VAE overlap, shared encode) to Future - Resolve stale Open Questions; keep torch.compile as Future item - Move 'Persistent model session' and 'Qwen Image Edit' from Future to implemented status with documented API - Update performance numbers with actual measured session+flash numbers - Document flash attention path alongside matmul fallback - Simplify data flow diagram, remove ComfyUI-only kontext_scale concept - Remove dead doc links (sdxl-generation.md)
This commit is contained in:
Michele Rossi
261
README.md
261
README.md
@@ -12,6 +12,10 @@ Built for **AMD Strix Halo** (Ryzen AI Max 395 Pro, Radeon 8060S, 128 GB unified
|
||||
memory). Also works on discrete AMD GPUs with ROCm. NVIDIA support is untested
|
||||
but should work if you swap the torch backend.
|
||||
|
||||
The 128 GB unified memory means VRAM is effectively unlimited (up to ~96 GB
|
||||
allocatable to GPU). The bottleneck is **GPU compute throughput**, not memory
|
||||
capacity — model offloading is pointless, everything stays resident.
|
||||
|
||||
## Install
|
||||
|
||||
```bash
|
||||
@@ -37,11 +41,13 @@ Symlink your ComfyUI models into `models/`:
|
||||
cd models
|
||||
ln -s /path/to/ComfyUI/models/checkpoints/novaAnimeXL_ilV190.safetensors .
|
||||
ln -s /path/to/ComfyUI/models/diffusion_models/qwen_image_edit_2509_fp8_e4m3fn.safetensors .
|
||||
ln -s /path/to/ComfyUI/models/text_encoders/qwen_2.5_vl_7b_fp8_scaled.safetensors .
|
||||
ln -s /path/to/ComfyUI/models/vae/qwen_image_vae.safetensors .
|
||||
ln -s /path/to/ComfyUI/models/loras/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors .
|
||||
```
|
||||
|
||||
The Qwen VAE and text encoder are downloaded automatically from HuggingFace Hub
|
||||
on first use (`Qwen/Qwen-Image` and `Qwen/Qwen2.5-VL-7B-Instruct`). No symlinks
|
||||
needed for those.
|
||||
|
||||
Or place the actual files there — the tool just reads whatever safetensors you
|
||||
point it at.
|
||||
|
||||
@@ -68,21 +74,10 @@ vnasset generate \
|
||||
| `--height` | `1024` | Image height |
|
||||
| `--steps` | `20` | Inference steps |
|
||||
| `--cfg` | `4.5` | CFG scale |
|
||||
| `--seed` | `0` | RNG seed (use `random` for random) |
|
||||
| `--seed` | `random` | RNG seed (integer or `random`) |
|
||||
| `--output` | `output.png` | Output path |
|
||||
| `--raw` | `false` | Disable Compel prompt weighting (fall back to plain diffusers encoding) |
|
||||
|
||||
When `--seed` is `random`, a random seed is generated and recorded in the
|
||||
metadata file.
|
||||
|
||||
### Output
|
||||
|
||||
Each generation produces:
|
||||
- `{output}.png` — the image
|
||||
- `{output}.json` — metadata (prompt, seed, model path, timing, resolution)
|
||||
|
||||
Directories in `--output` are created automatically.
|
||||
|
||||
### Edit (Qwen Image Edit)
|
||||
|
||||
```bash
|
||||
@@ -97,7 +92,7 @@ vnasset edit \
|
||||
|
||||
| Option | Default | Description |
|
||||
|--------|---------|-------------|
|
||||
| `--model` | (required) | Path to Qwen Image Edit `.safetensors` |
|
||||
| `--model` | (required) | Path to Qwen Image Edit `.safetensors` (FP8) |
|
||||
| `--input` | (required) | Input image to edit |
|
||||
| `--prompt` | (required) | Edit instruction |
|
||||
| `--steps` | `20` | Inference steps (`4` with Lightning LoRA) |
|
||||
@@ -106,36 +101,165 @@ vnasset edit \
|
||||
| `--lora` | (none) | Path to LoRA `.safetensors` |
|
||||
| `--output` | `output.png` | Output path |
|
||||
|
||||
## Current State
|
||||
**Turbo mode:** Use the Lightning 4-step LoRA with `--steps 4 --cfg 1.0` to cut
|
||||
inference time proportionally. The LoRA is fused into the transformer at load
|
||||
time, so there is no per-step LoRA overhead.
|
||||
|
||||
| Command | Status |
|
||||
|---------|--------|
|
||||
| `vnasset generate` | ✅ Working |
|
||||
| `vnasset edit` | ✅ Working |
|
||||
| `vnasset pipeline` | 🚧 Planned |
|
||||
### Output
|
||||
|
||||
Each generation produces:
|
||||
- `{output}.png` — the image
|
||||
- `{output}.json` — metadata (prompt, seed, model path, timing, resolution)
|
||||
|
||||
Directories in `--output` are created automatically.
|
||||
|
||||
### Session (persistent models)
|
||||
|
||||
For multi-call workflows, use `VnAssetsSession` to keep models loaded in GPU
|
||||
memory between operations. Models are loaded eagerly at construction:
|
||||
|
||||
```python
|
||||
from vnassets import VnAssetsSession
|
||||
|
||||
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")
|
||||
```
|
||||
|
||||
Either model can be omitted (`None`) for single-model sessions. Properties:
|
||||
- `vna.has_sdxl` / `vna.has_qwen` — check which models are loaded
|
||||
- `vna.load_time_s` — total session construction time
|
||||
- `vna.close()` — manual cleanup (automatic with `with`)
|
||||
|
||||
The standalone `vnasset generate` and `vnasset edit` CLI commands are thin
|
||||
wrappers around a one-shot session — same API, backwards compatible.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ VnAssetsSession │
|
||||
│ │
|
||||
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
|
||||
│ │ SDXL │ │ Qwen │ │ Qwen VL │ │ Qwen │ │
|
||||
│ │ UNet │ │ Transf. │ │ 7B TE │ │ VAE │ │
|
||||
│ │ (~3.5GB) │ │ (~20GB) │ │ (~14GB) │ │ (~1GB) │ │
|
||||
│ └──────────┘ └──────────┘ └──────────┘ └──────────┘ │
|
||||
│ │ │ │ │ │
|
||||
│ ▼ │ │ │ │
|
||||
│ ┌─────────┐ │ │ │ │
|
||||
│ │ Generate│ │ │ │ │
|
||||
│ │ │──────────┼───────────────┼──────────────┤ │
|
||||
│ │ │ base.png │ │ │ │
|
||||
│ └─────────┘ │ │ │ │
|
||||
│ │ ▼ ▼ ▼ │
|
||||
│ │ ┌──────────────────────────────────────┐ │
|
||||
│ └──────────► Edit Phase │ │
|
||||
│ │ base.png + prompts[] → variants[] │ │
|
||||
│ └──────────────────────────────────────┘ │
|
||||
│ │ │
|
||||
│ ▼ │
|
||||
│ ┌──────────────────┐ │
|
||||
│ │ base.png │ │
|
||||
│ │ happy.png │ │
|
||||
│ │ sad.png │ │
|
||||
│ │ angry.png │ │
|
||||
│ └──────────────────┘ │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
The Qwen transformer is loaded FP8 → BF16 at construction using
|
||||
`init_empty_weights` + incremental conversion to keep peak memory manageable
|
||||
(20B parameters: 20 GB FP8 on disk, ~40 GB BF16 at runtime). All models fit
|
||||
comfortably in 128 GB unified memory — no offloading, no swapping.
|
||||
|
||||
SDXL and Qwen use **separate VAEs** with different latent spaces. The SDXL
|
||||
checkpoint bundles its own VAE; Qwen uses `Qwen/Qwen-Image` VAE from
|
||||
HuggingFace. Both coexist in the same session without conflict.
|
||||
|
||||
## Data Flow
|
||||
|
||||
### Generate Phase
|
||||
|
||||
```
|
||||
prompt text ──► Compel (SDXL CLIPs) ──► conditioning ──┐
|
||||
├──► SDXL UNet (N steps) ──► latent ──► SDXL VAE decode ──► image
|
||||
noise + latent ─────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Edit Phase
|
||||
|
||||
```
|
||||
input image ──► Qwen VAE encode ──► latent ──┐
|
||||
│
|
||||
prompt text ──► Qwen VL 7B TE ──► conditioning┤
|
||||
├──► Qwen Transformer (N steps) ──► latent ──► Qwen VAE decode ──► image
|
||||
input image ──► Qwen VL 7B TE ──► visual tok.┘
|
||||
```
|
||||
|
||||
The text encoder handles both text conditioning and visual token encoding from
|
||||
the input image.
|
||||
|
||||
### Turbo vs Normal Mode
|
||||
|
||||
| Parameter | Normal | Turbo (Lightning LoRA) |
|
||||
|-----------|--------|------------------------|
|
||||
| Steps | 20 | 4 |
|
||||
| CFG | 4.0 | 1.0 |
|
||||
| LoRA | none | Lightning-4steps (fused at load) |
|
||||
| Sampler | Flow Match Euler | Flow Match Euler |
|
||||
|
||||
## Performance
|
||||
|
||||
Hardware: **AMD Strix Halo** (Ryzen AI Max 395 Pro, Radeon 8060S, 128 GB),
|
||||
bfloat16, `novaAnimeXL_ilV190.safetensors`.
|
||||
|
||||
### Generate (SDXL)
|
||||
|
||||
Radeon 8060S (Strix Halo iGPU), bfloat16, 1024×1024:
|
||||
- ~1s per step
|
||||
- ~20s for 20-step generation
|
||||
| Resolution | Steps | Load (s) | Inference (s) | Total (s) |
|
||||
|-----------|-------|----------|---------------|-----------|
|
||||
| 1024×1024 | 20 | ~2.5 | ~29 | ~31 |
|
||||
|
||||
Model loading adds ~5s cold-start overhead.
|
||||
Per-step breakdown (1024×1024):
|
||||
- Step 1 (warmup): ~1.3–1.6s
|
||||
- Steps 2–20 (steady): ~1.0–1.3s each
|
||||
- VAE decode included in final step
|
||||
|
||||
### Edit (Qwen Image Edit)
|
||||
Compel prompt weighting adds ~0.4s encoding overhead — negligible.
|
||||
|
||||
Radeon 8060S, bfloat16:
|
||||
- ~120s per step at 512×512
|
||||
- ~23s model loading
|
||||
### Edit (Qwen Image Edit, 4-step Lightning LoRA)
|
||||
|
||||
Larger resolutions scale proportionally. The SDPA math fallback in the text
|
||||
encoder's visual branch and the transformer's attention blocks is the main
|
||||
bottleneck.
|
||||
| Attention | Steps | Load (s) | Inference (s) | Total (s) |
|
||||
|-----------|-------|----------|---------------|-----------|
|
||||
| Matmul fallback | 4 | ~28 | ~87 | ~115 |
|
||||
| Flash attention | 4 | ~31 | ~53 | ~84 |
|
||||
|
||||
**Turbo mode:** Use the Lightning 4-step LoRA with `--steps 4 --cfg 1.0` to
|
||||
cut per-step time proportionally (4× fewer steps).
|
||||
Per-step breakdown (1024×1024, flash attention):
|
||||
- Step 1: ~0.3s (prefill/encoding)
|
||||
- Steps 2–4: ~1.6 → ~5.1 → ~7.0s (transformer + VAE)
|
||||
|
||||
Flash attention (`TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1`) cuts edit
|
||||
inference by **1.63×**. SDXL generate is unaffected (uses its own attention
|
||||
processor, not SDPA).
|
||||
|
||||
### Session (persistent models)
|
||||
|
||||
| Phase | Wall (s) | Notes |
|
||||
|-------|----------|-------|
|
||||
| Session load | ~28 | SDXL + Qwen transformer + VAE + TE + LoRA fuse |
|
||||
| Generate | ~31 | SDXL 20-step, 1024×1024 |
|
||||
| Edit (turbo) | ~87 / ~54 | Matmul / flash attention |
|
||||
| **Total (2 ops)** | **~146 / ~118** | One session load amortized |
|
||||
|
||||
With a session, each additional edit saves one model-load round trip (~28s).
|
||||
5 edits save 112s. Flash attention adds a further 1.6× multiplier on inference.
|
||||
|
||||
See [`docs/stats.md`](docs/stats.md) for detailed per-run breakdowns.
|
||||
|
||||
## Technical Notes
|
||||
|
||||
@@ -144,21 +268,43 @@ cut per-step time proportionally (4× fewer steps).
|
||||
`float16` causes GPU kernel crashes (segfault) on the Radeon 8060S. The tool
|
||||
uses `bfloat16` internally. This is transparent to the user.
|
||||
|
||||
### Custom attention
|
||||
### FP8 → BF16 conversion
|
||||
|
||||
The default PyTorch SDPA backends (flash attention, mem-efficient attention) are
|
||||
unstable on this AMD GPU. VNAsset uses a simple matmul-based attention
|
||||
implementation that avoids the SDPA dispatch entirely.
|
||||
The Qwen Image Edit model ships as FP8 (`fp8_e4m3fn`). It is converted to BF16
|
||||
at load time using `init_empty_weights` + incremental tensor conversion to keep
|
||||
peak memory manageable. RDNA 3.5 WMMA supports FP8 compute, but PyTorch ROCm FP8
|
||||
support is not yet mature enough to compute in FP8 — upcast to BF16 is the safe
|
||||
path and still fits in 128 GB.
|
||||
|
||||
### Attention backends
|
||||
|
||||
Two attention paths, selected automatically:
|
||||
|
||||
| Path | When | Performance |
|
||||
|------|------|-------------|
|
||||
| **Flash attention** | `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1` set | 1.63× faster edits |
|
||||
| **Matmul fallback** | Default (env var not set) | Stable, slower |
|
||||
|
||||
SDXL always uses a custom `AttnProcessor` (direct QKV matmul) regardless of the
|
||||
env var — its attention path is separate from the Qwen SDPA dispatch. The flash
|
||||
attention toggle only affects the Qwen transformer and text encoder.
|
||||
|
||||
```bash
|
||||
# Enable flash attention for the session
|
||||
TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 vnasset edit ...
|
||||
```
|
||||
|
||||
### ROCm torch version
|
||||
|
||||
Tested with `torch 2.11.0+rocm7.2`. Newer ROCm nightlies (2.13+, 2.14+) may
|
||||
cause GPU crashes. If you encounter segfaults, try matching this version.
|
||||
|
||||
## Documentation
|
||||
### LoRA loading
|
||||
|
||||
- **[SDXL Generation](docs/sdxl-generation.md)** — checkpoint loading, attention patches, prompt encoding, and generation pipeline details
|
||||
- **[ComfyUI Prompt Style Support](docs/comfyui-prompt-style.md)** — prompt weighting and BREAK syntax specification
|
||||
Lightning LoRA is loaded via `diffusers` native `load_lora_weights()` and fused
|
||||
into the transformer with `fuse_lora()` at session construction. The fusion
|
||||
takes ~1.5s and happens once — the turbo/normal switch is then just a
|
||||
steps+CFG change with no per-step LoRA overhead.
|
||||
|
||||
## Prompt Syntax
|
||||
|
||||
@@ -194,17 +340,36 @@ vnasset generate \
|
||||
|
||||
Use `--raw` to bypass weighting and fall back to plain diffusers encoding.
|
||||
|
||||
## Current State
|
||||
|
||||
| Feature | Status |
|
||||
|---------|--------|
|
||||
| `vnasset generate` | ✅ Working |
|
||||
| `vnasset edit` | ✅ Working |
|
||||
| `VnAssetsSession` (persistent models) | ✅ Working |
|
||||
| Compel prompt weighting + BREAK | ✅ Working |
|
||||
| Lightning LoRA fuse-at-load | ✅ Working |
|
||||
| Flash attention (experimental) | ✅ Working |
|
||||
| Metadata JSON output | ✅ Working |
|
||||
| `vnasset pipeline` (batch YAML config) | 🚧 Planned |
|
||||
| `vnasset serve` (daemon/HTTP API) | 🚧 Planned |
|
||||
| `torch.compile` on UNet | 🚧 Planned |
|
||||
| Batch edit loop (shared VAE encode) | 🚧 Planned |
|
||||
|
||||
## Future Improvements
|
||||
|
||||
- **Persistent model session** — keep models loaded between commands to
|
||||
eliminate the ~5s cold-start overhead per generation. A `vnasset serve`
|
||||
daemon or `vnasset batch` command.
|
||||
- **Pipeline batch mode** — `vnasset pipeline --config pipeline.yaml` for
|
||||
generate + multiple edits in one session from a YAML config file.
|
||||
- **`torch.compile` on UNet** — the UNet forward is identical each step; ROCm's
|
||||
`torch.compile` support is maturing and could cut per-step time in half.
|
||||
`torch.compile` support is maturing and could cut per-step time significantly.
|
||||
- **Shared encode optimization** — for N edit variants of the same input image,
|
||||
run VAE encode and VL visual token encoding once, then only text-encode and
|
||||
denoise per variant.
|
||||
- **Self-contained torch wheel** — bundle the known-working torch wheel file in
|
||||
the project (`wheels/torch-2.11.0+rocm7.2-cp312-cp312-linux_x86_64.whl`) so
|
||||
the install is reproducible without depending on PyTorch's nightly index
|
||||
availability or a ComfyUI installation.
|
||||
- **Qwen Image Edit support** — `vnasset edit` and `vnasset pipeline` for
|
||||
batch expression/outfit variant editing.
|
||||
|
||||
- **`vnasset serve`** — lightweight daemon with Unix socket or HTTP API for
|
||||
integrating VNAsset into external tools.
|
||||
- **Background removal** — green-screen trim and transparency pass for
|
||||
post-processing sprites.
|
||||
|
||||
Reference in New Issue
Block a user