RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

165
extensions_built_in/diffusion_models/example_model/README.md Normal file
View File

@ -0,0 +1,165 @@
# Example Model — a template for adding a new architecture to ai-toolkit
This folder is a complete, heavily commented template for wiring a brand-new
diffusion model into ai-toolkit. It assumes the worst (and most common) case:
**diffusers does not have your model**, so you vendor the network and a minimal
sampling pipeline yourself.
It is intentionally **not registered** — it never appears as a trainable arch.
It exists purely as a guide for people (and agents) adding image, editing,
video, or i2v models.
## File map
```
example/
├── README.md <- you are here
├── __init__.py <- exports ExampleModel (registration notes inside)
├── example_model.py <- the BaseModel subclass: every override documented
│ with exact inputs/outputs
└── src/ <- everything diffusers does NOT provide
├── model.py <- a minimal DiT with the gradient-checkpointing pattern
└── pipeline.py <- a minimal embeds-only flow-matching sampler
```
## How a model gets registered
1. `toolkit/util/get_model.py:get_all_models()` scans every package directly
under `extensions/` and `extensions_built_in/` for a module-level
`AI_TOOLKIT_MODELS` list.
2. For models in this folder, that list lives in
`extensions_built_in/diffusion_models/__init__.py` — import your class
there and append it to `AI_TOOLKIT_MODELS`.
(Alternatively, give your model its own folder under `extensions/` with its
own `AI_TOOLKIT_MODELS` list — see `extensions/z_image_pixel/`.)
3. The class attribute `arch` (e.g. `"example"`) is matched against
`model.arch` in the training config YAML to pick your class.
4. To expose it in the web UI, add an entry to
`ui/src/app/jobs/new/options.ts` (search for an existing arch like
`ideogram4` to copy the shape).
Minimal config YAML to train it:
```yaml
model:
arch: "example"
name_or_path: "/path/to/weights" # folder with transformer/, text_encoder/,
# tokenizer/, vae/
quantize: true # optional: qfloat8 the transformer
quantize_te: true # optional: qfloat8 the text encoder
train:
gradient_checkpointing: true
```
## Lifecycle — who calls what, in order
1. **Load** — `load_model()` builds the transformer, text encoder(s),
tokenizer(s), VAE and scheduler and stores them on `self`. Everything else
reads `self.model` / `self.vae` / `self.text_encoder`.
2. **Caching (optional)** — before training, the trainer may call
`encode_images()` per dataset image (latent cache) and
`get_prompt_embeds()` per caption (text-embed cache, saved via
`AdvancedPromptEmbeds.save`, one file per caption).
3. **Train step** (every step, see `extensions_built_in/sd_trainer/SDTrainer.py`):
1. clean latents come from the cache or `encode_images()`
2. noise + timestep are sampled; `add_noise()` (BaseModel) mixes them
3. `condition_noisy_latents(noisy_latents, batch)` — your hook to inject
control/reference conditioning
4. `get_noise_prediction(latent_model_input, timestep, text_embeddings)` —
the forward pass, under autograd
5. loss = MSE(prediction, `get_loss_target(noise=..., batch=...)`)
4. **Sampling previews** — `generate_images()` (BaseModel) encodes each sample
prompt with `get_prompt_embeds()`, then calls your
`get_generation_pipeline()` once and `generate_single_image(...)` per
prompt. Your pipeline only ever receives **embeds, never text**.
5. **Saving** — full fine-tunes go through `save_model()`. LoRA files are
written by the network code, with your
`convert_lora_weights_before_save/load()` mapping keys to the public
convention (usually the `diffusion_model.` prefix).
## Conventions to keep straight
- **Pixels** are `(B, 3, H, W)` in `[-1, 1]` (control tensors arrive in
`[0, 1]` — multiply by 2 and subtract 1 before encoding).
- **Latents** are `(B, C, h, w)`; video latents are `(B, C, frames, h, w)`.
- **Timesteps** cross the BaseModel API on a `0..1000` scale where 1000 is
pure noise. Convert to your model's native convention inside
`get_noise_prediction` — and watch for models whose native time runs the
other way (t=1 = clean); flip and/or negate there (ideogram4 does both).
- **Flow-matching target** in this codebase is `noise - clean`
(`get_loss_target`), i.e. the velocity pointing from data to noise.
- `self.model` / `self.transformer` / `self.unet` are aliases for the same
thing on BaseModel.
## AdvancedPromptEmbeds
`toolkit/advanced_prompt_embeds.py`. The flexible container for text
conditioning, preferred for all new models over the older `PromptEmbeds`:
- Every key holds a **list of tensors, one per batch item**
(`AdvancedPromptEmbeds(text_embeds=[t0, t1, ...])`). Store each item at its
natural length and pad to the batch max only at the model call
(`src/pipeline.py:pad_prompt_embeds`) — caches stay small and any prompts
can share a batch.
- Add as many keys as your model needs (`pooled_embeds`, image features, …).
- Keys that must not be dtype-cast (token ids, masks) go in
`embeds.frozen_dtype_keys`.
- CFG concat (`concat_prompt_embeds`), batch expansion, `.to()`, `.save()` /
`.load()` for the disk cache are all handled for you.
If you ever change what `get_prompt_embeds` produces, bump the
`text_embedding_space_version` property so stale on-disk caches invalidate.
## Gradient checkpointing
With `train.gradient_checkpointing: true`, `BaseSDTrainProcess` calls
`model.enable_gradient_checkpointing()` if it exists, else sets
`model.gradient_checkpointing = True`. Your network re-runs each block under
`torch.utils.checkpoint.checkpoint(..., use_reentrant=False)` when the flag is
set **and** `torch.is_grad_enabled()` is true — never gate on `self.training`.
See `src/model.py` for the full pattern and rationale.
## Adapting this template
### Editing / instruct model (image in, image out)
- In `condition_noisy_latents`, encode `batch.control_tensor`
(`(B, 3, H, W)` in `[0, 1]`) with the VAE and attach it to the noisy
latents — extra channels (`torch.cat(..., dim=1)`) or extra sequence tokens.
Slice the prediction back down in `get_noise_prediction` before returning.
Reference: `../flux_kontext/flux_kontext.py`.
- If the text encoder must *see* the control image (VL encoders), set
`self.encode_control_in_text_embeddings = True`; `get_prompt_embeds` then
receives `control_images`. Reference: `../qwen_image/qwen_image_edit.py`.
- Multiple reference images: `self.has_multiple_control_images = True`
(`batch.control_tensor_list`). Reference:
`../qwen_image/qwen_image_edit_plus.py`.
- In `generate_single_image`, load `gen_config.ctrl_img` (a file path) and run
the same conditioning for previews.
### Video model (t2v)
- Batches arrive as `(B, frames, 3, H, W)`; latents as
`(B, C, frames_latent, h, w)`. Override `encode_images`/`decode_latents`
for your video VAE (temporal compression means
`frames_latent = (frames - 1) // 4 + 1` for most VAEs).
- `gen_config.num_frames` drives previews; return a **list of PIL frames**
from `generate_single_image` and the harness saves a video.
- Reference: `../wan22/wan22_5b_model.py` and `../ltx2/`.
### Image-to-video (i2v)
- Same as video, plus first-frame conditioning: in `get_noise_prediction`
take frame 0 from `batch.tensor` (declare `batch` in your signature to
receive it), encode it, and merge it into the latent input. For previews do
the same with `gen_config.ctrl_img`.
- Reference: `../wan22/wan22_14b_i2v_model.py` and
`toolkit/models/wan21/wan_utils.py:add_first_frame_conditioning`.
### Other useful hooks (all on `toolkit/models/base_model.py:BaseModel`)
| Override | When you need it |
|---|---|
| `get_model_to_train()` | LoRA should attach to something other than `self.model` |
| `text_embedding_space_version` / `latent_space_version` | invalidate users' caches after a breaking change |
| `te_padding_side` | LLM text encoders that need left padding |
| `is_multistage`, `multistage_boundaries` | multi-expert models split by timestep range (`../wan22/wan22_14b_model.py`) |
| `load_training_adapter()` pattern | assistant LoRAs (de-distillation adapters), see `../z_image/z_image.py` |
| `get_latent_noise_from_latents()` | custom noise (default: `randn_like`) |
| `encode_audio()` | audio-conditioned models (`../ltx2/`) |

