# 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/`) |