From 2fa0bd87e6cd12f4d054346865687937b9f8df63 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 9 Oct 2026 05:13:39 -0600 Subject: [PATCH 1/4] Add support for Iris-3B --- README.md | 1 + .../diffusion_models/__init__.py | 2 + .../diffusion_models/iris3b/README.md | 59 + .../diffusion_models/iris3b/__init__.py | 3 + .../diffusion_models/iris3b/iris3b.py | 435 +++++++ .../diffusion_models/iris3b/src/__init__.py | 0 .../diffusion_models/iris3b/src/dit.py | 1136 +++++++++++++++++ .../diffusion_models/iris3b/src/pipeline.py | 154 +++ .../iris3b/src/text_encoder.py | 114 ++ extensions_built_in/diffusion_models/ui.tsx | 23 + .../inference_engine/latent_preview.py | 1 + testing/test_model_loading.py | 4 + toolkit/models/registry.py | 1 + toolkit/models/v2/text_encoders/qwen3_vl.py | 17 + 14 files changed, 1950 insertions(+) create mode 100644 extensions_built_in/diffusion_models/iris3b/README.md create mode 100644 extensions_built_in/diffusion_models/iris3b/__init__.py create mode 100644 extensions_built_in/diffusion_models/iris3b/iris3b.py create mode 100644 extensions_built_in/diffusion_models/iris3b/src/__init__.py create mode 100644 extensions_built_in/diffusion_models/iris3b/src/dit.py create mode 100644 extensions_built_in/diffusion_models/iris3b/src/pipeline.py create mode 100644 extensions_built_in/diffusion_models/iris3b/src/text_encoder.py diff --git a/README.md b/README.md index 1490eca..4809e40 100644 --- a/README.md +++ b/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) diff --git a/extensions_built_in/diffusion_models/__init__.py b/extensions_built_in/diffusion_models/__init__.py index 660f0e2..e94fa39 100644 --- a/extensions_built_in/diffusion_models/__init__.py +++ b/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, diff --git a/extensions_built_in/diffusion_models/iris3b/README.md b/extensions_built_in/diffusion_models/iris3b/README.md new file mode 100644 index 0000000..e2534fd --- /dev/null +++ b/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 ` 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. diff --git a/extensions_built_in/diffusion_models/iris3b/__init__.py b/extensions_built_in/diffusion_models/iris3b/__init__.py new file mode 100644 index 0000000..8d23bb1 --- /dev/null +++ b/extensions_built_in/diffusion_models/iris3b/__init__.py @@ -0,0 +1,3 @@ +from .iris3b import Iris3BModel + +__all__ = ["Iris3BModel"] diff --git a/extensions_built_in/diffusion_models/iris3b/iris3b.py b/extensions_built_in/diffusion_models/iris3b/iris3b.py new file mode 100644 index 0000000..355e0bb --- /dev/null +++ b/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: + ``/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 diff --git a/extensions_built_in/diffusion_models/iris3b/src/__init__.py b/extensions_built_in/diffusion_models/iris3b/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/extensions_built_in/diffusion_models/iris3b/src/dit.py b/extensions_built_in/diffusion_models/iris3b/src/dit.py new file mode 100644 index 0000000..29af73a --- /dev/null +++ b/extensions_built_in/diffusion_models/iris3b/src/dit.py @@ -0,0 +1,1136 @@ +"""Iris-3B pixel-space diffusion transformer, vendored from +https://github.com/speridlabs/iris-3b (``src/iris3b/models`` + ``src/iris3b/nn``). + +Two levels: + - patch stage: a hybrid trunk of ``dual_depth`` dual-stream MM-DiT blocks + followed by single-stream blocks over 16x16 patch tokens, text-conditioned + through a layerwise-attention adapter over stacked Qwen3-VL hidden layers; + - pixel stage: ``pixel.depth`` PiT blocks refine each patch's pixels, + conditioned on the patch stage's output tokens. +No VAE: input and output are RGB velocities in pixel space. + +Deviations from the reference, none of which change the state-dict layout: + - SDPA only (no fa3/fa4/cudnn backend switch); REPA feature capture dropped. + - Shared adaLN cores (``model.modulation = shared_*``) are evaluated ONCE per + forward in ``IrisDiT.forward`` and handed to the blocks, instead of each + block holding a hidden Python reference to the core Linear. The reference's + tuple trick bypasses module replacement (quantization / LoRA swap the + Linear under ``modulation_cores`` and the blocks would keep calling the old + one), so the core lives only in ``modulation_cores`` here. + - gradient checkpointing per block via ``enable_gradient_checkpointing()``. +""" + +import math +from dataclasses import asdict, dataclass, field, fields +from typing import Any, Dict, List, Optional + +import torch +import torch.nn.functional as F +from torch import nn +from torch.utils.checkpoint import checkpoint + +from toolkit.models.v2._mixin import OstrisModelMixin + + +# --------------------------------------------------------------------------- +# config +# --------------------------------------------------------------------------- + + +@dataclass +class IrisPixelConfig: + enabled: bool = True + depth: int = 4 + hidden_size: int = 16 + attn_hidden_size: int = 1280 + num_heads: int = 10 + mlp_ratio: float = 4.0 + modulation: str = "post" # pre | post + abs_pos_embed: bool = True + + +@dataclass +class IrisConfig: + """Mirror of the reference ``ModelConfig`` (defaults = Iris-3B). Keys the + reference config.yaml carries that do not affect the module (attn_backend, + repa_layer, adaln_zero_init) are accepted and ignored by ``from_dict``.""" + + block: str = "single_stream" # mmdit | single_stream + dual_depth: int = 8 + final_block_text: str = "keep" # keep | drop + hidden_size: int = 2560 + depth: int = 24 + num_heads: int = 20 + num_kv_heads: Optional[int] = 5 + gated_attention: bool = True + sandwich_norm: bool = True + patch_size: int = 16 + in_channels: int = 3 + mlp_ratio: float = 4.0 + qkv_bias: bool = False + qk_norm: bool = True + norm_eps: float = 1e-6 + modulation: str = "shared_bias" # per_block | shared_lowrank | shared_bias + modulation_rank: int = 64 + timestep_max_period: float = 10.0 + rope_theta: float = 10000.0 + rope_scale: float = 16.0 + rope_aspect: str = "isotropic" # square | isotropic + rope_frame_pairs: int = 0 + rope_frame_theta: float = 10.0 + text_rope: bool = True + text_rope_theta: float = 10000.0 + text_abs_pos_embed: bool = True + text_dim: int = 2560 + text_len: int = 300 + text_adapter: str = "lap_blocks2" # linear | blocks2 | lap_blocks2 + text_lap_num_layers: int = 12 + text_lap_num_heads: int = 32 + text_lap_mlp_ratio: float = 1.3 + pixel: IrisPixelConfig = field(default_factory=IrisPixelConfig) + + @classmethod + def from_dict(cls, raw: Optional[Dict[str, Any]]) -> "IrisConfig": + raw = dict(raw or {}) + pixel_raw = raw.pop("pixel", None) or {} + known = {f.name for f in fields(cls)} + kwargs = {k: v for k, v in raw.items() if k in known} + pixel_known = {f.name for f in fields(IrisPixelConfig)} + kwargs["pixel"] = IrisPixelConfig( + **{k: v for k, v in pixel_raw.items() if k in pixel_known} + ) + return cls(**kwargs) + + def to_dict(self) -> Dict[str, Any]: + return asdict(self) + + # BaseModel reads unet.config['in_channels'] / unet.config.patch_size + def __getitem__(self, key): + return getattr(self, key) + + def get(self, key, default=None): + return getattr(self, key, default) + + def validate(self) -> None: + if not 0 <= self.dual_depth <= self.depth: + raise ValueError( + f"dual_depth must be in [0, depth={self.depth}], got {self.dual_depth}" + ) + if self.final_block_text not in ("keep", "drop"): + raise ValueError(f"unknown final_block_text '{self.final_block_text}'") + if self.block not in ("mmdit", "single_stream"): + raise ValueError(f"unknown block '{self.block}'") + if self.modulation not in ("per_block", "shared_lowrank", "shared_bias"): + raise ValueError(f"unknown modulation '{self.modulation}'") + if self.text_adapter not in ("linear", "blocks2", "lap_blocks2"): + raise ValueError(f"unknown text_adapter '{self.text_adapter}'") + + +# --------------------------------------------------------------------------- +# primitives +# --------------------------------------------------------------------------- + + +class RMSNorm(nn.Module): + """RMS normalization computed in fp32, learned gain applied in input dtype.""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + normed = x.float() * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps) + return self.weight * normed.to(x.dtype) + + +class SwiGLU(nn.Module): + """Gated SiLU feed-forward with the 2/3 width rule, bias-free.""" + + def __init__(self, dim: int, mlp_ratio: float = 4.0): + super().__init__() + hidden = int(2 * int(dim * mlp_ratio) / 3) + self.w1 = nn.Linear(dim, hidden, bias=False) + self.w3 = nn.Linear(dim, hidden, bias=False) + self.w2 = nn.Linear(hidden, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +class GeluMLP(nn.Module): + def __init__(self, dim: int, mlp_ratio: float = 4.0): + super().__init__() + hidden = int(dim * mlp_ratio) + self.fc1 = nn.Linear(dim, hidden, bias=True) + self.fc2 = nn.Linear(hidden, dim, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc2(F.gelu(self.fc1(x))) + + +def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + return x * (1 + scale) + shift + + +# ---- rope ---- + + +def rope_2d( + head_dim: int, + height: int, + width: int, + theta: float = 10000.0, + scale: float = 16.0, + aspect: str = "square", + frame_pairs: int = 0, + frame_theta: float = 10.0, + frame_index: int = 0, +) -> torch.Tensor: + """Complex rotation factors ``[height*width, head_dim//2]`` for a row-major + grid. Coordinates are normalized to a fixed span ``[0, scale]`` so any grid + size covers the same angular range (``isotropic`` shares one step across + both axes so aspect ratio survives).""" + n_pairs = head_dim // 4 + if not 0 <= frame_pairs < n_pairs: + raise ValueError(f"frame_pairs must be in [0, {n_pairs}), got {frame_pairs}") + xy_pairs = n_pairs - frame_pairs + freqs = 1.0 / theta ** (torch.arange(0, head_dim, 4)[:n_pairs].float() / head_dim) + freqs = freqs[:xy_pairs] + if aspect == "square": + x_pos = torch.linspace(0, scale, width) + y_pos = torch.linspace(0, scale, height) + elif aspect == "isotropic": + step = scale / max(max(height, width) - 1, 1) + x_pos = torch.arange(width).float() * step + y_pos = torch.arange(height).float() * step + else: + raise ValueError(f"unknown rope aspect mode '{aspect}'") + x_ang = torch.outer(x_pos, freqs) + y_ang = torch.outer(y_pos, freqs) + x_cis = torch.polar(torch.ones_like(x_ang), x_ang) + y_cis = torch.polar(torch.ones_like(y_ang), y_ang) + x_grid = x_cis[None, :, :].expand(height, width, -1) + y_grid = y_cis[:, None, :].expand(height, width, -1) + cis = torch.stack([x_grid, y_grid], dim=-1).reshape(height * width, 2 * xy_pairs) + if frame_pairs: + n_slots = 2 * frame_pairs + f_freqs = 1.0 / frame_theta ** (torch.arange(n_slots).float() / n_slots) + f_cis = torch.polar(torch.ones(n_slots), float(frame_index) * f_freqs) + cis = torch.cat([cis, f_cis[None, :].expand(height * width, -1)], dim=-1) + return cis + + +def rope_1d(head_dim: int, length: int, theta: float = 10000.0) -> torch.Tensor: + freqs = 1.0 / theta ** (torch.arange(0, head_dim, 2)[: head_dim // 2].float() / head_dim) + ang = torch.outer(torch.arange(length).float(), freqs) + return torch.polar(torch.ones_like(ang), ang) + + +def apply_rope(q: torch.Tensor, k: torch.Tensor, freqs_cis: torch.Tensor): + """Rotate ``q``/``k`` of shape ``[B, N, H, head_dim]``; complex math in fp32.""" + q_c = torch.view_as_complex(q.float().reshape(*q.shape[:-1], -1, 2)) + k_c = torch.view_as_complex(k.float().reshape(*k.shape[:-1], -1, 2)) + fc = freqs_cis[None, :, None, :] + q_out = torch.view_as_real(q_c * fc).flatten(3) + k_out = torch.view_as_real(k_c * fc).flatten(3) + return q_out.to(q.dtype), k_out.to(k.dtype) + + +# ---- attention ---- + + +def scaled_dot_product(q, k, v, attn_mask=None, enable_gqa: bool = False): + """Attention over ``[B, heads, N, head_dim]`` tensors (SDPA).""" + return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, enable_gqa=enable_gqa) + + +class SelfAttention(nn.Module): + """Multi-head self-attention with per-head QK RMSNorm and RoPE.""" + + def __init__(self, dim, num_heads, qkv_bias=False, qk_norm=True, norm_eps=1e-6): + super().__init__() + if dim % num_heads: + raise ValueError(f"dim {dim} not divisible by heads {num_heads}") + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_norm = RMSNorm(self.head_dim, eps=norm_eps) if qk_norm else nn.Identity() + self.k_norm = RMSNorm(self.head_dim, eps=norm_eps) if qk_norm else nn.Identity() + self.proj = nn.Linear(dim, dim, bias=True) + + def forward(self, x, rope=None, attn_mask=None): + batch, n, dim = x.shape + q, k, v = self.qkv(x).reshape(batch, n, 3, self.num_heads, self.head_dim).unbind(2) + q, k = self.q_norm(q), self.k_norm(k) + if rope is not None: + q, k = apply_rope(q, k, rope) + q, k, v = (t.transpose(1, 2) for t in (q, k, v)) + out = scaled_dot_product(q, k, v, attn_mask=attn_mask) + return self.proj(out.transpose(1, 2).reshape(batch, n, dim)) + + +class JointAttention(nn.Module): + """MM-DiT joint attention: separate stream projections, one softmax over + the text-first concatenation. GQA unfuses QKV into Q/K/V linears.""" + + def __init__( + self, + dim, + num_heads, + qkv_bias=False, + qk_norm=True, + norm_eps=1e-6, + num_kv_heads=None, + text_out=True, + ): + super().__init__() + if dim % num_heads: + raise ValueError(f"dim {dim} not divisible by heads {num_heads}") + num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads + if num_kv_heads <= 0 or num_heads % num_kv_heads: + raise ValueError(f"query heads {num_heads} not divisible by KV heads {num_kv_heads}") + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_dim = dim // num_heads + if num_kv_heads == num_heads: + self.qkv_x = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.qkv_y = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_proj_x = self.k_proj_x = self.v_proj_x = None + self.q_proj_y = self.k_proj_y = self.v_proj_y = None + else: + kv_dim = num_kv_heads * self.head_dim + self.qkv_x = self.qkv_y = None + self.q_proj_x = nn.Linear(dim, dim, bias=qkv_bias) + self.k_proj_x = nn.Linear(dim, kv_dim, bias=qkv_bias) + self.v_proj_x = nn.Linear(dim, kv_dim, bias=qkv_bias) + self.q_proj_y = nn.Linear(dim, dim, bias=qkv_bias) + self.k_proj_y = nn.Linear(dim, kv_dim, bias=qkv_bias) + self.v_proj_y = nn.Linear(dim, kv_dim, bias=qkv_bias) + self.q_norm_x = RMSNorm(self.head_dim, eps=norm_eps) if qk_norm else nn.Identity() + self.k_norm_x = RMSNorm(self.head_dim, eps=norm_eps) if qk_norm else nn.Identity() + self.q_norm_y = RMSNorm(self.head_dim, eps=norm_eps) if qk_norm else nn.Identity() + self.k_norm_y = RMSNorm(self.head_dim, eps=norm_eps) if qk_norm else nn.Identity() + self.proj_x = nn.Linear(dim, dim, bias=True) + self.proj_y = nn.Linear(dim, dim, bias=True) if text_out else None + + def _project(self, tokens, stream): + batch, n_tokens, _ = tokens.shape + fused = self.qkv_x if stream == "x" else self.qkv_y + if fused is not None: + return fused(tokens).reshape(batch, n_tokens, 3, self.num_heads, self.head_dim).unbind(2) + if stream == "x": + q_proj, k_proj, v_proj = self.q_proj_x, self.k_proj_x, self.v_proj_x + else: + q_proj, k_proj, v_proj = self.q_proj_y, self.k_proj_y, self.v_proj_y + q = q_proj(tokens).reshape(batch, n_tokens, self.num_heads, self.head_dim) + k = k_proj(tokens).reshape(batch, n_tokens, self.num_kv_heads, self.head_dim) + v = v_proj(tokens).reshape(batch, n_tokens, self.num_kv_heads, self.head_dim) + return q, k, v + + def forward(self, x, y, rope_img, rope_txt=None, gate_x=None, gate_y=None): + batch, n_img, dim = x.shape + n_txt = y.shape[1] + qx, kx, vx = self._project(x, "x") + qy, ky, vy = self._project(y, "y") + qx, kx = self.q_norm_x(qx), self.k_norm_x(kx) + qy, ky = self.q_norm_y(qy), self.k_norm_y(ky) + qx, kx = apply_rope(qx, kx, rope_img) + if rope_txt is not None: + qy, ky = apply_rope(qy, ky, rope_txt) + q = torch.cat([qy, qx], dim=1).transpose(1, 2) + k = torch.cat([ky, kx], dim=1).transpose(1, 2) + v = torch.cat([vy, vx], dim=1).transpose(1, 2) + out = scaled_dot_product(q, k, v, enable_gqa=self.num_kv_heads != self.num_heads) + out = out.transpose(1, 2).reshape(batch, n_txt + n_img, dim) + out_x = out[:, n_txt:] + if gate_x is not None: + out_x = out_x * gate_x + if self.proj_y is None: + return self.proj_x(out_x), None + out_y = out[:, :n_txt] + if gate_y is not None: + out_y = out_y * gate_y + return self.proj_x(out_x), self.proj_y(out_y) + + +# ---- embedders ---- + + +class TimestepEmbedder(nn.Module): + """Sinusoidal features + 2-layer SiLU MLP. ``max_period`` is 10 (not + 10000): model time is the shifted flow level scaled to [0, 1000].""" + + def __init__(self, hidden_size: int, freq_dim: int = 256, max_period: float = 10.0): + super().__init__() + self.freq_dim = freq_dim + self.max_period = max_period + self.mlp = nn.Sequential( + nn.Linear(freq_dim, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + + def timestep_embedding(self, t: torch.Tensor) -> torch.Tensor: + n = self.freq_dim // 2 + k = torch.arange(n, dtype=torch.float32, device=t.device) + phase = torch.outer(t.float(), torch.exp(-math.log(self.max_period) * k / n)) + return torch.cat([phase.cos(), phase.sin()], dim=-1) + + def forward(self, t: torch.Tensor) -> torch.Tensor: + """``t``: [B] model time -> [B, 1, hidden].""" + batch = t.shape[0] + emb = self.timestep_embedding(t.reshape(-1)) + emb = self.mlp(emb.to(self.mlp[0].weight.dtype)) + return emb.reshape(batch, -1, emb.shape[-1]) + + +class PatchEmbedder(nn.Module): + def __init__(self, in_dim: int, hidden_size: int, norm: bool = False, norm_eps: float = 1e-6): + super().__init__() + self.proj = nn.Linear(in_dim, hidden_size, bias=True) + self.norm = RMSNorm(hidden_size, eps=norm_eps) if norm else nn.Identity() + + def forward(self, x): + return self.norm(self.proj(x)) + + +class TextEmbedder(nn.Module): + """Linear + RMSNorm text adapter (``text_adapter = linear``).""" + + def __init__(self, text_dim: int, hidden_size: int, norm_eps: float = 1e-6): + super().__init__() + self.proj = nn.Linear(text_dim, hidden_size, bias=True) + self.norm = RMSNorm(hidden_size, eps=norm_eps) + + def forward(self, y): + return self.norm(self.proj(y)) + + +class TextAdapterBlock(nn.Module): + """Unconditioned pre-norm block: ``x + attn(norm(x))``, ``x + swiglu(norm(x))``.""" + + def __init__(self, dim, num_heads, mlp_ratio=4.0, qkv_bias=False, qk_norm=True, norm_eps=1e-6): + super().__init__() + self.norm1 = RMSNorm(dim, eps=norm_eps) + self.attn = SelfAttention(dim, num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm, norm_eps=norm_eps) + self.norm2 = RMSNorm(dim, eps=norm_eps) + self.mlp = SwiGLU(dim, mlp_ratio) + + def forward(self, y, attn_mask=None): + y = y + self.attn(self.norm1(y), attn_mask=attn_mask) + return y + self.mlp(self.norm2(y)) + + +class LayerwiseAttentionBlock(nn.Module): + """Pre-norm block over one token's stack of encoder-layer states.""" + + def __init__(self, dim, num_heads, mlp_ratio=1.3, norm_eps=1e-6): + super().__init__() + hidden = int(dim * mlp_ratio) + self.norm1 = RMSNorm(dim, eps=norm_eps) + self.attn = SelfAttention(dim, num_heads, qkv_bias=False, qk_norm=False, norm_eps=norm_eps) + self.norm2 = RMSNorm(dim, eps=norm_eps) + self.mlp = nn.Sequential( + nn.Linear(dim, hidden, bias=True), + nn.SiLU(), + nn.Linear(hidden, dim, bias=True), + ) + + def forward(self, x): + x = x + self.attn(self.norm1(x)) + return x + self.mlp(self.norm2(x)) + + +class TransformerTextEmbedder(nn.Module): + """``Linear`` -> N unconditioned blocks -> RMSNorm. No RoPE (text positions + come from ``y_pos_embedding`` + the trunk's text RoPE). Pad KEYS are masked + inside its blocks; every query keeps its own position (mask OR diagonal) + so an all-pad row cannot produce NaN.""" + + def __init__( + self, + text_dim, + hidden_size, + num_blocks=2, + num_heads=24, + mlp_ratio=4.0, + qkv_bias=False, + qk_norm=True, + norm_eps=1e-6, + ): + super().__init__() + self.proj = nn.Linear(text_dim, hidden_size, bias=True) + self.blocks = nn.ModuleList( + [ + TextAdapterBlock( + hidden_size, num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, + qk_norm=qk_norm, norm_eps=norm_eps, + ) + for _ in range(num_blocks) + ] + ) + self.norm = RMSNorm(hidden_size, eps=norm_eps) + + def forward(self, y, mask): + if mask.shape != y.shape[:2]: + raise ValueError(f"text mask must be {tuple(y.shape[:2])}, got {tuple(mask.shape)}") + keep = mask.bool()[:, None, None, :] + eye = torch.eye(y.shape[1], dtype=torch.bool, device=y.device) + attn_mask = keep | eye + y = self.proj(y) + for block in self.blocks: + y = block(y, attn_mask=attn_mask) + return self.norm(y) + + +class LayerwiseTextEmbedder(nn.Module): + """Aggregate the stacked frozen-encoder layers per token (2 attention blocks + over the layer axis + learned pooling), then refine across tokens.""" + + def __init__( + self, + text_dim, + hidden_size, + num_layers, + layer_num_heads, + layer_mlp_ratio=1.3, + refiner_num_heads=24, + refiner_mlp_ratio=4.0, + qkv_bias=False, + qk_norm=True, + norm_eps=1e-6, + ): + super().__init__() + if num_layers <= 0: + raise ValueError(f"text_lap_num_layers must be positive, got {num_layers}") + self.text_dim = text_dim + self.num_layers = num_layers + self.layer_blocks = nn.ModuleList( + [ + LayerwiseAttentionBlock(text_dim, layer_num_heads, mlp_ratio=layer_mlp_ratio, norm_eps=norm_eps) + for _ in range(2) + ] + ) + self.layer_pool = nn.Linear(num_layers, 1, bias=True) + self.refiner = TransformerTextEmbedder( + text_dim, + hidden_size, + num_blocks=2, + num_heads=refiner_num_heads, + mlp_ratio=refiner_mlp_ratio, + qkv_bias=qkv_bias, + qk_norm=qk_norm, + norm_eps=norm_eps, + ) + + def forward(self, y, mask): + expected = (*mask.shape, self.num_layers, self.text_dim) + if tuple(y.shape) != expected: + raise ValueError(f"layerwise text states must be {expected}, got {tuple(y.shape)}") + batch, tokens, layers, dim = y.shape + y = y.reshape(batch * tokens, layers, dim) + for block in self.layer_blocks: + y = block(y) + y = self.layer_pool(y.transpose(1, 2)).squeeze(-1) + y = y.reshape(batch, tokens, dim) + return self.refiner(y, mask) + + +def _sincos_1d(embed_dim: int, pos: torch.Tensor) -> torch.Tensor: + omega = torch.arange(embed_dim // 2, dtype=torch.float64) / (embed_dim / 2.0) + omega = 1.0 / 10000.0**omega + out = torch.outer(pos.reshape(-1).double(), omega) + return torch.cat([torch.sin(out), torch.cos(out)], dim=1) + + +def sincos_pos_embed_2d(embed_dim: int, height: int, width: int) -> torch.Tensor: + """Fixed 2D sincos table ``[height*width, embed_dim]`` (MAE/DiT grid convention).""" + grid_h = torch.arange(height, dtype=torch.float32) + grid_w = torch.arange(width, dtype=torch.float32) + grid = torch.meshgrid(grid_w, grid_h, indexing="xy") + grid = torch.stack(grid, dim=0).reshape(2, 1, height, width) + emb = torch.cat([_sincos_1d(embed_dim // 2, grid[0]), _sincos_1d(embed_dim // 2, grid[1])], dim=1) + return emb.float() + + +class PixelEmbedder(nn.Module): + """Per-pixel linear embedding grouped into per-patch sequences: + ``[B, C, H, W] -> [B*(H/p)*(W/p), p*p, hidden]``.""" + + def __init__(self, in_channels: int, hidden_size: int, patch_size: int, abs_pos_embed: bool = True): + super().__init__() + self.patch_size = patch_size + self.abs_pos_embed = abs_pos_embed + self.proj = nn.Linear(in_channels, hidden_size, bias=True) + self._pos_cache: Dict[tuple, torch.Tensor] = {} + + def _pos(self, height, width, device, dtype): + key = (height, width, str(device), dtype) + table = self._pos_cache.get(key) + if table is None: + table = sincos_pos_embed_2d(self.proj.out_features, height, width) + table = table.reshape(height, width, -1).to(device=device, dtype=dtype) + self._pos_cache[key] = table + return table + + def forward(self, x): + batch, _, height, width = x.shape + p = self.patch_size + h_patches, w_patches = height // p, width // p + tokens = self.proj(x.permute(0, 2, 3, 1)) + if self.abs_pos_embed: + tokens = tokens + self._pos(height, width, tokens.device, tokens.dtype) + tokens = tokens.reshape(batch, h_patches, p, w_patches, p, -1) + tokens = tokens.permute(0, 1, 3, 2, 4, 5) + return tokens.reshape(batch * h_patches * w_patches, p * p, -1) + + +# ---- modulation ---- + + +class BlockModulation(nn.Module): + """One stream's adaLN projection for a patch block, same parameter names + as the reference for every mode: + + - ``per_block``: own ``weight``/``bias`` (the reference's ``nn.Linear(D, 6D)``) + - ``shared_bias``: ``bias`` only; output = shared core output + bias + - ``shared_lowrank``: ``down`` / ``adaln_up``; output = core + U(V(cond)) + + For the shared modes the core output ``shared`` ([B, 1, 6D]) is computed + once per forward by ``IrisDiT`` from ``modulation_cores`` and passed in. + """ + + def __init__(self, dim: int, mode: str, rank: int = 64): + super().__init__() + self.mode = mode + if mode == "per_block": + self.weight = nn.Parameter(torch.empty(6 * dim, dim)) + self.bias = nn.Parameter(torch.zeros(6 * dim)) + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + elif mode == "shared_bias": + self.bias = nn.Parameter(torch.zeros(6 * dim)) + elif mode == "shared_lowrank": + if rank <= 0: + raise ValueError(f"modulation_rank must be positive, got {rank}") + self.down = nn.Linear(dim, rank, bias=False) + self.adaln_up = nn.Linear(rank, 6 * dim, bias=False) + else: + raise ValueError(f"unknown modulation mode '{mode}'") + + def forward(self, cond: torch.Tensor, shared: Optional[torch.Tensor]) -> torch.Tensor: + if self.mode == "per_block": + return F.linear(cond, self.weight, self.bias) + if shared is None: + raise ValueError(f"{self.mode} modulation needs the shared core output") + if self.mode == "shared_bias": + return shared + self.bias + return shared + self.adaln_up(self.down(cond)) + + +# --------------------------------------------------------------------------- +# blocks +# --------------------------------------------------------------------------- + + +class MMDiTBlock(nn.Module): + """Dual-stream block: separate text/image weights, joint attention.""" + + def __init__( + self, + dim, + num_heads, + mlp_ratio=4.0, + qkv_bias=False, + qk_norm=True, + norm_eps=1e-6, + modulation="per_block", + modulation_rank=64, + num_kv_heads=None, + gated_attention=False, + sandwich_norm=False, + text_out=True, + ): + super().__init__() + self.norm_x1 = RMSNorm(dim, eps=norm_eps) + self.norm_x2 = RMSNorm(dim, eps=norm_eps) + self.norm_y1 = RMSNorm(dim, eps=norm_eps) + self.norm_y2 = RMSNorm(dim, eps=norm_eps) if text_out else None + self.attn = JointAttention( + dim, num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm, norm_eps=norm_eps, + num_kv_heads=num_kv_heads, text_out=text_out, + ) + self.attn_gate_x = nn.Linear(dim, dim, bias=False) if gated_attention else None + self.attn_gate_y = nn.Linear(dim, dim, bias=False) if gated_attention and text_out else None + self.attn_post_norm_x = RMSNorm(dim, eps=norm_eps) if sandwich_norm else nn.Identity() + self.attn_post_norm_y = ( + (RMSNorm(dim, eps=norm_eps) if sandwich_norm else nn.Identity()) if text_out else None + ) + self.mlp_x = SwiGLU(dim, mlp_ratio) + self.mlp_y = SwiGLU(dim, mlp_ratio) if text_out else None + self.mlp_post_norm_x = RMSNorm(dim, eps=norm_eps) if sandwich_norm else nn.Identity() + self.mlp_post_norm_y = ( + (RMSNorm(dim, eps=norm_eps) if sandwich_norm else nn.Identity()) if text_out else None + ) + self.adaln_img = BlockModulation(dim, modulation, modulation_rank) + self.adaln_txt = BlockModulation(dim, modulation, modulation_rank) + self.text_out = text_out + + def forward(self, x, y, cond, rope_img, rope_txt=None, mod_img=None, mod_txt=None): + (xs1, xc1, xg1, xs2, xc2, xg2) = self.adaln_img(cond, mod_img).chunk(6, dim=-1) + (ys1, yc1, yg1, ys2, yc2, yg2) = self.adaln_txt(cond, mod_txt).chunk(6, dim=-1) + hx = modulate(self.norm_x1(x), xs1, xc1) + hy = modulate(self.norm_y1(y), ys1, yc1) + attn_x, attn_y = self.attn( + hx, + hy, + rope_img, + rope_txt, + gate_x=None if self.attn_gate_x is None else torch.sigmoid(self.attn_gate_x(hx)), + gate_y=None if self.attn_gate_y is None else torch.sigmoid(self.attn_gate_y(hy)), + ) + x = x + xg1 * self.attn_post_norm_x(attn_x) + x = x + xg2 * self.mlp_post_norm_x(self.mlp_x(modulate(self.norm_x2(x), xs2, xc2))) + if not self.text_out: + return x, y + y = y + yg1 * self.attn_post_norm_y(attn_y) + y = y + yg2 * self.mlp_post_norm_y(self.mlp_y(modulate(self.norm_y2(y), ys2, yc2))) + return x, y + + +class SingleStreamBlock(nn.Module): + """Single-stream block: text and image share one set of weights.""" + + def __init__( + self, + dim, + num_heads, + mlp_ratio=4.0, + qkv_bias=False, + qk_norm=True, + norm_eps=1e-6, + modulation="per_block", + modulation_rank=64, + num_kv_heads=None, + gated_attention=False, + sandwich_norm=False, + text_out=True, + ): + super().__init__() + if dim % num_heads: + raise ValueError(f"dim {dim} not divisible by query heads {num_heads}") + num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads + if num_kv_heads <= 0 or num_heads % num_kv_heads: + raise ValueError(f"query heads {num_heads} not divisible by KV heads {num_kv_heads}") + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_dim = dim // num_heads + self.norm1 = RMSNorm(dim, eps=norm_eps) + self.norm2 = RMSNorm(dim, eps=norm_eps) + if num_kv_heads == num_heads: + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_proj = self.k_proj = self.v_proj = None + else: + kv_dim = num_kv_heads * self.head_dim + self.qkv = None + self.q_proj = nn.Linear(dim, dim, bias=qkv_bias) + self.k_proj = nn.Linear(dim, kv_dim, bias=qkv_bias) + self.v_proj = nn.Linear(dim, kv_dim, bias=qkv_bias) + self.q_norm = RMSNorm(self.head_dim, eps=norm_eps) if qk_norm else nn.Identity() + self.k_norm = RMSNorm(self.head_dim, eps=norm_eps) if qk_norm else nn.Identity() + self.attn_gate = nn.Linear(dim, dim, bias=False) if gated_attention else None + self.attn_proj = nn.Linear(dim, dim, bias=True) + self.attn_post_norm = RMSNorm(dim, eps=norm_eps) if sandwich_norm else nn.Identity() + self.mlp = SwiGLU(dim, mlp_ratio) + self.mlp_post_norm = RMSNorm(dim, eps=norm_eps) if sandwich_norm else nn.Identity() + self.adaln = BlockModulation(dim, modulation, modulation_rank) + self.text_out = text_out + + def _qkv(self, h, batch, n): + if self.qkv is not None: + return self.qkv(h).reshape(batch, n, 3, self.num_heads, self.head_dim).unbind(2) + q = self.q_proj(h).reshape(batch, n, self.num_heads, self.head_dim) + k = self.k_proj(h).reshape(batch, n, self.num_kv_heads, self.head_dim) + v = self.v_proj(h).reshape(batch, n, self.num_kv_heads, self.head_dim) + return q, k, v + + def forward(self, x, y, cond, rope_img, rope_txt=None, mod=None): + """Text and image rows share every weight, so each per-token op (norm, + modulation, projections, gate, MLP) runs per stream and only Q/K/V are + concatenated for the joint attention: same math as the reference's + cat-first layout, but no kernel iterates over a cat output, which + inductor cannot range-split under dynamic shapes (CantSplit).""" + batch, n_img, dim = x.shape + n_txt = y.shape[1] + (s1, c1, g1, s2, c2, g2) = self.adaln(cond, mod).chunk(6, dim=-1) + + hx = modulate(self.norm1(x), s1, c1) + hy = modulate(self.norm1(y), s1, c1) + qx, kx, vx = self._qkv(hx, batch, n_img) + qy, ky, vy = self._qkv(hy, batch, n_txt) + qx, kx = self.q_norm(qx), self.k_norm(kx) + qy, ky = self.q_norm(qy), self.k_norm(ky) + qx, kx = apply_rope(qx, kx, rope_img) + if rope_txt is not None: + qy, ky = apply_rope(qy, ky, rope_txt) + q = torch.cat([qy, qx], dim=1).transpose(1, 2) + k = torch.cat([ky, kx], dim=1).transpose(1, 2) + v = torch.cat([vy, vx], dim=1).transpose(1, 2) + attn = scaled_dot_product(q, k, v, enable_gqa=self.num_kv_heads != self.num_heads) + attn = attn.transpose(1, 2) # [B, n_txt + n_img, H, D] + attn_x = attn[:, n_txt:].reshape(batch, n_img, dim) + attn_y = attn[:, :n_txt].reshape(batch, n_txt, dim) if self.text_out else None + + def finish(tokens, h, a): + if self.attn_gate is not None: + a = a * torch.sigmoid(self.attn_gate(h)) + tokens = tokens + g1 * self.attn_post_norm(self.attn_proj(a)) + mlp = self.mlp(modulate(self.norm2(tokens), s2, c2)) + return tokens + g2 * self.mlp_post_norm(mlp) + + x = finish(x, hx, attn_x) + if not self.text_out: + # text supplied keys/values above; its rows are discarded downstream + return x, y + return x, finish(y, hy, attn_y) + + +class PiTBlock(nn.Module): + """Pixel-level transformer block over per-patch pixel sequences + ``[B*L, p*p, d_pix]`` conditioned on the patch's (timestep-fused) token. + ``post`` modulation: affine on branch outputs, chunk order (scale, shift).""" + + def __init__( + self, + hidden_size, + cond_dim, + pixels_per_patch, + attn_hidden_size, + num_heads, + mlp_ratio=4.0, + modulation="pre", + qk_norm=True, + norm_eps=1e-6, + ): + super().__init__() + if modulation not in ("pre", "post"): + raise ValueError(f"unknown PiT modulation '{modulation}'") + self.modulation = modulation + self.pixels_per_patch = pixels_per_patch + n_mod = 6 if modulation == "pre" else 4 + self.norm1 = RMSNorm(hidden_size, eps=norm_eps) + self.norm2 = RMSNorm(hidden_size, eps=norm_eps) + self.adaln = nn.Linear(cond_dim, n_mod * hidden_size * pixels_per_patch, bias=True) + self.compress = nn.Linear(pixels_per_patch * hidden_size, attn_hidden_size, bias=True) + self.expand = nn.Linear(attn_hidden_size, pixels_per_patch * hidden_size, bias=True) + self.attn = SelfAttention(attn_hidden_size, num_heads, qkv_bias=False, qk_norm=qk_norm, norm_eps=norm_eps) + self.mlp = GeluMLP(hidden_size, mlp_ratio) + + def _global_attn(self, pixels, rope, grid): + n_patches = grid[0] * grid[1] + batch = pixels.shape[0] // n_patches + compact = self.compress(pixels.reshape(batch * n_patches, -1)) + attn = self.attn(compact.reshape(batch, n_patches, -1), rope=rope) + return self.expand(attn.reshape(batch * n_patches, -1)).reshape_as(pixels) + + def forward(self, x, cond, rope, grid): + mods = self.adaln(cond).reshape(x.shape[0], self.pixels_per_patch, -1) + if self.modulation == "pre": + shift1, scale1, gate1, shift2, scale2, gate2 = mods.chunk(6, dim=-1) + h = modulate(self.norm1(x), shift1, scale1) + x = x + gate1 * self._global_attn(h, rope, grid) + x = x + gate2 * self.mlp(modulate(self.norm2(x), shift2, scale2)) + else: + scale1, shift1, scale2, shift2 = mods.chunk(4, dim=-1) + x = x + modulate(self._global_attn(self.norm1(x), rope, grid), shift1, scale1) + x = x + modulate(self.mlp(self.norm2(x)), shift2, scale2) + return x + + +class FinalLayer(nn.Module): + """Unmodulated head: RMSNorm + Linear.""" + + def __init__(self, hidden_size: int, out_dim: int, norm_eps: float = 1e-6): + super().__init__() + self.norm = RMSNorm(hidden_size, eps=norm_eps) + self.linear = nn.Linear(hidden_size, out_dim, bias=True) + + def forward(self, x): + return self.linear(self.norm(x)) + + +# --------------------------------------------------------------------------- +# model +# --------------------------------------------------------------------------- + + +class IrisDiT(nn.Module, OstrisModelMixin): + """Iris pixel-space DiT. ``forward(x, t, y, y_mask)`` -> velocity [B, C, H, W].""" + + def __init__(self, config): + super().__init__() + if isinstance(config, dict): + config = IrisConfig.from_dict(config) + config.validate() + self.config = config + cfg = config + self.gradient_checkpointing = False + p = cfg.patch_size + dim = cfg.hidden_size + + self.s_embedder = PatchEmbedder(p * p * cfg.in_channels, dim) + self.t_embedder = TimestepEmbedder(dim, max_period=cfg.timestep_max_period) + if cfg.text_adapter == "linear": + self.y_embedder = TextEmbedder(cfg.text_dim, dim, norm_eps=cfg.norm_eps) + elif cfg.text_adapter == "blocks2": + self.y_embedder = TransformerTextEmbedder( + cfg.text_dim, dim, num_blocks=2, num_heads=cfg.num_heads, mlp_ratio=cfg.mlp_ratio, + qkv_bias=cfg.qkv_bias, qk_norm=cfg.qk_norm, norm_eps=cfg.norm_eps, + ) + else: + self.y_embedder = LayerwiseTextEmbedder( + cfg.text_dim, + dim, + num_layers=cfg.text_lap_num_layers, + layer_num_heads=cfg.text_lap_num_heads, + layer_mlp_ratio=cfg.text_lap_mlp_ratio, + refiner_num_heads=cfg.num_heads, + refiner_mlp_ratio=cfg.mlp_ratio, + qkv_bias=cfg.qkv_bias, + qk_norm=cfg.qk_norm, + norm_eps=cfg.norm_eps, + ) + self._adapter_needs_mask = cfg.text_adapter != "linear" + self.y_pos_embedding = ( + nn.Parameter(torch.randn(1, cfg.text_len, dim)) if cfg.text_abs_pos_embed else None + ) + + # shared adaLN cores, one Linear(D, 6D) per stream, evaluated once per + # forward; the single-stream half of a hybrid trunk rides the dual + # half's image core (reference ``stream_aliases={"shared": "img"}``) + self.modulation_cores = nn.ModuleDict() + self._shared_modulation = cfg.modulation in ("shared_lowrank", "shared_bias") + self._stream_alias = {"shared": "img"} if cfg.dual_depth else {} + if self._shared_modulation: + streams = [] + if cfg.dual_depth: + streams += ["img", "txt"] + if cfg.dual_depth < cfg.depth: + streams.append(self._stream_alias.get("shared", "shared")) + for stream in dict.fromkeys(streams): + self.modulation_cores[f"adaln_{stream}"] = nn.Linear(dim, 6 * dim, bias=True) + + block_kwargs = dict( + mlp_ratio=cfg.mlp_ratio, + qkv_bias=cfg.qkv_bias, + qk_norm=cfg.qk_norm, + norm_eps=cfg.norm_eps, + modulation=cfg.modulation, + modulation_rank=cfg.modulation_rank, + num_kv_heads=cfg.num_kv_heads, + gated_attention=cfg.gated_attention, + sandwich_norm=cfg.sandwich_norm, + ) + tail_cls = MMDiTBlock if cfg.block == "mmdit" else SingleStreamBlock + final_text_out = cfg.final_block_text == "keep" + block_classes = [MMDiTBlock if i < cfg.dual_depth else tail_cls for i in range(cfg.depth)] + self.blocks = nn.ModuleList( + [ + block_cls(dim, cfg.num_heads, text_out=final_text_out or i < cfg.depth - 1, **block_kwargs) + for i, block_cls in enumerate(block_classes) + ] + ) + # recorded here rather than isinstance() at forward time: the trainer + # may replace entries with torch.compile wrappers (OptimizedModule) + self._block_is_dual = [block_cls is MMDiTBlock for block_cls in block_classes] + + if not cfg.pixel.enabled: + self.pixel_embedder = None + self.pixel_blocks = None + self.final_layer = FinalLayer(dim, p * p * cfg.in_channels, norm_eps=cfg.norm_eps) + else: + self.pixel_embedder = PixelEmbedder( + cfg.in_channels, cfg.pixel.hidden_size, p, abs_pos_embed=cfg.pixel.abs_pos_embed + ) + self.pixel_blocks = nn.ModuleList( + [ + PiTBlock( + cfg.pixel.hidden_size, + cond_dim=dim, + pixels_per_patch=p * p, + attn_hidden_size=cfg.pixel.attn_hidden_size, + num_heads=cfg.pixel.num_heads, + mlp_ratio=cfg.pixel.mlp_ratio, + modulation=cfg.pixel.modulation, + qk_norm=cfg.qk_norm, + norm_eps=cfg.norm_eps, + ) + for _ in range(cfg.pixel.depth) + ] + ) + self.final_layer = FinalLayer(cfg.pixel.hidden_size, cfg.in_channels, norm_eps=cfg.norm_eps) + + self._rope_img: Dict[tuple, torch.Tensor] = {} + self._rope_txt: Dict[tuple, torch.Tensor] = {} + self._rope_pix: Dict[tuple, torch.Tensor] = {} + + # -- toolkit hooks --------------------------------------------------------- + @property + def in_channels(self) -> int: + return self.config.in_channels + + @property + def device(self): + return self.s_embedder.proj.weight.device + + @property + def dtype(self): + return self.s_embedder.proj.weight.dtype + + def enable_gradient_checkpointing(self): + self.gradient_checkpointing = True + + def disable_gradient_checkpointing(self): + self.gradient_checkpointing = False + + @classmethod + def get_transformer_block_names(cls) -> Optional[List[str]]: + # streamed quantization over the trunk; the text adapter and pixel + # stage go through the whole-module "extras" pass (full names, so the + # exclude patterns below apply to them) + return ["blocks"] + + @classmethod + def get_quantization_exclude_modules(cls) -> Optional[List[str]]: + # embedders / heads / modulation cores feed every block or are tiny; + # the pixel stage is 16-wide per pixel and sets the final RGB values + return [ + "s_embedder*", + "t_embedder*", + "final_layer*", + "modulation_cores*", + "pixel_embedder*", + "pixel_blocks*", + "y_embedder.layer_pool*", + "adaln*", + "*.adaln*", + ] + + def get_offload_ignore_modules(self): + # tiny live state every block reads: keep resident + return [self.t_embedder, self.modulation_cores, self.s_embedder, self.final_layer] + + # -- rope caches ----------------------------------------------------------- + def _fetch_rope_img(self, grid, device): + key = (grid, str(device)) + if key not in self._rope_img: + cfg = self.config + head_dim = cfg.hidden_size // cfg.num_heads + self._rope_img[key] = rope_2d( + head_dim, grid[0], grid[1], theta=cfg.rope_theta, scale=cfg.rope_scale, + aspect=cfg.rope_aspect, frame_pairs=cfg.rope_frame_pairs, frame_theta=cfg.rope_frame_theta, + ).to(device) + return self._rope_img[key] + + def _fetch_rope_txt(self, length, device): + if not self.config.text_rope: + return None + key = (length, str(device)) + if key not in self._rope_txt: + head_dim = self.config.hidden_size // self.config.num_heads + self._rope_txt[key] = rope_1d(head_dim, length, theta=self.config.text_rope_theta).to(device) + return self._rope_txt[key] + + def _fetch_rope_pix(self, grid, device): + key = (grid, str(device)) + if key not in self._rope_pix: + cfg = self.config + head_dim = cfg.pixel.attn_hidden_size // cfg.pixel.num_heads + self._rope_pix[key] = rope_2d( + head_dim, grid[0], grid[1], theta=cfg.rope_theta, scale=cfg.rope_scale, aspect=cfg.rope_aspect, + ).to(device) + return self._rope_pix[key] + + def _run_block(self, block, *args): + if self.gradient_checkpointing and torch.is_grad_enabled(): + return checkpoint(block, *args, use_reentrant=False) + return block(*args) + + # -- forward --------------------------------------------------------------- + def forward( + self, + x: torch.Tensor, + t: torch.Tensor, + y: torch.Tensor, + y_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + x: noisy image [B, C, H, W], H and W divisible by patch_size. + t: model time in [0, 1000] (shifted flow level x 1000), [B]. + y: text states [B, T, text_dim]; for ``lap_blocks2`` the selected layer + stack [B, T, L, text_dim], or its flattened form [B, T, L*text_dim]. + y_mask: encoder attention mask [B, T] (1 = real token); required by the + masked adapters. + """ + cfg = self.config + p = cfg.patch_size + batch, _, height, width = x.shape + if height % p or width % p: + raise ValueError(f"input {height}x{width} is not divisible by patch_size {p}") + grid = (height // p, width // p) + n_patches = grid[0] * grid[1] + + if cfg.text_adapter == "lap_blocks2" and y.ndim == 3: + y = y.reshape(y.shape[0], y.shape[1], cfg.text_lap_num_layers, cfg.text_dim) + + patches = F.unfold(x, kernel_size=p, stride=p).transpose(1, 2) # [B, L, p*p*C] + s = self.s_embedder(patches) + t_emb = self.t_embedder(t) # [B, 1, D] + cond = F.silu(t_emb) + + y = y[:, : cfg.text_len] + if self._adapter_needs_mask: + if y_mask is None: + raise ValueError( + f"text_adapter='{cfg.text_adapter}' masks pad text positions and requires y_mask" + ) + y = self.y_embedder(y, y_mask[:, : cfg.text_len]) + else: + y = self.y_embedder(y) + if self.y_pos_embedding is not None: + y = y + self.y_pos_embedding[:, : y.shape[1]].to(y.dtype) + + rope_img = self._fetch_rope_img(grid, x.device) + rope_txt = self._fetch_rope_txt(y.shape[1], x.device) + + shared = {} + if self._shared_modulation: + shared = {key: core(cond) for key, core in self.modulation_cores.items()} + + def mod_for(stream): + if not self._shared_modulation: + return None + return shared[f"adaln_{self._stream_alias.get(stream, stream)}"] + + for block, is_dual in zip(self.blocks, self._block_is_dual): + if is_dual: + s, y = self._run_block(block, s, y, cond, rope_img, rope_txt, mod_for("img"), mod_for("txt")) + else: + s, y = self._run_block(block, s, y, cond, rope_img, rope_txt, mod_for("shared")) + + s = F.silu(t_emb + s) # timestep re-fused into every patch token + + if self.pixel_blocks is None: + out = self.final_layer(s) # [B, L, p*p*C] + folded = out.transpose(1, 2) + else: + s_cond = s.reshape(batch * n_patches, -1) + pixels = self.pixel_embedder(x) # [B*L, p*p, d_pix] + rope_pix = self._fetch_rope_pix(grid, x.device) + for block in self.pixel_blocks: + pixels = self._run_block(block, pixels, s_cond, rope_pix, grid) + out = self.final_layer(pixels) # [B*L, p*p, C] + folded = out.reshape(batch, n_patches, p * p, -1).permute(0, 3, 2, 1) + folded = folded.reshape(batch, -1, n_patches) + + return F.fold(folded, output_size=(height, width), kernel_size=p, stride=p) diff --git a/extensions_built_in/diffusion_models/iris3b/src/pipeline.py b/extensions_built_in/diffusion_models/iris3b/src/pipeline.py new file mode 100644 index 0000000..09dd7fe --- /dev/null +++ b/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] diff --git a/extensions_built_in/diffusion_models/iris3b/src/text_encoder.py b/extensions_built_in/diffusion_models/iris3b/src/text_encoder.py new file mode 100644 index 0000000..9f3473d --- /dev/null +++ b/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() diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index 9cee405..fb33d5a 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/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)", diff --git a/extensions_built_in/inference_engine/latent_preview.py b/extensions_built_in/inference_engine/latent_preview.py index a29f1ed..c184d6a 100644 --- a/extensions_built_in/inference_engine/latent_preview.py +++ b/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 } diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index d05af88..011ccc7 100644 --- a/testing/test_model_loading.py +++ b/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}, diff --git a/toolkit/models/registry.py b/toolkit/models/registry.py index a2357c2..3d94fc5 100644 --- a/toolkit/models/registry.py +++ b/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}}, diff --git a/toolkit/models/v2/text_encoders/qwen3_vl.py b/toolkit/models/v2/text_encoders/qwen3_vl.py index 2657bd1..5856ff9 100644 --- a/toolkit/models/v2/text_encoders/qwen3_vl.py +++ b/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.""" From 2fde764876d096f7bee8b5f5bd137b43aeee773a Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 9 Oct 2026 05:14:08 -0600 Subject: [PATCH 2/4] Version bump --- version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/version.py b/version.py index 8fd1283..dbc3e72 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.13.23" +VERSION = "0.13.24" From f370e7441d10b14818d88b9269c30a650f705fd3 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 9 Oct 2026 21:09:09 -0600 Subject: [PATCH 3/4] Add support for Qwen Image 2 Turbo with training adapter --- README.md | 2 + .../qwen_image_2/qwen_image_2.py | 125 ++++++++++++++- .../qwen_image_2/src/pipeline.py | 41 +++-- extensions_built_in/diffusion_models/ui.tsx | 143 ++++++++++++------ testing/test_model_loading.py | 12 ++ 5 files changed, 262 insertions(+), 61 deletions(-) diff --git a/README.md b/README.md index 4809e40..14823bc 100644 --- a/README.md +++ b/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) @@ -45,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) diff --git a/extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py b/extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py index 32e0770..d9cb84b 100644 --- a/extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py +++ b/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. diff --git a/extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py b/extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py index 710a554..0bd1d51 100644 --- a/extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py +++ b/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 diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index fb33d5a..f270a5d 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/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: ( +
+

+ This model'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. +

+

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

+

+ Changing this re-caches latents, and it cannot be combined with + a dataset using alpha_mask, which consumes the + alpha channel as a loss mask instead. +

+
+ ), + }, + }, +]; + 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: ( -
-

- This model'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. -

-

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

-

- Changing this re-caches latents, and it cannot be combined with - a dataset using alpha_mask, which consumes the - alpha channel as a loss mask instead. -

-
- ), - }, - }, + {}, + ], + }, + 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", diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index 011ccc7..4d6e585 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -189,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", From f7a1fb9a77e5418269b1e9eda4ace6d7c1b7bb47 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 9 Oct 2026 21:09:34 -0600 Subject: [PATCH 4/4] Version bump --- version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/version.py b/version.py index dbc3e72..3950142 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.13.24" +VERSION = "0.13.25"