12
extensions_built_in/diffusion_models/example_model/__init__.py Normal file
View File

@ -0,0 +1,12 @@
# This is a documentation-only TEMPLATE model. Start with README.md in this
# folder for the full guide to adding a new model architecture to ai-toolkit.
#
# It is intentionally NOT registered: the parent package
# (extensions_built_in/diffusion_models/__init__.py) does not import it, so it
# never shows up as a trainable arch. To register a real model, import its
# class there and append it to the AI_TOOLKIT_MODELS list. (Models can also
# live in their own folder under extensions/, which defines its own
# AI_TOOLKIT_MODELS list -- see extensions/z_image_pixel for a tiny example.)
from .example_model import ExampleModel
__all__ = ["ExampleModel"]

508
extensions_built_in/diffusion_models/example_model/example_model.py Normal file
View File

@ -0,0 +1,508 @@
"""ExampleModel -- a fully documented template for adding a new model to ai-toolkit.
Read README.md in this folder first for the big picture (lifecycle, data flow,
registration, and how to adapt this template into an edit / video / i2v model).
Every override below documents:
- WHEN ai-toolkit calls it
- WHAT comes in (shapes, dtypes, scales)
- WHAT must come out
The model itself is a made-up flow-matching DiT whose architecture lives in
./src/model.py and whose preview sampler lives in ./src/pipeline.py, simulating
the common case where diffusers does not ship your model and you vendor both.
"""
import os
from typing import List, Optional
import torch
import yaml
from safetensors.torch import load_file, save_file
from diffusers import AutoencoderKL
from transformers import AutoTokenizer, AutoModel
from optimum.quanto import freeze
from toolkit.accelerator import unwrap_model
from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
from toolkit.basic import flush
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from toolkit.models.base_model import BaseModel
from toolkit.samplers.custom_flowmatch_sampler import (
CustomFlowMatchEulerDiscreteScheduler,
)
from toolkit.util.quantize import quantize, get_qtype, quantize_model
from .src.model import ExampleTransformer2DModel
from .src.pipeline import ExamplePipeline, pad_prompt_embeds
# Config for the training/sampling noise scheduler. ai-toolkit's flow-matching
# models all use CustomFlowMatchEulerDiscreteScheduler; ``shift`` warps the
# timestep distribution toward the high-noise end (bigger = more high-noise
# steps, typical for high-resolution models).
scheduler_config = {
"num_train_timesteps": 1000,
"use_dynamic_shifting": False,
"shift": 3.0,
}
class ExampleModel(BaseModel):
# ``arch`` is the unique id that ties everything together:
# - ``model.arch: "example"`` in the training config YAML selects this class
# (resolved by toolkit/util/get_model.py:get_model_class)
# - it is the default cache key for text-embedding / latent caches
arch = "example"
def __init__(
self,
device, # "cuda:0" etc.
model_config: ModelConfig, # the parsed ``model:`` section of the YAML
dtype="bf16",
custom_pipeline=None,
noise_scheduler=None,
**kwargs,
):
super().__init__(
device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
)
# --- flags the rest of the toolkit reads ---
# flow matching (velocity prediction) vs ddpm-style epsilon prediction
self.is_flow_matching = True
# transformer (DiT) vs unet: affects LoRA naming ("transformer." prefix)
self.is_transformer = True
# Class names of modules whose Linear layers get LoRA'd. Matched against
# type(module).__name__, so this must equal the class name in src/model.py.
self.target_lora_modules = ["ExampleTransformer2DModel"]
# --- values used by our own overrides below ---
self.patch_size = 2 # transformer patch size (latent px per token)
self.vae_scale_factor = 8 # pixels per latent px (8x downsampling VAE)
# hard cap on prompt token length (truncation only -- embeds are stored
# per-sample at natural length, see get_prompt_embeds)
self.max_text_length = 512
# Other flags you may need (all default False, set in BaseModel.__init__):
# self.encode_control_in_text_embeddings = True
# -> get_prompt_embeds receives control_images (vision-language TEs
# that look at the control image, e.g. qwen_image_edit)
# self.has_multiple_control_images = True
# -> control images arrive as a list (qwen_image_edit_plus)
# self.use_raw_control_images = True
# -> control images are not resized to match the target image
# self.is_multistage = True
# -> model has multiple experts trained on timestep ranges (wan22 14b)
@staticmethod
def get_train_scheduler():
"""Build the noise scheduler used for BOTH training and sampling.
Called when loading the model, and again by the pipeline for every
preview run (a fresh instance, because scheduler state is mutable).
"""
return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
def get_bucket_divisibility(self):
"""Pixel multiple that dataset resolution buckets must snap to.
The data loader crops every image so width/height are divisible by
this. Latents are 1/8 the pixel size (VAE) and the transformer eats
2x2 latent patches, so pixels must be divisible by 8 * 2 = 16.
"""
return self.vae_scale_factor * self.patch_size
# ------------------------------------------------------------------
# Loading
# ------------------------------------------------------------------
def load_model(self):
"""Load every component and store them on ``self``.
Called once at startup. ``self.model_config`` is the ``model:`` section
of the training YAML; the fields used here:
- name_or_path: local folder (or HF repo) with the weights
- quantize / qtype: quantize the transformer (e.g. "qfloat8")
- quantize_te / qtype_te: quantize the text encoder
- low_vram: keep big components on CPU; your other overrides then
move them to GPU on demand (see the device checks below)
MUST set, before returning:
self.model the trainable denoiser (transformer/unet)
self.vae the (frozen) VAE
self.text_encoder one module or a list of modules (frozen unless
training the TE)
self.tokenizer one tokenizer or a list, parallel to text_encoder
self.noise_scheduler from get_train_scheduler()
self.pipeline anything generate_single_image can use
"""
dtype = self.torch_dtype
self.print_and_status_update("Loading Example model")
# Expected layout (diffusers-style folder):
# <name_or_path>/transformer/model.safetensors
# <name_or_path>/text_encoder/ + /tokenizer/ (transformers format)
# <name_or_path>/vae/ (diffusers AutoencoderKL)
model_path = self.model_config.name_or_path
# --- transformer (the custom model from src/) ---
self.print_and_status_update("Loading transformer")
# Instantiate on the meta device (no RAM used), then materialize the
# real tensors straight from the checkpoint with assign=True. This
# avoids allocating the model twice. If your model has non-persistent
# buffers, rebuild them after this (see ideogram4.py for an example).
with torch.device("meta"):
transformer = ExampleTransformer2DModel()
state_dict = load_file(
os.path.join(model_path, "transformer", "model.safetensors")
)
state_dict = {k: v.to(dtype) for k, v in state_dict.items()}
transformer.load_state_dict(state_dict, assign=True)
del state_dict
flush() # gc + empty cuda cache; call it after dropping anything big
if self.model_config.quantize:
# quantize_model handles qtype selection, exclusions and device
# juggling, and leaves the model on CPU
self.print_and_status_update("Quantizing transformer")
quantize_model(self, transformer)
flush()
if self.model_config.low_vram:
# leave it on CPU; get_noise_prediction moves it over when needed
transformer.to("cpu")
else:
transformer.to(self.device_torch, dtype=dtype)
flush()
# For partial layer offloading support see MemoryManager.attach usage
# in ../ideogram4/ideogram4.py or ../z_image/z_image.py.
# --- text encoder + tokenizer (stock transformers model) ---
self.print_and_status_update("Loading text encoder")
tokenizer = AutoTokenizer.from_pretrained(model_path, subfolder="tokenizer")
text_encoder = AutoModel.from_pretrained(
model_path, subfolder="text_encoder", torch_dtype=dtype
)
text_encoder.to(self.te_device_torch)
# the TE is frozen here; only set requires_grad if you train it
text_encoder.eval()
text_encoder.requires_grad_(False)
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing text encoder")
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
# --- VAE ---
self.print_and_status_update("Loading VAE")
vae = AutoencoderKL.from_pretrained(model_path, subfolder="vae")
vae.to(self.vae_device_torch, dtype=self.vae_torch_dtype)
vae.eval()
vae.requires_grad_(False)
flush()
# --- scheduler + store everything ---
self.noise_scheduler = ExampleModel.get_train_scheduler()
self.vae = vae
self.text_encoder = text_encoder # could be a list for multi-TE models
self.tokenizer = tokenizer # parallel list if multiple TEs
self.model = transformer # aliased as self.transformer / self.unet
self.pipeline = ExamplePipeline(self)
self.print_and_status_update("Model Loaded")
# ------------------------------------------------------------------
# Sampling (training previews)
# ------------------------------------------------------------------
def get_generation_pipeline(self):
"""Return a fresh pipeline for a round of preview sampling.
Called once per sampling round by BaseModel.generate_images. Our
pipeline holds no state, so a new lightweight wrapper is enough.
"""
return ExamplePipeline(self)
def generate_single_image(
self,
pipeline: ExamplePipeline,
gen_config: GenerateImageConfig, # one sample_prompts entry: width,
# height, seed, num_inference_steps,
# guidance_scale, ctrl_img, num_frames...
conditional_embeds: AdvancedPromptEmbeds, # already-encoded prompt
unconditional_embeds: AdvancedPromptEmbeds, # already-encoded negative prompt
generator: torch.Generator, # seeded with gen_config.seed
extra: dict, # adapter kwargs (controlnet etc.)
):
"""Render ONE preview image.
The harness (BaseModel.generate_images) has already encoded the
prompts with get_prompt_embeds -- the pipeline never sees text.
Returns a PIL.Image (or for video models a list of PIL frames).
"""
# low_vram: components may be parked on CPU between steps
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
# snap requested size to the model's divisibility
sc = self.get_bucket_divisibility()
gen_config.width = int(gen_config.width // sc * sc)
gen_config.height = int(gen_config.height // sc * sc)
img = pipeline(
conditional_embeds=conditional_embeds,
unconditional_embeds=unconditional_embeds,
height=gen_config.height,
width=gen_config.width,
num_inference_steps=gen_config.num_inference_steps,
guidance_scale=gen_config.guidance_scale,
latents=gen_config.latents, # usually None; pre-made noise if set
generator=generator,
)[0]
return img
# ------------------------------------------------------------------
# Training hooks
# ------------------------------------------------------------------
def get_noise_prediction(
self,
latent_model_input: torch.Tensor,
timestep: torch.Tensor,
text_embeddings: AdvancedPromptEmbeds,
**kwargs,
):
"""The actual forward pass of the denoiser. Called every train step
(with grads) via BaseModel.predict_noise, and also by some adapters.
in:
latent_model_input (B, C, h, w) noisy latents: the output of
add_noise(clean_latents, noise, timestep), after
condition_noisy_latents (channel-concat models
would see extra channels here).
For video models this is (B, C, frames, h, w).
timestep (B,) float on the 0..1000 scale, 1000 = pure noise
text_embeddings AdvancedPromptEmbeds for the batch; every key you
stored in get_prompt_embeds holds a list of B
tensors (cached per-sample embeds are expanded /
concatenated for you)
**kwargs may include ``batch`` (DataLoaderBatchDTO),
guidance_embedding_scale, adapter residuals, ...
only passed if your signature declares them
out:
(B, C, h, w) the model prediction. For flow matching that is the
velocity in the same convention as get_loss_target (here:
noise - clean). Shape must match the TARGET latents -- if you
concatenated control channels/tokens in, slice them off before
returning (see ../flux_kontext/flux_kontext.py).
"""
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
# toolkit timestep (0..1000) -> our model's flow time in [0, 1].
# WATCH OUT: every model has its own time convention. If the original
# repo uses t=1 for clean images, flip it here (see
# ../ideogram4/src/pipeline.py predict_velocity for an example).
t01 = timestep.to(self.device_torch, dtype=torch.float32) / 1000.0
# per-sample embed lists -> padded batch tensor + attention mask
llm_features, text_mask = pad_prompt_embeds(
text_embeddings.text_embeds, self.device_torch, self.torch_dtype
)
noise_pred = self.model(
hidden_states=latent_model_input.to(self.device_torch, self.torch_dtype),
timestep=t01,
encoder_hidden_states=llm_features,
attention_mask=text_mask,
)
return noise_pred
def get_prompt_embeds(self, prompt) -> AdvancedPromptEmbeds:
"""Encode prompt text into whatever conditioning the model eats.
Called for dataset captions (optionally cached to disk per caption),
for sample prompts, and for the empty string (unconditional).
in: prompt a str or list[str]
out: AdvancedPromptEmbeds. Each key holds a LIST of tensors, one per
prompt, each at its natural (unpadded) length. Padding to the
batch max is deferred to get_noise_prediction / the pipeline,
which keeps caches small and lets any prompts share a batch.
You can store any number of keys (pooled embeds, image features,
...). If a key must keep its dtype when everything else is cast
(masks, token ids), list it in ``embeds.frozen_dtype_keys``.
NOTE: if you change how embeddings are computed after release, bump
``text_embedding_space_version`` (a property on BaseModel) to
invalidate users' on-disk caches.
"""
if isinstance(prompt, str):
prompt = [prompt]
# low_vram support: TE might be parked on CPU
if self.text_encoder.device == torch.device("cpu"):
self.text_encoder.to(self.device_torch)
embeds_list = []
for p in prompt:
tokens = self.tokenizer(
p,
truncation=True,
max_length=self.max_text_length,
return_tensors="pt",
).to(self.text_encoder.device)
# no padding: encode each prompt at its own length
with torch.no_grad():
output = self.text_encoder(**tokens, output_hidden_states=True)
# (L, D) -- drop the batch dim, one tensor per prompt
embeds_list.append(output.last_hidden_state[0].to(self.torch_dtype))
return AdvancedPromptEmbeds(text_embeds=embeds_list)
def get_loss_target(self, *args, **kwargs):
"""The ground-truth tensor the prediction is MSE'd against.
kwargs: noise (B, C, h, w), batch (DataLoaderBatchDTO with .latents =
the clean latents), timesteps. For flow matching the velocity target
is noise - clean. Must be detached.
"""
noise = kwargs.get("noise")
batch = kwargs.get("batch")
return (noise - batch.latents).detach()
def condition_noisy_latents(
self, latents: torch.Tensor, batch
) -> torch.Tensor:
"""Optional hook: modify noisy latents before the model sees them.
Called every train step right after noise is added. This is THE hook
for editing / inpainting / i2v models that feed reference latents in
alongside the noisy target (the reference is concatenated here, then
consumed -- and sliced off the prediction -- in get_noise_prediction).
in: latents (B, C, h, w) noisy latents
batch DataLoaderBatchDTO -- batch.control_tensor holds the
control image(s) as (B, 3, H, W) in [0, 1] when the
dataset config has a control_path
out: latents, conditioned (return .detach()'d -- no grads here)
This base text-to-image model needs nothing, so it passes through.
Real examples: ../flux_kontext/flux_kontext.py (concat control latents
as extra tokens), ../qwen_image/qwen_image_edit.py.
"""
return latents
# ------------------------------------------------------------------
# VAE encode / decode
# ------------------------------------------------------------------
# BaseModel.encode_images / decode_latents already handle a diffusers
# AutoencoderKL (scaling_factor / shift_factor) and would work unchanged
# for this model. They are overridden here anyway to document the
# contract, since custom VAEs (or latent normalization, patchified
# latents, video VAEs...) usually need it.
def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None):
"""Pixels -> latents. Used for latent caching and for control images.
in: image_list list of (3, H, W) tensors -- or a (B, 3, H, W) batch --
with values in [-1, 1], already crop/bucket-sized
out: (B, C, h, w) latents, normalized the way the transformer expects
(for AutoencoderKL: (z - shift_factor) * scaling_factor)
"""
if device is None:
device = self.vae_device_torch
if dtype is None:
dtype = self.vae_torch_dtype
if self.vae.device == torch.device("cpu"):
self.vae.to(self.vae_device_torch)
if isinstance(image_list, list):
images = torch.stack(image_list, dim=0)
else:
images = image_list
images = images.to(device, dtype=dtype)
latents = self.vae.encode(images).latent_dist.sample()
shift = self.vae.config["shift_factor"] or 0
latents = (latents - shift) * self.vae.config["scaling_factor"]
return latents.to(device, dtype=dtype)
def decode_latents(self, latents: torch.Tensor, device=None, dtype=None):
"""Latents -> pixels. Used when rendering previews.
in: (B, C, h, w) latents in the normalized space encode_images produces
out: (B, 3, H, W) images in [-1, 1]
"""
if device is None:
device = self.vae_device_torch
if dtype is None:
dtype = self.vae_torch_dtype
if self.vae.device == torch.device("cpu"):
self.vae.to(self.vae_device_torch)
latents = latents.to(device, dtype=dtype)
shift = self.vae.config["shift_factor"] or 0
latents = latents / self.vae.config["scaling_factor"] + shift
return self.vae.decode(latents).sample
# ------------------------------------------------------------------
# Saving / bookkeeping
# ------------------------------------------------------------------
def get_model_has_grad(self):
"""True only if the base denoiser weights themselves require grad
(full fine-tune). LoRA training: False. Used to save/restore device
and grad state around sampling."""
return False
def get_te_has_grad(self):
"""Same as above for the text encoder."""
return False
def save_model(self, output_path, meta, save_dtype):
"""Save the FULL model (fine-tune checkpoints; LoRA saving is handled
elsewhere and only consults convert_lora_weights_before_save).
``output_path`` is a directory (no extension). Save in whatever layout
load_model can read back; include aitk_meta.yaml for provenance.
"""
transformer: ExampleTransformer2DModel = unwrap_model(self.model)
os.makedirs(os.path.join(output_path, "transformer"), exist_ok=True)
state_dict = {
k: v.clone().to("cpu", dtype=save_dtype)
for k, v in transformer.state_dict().items()
}
save_file(
state_dict, os.path.join(output_path, "transformer", "model.safetensors")
)
with open(os.path.join(output_path, "aitk_meta.yaml"), "w") as f:
yaml.dump(meta, f)
def get_base_model_version(self):
"""Free-form version string written into LoRA metadata so other tools
can identify the base model family."""
return "example.1"
def get_transformer_block_names(self) -> Optional[List[str]]:
"""Attribute name(s) on self.model that hold the repeated transformer
blocks (a ModuleList). Used for LoRA block targeting; must match the
attribute in src/model.py."""
return ["blocks"]
def convert_lora_weights_before_save(self, state_dict):
"""Map internal LoRA keys to the ecosystem-standard naming right before
the .safetensors is written. Most modern models ship LoRAs with a
``diffusion_model.`` prefix (ComfyUI convention); internally ai-toolkit
uses ``transformer.``."""
return {
k.replace("transformer.", "diffusion_model."): v
for k, v in state_dict.items()
}
def convert_lora_weights_before_load(self, state_dict):
"""Inverse of the above, applied when resuming from a saved LoRA."""
return {
k.replace("diffusion_model.", "transformer."): v
for k, v in state_dict.items()
}

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]