523 lines
23 KiB
Python
523 lines
23 KiB
Python
"""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"
|
|
|
|
# ALL NEW MODELS should set this to False. ``BaseModel`` defaults it to True
|
|
# only for backwards-compatibility with already-released LoKr checkpoints; the
|
|
# newer LoKr weight format is the correct one for any new architecture.
|
|
use_old_lokr_format = False
|
|
|
|
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.
|
|
|
|
Each per-prompt tensor MUST be 2D ``(L, D)`` -- BaseModel infers the
|
|
text batch size from the list and only treats it as one-per-prompt
|
|
when the tensors are 2D; a 3D per-prompt tensor is misread as an
|
|
already-batched ``(B, L, D)`` and training fails with a latents-vs-
|
|
text batch-size mismatch. If your conditioning has an extra axis
|
|
(e.g. N stacked encoder layers -> ``(L, N, D)``), flatten it here
|
|
(``(L, N*D)``) and restore it (``reshape(B, Lt, N, D)``) at the
|
|
model call.
|
|
|
|
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()
|
|
}
|