diff --git a/extensions_built_in/diffusion_models/example_model/README.md b/extensions_built_in/diffusion_models/example_model/README.md
new file mode 100644
index 00000000..db791330
--- /dev/null
+++ b/extensions_built_in/diffusion_models/example_model/README.md
@@ -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/`) |
diff --git a/extensions_built_in/diffusion_models/example_model/__init__.py b/extensions_built_in/diffusion_models/example_model/__init__.py
new file mode 100644
index 00000000..0fd9402a
--- /dev/null
+++ b/extensions_built_in/diffusion_models/example_model/__init__.py
@@ -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"]
diff --git a/extensions_built_in/diffusion_models/example_model/example_model.py b/extensions_built_in/diffusion_models/example_model/example_model.py
new file mode 100644
index 00000000..46e224c5
--- /dev/null
+++ b/extensions_built_in/diffusion_models/example_model/example_model.py
@@ -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()
+        }
diff --git a/extensions_built_in/diffusion_models/example_model/src/__init__.py b/extensions_built_in/diffusion_models/example_model/src/__init__.py
new file mode 100644
index 00000000..38747395
--- /dev/null
+++ b/extensions_built_in/diffusion_models/example_model/src/__init__.py
@@ -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
diff --git a/extensions_built_in/diffusion_models/example_model/src/model.py b/extensions_built_in/diffusion_models/example_model/src/model.py
new file mode 100644
index 00000000..e3ad7264
--- /dev/null
+++ b/extensions_built_in/diffusion_models/example_model/src/model.py
@@ -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
diff --git a/extensions_built_in/diffusion_models/example_model/src/pipeline.py b/extensions_built_in/diffusion_models/example_model/src/pipeline.py
new file mode 100644
index 00000000..30416e36
--- /dev/null
+++ b/extensions_built_in/diffusion_models/example_model/src/pipeline.py
@@ -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]