Add support for Iris-3B

README.md

@@ -32,6 +32,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)

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
@@ -56,6 +57,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/ui.tsx

@@ -1888,6 +1888,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

@@ -50,6 +50,7 @@ ARCH_LATENT_FORMAT = {
     "anima": "Wan21",
     "cogview4": "Flux",
     "prx_pixel": None,  # pixel space
+    "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},

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."""