Add vnasset SDXL generate command

ROCm-safe bfloat16 inference with custom matmul attention.
Automatic output directories, random seed, timing metadata.
This commit is contained in:
Michele Rossi
2026-07-06 16:29:38 +02:00
commit 06fba9c234
7 changed files with 360 additions and 0 deletions

7
.gitignore vendored Normal file
View File

@@ -0,0 +1,7 @@
.venv/
__pycache__/
*.pyc
*.egg-info/
models/
dist/
build/

146
README.md Normal file
View File

@@ -0,0 +1,146 @@
# VNAsset
Fast CLI pipeline for visual novel image asset generation.
Drop-in replacement for the ComfyUI workflow loop: generate base character sprites
with SDXL, then batch-edit variants (expressions, outfits) with Qwen Image Edit —
all in one warm session, no node-graph overhead.
## Hardware
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.
## Install
```bash
git clone <repo> vnassets
cd vnassets
# Create venv (Python 3.12 required for ROCm torch compatibility)
python3.12 -m venv .venv
source .venv/bin/activate
# Install ROCm PyTorch (adjust index URL for your ROCm version)
pip install torch --index-url https://download.pytorch.org/whl/rocm7.2
# Install the rest
pip install -e .
```
### Models
Symlink your ComfyUI models into `models/`:
```bash
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 .
```
Or place the actual files there — the tool just reads whatever safetensors you
point it at.
## Usage
### Generate (SDXL text-to-image)
```bash
vnasset generate \
--checkpoint models/novaAnimeXL_ilV190.safetensors \
--prompt "1girl, solo, red hair, glasses, blue eyes, white crop top, standing, portrait" \
--negative-prompt "deformed, ugly, bad quality, lowres" \
--steps 20 \
--seed 42 \
--output output/character_base.png
```
| Option | Default | Description |
|--------|---------|-------------|
| `--checkpoint` | (required) | Path to SDXL `.safetensors` |
| `--prompt` | (required) | Positive prompt |
| `--negative-prompt` | `""` | Negative prompt |
| `--width` | `1024` | Image width |
| `--height` | `1024` | Image height |
| `--steps` | `20` | Inference steps |
| `--cfg` | `4.5` | CFG scale |
| `--seed` | `0` | RNG seed (use `random` for random) |
| `--output` | `output.png` | Output path |
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.
## Current State
| Command | Status |
|---------|--------|
| `vnasset generate` | ✅ Working |
| `vnasset edit` | 🚧 Planned |
| `vnasset pipeline` | 🚧 Planned |
## Performance
Radeon 8060S (Strix Halo iGPU), bfloat16, 1024×1024:
- ~1s per step
- ~20s for 20-step generation
Model loading adds ~5s cold-start overhead.
## Technical Notes
### bfloat16 required on RDNA 3.5
`float16` causes GPU kernel crashes (segfault) on the Radeon 8060S. The tool
uses `bfloat16` internally. This is transparent to the user.
### Custom attention
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.
### 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.
## Prompt Compatibility
VNAsset uses standard diffusers SDXL encoding, which is equivalent to ComfyUI's
`BNK_CLIPTextEncodeAdvanced` with `token_normalization=none` and
`weight_interpretation=comfy` for plain comma-separated prompts.
ComfyUI-specific syntax is **not currently supported**:
- `(word:1.2)` — prompt weighting
- `BREAK` — conditioning chunking
If your prompts rely on these, you'll get different output than the ComfyUI
workflow. compel integration is planned for later.
## 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.
- **`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.
- **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.
- **compel prompt weighting** — support `(word:weight)` and `BREAK` syntax for
parity with ComfyUI prompt encoding.

24
pyproject.toml Normal file
View File

@@ -0,0 +1,24 @@
[build-system]
requires = ["setuptools>=64"]
build-backend = "setuptools.build_meta"
[project]
name = "vnassets"
version = "0.1.0"
requires-python = ">=3.12"
dependencies = [
"torch",
"diffusers",
"transformers",
"accelerate",
"safetensors",
"pillow",
"pyyaml",
"click",
]
[tool.setuptools.packages.find]
exclude = ["models*"]
[project.scripts]
vnasset = "vnassets.cli:main"

0
vnassets/__init__.py Normal file
View File

61
vnassets/attention.py Normal file
View File

@@ -0,0 +1,61 @@
"""Custom attention processor that avoids SDPA on ROCm."""
import torch
import torch.nn.functional as F
from diffusers.models.attention_processor import Attention
def simple_attention_forward(
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""Simple QKV attention using raw matmul, bypassing SDPA dispatch."""
batch_size, seq_len, _ = hidden_states.shape
inner_dim = attn.inner_dim if hasattr(attn, 'inner_dim') else hidden_states.shape[-1]
# Project to Q, K, V
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states if encoder_hidden_states is None else encoder_hidden_states)
value = attn.to_v(hidden_states if encoder_hidden_states is None else encoder_hidden_states)
# Reshape to (batch, heads, seq, head_dim)
head_dim = inner_dim // attn.heads
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
# Attention: softmax(Q @ K^T / sqrt(d)) @ V
scale = head_dim ** -0.5
attn_weights = query @ key.transpose(-2, -1) * scale
attn_weights = F.softmax(attn_weights, dim=-1).to(query.dtype)
hidden_states = attn_weights @ value
# Reshape back
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
# Output projection
hidden_states = attn.to_out[0](hidden_states)
if len(attn.to_out) > 1:
hidden_states = attn.to_out[1](hidden_states) # dropout
return hidden_states
def patch_unet_attention(unet):
"""Replace all attention processors with simple matmul-based attention."""
from diffusers.models.attention_processor import AttnProcessor
class SimpleAttnProcessor(AttnProcessor):
def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
return simple_attention_forward(attn, hidden_states, encoder_hidden_states, attention_mask, **kwargs)
processor = SimpleAttnProcessor()
unet.set_attn_processor(processor)

48
vnassets/cli.py Normal file
View File

@@ -0,0 +1,48 @@
"""VNAsset — Visual Novel Asset Pipeline CLI."""
import click
from .generate import generate
def _parse_seed(ctx, param, value):
if value is None or value == "":
return None
if value.lower() == "random":
return None
try:
return int(value)
except ValueError:
raise click.BadParameter("seed must be an integer or 'random'")
@click.group()
def main():
"""VNAsset — fast CLI pipeline for visual novel image assets."""
@main.command()
@click.option("--checkpoint", required=True, help="Path to SDXL checkpoint (.safetensors)")
@click.option("--prompt", required=True, help="Positive prompt")
@click.option("--negative-prompt", default="", help="Negative prompt")
@click.option("--width", default=1024, type=int)
@click.option("--height", default=1024, type=int)
@click.option("--steps", default=20, type=int)
@click.option("--cfg", default=4.5, type=float)
@click.option(
"--seed", default="random", callback=_parse_seed,
help="RNG seed (integer or 'random')",
)
@click.option("--output", default="output.png", help="Output image path")
def generate_cmd(checkpoint, prompt, negative_prompt, width, height, steps, cfg, seed, output):
"""Generate an image from an SDXL checkpoint."""
generate(
checkpoint_path=checkpoint,
prompt=prompt,
negative_prompt=negative_prompt,
width=width,
height=height,
steps=steps,
cfg=cfg,
seed=seed,
output_path=output,
)

74
vnassets/generate.py Normal file
View File

@@ -0,0 +1,74 @@
"""SDXL text-to-image generation."""
import json
import random
import time
from pathlib import Path
import torch
from diffusers import StableDiffusionXLPipeline
from .attention import patch_unet_attention
def generate(
checkpoint_path: str,
prompt: str,
negative_prompt: str = "",
width: int = 1024,
height: int = 1024,
steps: int = 20,
cfg: float = 4.5,
seed: int | None = None,
output_path: str = "output.png",
) -> None:
device = "cuda" if torch.cuda.is_available() else "cpu"
# bfloat16 avoids ROCm kernel crashes on RDNA 3.5; float16 segfaults
dtype = torch.bfloat16
if seed is None:
seed = random.randint(0, 2**32 - 1)
output = Path(output_path)
output.parent.mkdir(parents=True, exist_ok=True)
t0 = time.perf_counter()
pipe = StableDiffusionXLPipeline.from_single_file(
checkpoint_path,
torch_dtype=dtype,
)
pipe.to(device)
patch_unet_attention(pipe.unet)
t_load = time.perf_counter() - t0
generator = torch.Generator(device=device).manual_seed(seed)
t1 = time.perf_counter()
image = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
width=width,
height=height,
num_inference_steps=steps,
guidance_scale=cfg,
generator=generator,
).images[0]
t_infer = time.perf_counter() - t1
image.save(output)
print(f"Saved {output}")
meta_path = output.with_suffix(".json")
meta = {
"checkpoint": str(Path(checkpoint_path).resolve()),
"prompt": prompt,
"negative_prompt": negative_prompt,
"width": width,
"height": height,
"steps": steps,
"cfg": cfg,
"seed": seed,
"load_time_s": round(t_load, 2),
"inference_time_s": round(t_infer, 2),
}
meta_path.write_text(json.dumps(meta, indent=2))
print(f"Saved {meta_path}")