495 lines
20 KiB
Markdown
495 lines
20 KiB
Markdown
# 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
|