diff --git a/.gitignore b/.gitignore index c35da88..47bcd88 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,4 @@ __pycache__/ models/ dist/ build/ +output/ diff --git a/README.md b/README.md index 35e06fd..a7e0bbe 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/docs/sdxl-generation.md b/docs/sdxl-generation.md new file mode 100644 index 0000000..7d28aca --- /dev/null +++ b/docs/sdxl-generation.md @@ -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 → [, 1, girl, ,, red, hair, ] ← 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 diff --git a/vnassets/attention.py b/vnassets/attention.py index fc46e8b..95c21da 100644 --- a/vnassets/attention.py +++ b/vnassets/attention.py @@ -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 diff --git a/vnassets/edit.py b/vnassets/edit.py index bd67ce9..6d89384 100644 --- a/vnassets/edit.py +++ b/vnassets/edit.py @@ -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), }