Merge in from main

README.md

@@ -18,6 +18,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
 - [Qwen/Qwen-Image](https://huggingface.co/Qwen/Qwen-Image) (Qwen-Image)
 - [Qwen/Qwen-Image-2512](https://huggingface.co/Qwen/Qwen-Image-2512) (Qwen-Image-2512)
 - [Qwen/Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1) (Qwen-Image-2.1)
+- [Qwen/Qwen-Image-2.1-Turbo](https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo) (Qwen-Image-2.1 Turbo)
 - [HiDream-ai/HiDream-I1-Full](https://huggingface.co/HiDream-ai/HiDream-I1-Full) (HiDream I1)
 - [OmniGen2/OmniGen2](https://huggingface.co/OmniGen2/OmniGen2) (OmniGen2)
 - [Tongyi-MAI/Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) (Z-Image Turbo)
@@ -32,6 +33,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
 - [HiDream-ai/HiDream-O1-Image](https://huggingface.co/HiDream-ai/HiDream-O1-Image) (HiDream O1)
 - [ideogram-ai/ideogram-4-fp8](https://huggingface.co/ideogram-ai/ideogram-4-fp8) (Ideogram 4 FP8)
 - [Photoroom/prxpixel-t2i](https://huggingface.co/Photoroom/prxpixel-t2i) (PRXPixel)
+- [speridlabs/iris-3b](https://huggingface.co/speridlabs/iris-3b) (Iris-3B)
 - [circlestone-labs/Anima-Base-v1.0-Diffusers](https://huggingface.co/circlestone-labs/Anima-Base-v1.0-Diffusers) (Anima)
 - [krea/Krea-2-Raw](https://huggingface.co/krea/Krea-2-Raw) (Krea 2)
 - [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo) (Krea 2 Turbo)
@@ -44,6 +46,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
 - [Qwen/Qwen-Image-Edit-2509](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) (Qwen-Image-Edit-2509)
 - [Qwen/Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) (Qwen-Image-Edit-2511)
 - [Qwen/Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1) (Qwen-Image-2.1) - one model for both; it edits when your dataset has control images
+- [Qwen/Qwen-Image-2.1-Turbo](https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo) (Qwen-Image-2.1 Turbo) - same, with the 8-step turbo weights
 - [HiDream-ai/HiDream-E1-1](https://huggingface.co/HiDream-ai/HiDream-E1-1) (HiDream E1)
 - [Boogu/Boogu-Image-0.1-Edit](https://huggingface.co/Boogu/Boogu-Image-0.1-Edit) (Boogu Image Edit)
 - [krea/Krea-2-Raw](https://huggingface.co/krea/Krea-2-Raw) (Krea 2 Edit Training)

extensions_built_in/diffusion_models/__init__.py

@@ -19,6 +19,7 @@ from .anima import AnimaModel
 from .ideogram4 import Ideogram4Model
 from .prx_pixel_t2i import PRXPixelT2IModel
 from .krea2 import Krea2Model
+from .iris3b import Iris3BModel
 from .boogu_image import BooguImageModel, BooguImageEditModel
 from .mageflow import MageFlowModel, MageFlowEditModel
 from .minimax_h3 import MinimaxH3Model, MinimaxH3Ref2VAModel, MinimaxH3FastModel
@@ -57,6 +58,7 @@ AI_TOOLKIT_MODELS = [
     Ideogram4Model,
     PRXPixelT2IModel,
     Krea2Model,
+    Iris3BModel,
     BooguImageModel,
     BooguImageEditModel,
     MageFlowModel,

extensions_built_in/diffusion_models/iris3b/README.md

@@ -0,0 +1,59 @@
+# Iris-3B (Sperid Labs, pixel-space text-to-image)
+
+ai-toolkit integration for [`speridlabs/iris-3b`](https://huggingface.co/speridlabs/iris-3b)
+([arXiv:2610.09450](https://arxiv.org/abs/2610.09450)), a 3B-parameter diffusion
+transformer that generates directly in pixel space. Apache 2.0.
+
+The architecture is vendored in [src/dit.py](src/dit.py) (from the reference
+[speridlabs/iris-3b](https://github.com/speridlabs/iris-3b) repo), the prompt
+encoder in [src/text_encoder.py](src/text_encoder.py), and the preview sampler
+(the reference FlowDPM-Solver++) in [src/pipeline.py](src/pipeline.py).
+
+## What is different about this model
+
+| Property | What it means | How it's handled |
+|---|---|---|
+| **Pixel space** | No VAE — the transformer denoises raw RGB (`in_channels=3`, `patch_size=16`) | `FakeVAE` (identity); "latents" are the image in `[-1, 1]`; buckets snap to 16 px |
+| **v-prediction, shift 4** | Rectified flow, velocity = `noise - clean`; the network is conditioned on the *shifted* noise level × 1000 | Toolkit convention already matches, so the toolkit `timestep` feeds the model as-is. Use `train.timestep_type: shift` (the UI default) to train on the reference schedule |
+| **Stacked text layers** | 12 Qwen3-VL-4B hidden layers per token, fixed 300-token chat template, pooled by a learned layer-attention adapter inside the DiT | Prompts are always padded to 300 tokens (the trunk attends to pad positions, so the pad count is part of the trained distribution); embeds are stored as `(300, 12*2560)` |
+| **Hybrid trunk** | 8 dual-stream MM-DiT blocks + 16 single-stream blocks, GQA 20/5 heads, gated attention, sandwich RMSNorm, shared-bias adaLN | Shared adaLN cores live in `modulation_cores` and are evaluated once per forward |
+| **Pixel head** | 4 PiT blocks decode each 16×16 patch back to pixels | Kept in full precision when quantizing |
+
+## Text encoder
+
+Qwen3-VL-4B-Instruct, frozen. Config and tokenizer come from
+`Qwen/Qwen3-VL-4B-Instruct`; the weights come from the ComfyUI repack
+`Comfy-Org/Krea-2/text_encoders/qwen3vl_4b_bf16.safetensors` (the same file
+Krea 2 uses). A copy already under `MODELS_PATH/text_encoders/` is used in
+place; otherwise it is downloaded there. Point `model.te_name_or_path` at
+another `.safetensors` file, a transformers folder, or a hub repo to override.
+
+## Train it
+
+```yaml
+model:
+  arch: "iris3b"
+  name_or_path: "speridlabs/iris-3b"   # hub repo, or a dir with model.safetensors + config.yaml
+  quantize_te: true                     # optional: quantize the 4B text encoder
+  # te_name_or_path: "Qwen/Qwen3-VL-4B-Instruct"   # load the vendor repo instead of the repack
+train:
+  timestep_type: "shift"                # reference training schedule (shift 4.0)
+  gradient_checkpointing: true
+sample:
+  guidance_scale: 3.0
+  sample_steps: 30                      # DPM-Solver++ order 2; the paper uses 100
+```
+
+`model_kwargs`: `flow_shift` (override the checkpoint's shift), `sample_order`
+(DPM-Solver++ order, default 2), `cfg_interval` (default `[0, 1]`),
+`checkpoint_filename` (file to pick inside a directory / repo),
+`model_config` (architecture overrides merged over the checkpoint's `config.yaml`).
+
+Full fine-tunes save as `model.safetensors` + `config.yaml`, the reference
+export layout, so `scripts/sample.py --checkpoint <dir>` from the Iris repo
+loads them directly. LoRA keys use the `diffusion_model.` prefix.
+
+The `depth/` and `upscaler/` fine-tunes in the hub repo share this architecture
+but are not text-to-image models; they load (`name_or_path:
+speridlabs/iris-3b/depth/model.safetensors`) but their input conditioning is
+not implemented here.

extensions_built_in/diffusion_models/iris3b/__init__.py

@@ -0,0 +1,3 @@
+from .iris3b import Iris3BModel
+
+__all__ = ["Iris3BModel"]

extensions_built_in/diffusion_models/iris3b/iris3b.py

@@ -0,0 +1,435 @@
+"""Iris3BModel -- Sperid Labs' Iris-3B pixel-space text-to-image DiT
+(https://huggingface.co/speridlabs/iris-3b, arXiv:2610.09450) in ai-toolkit.
+
+  - **Pixel space, no VAE.** The transformer denoises raw RGB (``in_channels=3``,
+    ``patch_size=16``) through a 24-block hybrid trunk (8 dual-stream MM-DiT +
+    16 single-stream) and a 4-block per-pixel PiT head. A ``FakeVAE`` (identity)
+    makes the toolkit's "latents" the image in [-1, 1]; buckets snap to 16 px.
+    Architecture vendored in ``src/dit.py``.
+  - **v-prediction rectified flow, shift 4.** Velocity = noise - clean on
+    ``x_t = (1 - t) * clean + t * noise`` -- ai-toolkit's own convention, so no
+    flip/negation. The model is conditioned on *shifted* model time
+    ``1000 * sigma'``; the toolkit's ``timestep_type: shift`` with the scheduler
+    ``shift`` reproduces the reference training schedule and the toolkit
+    ``timestep`` feeds the model as-is.
+  - **Qwen3-VL-4B-Instruct text encoder, 12 stacked hidden layers.** Prompts go
+    through the Iris chat template and are padded to a fixed 300 tokens
+    (``src/text_encoder.py``); the DiT's layerwise-attention adapter pools the
+    layer stack. Weights come from the ComfyUI repack
+    (``Comfy-Org/Krea-2/text_encoders/qwen3vl_4b_bf16.safetensors``, the same
+    file Krea 2 uses): a local copy under ``MODELS_PATH/text_encoders`` is used
+    in place, otherwise it is downloaded there. ``model.te_name_or_path`` points
+    at any other file, transformers folder or hub repo instead.
+  - **Previews** use the reference FlowDPM-Solver++ (order 2) sampler
+    (``src/pipeline.py``); 30 steps at CFG 3 is a good default (the paper uses
+    100).
+
+``name_or_path`` is the hub repo (``speridlabs/iris-3b``), a local export dir
+(``model.safetensors`` + ``config.yaml``), or a ``.safetensors`` file. The
+checkpoint's ``config.yaml`` (model / text_encoder / flow sections) drives the
+architecture, the text layers and the flow shift; the Iris-3B defaults apply
+when it is absent. Full fine-tunes save back in that same layout, so
+``scripts/sample.py`` from the reference repo loads them directly.
+"""
+
+import os
+from typing import List, Optional
+
+import torch
+import yaml
+import huggingface_hub
+from huggingface_hub.errors import EntryNotFoundError
+
+from toolkit.accelerator import unwrap_model
+from toolkit.basic import flush
+from toolkit.config_modules import GenerateImageConfig, ModelConfig
+from toolkit.metadata import get_meta_for_safetensors
+from toolkit.models.base_model import BaseModel
+from toolkit.models.FakeVAE import FakeVAE
+from toolkit.models.v2.resolver import resolve_component_file
+from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder
+from toolkit.prompt_utils import PromptEmbeds
+from toolkit.samplers.custom_flowmatch_sampler import (
+    CustomFlowMatchEulerDiscreteScheduler,
+)
+
+from .src.dit import IrisConfig, IrisDiT
+from .src.pipeline import Iris3BPipeline
+from .src.text_encoder import (
+    DEFAULT_HIDDEN_LAYERS,
+    DEFAULT_TEXT_LEN,
+    IrisPromptEncoder,
+)
+
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+
+# Reference flow settings (config.yaml ``flow`` section of the release).
+DEFAULT_FLOW_SHIFT = 4.0
+DEFAULT_NUM_TRAIN_TIMESTEPS = 1000
+DEFAULT_X_PRED_SIGMA_MIN = 0.05
+
+# Qwen3-VL-4B-Instruct: config + tokenizer (tiny files) from the vendor repo,
+# weights from the ComfyUI repack shared with Krea 2.
+QWEN3_VL_REPO = "Qwen/Qwen3-VL-4B-Instruct"
+DEFAULT_TE_PATH = "Comfy-Org/Krea-2/text_encoders/qwen3vl_4b_bf16.safetensors"
+
+
+def _resolve_checkpoint(name_or_path: str, filename: Optional[str] = None):
+    """``(weights_path, config_yaml_path_or_None)`` for a local file, a local
+    export dir, a hub repo id, or ``org/repo/sub/file.safetensors``."""
+    if os.path.isfile(name_or_path):
+        cfg = os.path.join(os.path.dirname(name_or_path), "config.yaml")
+        return name_or_path, (cfg if os.path.isfile(cfg) else None)
+    if os.path.isdir(name_or_path):
+        weights = os.path.join(name_or_path, filename or "model.safetensors")
+        if not os.path.isfile(weights):
+            candidates = [f for f in os.listdir(name_or_path) if f.endswith(".safetensors")]
+            if len(candidates) != 1:
+                raise FileNotFoundError(
+                    f"Could not pick an Iris checkpoint in {name_or_path}: found {candidates}. "
+                    "Set model.model_kwargs.checkpoint_filename."
+                )
+            weights = os.path.join(name_or_path, candidates[0])
+        cfg = os.path.join(name_or_path, "config.yaml")
+        return weights, (cfg if os.path.isfile(cfg) else None)
+
+    parts = name_or_path.split("/")
+    if name_or_path.endswith(".safetensors") and len(parts) >= 3:
+        repo_id, rel = "/".join(parts[:2]), "/".join(parts[2:])
+    elif len(parts) == 2:
+        repo_id, rel = name_or_path, (filename or "model.safetensors")
+    else:
+        raise FileNotFoundError(
+            f"Iris checkpoint {name_or_path!r} is not a local file/dir, a hub repo id, or "
+            "'org/repo/path/file.safetensors'."
+        )
+    weights = huggingface_hub.hf_hub_download(repo_id=repo_id, filename=rel, token=HF_TOKEN)
+    cfg_rel = os.path.join(os.path.dirname(rel), "config.yaml") if os.path.dirname(rel) else "config.yaml"
+    try:
+        cfg = huggingface_hub.hf_hub_download(repo_id=repo_id, filename=cfg_rel, token=HF_TOKEN)
+    except EntryNotFoundError:
+        cfg = None
+    return weights, cfg
+
+
+class Iris3BModel(BaseModel):
+    arch = "iris3b"
+
+    def __init__(
+        self,
+        device,
+        model_config: ModelConfig,
+        dtype="bf16",
+        custom_pipeline=None,
+        noise_scheduler=None,
+        **kwargs,
+    ):
+        super().__init__(device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs)
+        self.is_flow_matching = True
+        self.is_transformer = True
+        # LoRA targets, matched by class name: the trunk blocks only. The
+        # transformer_only filter is a substring match on "blocks", which also
+        # catches pixel_blocks (per-pixel PiT head) and y_embedder.*.blocks (text
+        # adapter); scoping the target classes keeps those frozen.
+        self.target_lora_modules = ["MMDiTBlock", "SingleStreamBlock"]
+
+        self.patch_size = 16
+        self.vae_scale_factor = 1  # pixel space
+        # flow / text settings; refreshed from the checkpoint's config.yaml at load
+        self.flow_shift = float(self.model_config.model_kwargs.get("flow_shift", DEFAULT_FLOW_SHIFT))
+        self.num_train_timesteps = DEFAULT_NUM_TRAIN_TIMESTEPS
+        self.prediction = "v"
+        self.x_pred_sigma_min = DEFAULT_X_PRED_SIGMA_MIN
+        self.text_len = DEFAULT_TEXT_LEN
+        self.text_hidden_layers = DEFAULT_HIDDEN_LAYERS
+        self.te_config_repo = QWEN3_VL_REPO
+        self.iris_raw_config: dict = {}
+        self.prompt_encoder: Optional[IrisPromptEncoder] = None
+
+    @staticmethod
+    def get_train_scheduler(
+        shift: float = DEFAULT_FLOW_SHIFT,
+        num_train_timesteps: int = DEFAULT_NUM_TRAIN_TIMESTEPS,
+    ):
+        # called on the class by the trainer / generate process (reference
+        # defaults); load_model rebuilds it with the checkpoint's own shift
+        return CustomFlowMatchEulerDiscreteScheduler(
+            num_train_timesteps=num_train_timesteps,
+            shift=shift,
+            use_dynamic_shifting=False,
+        )
+
+    def get_bucket_divisibility(self):
+        return self.vae_scale_factor * self.patch_size
+
+    # ------------------------------------------------------------------
+    # Loading
+    # ------------------------------------------------------------------
+    def _read_checkpoint_config(self, cfg_path: Optional[str]) -> IrisConfig:
+        raw = {}
+        if cfg_path is not None:
+            with open(cfg_path, "r") as f:
+                raw = yaml.safe_load(f) or {}
+        self.iris_raw_config = raw
+
+        flow = raw.get("flow", {}) or {}
+        if "flow_shift" not in self.model_config.model_kwargs:
+            self.flow_shift = float(flow.get("shift", DEFAULT_FLOW_SHIFT))
+        self.num_train_timesteps = int(flow.get("num_train_timesteps", DEFAULT_NUM_TRAIN_TIMESTEPS))
+        self.prediction = str(flow.get("prediction", "v"))
+        if self.prediction not in ("v", "x"):
+            raise ValueError(f"unknown flow.prediction '{self.prediction}' (v | x)")
+        self.x_pred_sigma_min = float(flow.get("x_pred_sigma_min", DEFAULT_X_PRED_SIGMA_MIN))
+
+        te = raw.get("text_encoder", {}) or {}
+        self.text_len = int(te.get("max_length", DEFAULT_TEXT_LEN))
+        layers = te.get("hidden_layers", None)
+        self.text_hidden_layers = tuple(int(i) for i in layers) if layers else DEFAULT_HIDDEN_LAYERS
+        self.te_config_repo = str(te.get("pretrained", QWEN3_VL_REPO))
+
+        model_raw = dict(raw.get("model", {}) or {})
+        model_raw.update(self.model_config.model_kwargs.get("model_config", {}) or {})
+        config = IrisConfig.from_dict(model_raw)
+        if config.text_adapter == "lap_blocks2" and config.text_lap_num_layers != len(self.text_hidden_layers):
+            raise ValueError(
+                f"model.text_lap_num_layers={config.text_lap_num_layers} but text_encoder.hidden_layers "
+                f"has {len(self.text_hidden_layers)} entries"
+            )
+        if config.text_len != self.text_len:
+            raise ValueError(
+                f"model.text_len={config.text_len} but text_encoder.max_length={self.text_len}"
+            )
+        return config
+
+    def _load_transformer(self) -> IrisDiT:
+        self.print_and_status_update("Loading Iris-3B transformer")
+        weights, cfg_path = _resolve_checkpoint(
+            self.model_config.name_or_path,
+            self.model_config.model_kwargs.get("checkpoint_filename", None),
+        )
+        config = self._read_checkpoint_config(cfg_path)
+        self.patch_size = config.patch_size
+        self.print_and_status_update(f"  - weights: {weights}")
+        # fp32 release weights -> load dtype; quantize / offload / place per model_config
+        transformer = IrisDiT.load(
+            weights,
+            config=config,
+            **self.component_load_kwargs("transformer"),
+        )
+        flush()
+        return transformer
+
+    def _load_text_encoder(self):
+        dtype = self.torch_dtype
+        te_path = self.model_config.te_name_or_path or DEFAULT_TE_PATH
+        tokenizer = Qwen3VLTextEncoder.load_tokenizer(self.te_config_repo, subfolder="", token=HF_TOKEN)
+
+        if os.path.isdir(te_path) or not te_path.endswith(".safetensors"):
+            # transformers folder or hub repo: load exactly what it names
+            self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}")
+            text_encoder = Qwen3VLTextEncoder.load_model(
+                te_path, dtype=dtype, subfolder="", use_comfy_weights=False, token=HF_TOKEN
+            )
+        else:
+            # single-file ComfyUI checkpoint: local copy in MODELS_PATH/text_encoders
+            # wins, otherwise downloaded there; config from the vendor repo
+            te_file = resolve_component_file(
+                te_path,
+                "text_encoders",
+                component="Qwen3-VL text encoder",
+                hf_token=HF_TOKEN,
+                status_fn=self.print_and_status_update,
+            )
+            self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_file}")
+            text_encoder = Qwen3VLTextEncoder.load_model(
+                te_file,
+                dtype=dtype,
+                config_path=self.te_config_repo,
+                subfolder="",
+                use_comfy_weights=False,
+            )
+        # text-only conditioning: the vision tower is dead weight
+        text_encoder.drop_vision_tower()
+        text_encoder.eval()
+        text_encoder.requires_grad_(False)
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
+        flush()
+        return tokenizer, text_encoder
+
+    def load_model(self):
+        transformer = self._load_transformer()
+        tokenizer, text_encoder = self._load_text_encoder()
+
+        self.print_and_status_update("Preparing pixel-space VAE (identity)")
+        vae = FakeVAE(scaling_factor=1.0)
+        vae.to(self.vae_device_torch, dtype=self.vae_torch_dtype)
+
+        self.noise_scheduler = Iris3BModel.get_train_scheduler(self.flow_shift, self.num_train_timesteps)
+        self.vae = vae
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.prompt_encoder = IrisPromptEncoder(
+            tokenizer, text_len=self.text_len, hidden_layers=self.text_hidden_layers
+        )
+        self.model = transformer
+        self.pipeline = Iris3BPipeline(self)
+        self.print_and_status_update("Model Loaded")
+
+    # ------------------------------------------------------------------
+    # Sampling (training previews)
+    # ------------------------------------------------------------------
+    def get_generation_pipeline(self):
+        return Iris3BPipeline(self)
+
+    def generate_single_image(
+        self,
+        pipeline: Iris3BPipeline,
+        gen_config: GenerateImageConfig,
+        conditional_embeds: PromptEmbeds,
+        unconditional_embeds: PromptEmbeds,
+        generator: torch.Generator,
+        extra: dict,
+    ):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+
+        sc = self.get_bucket_divisibility()
+        gen_config.width = int(gen_config.width // sc * sc)
+        gen_config.height = int(gen_config.height // sc * sc)
+
+        mkw = self.model_config.model_kwargs
+        cfg_interval = mkw.get("cfg_interval", (0.0, 1.0))
+        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,
+            generator=generator,
+            order=int(mkw.get("sample_order", 2)),
+            cfg_interval=(float(cfg_interval[0]), float(cfg_interval[1])),
+        )[0]
+        return img
+
+    # ------------------------------------------------------------------
+    # Training hooks
+    # ------------------------------------------------------------------
+    def model_velocity(
+        self,
+        x_t: torch.Tensor,
+        t_model: torch.Tensor,
+        text_feats: torch.Tensor,
+        text_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run the DiT and return the flow velocity ``noise - clean``. ``t_model``
+        is Iris model time in [0, 1000] (= the toolkit timestep). An
+        x-prediction checkpoint is converted via ``(x_t - x0) / max(sigma, s_min)``."""
+        out = self.model(x_t, t_model, text_feats, y_mask=text_mask)
+        if self.prediction == "x":
+            sigma = (t_model.float() / self.num_train_timesteps).view(-1, 1, 1, 1).to(out.dtype)
+            out = (x_t - out) / sigma.clamp(min=self.x_pred_sigma_min)
+        return out
+
+    def get_noise_prediction(
+        self,
+        latent_model_input: torch.Tensor,  # (B, 3, H, W) pixels in [-1, 1] + noise
+        timestep: torch.Tensor,  # 0..1000 scale, 1000 = pure noise
+        text_embeddings: PromptEmbeds,
+        **kwargs,
+    ):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+
+        # toolkit timestep IS Iris model time (1000 * noise level); with
+        # timestep_type=shift that is the shifted sigma the reference trained on
+        t_model = timestep.to(self.device_torch, dtype=torch.float32)
+        if t_model.dim() == 0:
+            t_model = t_model.unsqueeze(0)
+        if t_model.shape[0] != latent_model_input.shape[0]:
+            t_model = t_model.expand(latent_model_input.shape[0])
+
+        feats = text_embeddings.text_embeds.to(self.device_torch, self.torch_dtype)
+        mask = getattr(text_embeddings, "attention_mask", None)
+        if mask is None:
+            raise ValueError("Iris prompt embeds must carry an attention_mask")
+        mask = mask.to(self.device_torch)
+
+        return self.model_velocity(
+            latent_model_input.to(self.device_torch, self.torch_dtype), t_model, feats, mask
+        )
+
+    def get_prompt_embeds(self, prompt, **kwargs) -> PromptEmbeds:
+        """Fixed-length (text_len) stacked Qwen3-VL layer features.
+
+        ``text_embeds`` is ``(B, text_len, num_layers * 2560)`` -- the layer axis
+        is flattened so cached embeds batch/concat like ordinary (B, L, D)
+        embeddings; the DiT restores it. ``attention_mask`` is ``(B, text_len)``
+        bool. The empty prompt encodes to the reference CFG null."""
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        if self.text_encoder.device == torch.device("cpu"):
+            self.text_encoder.to(self.device_torch)
+        embeds, mask = self.prompt_encoder(self.text_encoder, list(prompt), dtype=self.torch_dtype)
+        pe = PromptEmbeds(embeds)
+        pe.attention_mask = mask
+        return pe
+
+    def condition_noisy_latents(self, latents: torch.Tensor, batch):
+        return latents
+
+    def get_model_has_grad(self):
+        return False
+
+    def get_te_has_grad(self):
+        return False
+
+    # ------------------------------------------------------------------
+    # Saving / bookkeeping
+    # ------------------------------------------------------------------
+    def save_model(self, output_path, meta, save_dtype):
+        """Full fine-tune save in the reference export layout:
+        ``<output_path>/model.safetensors`` + ``config.yaml``."""
+        os.makedirs(output_path, exist_ok=True)
+        transformer: IrisDiT = unwrap_model(self.model)
+        transformer.save_model(
+            os.path.join(output_path, "model.safetensors"),
+            dtype=save_dtype,
+            metadata=get_meta_for_safetensors(meta, name="iris3b"),
+        )
+        raw = dict(self.iris_raw_config)
+        raw["model"] = transformer.config.to_dict()
+        te = dict(raw.get("text_encoder", {}) or {})
+        te.update(
+            {
+                "name": "qwen3_vl",
+                "pretrained": self.te_config_repo,
+                "dim": transformer.config.text_dim,
+                "max_length": self.text_len,
+                "hidden_layers": list(self.text_hidden_layers),
+            }
+        )
+        raw["text_encoder"] = te
+        flow = dict(raw.get("flow", {}) or {})
+        flow.update(
+            {
+                "num_train_timesteps": self.num_train_timesteps,
+                "shift": self.flow_shift,
+                "prediction": self.prediction,
+                "x_pred_sigma_min": self.x_pred_sigma_min,
+            }
+        )
+        raw["flow"] = flow
+        with open(os.path.join(output_path, "config.yaml"), "w") as f:
+            yaml.safe_dump(raw, f, sort_keys=False)
+        with open(os.path.join(output_path, "aitk_meta.yaml"), "w") as f:
+            yaml.dump(meta, f)
+
+    def get_base_model_version(self):
+        return "iris3b"
+
+    def get_transformer_block_names(self) -> Optional[List[str]]:
+        return ["blocks"]
+
+    lora_keys_use_comfy_prefix = True

extensions_built_in/diffusion_models/iris3b/src/pipeline.py

@@ -0,0 +1,154 @@
+"""Preview sampler for Iris-3B: FlowDPM-Solver++ (multistep, order 2) on the
+shifted rectified-flow grid, ported from the reference ``iris3b/flow/solver.py``.
+
+Data-prediction formulation with ``alpha(t) = 1 - t``, ``sigma(t) = t``,
+``lambda(t) = log((1 - t) / t)``. The model predicts velocity
+(``x0_hat = x - t * v``). The time grid is
+``sigma = 1 - linspace(1, 0.001, steps + 1)`` remapped by the training shift,
+descending and ending exactly at t = 0 so the last (first-order) update is an
+exact projection onto the predicted clean image. NFE == steps.
+
+CFG is applied to the raw velocity (``uncond + s * (cond - uncond)``) and gated
+to model times strictly inside ``cfg_interval``. Pixel space: the integrator
+state IS the image in [-1, 1]; "decode" is a clamp + uint8 cast.
+"""
+
+import math
+from typing import List, Optional, Tuple
+
+import torch
+from PIL import Image
+from diffusers.utils.torch_utils import randn_tensor
+
+
+def shift_sigma(sigma: torch.Tensor, shift: float) -> torch.Tensor:
+    return shift * sigma / (1 + (shift - 1) * sigma)
+
+
+def time_grid(steps: int, shift: float) -> List[float]:
+    sigma = 1.0 - torch.linspace(1.0, 0.001, steps + 1, dtype=torch.float64)
+    return shift_sigma(sigma, shift).flip(0).tolist()  # descending, last value exactly 0
+
+
+def _lambda(t: float) -> float:
+    if t <= 0.0:
+        return math.inf
+    return math.log((1.0 - t) / t)
+
+
+def _first_order(x, s: float, t: float, x0):
+    h = _lambda(t) - _lambda(s)
+    phi1 = math.expm1(-h)
+    return (t / s) * x - (1.0 - t) * phi1 * x0
+
+
+def _second_order(x, prev1, prev0, t: float):
+    (s1, x0_1), (s0, x0_0) = prev1, prev0
+    lam_t, lam_0, lam_1 = _lambda(t), _lambda(s0), _lambda(s1)
+    h, h0 = lam_t - lam_0, lam_0 - lam_1
+    r0 = h0 / h
+    d = (x0_0 - x0_1) / r0
+    phi1 = math.expm1(-h)
+    return (t / s0) * x - (1.0 - t) * phi1 * x0_0 - 0.5 * (1.0 - t) * phi1 * d
+
+
+class Iris3BPipeline:
+    """Lightweight sampler used by ai-toolkit's preview generation."""
+
+    def __init__(self, model):
+        # model: the Iris3BModel holder (transformer, devices, flow settings)
+        self.model = model
+
+    @property
+    def device(self):
+        return self.model.device_torch
+
+    def to(self, *args, **kwargs):
+        return self
+
+    def set_progress_bar_config(self, **kwargs):
+        pass
+
+    @staticmethod
+    def _feats(embeds, device, dtype):
+        feats = embeds.text_embeds.to(device, dtype=dtype)
+        mask = getattr(embeds, "attention_mask", None)
+        if mask is None:
+            raise ValueError("Iris prompt embeds must carry an attention_mask")
+        return feats, mask.to(device)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        conditional_embeds,
+        unconditional_embeds,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = 30,
+        guidance_scale: float = 3.0,
+        latents: Optional[torch.Tensor] = None,
+        generator: Optional[torch.Generator] = None,
+        order: int = 2,
+        cfg_interval: Tuple[float, float] = (0.0, 1.0),
+        shift: Optional[float] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        model = self.model
+        device = model.device_torch
+        dtype = model.torch_dtype
+        transformer = model.transformer
+        shift = model.flow_shift if shift is None else shift
+        num_train_timesteps = model.num_train_timesteps
+
+        do_cfg = unconditional_embeds is not None and guidance_scale != 1.0
+        lo, hi = cfg_interval
+
+        if latents is None:
+            shape = (1, transformer.in_channels, height, width)
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        x = latents.to(device, dtype=torch.float32)
+        batch = x.shape[0]
+
+        cond_feats, cond_mask = self._feats(conditional_embeds, device, dtype)
+        if do_cfg:
+            uncond_feats, uncond_mask = self._feats(unconditional_embeds, device, dtype)
+            cfg_feats = torch.cat([uncond_feats, cond_feats], dim=0)
+            cfg_mask = torch.cat([uncond_mask, cond_mask], dim=0)
+
+        def velocity(x_t: torch.Tensor, t: float) -> torch.Tensor:
+            t_model = torch.full((batch,), t * num_train_timesteps, device=device, dtype=torch.float32)
+            if do_cfg and lo < t < hi:
+                out = model.model_velocity(
+                    torch.cat([x_t, x_t], dim=0).to(dtype),
+                    torch.cat([t_model, t_model], dim=0),
+                    cfg_feats,
+                    cfg_mask,
+                )
+                out_uncond, out_cond = out.float().chunk(2, dim=0)
+                return out_uncond + guidance_scale * (out_cond - out_uncond)
+            return model.model_velocity(x_t.to(dtype), t_model, cond_feats, cond_mask).float()
+
+        grid = time_grid(num_inference_steps, shift)
+        history: List[Tuple[float, torch.Tensor]] = []
+        emit = getattr(model, "_emit_sample_step", None)
+        for i in range(1, num_inference_steps + 1):
+            s, t = grid[i - 1], grid[i]
+            x0 = x - s * velocity(x, s)
+            if emit is not None and getattr(model, "sample_step_hook", None) is not None:
+                emit(x0, i - 1, num_inference_steps)
+            history.append((s, x0))
+            # order ramps up over the first steps and back down at the end so
+            # the terminal t=0 update is first-order: an exact x <- x0 projection
+            step_order = min(i, order, num_inference_steps + 1 - i)
+            if step_order == 1:
+                x = _first_order(x, s, t, x0)
+            else:
+                x = _second_order(x, history[-2], history[-1], t)
+            if len(history) > 2:
+                history.pop(0)
+
+        images = model.decode_latents(x, device=device, dtype=torch.float32)
+        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]

extensions_built_in/diffusion_models/iris3b/src/text_encoder.py

@@ -0,0 +1,114 @@
+"""Qwen3-VL prompt encoding for Iris-3B, ported from the reference
+``iris3b/text/qwen3_vl.py``.
+
+The conditioning is a fixed-length (``text_len`` = 300) stack of 12 Qwen3-VL
+hidden layers per token. The prompt is wrapped in a chat template: the system
+prefix is fed through the encoder as context but its hidden states are sliced
+off; the caption is truncated to ``text_len - len(suffix)`` tokens BEFORE the
+assistant-turn suffix is appended, so the suffix can never be the part that is
+cut; the rest is right-padded to ``text_len`` and the pad positions are zeroed.
+
+Fixed-length padding is load-bearing, not a convenience: the trunk's joint
+attention sees every text position unmasked (only the text adapter masks pad
+keys), so the number of pad tokens is part of the distribution the model was
+trained on. Always encode to the full ``text_len``.
+"""
+
+from typing import List, Tuple
+
+import torch
+
+# Chat template around every prompt (identical to the reference and to Krea 2).
+PROMPT_PREFIX = (
+    "<|im_start|>system\n"
+    "Describe the image by detailing the color, shape, size, texture, quantity, text, spatial "
+    "relationships of the objects and background:<|im_end|>\n"
+    "<|im_start|>user\n"
+)
+PROMPT_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"
+
+# 1-based decoder layers whose hidden states are stacked (hidden_states[i] is
+# the output of layer i; index 0 is the embedding output).
+DEFAULT_HIDDEN_LAYERS: Tuple[int, ...] = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)
+DEFAULT_TEXT_LEN = 300
+
+
+class IrisPromptEncoder:
+    """Tokenizes with the Iris template and runs the frozen Qwen3-VL language
+    stack. ``__call__`` returns ``(features, mask)`` with features
+    ``(B, text_len, num_layers * hidden)`` (layer axis flattened so the toolkit
+    treats it as an ordinary (B, L, D) embedding) and ``mask`` ``(B, text_len)``
+    bool, True at real tokens."""
+
+    def __init__(
+        self,
+        tokenizer,
+        text_len: int = DEFAULT_TEXT_LEN,
+        hidden_layers: Tuple[int, ...] = DEFAULT_HIDDEN_LAYERS,
+    ):
+        self.tokenizer = tokenizer
+        self.text_len = int(text_len)
+        self.hidden_layers = tuple(int(i) for i in hidden_layers)
+        prefix_ids = tokenizer.encode(PROMPT_PREFIX, add_special_tokens=False)
+        suffix_ids = tokenizer.encode(PROMPT_SUFFIX, add_special_tokens=False)
+        if not prefix_ids or not suffix_ids:
+            raise ValueError("Iris prompt template produced an empty token prefix or suffix")
+        self.prefix_ids = torch.tensor(prefix_ids, dtype=torch.long)
+        self.suffix_ids = torch.tensor(suffix_ids, dtype=torch.long)
+        self.caption_budget = self.text_len - len(suffix_ids)
+        if self.caption_budget < 1:
+            raise ValueError(f"text_len={self.text_len} leaves no room for a caption")
+        pad_id = tokenizer.pad_token_id
+        if pad_id is None:
+            pad_id = tokenizer.eos_token_id
+        if pad_id is None:
+            raise ValueError("tokenizer exposes neither pad_token_id nor eos_token_id")
+        self.pad_id = int(pad_id)
+
+    def tokenize(self, prompts: List[str]) -> Tuple[torch.Tensor, torch.Tensor]:
+        """``(input_ids, attention_mask)`` of shape ``(B, len(prefix) + text_len)``."""
+        caption_ids = self.tokenizer(prompts, add_special_tokens=False)["input_ids"]
+        start = len(self.prefix_ids)
+        stop = start + self.text_len
+        input_ids = torch.full((len(prompts), stop), self.pad_id, dtype=torch.long)
+        attention_mask = torch.zeros((len(prompts), stop), dtype=torch.long)
+        input_ids[:, :start] = self.prefix_ids
+        n_suffix = len(self.suffix_ids)
+        for row, ids in enumerate(caption_ids):
+            n = min(len(ids), self.caption_budget)
+            if n:
+                input_ids[row, start : start + n] = torch.tensor(ids[:n], dtype=torch.long)
+            input_ids[row, start + n : start + n + n_suffix] = self.suffix_ids
+            attention_mask[row, : start + n + n_suffix] = 1
+        return input_ids, attention_mask
+
+    @torch.no_grad()
+    def __call__(self, text_encoder, prompts: List[str], dtype=None) -> Tuple[torch.Tensor, torch.Tensor]:
+        # the language stack alone: Qwen3VLForConditionalGeneration.model is the
+        # inner Qwen3VLModel (skips the lm_head logits over 150k vocab)
+        decoder = getattr(text_encoder, "model", text_encoder)
+        device = text_encoder.device
+        input_ids, attention_mask = self.tokenize(prompts)
+        input_ids = input_ids.to(device)
+        attention_mask = attention_mask.to(device)
+        output = decoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            use_cache=False,
+            return_dict=True,
+            output_hidden_states=True,
+        )
+        hidden_states = output.hidden_states
+        if hidden_states is None or len(hidden_states) <= self.hidden_layers[-1]:
+            raise RuntimeError("Qwen3-VL decoder did not return the requested hidden layers")
+        start = len(self.prefix_ids)
+        stop = start + self.text_len
+        # (B, T, L, D), pads zeroed; with layer offloading the parameters'
+        # device is not the compute device, so follow the hidden states
+        embeds = torch.stack([hidden_states[layer][:, start:stop] for layer in self.hidden_layers], dim=2)
+        mask = attention_mask[:, start:stop].to(embeds.device)
+        embeds = embeds * mask[:, :, None, None].to(embeds.dtype)
+        if dtype is not None:
+            embeds = embeds.to(dtype)
+        embeds = embeds.reshape(embeds.shape[0], embeds.shape[1], -1)
+        return embeds, mask.bool()

extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py

@@ -19,19 +19,30 @@ see the sampled timestep.
 
 Flow-matching convention matches ai-toolkit (t=1 noise -> t=0 clean, target =
 noise - clean), so `get_noise_prediction` does no time flip or negation.
+
+Qwen-Image-2.1-Turbo is the same architecture distilled to 8 steps at CFG 1.
+Point `name_or_path` at its weight file in the repack
+(`Comfy-Org/Qwen-Image-2.1/diffusion_models/qwen_image_2.1_turbo_int8_convrot.safetensors`);
+`model_kwargs.turbo` samples with the schedule the checkpoint ships. Training
+it goes through a training adapter
+(`assistant_lora_path`): a LoRA that is live on the DiT while training and
+off while sampling, so samples show the turbo model plus what was trained.
 """
 
 import os
 from typing import TYPE_CHECKING, List, Optional
 
+import huggingface_hub
 import numpy as np
 import torch
 from PIL import Image
+from safetensors.torch import load_file
 
 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.config_modules import GenerateImageConfig, ModelConfig, NetworkConfig
+from toolkit.lora_special import LoRASpecialNetwork
 from toolkit.metadata import get_meta_for_safetensors
 from toolkit.models.base_model import BaseModel
 from toolkit.samplers.custom_flowmatch_sampler import (
@@ -80,6 +91,16 @@ scheduler_config = {
 COMFY_REPO = "Comfy-Org/Qwen-Image-2.1"
 BASE_REPO = "Qwen/Qwen-Image-2.1"
 
+# the turbo checkpoint is sampled with the schedule it ships
+# (TURBO_SAMPLE_SIGMAS in src/pipeline.py) as-is: no resolution shift, no
+# terminal stretch; the step count is fixed
+turbo_scheduler_config = {
+    **scheduler_config,
+    "shift": 1.0,
+    "shift_terminal": None,
+    "use_dynamic_shifting": False,
+}
+
 # decode above this many output pixels goes through the VAE's tiled path
 TILE_DECODE_ABOVE_PIXELS = 1024 * 1024
 
@@ -141,6 +162,15 @@ class QwenImage2Model(BaseModel):
     def get_train_scheduler():
         return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
 
+    @property
+    def is_turbo(self) -> bool:
+        return bool(self.model_config.model_kwargs.get("turbo", False))
+
+    def get_sample_scheduler(self):
+        if self.is_turbo:
+            return CustomFlowMatchEulerDiscreteScheduler(**turbo_scheduler_config)
+        return self.get_train_scheduler()
+
     def get_bucket_divisibility(self):
         # 16 for the VAE, 2 more because the DiT groups target latent tokens
         # into 2x2 blocks, one per vision slot
@@ -165,12 +195,35 @@ class QwenImage2Model(BaseModel):
             # a local full checkpoint supplies its own text encoder / vae
             base_model_path = model_path
 
+        load_kwargs = self.component_load_kwargs("transformer")
         self.print_and_status_update("Loading transformer")
-        transformer = QwenImage21Transformer2DModel.load(
-            model_path,
-            config_path=base_model_path,
-            **self.component_load_kwargs("transformer"),
-        )
+        if self.model_config.assistant_lora_path is None:
+            transformer = QwenImage21Transformer2DModel.load(
+                model_path, config_path=base_model_path, **load_kwargs
+            )
+        else:
+            # the adapter stays a live LoRA on the quantized linears, attached
+            # after quantization (it must wrap the quantized forward) and
+            # before layer offloading (the offloader routes through it)
+            offload = load_kwargs.pop("offload", 0.0)
+            transformer = QwenImage21Transformer2DModel.load_model(
+                model_path,
+                dtype=load_kwargs["dtype"],
+                config_path=base_model_path,
+                qtype=(load_kwargs.get("qtype") or "").split("|", 1)[0] or None,
+                quantize_on_load=False,
+            )
+            transformer.aitk_post_load(offload=0.0, **load_kwargs)
+            self.load_training_adapter(transformer)
+            if offload and offload > 0:
+                from toolkit.memory_management import MemoryManager
+
+                MemoryManager.attach(
+                    transformer,
+                    torch.device(load_kwargs["quantize_device"]),
+                    offload_percent=offload,
+                    ignore_modules=list(transformer.get_offload_ignore_modules() or []),
+                )
         flush()
 
         self.print_and_status_update("Loading text encoder")
@@ -204,6 +257,66 @@ class QwenImage2Model(BaseModel):
         self.pipeline = QwenImage21Pipeline(self)
         self.print_and_status_update("Model Loaded")
 
+    def load_training_adapter(self, transformer: QwenImage21Transformer2DModel):
+        """Attach the training adapter (a LoRA on the DiT) as a live assistant:
+        on at 1.0 while training, off while sampling so samples show the turbo
+        model. It is never merged: the weights are int8, whose grid swallows a
+        delta this small on a requantize."""
+        self.print_and_status_update("Loading assistant LoRA")
+        lora_path = self.model_config.assistant_lora_path
+        if not os.path.exists(lora_path):
+            # assume it is a hub path
+            lora_splits = lora_path.split("/")
+            if len(lora_splits) != 3:
+                raise ValueError(
+                    f"Assistant LoRA path {lora_path} is not a valid local path or hub path."
+                )
+            repo_id = "/".join(lora_splits[:2])
+            filename = lora_splits[2]
+            try:
+                lora_path = huggingface_hub.hf_hub_download(repo_id=repo_id, filename=filename)
+                self.model_config.assistant_lora_path = lora_path
+            except Exception as e:
+                raise ValueError(f"Failed to download assistant LoRA from {lora_path}: {e}")
+        lora_state_dict = load_file(lora_path)
+        dim_key = next(k for k in lora_state_dict if k.endswith("lora_A.weight"))
+        dim = int(lora_state_dict[dim_key].shape[0])
+        lora_state_dict = {
+            key.replace("diffusion_model.", "transformer."): value
+            for key, value in lora_state_dict.items()
+        }
+
+        network_config = NetworkConfig(type="lora", linear=dim, linear_alpha=dim, transformer_only=True)
+        LoRASpecialNetwork.LORA_PREFIX_UNET = "lora_transformer"
+        network = LoRASpecialNetwork(
+            text_encoder=None,
+            unet=transformer,
+            lora_dim=network_config.linear,
+            multiplier=1.0,
+            alpha=network_config.linear_alpha,
+            train_unet=True,
+            train_text_encoder=False,
+            network_config=network_config,
+            network_type=network_config.type,
+            transformer_only=network_config.transformer_only,
+            is_transformer=True,
+            target_lin_modules=self.target_lora_modules,
+            is_assistant_adapter=True,
+            is_ara=True,
+            # transformer_only filters to the block list only through the holder
+            base_model=self,
+        )
+        network.apply_to(None, transformer, apply_text_encoder=False, apply_unet=True)
+        network.force_to(self.device_torch, dtype=self.torch_dtype)
+        network.eval()
+        network.requires_grad_(False)
+        network._update_torch_multiplier()
+        network.load_weights(lora_state_dict)
+        self.assistant_lora: LoRASpecialNetwork = network
+        self.assistant_lora.multiplier = 1.0
+        self.assistant_lora.is_active = True
+        self.invert_assistant_lora = False
+
     # ------------------------------------------------------------------
     # VAE. The latents are RGBA. Images without alpha get an opaque one on
     # encode, and decode drops it again unless RGBA output is on.

extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py

@@ -303,6 +303,21 @@ def run_transformer(
     return unpack_latents(output[:, -target_tokens:], height, width)
 
 
+# Qwen-Image-2.1-Turbo's 8-step schedule, from `sample_sigmas` in the
+# checkpoint's model_index.json (Qwen/Qwen-Image-2.1-Turbo); the final step
+# runs from the last sigma to 0
+TURBO_SAMPLE_SIGMAS = [
+    1.0,
+    0.978453,
+    0.95418,
+    0.926626,
+    0.89508,
+    0.845148,
+    0.704534,
+    0.414568,
+]
+
+
 def calculate_shift(
     image_seq_len: int,
     base_seq_len: int = 256,
@@ -367,16 +382,22 @@ class QwenImage21Pipeline:
             )
         latents = latents.to(device, dtype=dtype)
 
-        scheduler = model.get_train_scheduler()
-        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
-        mu = calculate_shift(
-            latent_height * latent_width,
-            scheduler.config.get("base_image_seq_len", 256),
-            scheduler.config.get("max_image_seq_len", 8192),
-            scheduler.config.get("base_shift", 0.5),
-            scheduler.config.get("max_shift", 0.9),
-        )
-        scheduler.set_timesteps(sigmas=sigmas, device=device, mu=mu)
+        scheduler = model.get_sample_scheduler()
+        if model.is_turbo:
+            # the turbo checkpoint's own 8-step schedule; num_inference_steps
+            # does not override it (other schedules are unevaluated)
+            sigmas = np.array(TURBO_SAMPLE_SIGMAS, dtype=np.float32)
+            scheduler.set_timesteps(sigmas=sigmas, device=device)
+        else:
+            sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+            mu = calculate_shift(
+                latent_height * latent_width,
+                scheduler.config.get("base_image_seq_len", 256),
+                scheduler.config.get("max_image_seq_len", 8192),
+                scheduler.config.get("base_shift", 0.5),
+                scheduler.config.get("max_shift", 0.9),
+            )
+            scheduler.set_timesteps(sigmas=sigmas, device=device, mu=mu)
         scheduler.set_begin_index(0)
 
         # 2.1 is meant to be sampled without guidance; a scale of 1 skips it

extensions_built_in/diffusion_models/ui.tsx

@@ -2,7 +2,7 @@
 // registers in AI_TOOLKIT_MODELS. Loaded at runtime by the UI, not bundled:
 // see ui/src/extensions/README.md for the convention and the allowed imports.
 import Link from "next/link";
-import type { ModelArch } from "@/app/jobs/new/options";
+import type { CustomModelOption, ModelArch } from "@/app/jobs/new/options";
 import type { JobConfig } from "@/types";
 import {
   defaultSampleConfig,
@@ -12,6 +12,55 @@ import {
 const defaultNameOrPath = "";
 const defaultLinearRank = 32;
 
+// Qwen-Image-2.1 and its turbo distill share the RGBA VAE, so both entries
+// carry this option
+const qwenImage21TransparencyOption: CustomModelOption[] = [
+  {
+    type: "checkbox",
+    label: "Transparency (RGBA)",
+    getValue: (config: JobConfig) =>
+      config?.config?.process?.[0]?.model?.model_kwargs?.rgba ?? false,
+    onChange: (
+      value: boolean,
+      config: JobConfig,
+      setJobConfig: (value: any, key: string) => void,
+    ) => {
+      const kwargs = {
+        ...(config?.config?.process?.[0]?.model?.model_kwargs ?? {}),
+      };
+      if (value) {
+        kwargs.rgba = true;
+      } else {
+        delete kwargs.rgba;
+      }
+      setJobConfig(kwargs, "config.process[0].model.model_kwargs");
+    },
+    doc: {
+      title: "Transparency (RGBA)",
+      description: (
+        <div className="space-y-2">
+          <p>
+            This model&apos;s VAE is natively RGBA, so alpha can be carried
+            end to end. When on, dataset images and reference images load
+            with their alpha channel, the VAE encodes all four channels, and
+            samples are saved as PNGs with their transparency intact.
+          </p>
+          <p>
+            Images with no alpha of their own get a fully opaque one, so a
+            mixed dataset is fine. Turn this off to train and sample flat
+            RGB: alpha is dropped on the way in and added back as opaque.
+          </p>
+          <p>
+            Changing this re-caches latents, and it cannot be combined with
+            a dataset using <code>alpha_mask</code>, which consumes the
+            alpha channel as a loss mask instead.
+          </p>
+        </div>
+      ),
+    },
+  },
+];
+
 export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
   {
     name: "anima",
@@ -573,52 +622,56 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
       "model.low_vram",
       "model.layer_offloading",
     ],
-    customModelSelectOptions: [
-      {
-        type: "checkbox",
-        label: "Transparency (RGBA)",
-        getValue: (config: JobConfig) =>
-          config?.config?.process?.[0]?.model?.model_kwargs?.rgba ?? false,
-        onChange: (
-          value: boolean,
-          config: JobConfig,
-          setJobConfig: (value: any, key: string) => void,
-        ) => {
-          const kwargs = {
-            ...(config?.config?.process?.[0]?.model?.model_kwargs ?? {}),
-          };
-          if (value) {
-            kwargs.rgba = true;
-          } else {
-            delete kwargs.rgba;
-          }
-          setJobConfig(kwargs, "config.process[0].model.model_kwargs");
+    customModelSelectOptions: qwenImage21TransparencyOption,
+  },
+  {
+    name: "qwen_image_2:turbo",
+    label: "Qwen-Image-2.1 Turbo (w/ Training Adapter)",
+    generateNameOverride: "Qwen-Image-2.1 Turbo",
+    group: "image",
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      // the turbo weight file in the repack (text encoder and VAE come from
+      // the same repack); model_kwargs.turbo selects the turbo schedule
+      "config.process[0].model.name_or_path": [
+        "Comfy-Org/Qwen-Image-2.1/diffusion_models/qwen_image_2.1_turbo_int8_convrot.safetensors",
+        defaultNameOrPath,
+      ],
+      "config.process[0].model.quantize": [true, false],
+      "config.process[0].model.quantize_te": [true, false],
+      "config.process[0].model.low_vram": [true, false],
+      "config.process[0].train.unload_text_encoder": [false, false],
+      "config.process[0].sample.sampler": ["flowmatch", "flowmatch"],
+      "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"],
+      "config.process[0].train.timestep_type": ["shift", "sigmoid"],
+      "config.process[0].model.qtype": ["convrot8", "qfloat8"],
+      "config.process[0].model.qtype_te": ["convrot8", "qfloat8"],
+      // training adapter: live on the DiT while training, off while sampling
+      "config.process[0].model.assistant_lora_path": [
+        'ostris/qwen_image_2_turbo_training_adapter/qwen_image_2_turbo_training_adapter_v1.safetensors',
+        undefined
+      ],
+      // the turbo checkpoint samples with its own fixed 8-step schedule at
+      // CFG 1; the step count here is informational
+      "config.process[0].sample.guidance_scale": [1.0, 4.0],
+      "config.process[0].sample.sample_steps": [8, 25],
+      "config.process[0].model.model_kwargs": [
+        {
+          turbo: true,
+          rgba: false,
         },
-        doc: {
-          title: "Transparency (RGBA)",
-          description: (
-            <div className="space-y-2">
-              <p>
-                This model&apos;s VAE is natively RGBA, so alpha can be carried
-                end to end. When on, dataset images and reference images load
-                with their alpha channel, the VAE encodes all four channels, and
-                samples are saved as PNGs with their transparency intact.
-              </p>
-              <p>
-                Images with no alpha of their own get a fully opaque one, so a
-                mixed dataset is fine. Turn this off to train and sample flat
-                RGB: alpha is dropped on the way in and added back as opaque.
-              </p>
-              <p>
-                Changing this re-caches latents, and it cannot be combined with
-                a dataset using <code>alpha_mask</code>, which consumes the
-                alpha channel as a loss mask instead.
-              </p>
-            </div>
-          ),
-        },
-      },
+        {},
+      ],
+    },
+    disableSections: ["network.conv", "train.unload_text_encoder"],
+    additionalSections: [
+      "datasets.multi_control_paths",
+      "sample.multi_ctrl_imgs",
+      "model.low_vram",
+      "model.layer_offloading",
+      "model.assistant_lora_path",
     ],
+    customModelSelectOptions: qwenImage21TransparencyOption,
   },
   {
     name: "ming_image",
@@ -1920,6 +1973,29 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
     disableSections: ["network.conv"],
     additionalSections: ["model.low_vram", "model.layer_offloading"],
   },
+  {
+    name: "iris3b",
+    label: "Iris-3B (pixel space)",
+    group: "image",
+    defaults: {
+      "config.process[0].model.name_or_path": [
+        "speridlabs/iris-3b",
+        defaultNameOrPath,
+      ],
+      "config.process[0].model.quantize": [true, false],
+      "config.process[0].model.quantize_te": [true, false],
+      "config.process[0].model.qtype": ["convrot8", "qfloat8"],
+      "config.process[0].model.qtype_te": ["convrot8", "qfloat8"],
+      "config.process[0].train.timestep_type": ["shift", "sigmoid"],
+      "config.process[0].sample.guidance_scale": [3, 4],
+      "config.process[0].sample.sample_steps": [30, 25],
+      "config.process[0].network.conv": [undefined, 16],
+      "config.process[0].network.conv_alpha": [undefined, 16],
+      "config.process[0].model.low_vram": [true, false],
+    },
+    disableSections: ["network.conv"],
+    additionalSections: ["model.low_vram", "model.layer_offloading"],
+  },
   {
     name: "krea2",
     label: "Krea 2 (raw)",

extensions_built_in/inference_engine/latent_preview.py

@@ -51,6 +51,7 @@ ARCH_LATENT_FORMAT = {
     "cogview4": "Flux",
     "prx_pixel": None,  # pixel space
     "flux3": "Flux3Image",
+    "iris3b": None,  # pixel space
 }
 
 

testing/test_model_loading.py

@@ -160,6 +160,10 @@ MODEL_TESTS = {
         "model": {"name_or_path": "Photoroom/prxpixel-t2i", "quantize_te": True},
         "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0},
     },
+    "iris3b": {
+        "model": {"name_or_path": "speridlabs/iris-3b", "quantize_te": True},
+        "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 3.0},
+    },
     "zeta_chroma": {
         "model": {"name_or_path": "lodestones/Zeta-Chroma/zeta-chroma-base-x0-pixel-dino-distance.safetensors", "extras_name_or_path": "Tongyi-MAI/Z-Image-Turbo", "quantize": True, "quantize_te": True},
         "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0},
@@ -185,6 +189,18 @@ MODEL_TESTS = {
         "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 1.0},
         "needs_control_image": True,
     },
+    "qwen_image_2:turbo": {
+        # model_kwargs.turbo samples with the checkpoint's fixed 8-step
+        # schedule at CFG 1
+        "model": {
+            "name_or_path": "Comfy-Org/Qwen-Image-2.1/diffusion_models/qwen_image_2.1_turbo_int8_convrot.safetensors",
+            "quantize": True,
+            "quantize_te": True,
+            "model_kwargs": {"turbo": True},
+        },
+        "sample": {**IMG, "num_inference_steps": 8, "guidance_scale": 1.0},
+        "needs_control_image": True,
+    },
     "ming_image": {
         "model": {
             "name_or_path": "Comfy-Org/Ming-Image",

toolkit/models/registry.py

@@ -44,6 +44,7 @@ ARCH_REGISTRY: Dict[str, dict] = {
     "flux2_klein_4b": {"modality": "image", "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-4B", "quantize_te": True}, "sample": dict(IMG)},
     "flux2_klein_9b": {"modality": "image", "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-9B", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
     "prx_pixel": {"modality": "image", "model": {"name_or_path": "Photoroom/prxpixel-t2i", "quantize_te": True}, "sample": dict(IMG)},
+    "iris3b": {"modality": "image", "model": {"name_or_path": "speridlabs/iris-3b", "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 30, "guidance_scale": 3.0}},
     "f-lite": {"modality": "image", "model": {"name_or_path": "Freepik/F-Lite", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
     "sd1": {"modality": "image", "model": {"name_or_path": "stable-diffusion-v1-5/stable-diffusion-v1-5"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 7.5}},
     "sdxl": {"modality": "image", "model": {"name_or_path": "stabilityai/stable-diffusion-xl-base-1.0"}, "sample": {**IMG, "guidance_scale": 6.0}},

toolkit/models/v2/text_encoders/qwen3_vl.py

@@ -44,6 +44,23 @@ class Qwen3VLTextEncoder(Qwen3VLForConditionalGeneration, OstrisTransformersMixi
     def get_transformer_block_names(cls):
         return ["model.language_model.layers"]
 
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        """Single-file (ComfyUI repack) checkpoints: accept the older comfy key
+        layout (``model.layers.*`` / ``visual.*``) next to the transformers one,
+        and materialize the tied ``lm_head`` the files omit."""
+        out = {}
+        for key, value in state_dict.items():
+            if key.startswith("model.layers.") or key in ("model.embed_tokens.weight", "model.norm.weight"):
+                key = "model.language_model." + key[len("model.") :]
+            elif key.startswith("visual."):
+                key = "model." + key
+            out[key] = value
+        embed = out.get("model.language_model.embed_tokens.weight")
+        if "lm_head.weight" not in out and embed is not None:
+            out["lm_head.weight"] = embed
+        return out
+
     def drop_vision_tower(self):
         """Text-only conditioning: the vision tower is dead weight — drop it to
         free VRAM and skip loading its (bf16-slow) Conv3d patch_embed."""

version.py

@@ -1 +1 @@
-VERSION = "0.13.23"
+VERSION = "0.13.25"