github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Added an example model with docs so people and agents can add models easier.

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-13 08:17:06 -06:00
parent 2f91db8363
commit a01c83073a
6 changed files with 1112 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
extensions_built_in/diffusion_models/example_model/src/__init__.py Normal file
View File

@@ -0,0 +1,4 @@
# Everything diffusers does NOT provide for your model lives in src/:
# the network architecture and a minimal sampling pipeline.
from .model import ExampleTransformer2DModel
from .pipeline import ExamplePipeline, pad_prompt_embeds

265
extensions_built_in/diffusion_models/example_model/src/model.py Normal file
View File

@@ -0,0 +1,265 @@
"""A minimal diffusion transformer (DiT) used by the example model extension.
This file stands in for the situation where diffusers does NOT have your model.
You vendor the architecture yourself inside your extension's ``src/`` folder and
load the weights manually in your model class (see ``../example_model.py``).
The architecture here is intentionally tiny and boring:
latents (B, C, h, w)
-> patchify with a strided conv (B, N_img, hidden)
text embeds (B, L, text_dim)
-> linear projection (B, L, hidden)
concat [text | image] into one joint sequence (B, L + N_img, hidden)
-> N transformer blocks (self attention + mlp, adaLN-zero
modulated by the timestep embedding)
-> final modulated norm + linear
take only the image tokens and unpatchify back to (B, C, h, w)
Real models add RoPE position embeddings, fancier attention, guidance
embeddings, etc. For real-world reference implementations in this repo see:
- ../../chroma/src/model.py (flux-style double/single stream blocks)
- ../../ernie_image/transformer.py (diffusers ModelMixin based)
- ../../ideogram4/src/transformer.py (packed single-sequence model)
GRADIENT CHECKPOINTING
======================
ai-toolkit enables gradient checkpointing on your model from
``jobs/process/BaseSDTrainProcess.py`` which does, in order of preference:
if hasattr(unet, 'enable_gradient_checkpointing'):
unet.enable_gradient_checkpointing()
elif hasattr(unet, 'gradient_checkpointing'):
unet.gradient_checkpointing = True
So a custom model only needs:
1. a ``self.gradient_checkpointing`` flag (default False)
2. (optionally) an ``enable_gradient_checkpointing()`` method
3. to wrap each transformer block call in ``torch.utils.checkpoint.checkpoint``
when the flag is set AND grads are enabled.
IMPORTANT: gate on ``torch.is_grad_enabled()``, NOT on ``self.training``.
Sampling runs under ``torch.no_grad()`` where checkpointing is pure overhead,
and some training setups (e.g. certain adapters) run the module in eval mode
while still needing gradients. ``torch.is_grad_enabled()`` handles both.
"""
import math
import torch
import torch.nn.functional as F
from torch import nn
from torch.utils.checkpoint import checkpoint
def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
"""Standard sinusoidal embedding.
in: t (B,) float tensor, the flow-matching time in [0, 1] (1 = pure noise)
out: emb (B, dim)
We scale t by 1000 before embedding so the sinusoids get a useful range,
the same trick flux and friends use.
"""
t = t.float() * 1000.0
half = dim // 2
freqs = torch.exp(
-math.log(max_period) * torch.arange(half, dtype=torch.float32, device=t.device) / half
)
args = t[:, None] * freqs[None]
return torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
class ExampleTransformerBlock(nn.Module):
"""One DiT block: adaLN-zero modulated self-attention + MLP.
in: x (B, S, hidden) the joint [text | image] token sequence
temb (B, hidden) the timestep embedding
attn_mask (B, 1, 1, S) bool, True = attend, False = padding
out: x (B, S, hidden)
"""
def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0):
super().__init__()
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.qkv = nn.Linear(hidden_size, hidden_size * 3)
self.proj = nn.Linear(hidden_size, hidden_size)
self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
mlp_hidden = int(hidden_size * mlp_ratio)
self.mlp = nn.Sequential(
nn.Linear(hidden_size, mlp_hidden),
nn.GELU(approximate="tanh"),
nn.Linear(mlp_hidden, hidden_size),
)
# adaLN-zero: timestep embedding -> shift/scale/gate for attn and mlp.
# Zero-init so the block starts as identity (standard DiT trick).
self.adaLN_modulation = nn.Sequential(
nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size)
)
nn.init.zeros_(self.adaLN_modulation[-1].weight)
nn.init.zeros_(self.adaLN_modulation[-1].bias)
def forward(self, x: torch.Tensor, temb: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
b, s, d = x.shape
shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = (
self.adaLN_modulation(temb).unsqueeze(1).chunk(6, dim=-1)
) # each (B, 1, hidden), broadcasts over the sequence
# --- attention ---
h = self.norm1(x) * (1 + scale_a) + shift_a
q, k, v = self.qkv(h).chunk(3, dim=-1)
q = q.view(b, s, self.num_heads, self.head_dim).transpose(1, 2)
k = k.view(b, s, self.num_heads, self.head_dim).transpose(1, 2)
v = v.view(b, s, self.num_heads, self.head_dim).transpose(1, 2)
h = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
h = h.transpose(1, 2).reshape(b, s, d)
x = x + gate_a * self.proj(h)
# --- mlp ---
h = self.norm2(x) * (1 + scale_m) + shift_m
x = x + gate_m * self.mlp(h)
return x
class ExampleTransformer2DModel(nn.Module):
"""The denoiser. Plain ``nn.Module`` on purpose.
You could also subclass ``diffusers.ModelMixin``/``ConfigMixin`` (see
../../ernie_image/transformer.py) to get ``save_pretrained``,
``_gradient_checkpointing_func`` etc. for free, but a plain module shows
exactly what ai-toolkit actually requires, which is very little:
- a forward pass
- ``device`` / ``dtype`` properties (BaseModel reads ``self.model.device``
and ``self.model.dtype`` in a few places, e.g. save_device_state)
- the gradient checkpointing flag described in the module docstring
NOTE: the class NAME matters. ``ExampleModel.target_lora_modules`` lists
"ExampleTransformer2DModel" -- that string is matched against module class
names when deciding where to attach LoRA layers.
"""
def __init__(
self,
in_channels: int = 16, # VAE latent channels
out_channels: int = 16, # predicted velocity has the same channels
patch_size: int = 2, # latent pixels per token side
hidden_size: int = 1024,
num_heads: int = 16,
num_layers: int = 12,
text_dim: int = 2048, # width of the text encoder hidden states
):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.patch_size = patch_size
self.hidden_size = hidden_size
# latent (B, C, h, w) -> image tokens (B, N_img, hidden)
self.x_embedder = nn.Conv2d(
in_channels, hidden_size, kernel_size=patch_size, stride=patch_size
)
# text encoder hidden states -> model width
self.text_proj = nn.Linear(text_dim, hidden_size)
# sinusoidal timestep embedding -> mlp
self.t_embedder = nn.Sequential(
nn.Linear(hidden_size, hidden_size),
nn.SiLU(),
nn.Linear(hidden_size, hidden_size),
)
# ``blocks`` is the repeated-layer ModuleList. The attribute name is
# what ExampleModel.get_transformer_block_names() returns, which the
# LoRA code uses for block targeting / "transformer only" training.
self.blocks = nn.ModuleList(
[
ExampleTransformerBlock(hidden_size, num_heads)
for _ in range(num_layers)
]
)
# final adaLN + projection back to patch pixels, zero-init so the
# untrained model predicts zeros.
self.norm_out = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.adaLN_out = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size))
self.proj_out = nn.Linear(hidden_size, patch_size * patch_size * out_channels)
nn.init.zeros_(self.adaLN_out[-1].weight)
nn.init.zeros_(self.adaLN_out[-1].bias)
nn.init.zeros_(self.proj_out.weight)
nn.init.zeros_(self.proj_out.bias)
# gradient checkpointing flag, flipped on by the trainer (see module
# docstring). Off by default so inference pays no cost.
self.gradient_checkpointing = False
# the trainer prefers this method if it exists
def enable_gradient_checkpointing(self, enable: bool = True):
self.gradient_checkpointing = enable
def disable_gradient_checkpointing(self):
self.gradient_checkpointing = False
@property
def device(self):
return next(self.parameters()).device
@property
def dtype(self):
return next(self.parameters()).dtype
def forward(
self,
hidden_states: torch.Tensor, # (B, C, h, w) noisy latents
timestep: torch.Tensor, # (B,) flow time in [0, 1], 1 = pure noise
encoder_hidden_states: torch.Tensor, # (B, L, text_dim) padded text features
attention_mask: torch.Tensor, # (B, L) 1 = real text token, 0 = padding
) -> torch.Tensor:
"""Predict the flow-matching velocity.
out: (B, C, h, w) velocity in the ai-toolkit convention
(noise - clean), matching ExampleModel.get_loss_target().
"""
b, c, h, w = hidden_states.shape
p = self.patch_size
gh, gw = h // p, w // p
n_img = gh * gw
# tokens
img = self.x_embedder(hidden_states) # (B, hidden, gh, gw)
img = img.flatten(2).transpose(1, 2) # (B, N_img, hidden)
txt = self.text_proj(encoder_hidden_states) # (B, L, hidden)
x = torch.cat([txt, img], dim=1) # (B, L + N_img, hidden)
# timestep conditioning
temb = self.t_embedder(timestep_embedding(timestep, self.hidden_size))
temb = temb.to(x.dtype)
# joint attention mask: text padding is masked out, image tokens and
# real text tokens attend everywhere. (B, 1, 1, S) bool for sdpa.
img_mask = torch.ones(b, n_img, dtype=torch.bool, device=x.device)
attn_mask = torch.cat([attention_mask.bool(), img_mask], dim=1)
attn_mask = attn_mask[:, None, None, :]
for block in self.blocks:
if torch.is_grad_enabled() and self.gradient_checkpointing:
# Recompute this block's activations during backward instead
# of storing them -- trades compute for a big VRAM saving.
# use_reentrant=False is the modern, correct variant.
x = checkpoint(block, x, temb, attn_mask, use_reentrant=False)
else:
x = block(x, temb, attn_mask)
# final modulation + project, keep only the image tokens
shift, scale = self.adaLN_out(temb).unsqueeze(1).chunk(2, dim=-1)
x = self.norm_out(x) * (1 + scale) + shift
x = self.proj_out(x)[:, -n_img:] # (B, N_img, p*p*C)
# unpatchify back to the latent layout
x = x.view(b, gh, gw, p, p, self.out_channels)
x = x.permute(0, 5, 1, 3, 2, 4).reshape(b, self.out_channels, h, w)
return x

