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:
7
.gitignore
vendored
Normal file
7
.gitignore
vendored
Normal file
@@ -0,0 +1,7 @@
|
||||
.venv/
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.egg-info/
|
||||
models/
|
||||
dist/
|
||||
build/
|
||||
146
README.md
Normal file
146
README.md
Normal 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
24
pyproject.toml
Normal 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
0
vnassets/__init__.py
Normal file
61
vnassets/attention.py
Normal file
61
vnassets/attention.py
Normal 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
48
vnassets/cli.py
Normal 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
74
vnassets/generate.py
Normal 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}")
|
||||
Reference in New Issue
Block a user