Buckets:
| # 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. | |
| - **`use_old_lokr_format = False`** — set this class attribute on every NEW | |
| model. `BaseModel` defaults it to `True` purely for backwards-compatibility | |
| with LoKr checkpoints trained before the format change; all new architectures | |
| should use the new LoKr format. (Plain LoRA training is unaffected — this only | |
| matters for `network.type: "lokr"`.) | |
| ## 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. | |
| - **Keep each per-item tensor 2D `(L, D)`.** This is a hard requirement, not a | |
| convention: `BaseModel.predict_noise` infers the text batch size from the | |
| embed list, and it only counts the list as one-per-item when each tensor is | |
| 2D (`len(text_embeds[0].shape) == 2`). A 3D per-item tensor is read as an | |
| already-batched `(B, L, D)` and its *first axis* is taken as the batch size — | |
| so a single 3D prompt of length `L` looks like a batch of `L`, and training | |
| dies with *"Batch size of latents must be the same or half the batch size of | |
| text embeddings."* If your conditioning has an extra axis (e.g. a stack of N | |
| encoder layers, giving `(L, N, D)`), **flatten it into the feature axis** | |
| (`(L, N*D)`) in `get_prompt_embeds` and **restore it** (`reshape(B, Lt, N, D)`) | |
| in `get_noise_prediction` / the pipeline, right before the model call. | |
| - 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. | |
| ## Quantization | |
| With `quantize: true`, `quantize_model` swaps every `nn.Linear` for an | |
| `optimum.quanto` quantized one. Their matmul kernel **only accepts 2D or 3D | |
| activations** (`assert activations.ndim in (2, 3)`) — a `Linear` you feed a 4D | |
| tensor works fine in bf16 but throws once quantized. If your network applies a | |
| `Linear` over a 4D tensor (e.g. projecting a `(B, L, D, N)` layer axis), | |
| reshape to 3D for the call and back afterwards. | |
| Also watch out for **slow bf16 kernels on vendored components**: `Conv3d` has no | |
| fast cuDNN bf16 path (it falls back to a slow one). If a frozen sub-model carries | |
| a `Conv3d` you don't actually run — e.g. a vision tower's patch embed on a VL | |
| text encoder — drop it (`text_encoder.model.visual = None`) to skip loading it; | |
| if you must run one, consider running that component in fp16/fp32. | |
| ## Attention backends (don't force flash-attn) | |
| Reference repos very often hard-code an attention kernel — `flash_attn`, | |
| xformers, sage — and import it at module top level. **Do not carry that | |
| requirement over.** ai-toolkit has to import and load your model on machines | |
| where that package isn't installed (CPU boxes, headless CI, plain installs), so | |
| a top-level `from flash_attn import ...` turns "load the model" into an | |
| `ImportError`. | |
| The rule: | |
| - **Default to torch's built-in `F.scaled_dot_product_attention`** (the | |
| "native" backend). It needs no extra dependency, runs on CPU and CUDA, and | |
| already dispatches to a fused/flash kernel on supported hardware. `src/model.py` | |
| does exactly this. | |
| - **Make any other kernel OPTIONAL**, selected at runtime — never required at | |
| import. The clean pattern: | |
| 1. Guard the import so a missing package is a flag, not a crash: | |
| ```python | |
| try: | |
| from flash_attn import flash_attn_varlen_func | |
| _FLASH_ATTN_AVAILABLE = True | |
| except ImportError: | |
| flash_attn_varlen_func = None | |
| _FLASH_ATTN_AVAILABLE = False | |
| ``` | |
| 2. Give each attention module an `attention_backend` flag (default | |
| `"native"`) and **branch inside its forward** — `"flash"` runs the flash | |
| kernel, anything else runs SDPA. | |
| 3. Expose a `set_attention_backend("native"|"flash")` on the parent model | |
| that validates the name, raises a clear error if `"flash"` is requested | |
| while `_FLASH_ATTN_AVAILABLE` is `False`, and propagates the flag to every | |
| attention module. | |
| 4. Wire it to a config knob so it stays opt-in, e.g. | |
| `model_kwargs.attention_backend: "flash"` read in `load_model`. | |
| Branch on a per-module **flag**, don't swap the processor/module instance: | |
| attention modules that own trained q/k/v weights (joint/dual-stream blocks) | |
| would lose those weights if you replaced them with a different instance. | |
| Worked implementations to copy: `../ideogram4/src/transformer.py` | |
| (`set_attention_backend`, native+flash in one `Attention.forward`) and | |
| `../boogu_image/src/attention_processor.py` (guarded import, per-processor | |
| `attention_backend` flag, flash varlen branch alongside SDPA). | |
| ## 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/`) | | |
Xet Storage Details
- Size:
- 12.7 kB
- Xet hash:
- 5b44154a0ba16b2349610a50445ad7359e4c12519e1f6cc14607f66660e7241a
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.