Improve doc, cleanup

.gitignore

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

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.

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 → [<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

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."""
+    """Patch Qwen transformer attention.
-    from diffusers.models.attention_dispatch import dispatch_attention_fn
+
-    # Monkey-patch the dispatch function
+    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

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),
     }