158
extensions_built_in/diffusion_models/example_model/src/pipeline.py Normal file
View File

@@ -0,0 +1,158 @@
"""A minimal sampling pipeline for the example model.
ai-toolkit only uses your pipeline to render preview/sample images during
training (see BaseModel.generate_images -> ExampleModel.generate_single_image).
It does NOT need to be a diffusers DiffusionPipeline, and because ai-toolkit
always encodes the prompts itself (so it can cache embeds, apply trigger words,
run adapters, etc.) the pipeline never sees raw prompt strings -- only
already-encoded ``AdvancedPromptEmbeds``.
So all a pipeline has to do is:
1. make starting noise
2. loop the scheduler over timesteps, calling the transformer
3. apply classifier-free guidance (cond vs uncond prediction)
4. decode the final latents with the VAE and return PIL images
The pattern of passing the whole BaseModel instance into the pipeline (rather
than individual components) is borrowed from ../../ideogram4/src/pipeline.py.
It keeps the pipeline tiny because it can reuse the model's scheduler factory,
``decode_latents`` and device/dtype bookkeeping.
"""
from typing import List, Optional
import torch
from PIL import Image
from diffusers.utils.torch_utils import randn_tensor
def pad_prompt_embeds(
embeds_list: List[torch.Tensor],
device: torch.device,
dtype: torch.dtype,
):
"""Right-pad a list of per-sample text features into one batch tensor.
in: embeds_list list (len B) of (L_i, D) tensors -- this is exactly what
``AdvancedPromptEmbeds.text_embeds`` holds: one tensor per
batch item, each at its own natural length.
out: features (B, L_max, D) zero-padded on the right
mask (B, L_max) long, 1 = real token, 0 = padding
Storing embeds unpadded per item and only padding at the model call is the
preferred pattern: cached embeds stay small, and items of very different
prompt lengths can share a batch.
"""
lengths = [e.shape[0] for e in embeds_list]
max_len = max(lengths)
dim = embeds_list[0].shape[-1]
batch_size = len(embeds_list)
features = torch.zeros(batch_size, max_len, dim, device=device, dtype=dtype)
mask = torch.zeros(batch_size, max_len, dtype=torch.long, device=device)
for i, e in enumerate(embeds_list):
n = e.shape[0]
features[i, :n] = e.to(device, dtype)
mask[i, :n] = 1
return features, mask
class ExamplePipeline:
"""Lightweight flow-matching sampler used for training previews."""
def __init__(self, model):
# ``model`` is the ExampleModel (a BaseModel subclass), giving us
# access to model.transformer, model.vae, model.decode_latents, etc.
self.model = model
@property
def device(self):
return self.model.device_torch
def to(self, *args, **kwargs):
# BaseModel.generate_images may call pipeline.to(device); we manage
# devices through the model itself, so this is a no-op.
return self
def set_progress_bar_config(self, **kwargs):
# called by the sampler harness (inside a try/except, so optional);
# diffusers pipelines use it to silence tqdm. Nothing to do here.
pass
@torch.no_grad()
def __call__(
self,
# AdvancedPromptEmbeds with key ``text_embeds`` (list of (L, D) tensors)
conditional_embeds,
unconditional_embeds,
height: int = 1024,
width: int = 1024,
num_inference_steps: int = 25,
guidance_scale: float = 4.0,
latents: Optional[torch.Tensor] = None, # pre-made noise, usually None
generator: Optional[torch.Generator] = None, # seeded RNG for reproducible samples
**kwargs,
) -> List[Image.Image]:
model = self.model
device = model.device_torch
dtype = model.torch_dtype
transformer = model.transformer
# Always sample with a FRESH scheduler. The training scheduler is
# stateful; mutating it mid-training would corrupt the train step.
scheduler = model.get_train_scheduler()
scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = scheduler.timesteps # 1000 -> 0 scale
# pixel size -> latent size (VAE downsample only; the transformer
# patchifies internally so latents stay unpacked here)
gh = height // model.vae_scale_factor
gw = width // model.vae_scale_factor
do_cfg = unconditional_embeds is not None and guidance_scale != 1.0
# 1. starting noise (keep it float32; cast per model call)
if latents is None:
shape = (1, transformer.in_channels, gh, gw)
latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
latents = latents.to(device, dtype=torch.float32)
# 2. pad the per-item embed lists into batch tensors once, up front
cond_feats, cond_mask = pad_prompt_embeds(conditional_embeds.text_embeds, device, dtype)
if do_cfg:
uncond_feats, uncond_mask = pad_prompt_embeds(unconditional_embeds.text_embeds, device, dtype)
# 3. denoising loop
for t in timesteps:
# scheduler timesteps are on a 0-1000 scale; the transformer wants
# flow time in [0, 1] with 1 = pure noise
t01 = (t / 1000.0).to(device).expand(latents.shape[0])
v_cond = transformer(
hidden_states=latents.to(dtype),
timestep=t01,
encoder_hidden_states=cond_feats,
attention_mask=cond_mask,
)
if do_cfg:
v_uncond = transformer(
hidden_states=latents.to(dtype),
timestep=t01,
encoder_hidden_states=uncond_feats,
attention_mask=uncond_mask,
)
# classifier-free guidance: push the prediction away from the
# unconditional (negative prompt) direction
v = v_uncond + guidance_scale * (v_cond - v_uncond)
else:
v = v_cond
latents = scheduler.step(v.to(torch.float32), t, latents, return_dict=False)[0]
# 4. decode latents -> images in [-1, 1] -> uint8 PIL
images = model.decode_latents(latents, device=device, dtype=dtype)
images = images.float().clamp(-1.0, 1.0)
images = ((images + 1.0) * 127.5).round().to(torch.uint8)
images = images.permute(0, 2, 3, 1).cpu().numpy()
return [Image.fromarray(arr) for arr in images]