jmp-games/vnassets
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve doc, cleanup

Browse Source
This commit is contained in:
Michele Rossi
2026-07-08 09:13:46 +02:00
parent 9564202e6d
commit e7cde842b3
5 changed files with 528 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@@ -5,3 +5,4 @@ __pycache__/
models/
dist/
build/
output/

5
README.md
View File

@@ -155,6 +155,11 @@ implementation that avoids the SDPA dispatch entirely.
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
- **[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
## Prompt Syntax
VNAsset supports **ComfyUI-style prompt weighting** via the `compel` library.

494
docs/sdxl-generation.md Normal file
View File

@@ -0,0 +1,494 @@
# SDXL Generation
## Overview
VNAsset uses **Stable Diffusion XL (SDXL)** for the base sprite generation phase.
SDXL is loaded from a single `.safetensors` checkpoint file (the same format
used by ComfyUI and Automatic1111) and run via HuggingFace `diffusers` with
custom patches for AMD GPU stability.
No node graph, no serialization — just direct PyTorch forward calls through the
SDXL UNet, VAE, and CLIP text encoders.
---
## Architecture: How the Three Components Work
SDXL is a latent diffusion model. It generates images by reversing a
noise-adding process in a compressed latent space, guided by text conditioning.
Three neural networks cooperate to do this:
### CLIP Text Encoders
**What they do:** Turn a text prompt into a numeric tensor the UNet can
understand.
SDXL uses **two** CLIP encoders, not one:
| Encoder | Architecture | Tokenizer | Output shape per token | Pooled output? |
|----------|-------------|-----------|----------------------|----------------|
| CLIP-L | ViT-L/14 | `tokenizer` | 768-d | No |
| OpenCLIP-G | ViT-bigG/14 | `tokenizer_2` | 1280-d | **Yes** — 1280-d vector |
**Why two?** CLIP-L was trained on proprietary OpenAI data; OpenCLIP-G was
trained on LAION-2B, a public dataset. They capture complementary semantic
information. Using both improves prompt adherence and image quality.
**Token sequence encoding (both encoders):**
```
"1girl, red hair"
│
▼
Tokenizer → [<s>, 1, girl, ,, red, hair, </s>] ← token IDs with BOS/EOS
│
▼
Token embedding lookup → 7 × 768 matrix (CLIP-L) / 7 × 1280 (OpenCLIP-G)
│
▼
Transformer encoder (self-attention over sequence) → contextualized embeddings
│
▼
Output: [1, 77, 768] (CLIP-L) + [1, 77, 1280] (OpenCLIP-G)
```
Each encoder takes its token sequence and runs it through a stack of
Transformer blocks. Self-attention lets each token attend to every other token,
so `hair` gets context from `red` and `girl`.
**Pooled embedding (OpenCLIP-G only):**
OpenCLIP-G produces a **second** output — a single 1280-d vector (the pooled
representation) that summarizes the entire prompt. This is concatenated with
the timestep embedding inside the UNet and modulates its cross-attention
blocks, giving a global "this is what the whole prompt means" signal.
**SDXL concatenation:**
The two encoder outputs are concatenated along the feature dimension:
```
CLIP-L output: [1, 77, 768]
OpenCLIP-G output: [1, 77, 1280]
│
▼
Concatenated: [1, 77, 2048] ← what the UNet receives
Pooled: [1, 1280] ← separate global conditioning
```
The UNet's cross-attention layers have 2048-d key/value projection weights
to match this concatenated dimension.
---
### VAE (Variational Autoencoder)
**What it does:** Compresses images into a compact latent code, and
decompresses latents back into pixels. This is what makes SDXL a *latent*
diffusion model: the expensive diffusion math happens in a smaller space.
**Why compress?** A 1024×1024 RGB image is `3 × 1024 × 1024 = 3,145,728`
values. The VAE compresses this by a factor of **8× spatially** and expands
the channel count from 3 to 4:
```
Input image: [B, 3, 1024, 1024] → 3,145,728 values
Latent: [B, 4, 128, 128 ] → 65,536 values (48× smaller)
```
Denoising 65k values instead of 3.1M is dramatically cheaper in both memory
and compute — roughly **48× cheaper** for the UNet.
**Architecture:**
The SDXL VAE is a convolutional autoencoder:
```
Encoder Decoder
image ──► [Conv → ResBlock → Downsample] × 4 ──► latent ──► [Conv → ResBlock → Upsample] × 4 ──► image
(each stage halves spatial dims) (each stage doubles spatial dims)
```
Each downsampling stage uses stride-2 convolutions to halve the spatial
resolution. The decoder mirrors this with nearest-neighbor upsampling followed
by convolutions. Residual blocks provide gradient flow through the deep stack.
The VAE is **frozen during diffusion training**. It's pre-trained separately
(on reconstruction + KL regularization), then treated as a fixed
encoder/decoder. The diffusion model only ever sees latents.
**In the pipeline:**
- **Training time:** VAE encodes real images into latents, noise is added,
UNet learns to denoise.
- **Inference time (VNAsset):** VAE is used only at the end — the UNet
denoiser starts from pure noise in latent space, no encoding step needed.
The decoder converts the cleaned latent back to pixels.
---
### UNet
**What it does:** Predicts and removes noise from a latent, one small step at
a time. This is the core of the diffusion process.
**The diffusion idea:**
1. Start with a clean latent `z_0` (what you want).
2. Add Gaussian noise over many small steps, producing `z_t` at timestep `t`.
3. The UNet is trained to predict the noise that was added, given `z_t` and `t`.
At inference, you start from pure Gaussian noise `z_T` and repeatedly apply
the UNet's prediction to step toward `z_0`:
```
z_T (pure noise) ──► z_{T-1} ──► z_{T-2} ──► ... ──► z_0 (clean latent)
↑ ↑ ↑ ↑
UNet UNet UNet UNet
```
Each arrow above is one denoising step. With 20 steps and 1024×1024
resolution, that's 20 UNet forward passes (hence ~20 seconds at ~1 s/step).
**UNet architecture:**
The SDXL UNet has a U-shaped structure — an encoder (down) path and a decoder
(up) path connected by skip connections:
```
┌─────────────┐
latent ──► Conv │ DownBlock │
[4,128,128] │ 320 ch │──── skip ──────────────────────────┐
└──────┬──────┘ │
│ down×2 │
┌──────▼──────┐ │
│ DownBlock │ │
│ 640 ch │──── skip ────────────────────┐ │
└──────┬──────┘ │ │
│ down×2 │ │
┌──────▼──────┐ │ │
│ DownBlock │ │ │
│ 1280 ch │──── skip ──────────────┐ │ │
└──────┬──────┘ │ │ │
│ down×2 │ │ │
┌──────▼──────┐ │ │ │
│ MidBlock │ │ │ │
│ 1280 ch │ │ │ │
└──────┬──────┘ │ │ │
│ up×2 │ │ │
┌──────▼──────┐ │ │ │
│ UpBlock │──── skip ──────────────┘ │ │
│ 1280 ch │ │ │
└──────┬──────┘ │ │
│ up×2 │ │
┌──────▼──────┐ │ │
│ UpBlock │──── skip ────────────────────┘ │
│ 640 ch │ │
└──────┬──────┘ │
│ up×2 │
┌──────▼──────┐ │
│ UpBlock │──── skip ──────────────────────────┘
│ 320 ch │
└──────┬──────┘
│
┌──────▼──────┐
│ Conv │
│ out: 4 ch │
└─────────────┘
│
▼
predicted noise
(same shape as latent)
```
**What's inside each block:**
Each DownBlock and UpBlock is a stack of:
- **ResBlocks** (Residual blocks): Convolution layers with skip connections
that process the feature map.
- **Transformer blocks** (in the last two stages): Self-attention + cross-attention.
Cross-attention is where text conditioning enters — the feature map
(as queries) attends to the CLIP embeddings (as keys/values).
**Cross-attention: how text controls generation:**
```
Q = Linear(feature_pixels) # "what is each pixel looking for?"
K = Linear(clip_embeddings) # "what do the words represent?"
V = Linear(clip_embeddings) # "what information do the words carry?"
attention_weights = softmax(Q @ K^T / sqrt(d_k))
output = attention_weights @ V
```
Every spatial position in the latent attends to every token in the prompt.
This is how `red hair` ends up controlling the hair region — the pixels that
activate for the hair region will learn to attend strongly to the `red` and
`hair` token embeddings.
**Timestep conditioning:**
The UNet also receives the current timestep `t`. It's embedded via a
sinusoidal encoding and fed through MLPs into every ResBlock as a scale/shift
modulation (similar to adaptive group normalization). This tells the network
*how much* noise to expect — at early timesteps (high noise), the UNet makes
large structural changes; at late timesteps (low noise), it refines fine
details.
**Classifier-free guidance (CFG):**
During inference, the UNet runs **twice** per step: once with the positive
prompt, once with the negative prompt (or an empty prompt). The two noise
predictions are combined:
```
predicted_noise = neg_noise + cfg_scale × (pos_noise - neg_noise)
```
A higher CFG scale (e.g. 7–10) pushes the result harder toward the positive
prompt, away from the negative. This improves prompt adherence but at the cost
of reduced diversity and, at extreme values, artifacts. SDXL typically works
well at 4.5–7; VNAsset defaults to 4.5.
**SDXL UNet size:** ~2.6B parameters. On the Radeon 8060S this occupies
~3.5 GB in bfloat16.
---
## Pipeline Flow
```
prompt text ──► SDXL CLIP encoders (CLIP-L + OpenCLIP-G) ──► conditioning ──┐
│
seed ──► Generator ──► noise ──► empty latent ─────────────────────────────┤
│
├──► SDXL UNet
│ (Euler scheduler,
│ N steps, CFG)
│ │
▼
latent
│
▼
SDXL VAE decode
│
▼
output.png
```
### Step by step
1. **Load checkpoint.** `StableDiffusionXLPipeline.from_single_file()` loads the
`.safetensors` into the `diffusers` SDXL pipeline object (UNet, VAE, CLIP-L
text encoder, OpenCLIP-G text encoder, scheduler).
2. **Move to GPU, set dtype.** Pipeline moves to `cuda` (ROCm HIP), with
`torch.bfloat16`. `float16` is avoided because it causes GPU kernel segfaults
on RDNA 3.5 (Radeon 8060S).
3. **Patch attention.** All UNet attention processors are replaced with a simple
matmul-based implementation that bypasses PyTorch's unstable SDPA dispatch on
AMD GPUs (see [Custom Attention Patches](#custom-attention-patches)).
4. **Encode prompts.** If Compel weighting is enabled (the default), the prompt
and negative prompt are parsed for `(word:weight)` and `BREAK` syntax,
translated through the dual CLIP encoders, and returned as pre-computed
embedding tensors. If `--raw` is set, the plain string path through
`pipe.__call__()` is used instead.
5. **Generate.** The UNet denoises a random latent guided by the text
conditioning for the requested number of steps. Euler ancestral scheduling
is used; the CFG scale balances prompt adherence (higher = stronger prompt
alignment).
6. **Decode.** The SDXL VAE decodes the final latent into a 1024×1024 (or
custom resolution) RGB image.
7. **Save outputs.** The image is written as PNG. A sidecar JSON file records
metadata (prompt, seed, timing, model path, resolution).
8. **Cleanup.** Pipeline is deleted and `torch.cuda.empty_cache()` is called to
free GPU memory.
---
## Checkpoint Compatibility
VNAsset loads **any** SDXL `.safetensors` checkpoint. It uses
`StableDiffusionXLPipeline.from_single_file()`, which handles:
- Standard SDXL checkpoints (e.g., `sd_xl_base_1.0.safetensors`)
- Fine-tuned checkpoints (e.g., `novaAnimeXL_ilV190.safetensors`)
- Checkpoints with baked VAE
- Checkpoints with separate VAE (the pipeline detects and loads what's present)
The checkpoint is specified via `--checkpoint`:
```bash
vnasset generate \
--checkpoint models/novaAnimeXL_ilV190.safetensors \
--prompt "1girl, solo, red hair, blue eyes" \
--output output/character.png
```
---
## Prompt Encoding
### Default: Compel weighting (ComfyUI-compatible)
By default, VNAsset uses the [`compel`](https://github.com/damian0815/compel)
library to support ComfyUI-style prompt syntax:
| Syntax | Effect |
|--------|--------|
| `(word)` | Boost ×1.1 |
| `(word:1.5)` | Boost ×1.5 |
| `(word:0.6)` | De-emphasize ×0.6 |
| `[word]` | De-emphasize ×0.9 |
| `BREAK` | Split into independent conditioning chunks |
Compel applies weights at the embedding tensor level — it multiplies token
embeddings by the specified weight before passing them to the UNet. Both CLIP
encoders (CLIP-L and OpenCLIP-G) are handled, including the pooled embedding
from OpenCLIP-G.
For details on syntax and the underlying mechanism, see
[`docs/comfyui-prompt-style.md`](comfyui-prompt-style.md).
### Raw mode (`--raw`)
When `--raw` is passed, Compel is bypassed entirely. The prompt string is sent
directly to `pipe(prompt=..., negative_prompt=...)`, using diffusers' built-in
CLIP encoding without any weighting. This is useful for:
- Prompts that contain literal parentheses (no escaping needed)
- Debugging — comparing weighted vs unweighted output
- Situations where Compel's overhead is undesirable
```bash
vnasset generate \
--checkpoint models/novaAnimeXL_ilV190.safetensors \
--prompt "1girl, red hair, blue eyes" \
--raw \
--output output/character.png
```
---
## Custom Attention Patches
PyTorch's default SDPA backends (flash attention, mem-efficient attention) are
unstable on AMD RDNA 3.5 GPUs under ROCm — they can produce NaN outputs or
segfault. VNAsset replaces the UNet's attention processor with a manual
matmul-based implementation.
### What gets patched
Every `Attention` module in the SDXL UNet is given a `SimpleAttnProcessor`:
```
AttnProcessor (default, dispatches to SDPA)
│
▼
SimpleAttnProcessor (manual Q·K^T·V with softmax)
```
### What the custom attention does
```python
Q, K, V = Linear(hidden), Linear(hidden), Linear(hidden)
scores = Q @ K^T / sqrt(d_k)
weights = softmax(scores)
output = weights @ V
```
No fused kernel dispatch, no flash attention, no mem-efficient attention —
just straightforward matmul + softmax. This is slower than fused attention
but stable on ROCm.
### When it's not needed
If the environment variable `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1` is set,
ROCm's experimental flash attention backend is used instead and no patch is
applied. This affects only the Qwen Image Edit transformer — the SDXL UNet
always gets patched regardless of this flag.
### Implementation
The function `patch_unet_attention()` in [`vnassets/attention.py`](../vnassets/attention.py)
iterates over all attention layers and swaps in the custom processor:
```python
from vnassets.attention import patch_unet_attention
pipe = StableDiffusionXLPipeline.from_single_file(checkpoint_path, torch_dtype=dtype)
pipe.to(device)
patch_unet_attention(pipe.unet) # ← replaces all attention processors
```
---
## Dependencies
| Package | Role |
|---------|------|
| `diffusers` | SDXL pipeline (UNet, VAE, scheduler) |
| `compel` | ComfyUI-style prompt weighting |
| `torch` (ROCm) | GPU compute via HIP backend |
| `safetensors` | Checkpoint file format |
All are declared in `pyproject.toml` and installed with `pip install -e .`.
---
## Code Locations
| Component | File |
|-----------|------|
| Generation entry point | [`vnassets/generate.py`](../vnassets/generate.py) — `generate()` |
| CLI command binding | [`vnassets/cli.py`](../vnassets/cli.py) — `generate_cmd()` |
| Attention patches | [`vnassets/attention.py`](../vnassets/attention.py) — `patch_unet_attention()`, `simple_attention_forward()` |
| Prompt weighting | [`vnassets/prompt.py`](../vnassets/prompt.py) — `build_compel()`, `encode_prompts()` |
---
## Performance
All measurements on Radeon 8060S (Strix Halo iGPU), bfloat16, 1024×1024.
| Phase | Time |
|-------|------|
| Model loading (checkpoint → VRAM) | ~5 s |
| Per inference step | ~1 s |
| 20-step generation (total inference) | ~20 s |
| VAE decode | ~1 s |
The model is freshly loaded and then torn down each invocation. The planned
`vnasset pipeline` / `vnasset serve` will keep models resident across
generations to eliminate the ~5 s cold-start overhead.
---
## Parameter Reference
| Parameter | CLI flag | Type | Default | Description |
|-----------|----------|------|---------|-------------|
| Checkpoint | `--checkpoint` | path | *(required)* | Path to SDXL `.safetensors` |
| Prompt | `--prompt` | string | *(required)* | Text prompt (supports Compel weighting) |
| Negative prompt | `--negative-prompt` | string | `""` | Negative prompt |
| Width | `--width` | int | `1024` | Output image width |
| Height | `--height` | int | `1024` | Output image height |
| Steps | `--steps` | int | `20` | Denoising steps |
| CFG scale | `--cfg` | float | `4.5` | Classifier-free guidance scale |
| Seed | `--seed` | int/random | `random` | RNG seed |
| Output path | `--output` | path | `output.png` | Output PNG path |
| Raw mode | `--raw` | flag | `false` | Bypass Compel weighting |
---
## See Also
- [ComfyUI Prompt Style Support](comfyui-prompt-style.md) — prompt weighting and BREAK syntax details
- [TECH_SPEC.md](../TECH_SPEC.md) — full pipeline architecture and roadmap
- [README.md](../README.md) — usage examples and install guide

18
vnassets/attention.py
View File

@@ -69,11 +69,21 @@ def patch_unet_attention(unet):
def patch_qwen_transformer(transformer):
"""Patch Qwen transformer to use matmul attention instead of SDPA."""
from diffusers.models.attention_dispatch import dispatch_attention_fn
# Monkey-patch the dispatch function
"""Patch Qwen transformer attention.
If the ROCm experimental flash attention env var is set
(TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1), the default SDPA dispatch
will use flash attention — no patch needed.
Otherwise, fall back to a simple matmul-based attention that avoids
the unstable SDPA math backend on AMD GPUs.
"""
import os
if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL") == "1":
return # flash attention available, no patch needed
# Monkey-patch the dispatch function with matmul fallback
import diffusers.models.attention_dispatch as ad
ad.dispatch_attention_fn = _matmul_attention
# Also patch the module that imports it
import diffusers.models.transformers.transformer_qwenimage as tq
tq.dispatch_attention_fn = _matmul_attention

14
vnassets/edit.py
View File

@@ -61,6 +61,7 @@ def edit(
cfg: float = 4.0,
seed: int | None = None,
output_path: str = "output.png",
lora_path: str | None = None,
) -> None:
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16
@@ -93,6 +94,16 @@ def edit(
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")
@@ -125,6 +136,9 @@ def edit(
"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),
}