Add support for Qwen Image 2.1

extensions_built_in/diffusion_models/__init__.py

@@ -5,6 +5,7 @@ from .omnigen2 import OmniGen2Model
 from .flux_kontext import FluxKontextModel
 from .wan22 import Wan225bModel, Wan2214bModel, Wan2214bI2VModel
 from .qwen_image import QwenImageModel, QwenImageEditModel, QwenImageEditPlusModel
+from .qwen_image_2 import QwenImage2Model
 from .flux2 import Flux2Model, Flux2Klein4BModel, Flux2Klein9BModel
 from .z_image import ZImageModel
 from .ltx2 import LTX2Model, LTX23Model, LTX25Model
@@ -36,6 +37,7 @@ AI_TOOLKIT_MODELS = [
     QwenImageModel,
     QwenImageEditModel,
     QwenImageEditPlusModel,
+    QwenImage2Model,
     Flux2Model,
     ZImageModel,
     LTX2Model,

extensions_built_in/diffusion_models/qwen_image_2/__init__.py

@@ -0,0 +1,3 @@
+from .qwen_image_2 import QwenImage2Model
+
+__all__ = ["QwenImage2Model"]

extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py

@@ -0,0 +1,570 @@
+"""Qwen-Image 2.1 for ai-toolkit.
+
+One checkpoint does text-to-image and editing, so there is one arch: give it
+reference images and it edits, give it none and it is plain T2I.
+
+The pieces, all pulled from the Comfy-Org repack into the toolkit's ComfyUI
+models folder:
+
+  - transformer: 32-layer single-stream block-causal DiT (`src/transformer.py`),
+  - text encoder: Qwen3-VL-8B (`src/text_encoder.py`), read at the last decoder
+    layer BEFORE its final RMSNorm,
+  - VAE: RGBA, 16x spatial, 64 latent channels (`src/vae.py`).
+
+Text and images share one sequence: the Qwen3-VL encoder reserves a
+`<|image_pad|>` slot per reference image and the DiT drops that image's VAE
+latents into the slot, four latent tokens per slot. Reference tokens and text
+are modulated from t=0 (`causal_condition`), so only the target image's tokens
+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.
+"""
+
+import os
+from typing import TYPE_CHECKING, List, Optional
+
+import numpy as np
+import torch
+from PIL import Image
+
+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.metadata import get_meta_for_safetensors
+from toolkit.models.base_model import BaseModel
+from toolkit.samplers.custom_flowmatch_sampler import (
+    CustomFlowMatchEulerDiscreteScheduler,
+)
+
+from .src.pipeline import (
+    QwenImage21Pipeline,
+    QwenImage21PromptEncoder,
+    VAE_SCALE_FACTOR,
+    VISION_TOKEN_PIXELS,
+    pack_latents,
+    pad_prompt_batch,
+    prepare_condition_image,
+    run_transformer,
+    tensor_to_pil,
+)
+from .src.text_encoder import QwenImage21TextEncoder
+from .src.transformer import QwenImage21Transformer2DModel
+from .src.vae import AutoencoderKLQwenImage21
+
+if TYPE_CHECKING:
+    from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
+
+
+# matches scheduler/scheduler_config.json in the base repo
+scheduler_config = {
+    "base_image_seq_len": 256,
+    "base_shift": 0.5,
+    "invert_sigmas": False,
+    "max_image_seq_len": 8192,
+    "max_shift": 0.9,
+    "num_train_timesteps": 1000,
+    "shift": 1.0,
+    "shift_terminal": 0.02,
+    "stochastic_sampling": False,
+    "time_shift_type": "exponential",
+    "use_beta_sigmas": False,
+    "use_dynamic_shifting": True,
+    "use_exponential_sigmas": False,
+    "use_karras_sigmas": False,
+}
+
+# the Comfy-Org repack is the weight source; the original repo supplies the
+# configs and the processor, which the repack does not carry
+COMFY_REPO = "Comfy-Org/Qwen-Image-2.1"
+BASE_REPO = "Qwen/Qwen-Image-2.1"
+
+# decode above this many output pixels goes through the VAE's tiled path
+TILE_DECODE_ABOVE_PIXELS = 1024 * 1024
+
+
+def _drop_repeats(images):
+    """Drop a reference that repeats the one before it.
+
+    `GenerateImageConfig.ctrl_img_1` defaults to `ctrl_img`, so the sampler
+    hands the same reference over twice when only one was configured. Here that
+    is not merely wasteful: the prompt would reserve slots for two images while
+    only one set of latents arrives.
+    """
+    kept = []
+    for image in images:
+        previous = kept[-1] if kept else None
+        if (
+            previous is not None
+            and previous.shape == image.shape
+            and torch.equal(previous, image)
+        ):
+            continue
+        kept.append(image)
+    return kept
+
+
+class QwenImage2Model(BaseModel):
+    arch = "qwen_image_2"
+
+    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
+        self.target_lora_modules = ["QwenImage21Transformer2DModel"]
+        self.vae_scale_factor = VAE_SCALE_FACTOR
+        self.prompt_encoder: Optional[QwenImage21PromptEncoder] = None
+
+        # Editing is not a separate model here, it is what happens when the
+        # dataset has a control path: reference images ride into the text
+        # embeddings as vision tokens and their latents into the sequence.
+        self.encode_control_in_text_embeddings = True
+        self.has_multiple_control_images = True
+        # References are bucket-resized to the target size by the dataloader.
+        # The DiT reads the image-slot layout from row 0 of the batch, so every
+        # sample must contribute the same slot count -- raw (per-image aspect)
+        # references would break that.
+        self.use_raw_control_images = False
+
+    @staticmethod
+    def get_train_scheduler():
+        return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
+
+    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
+        return VISION_TOKEN_PIXELS
+
+    # ------------------------------------------------------------------
+    # Loading
+    # ------------------------------------------------------------------
+    def load_model(self):
+        dtype = self.torch_dtype
+        self.print_and_status_update("Loading Qwen-Image 2.1 model")
+        model_path = self.model_config.name_or_path
+        base_model_path = self.model_config.extras_name_or_path
+
+        if base_model_path == model_path and not os.path.isdir(base_model_path):
+            # extras default to name_or_path, which is the comfy repack (or a
+            # single file); neither carries the configs or the processor
+            base_model_path = BASE_REPO
+        elif os.path.isdir(model_path) and os.path.isdir(
+            os.path.join(model_path, "text_encoder")
+        ):
+            # a local full checkpoint supplies its own text encoder / vae
+            base_model_path = model_path
+
+        self.print_and_status_update("Loading transformer")
+        transformer = QwenImage21Transformer2DModel.load(
+            model_path,
+            config_path=base_model_path,
+            **self.component_load_kwargs("transformer"),
+        )
+        flush()
+
+        self.print_and_status_update("Loading text encoder")
+        processor = QwenImage21TextEncoder.load_processor(base_model_path)
+        text_encoder = QwenImage21TextEncoder.load_model(
+            base_model_path, dtype=dtype, subfolder="text_encoder"
+        )
+        # the vision tower stays: any prompt may carry reference images. bf16
+        # Conv3d has no fast kernel, the equivalent GEMM does
+        text_encoder.patch_vision_patch_embed()
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
+        text_encoder.requires_grad_(False)
+        text_encoder.eval()
+        flush()
+
+        self.print_and_status_update("Loading VAE")
+        vae = AutoencoderKLQwenImage21.load(
+            base_model_path, **self.component_load_kwargs("vae")
+        )
+        vae.requires_grad_(False)
+        vae.eval()
+
+        self.noise_scheduler = QwenImage2Model.get_train_scheduler()
+
+        self.vae = vae
+        self.text_encoder = [text_encoder]
+        self.tokenizer = [processor.tokenizer]
+        self.processor = processor
+        self.model = transformer
+        self.prompt_encoder = QwenImage21PromptEncoder(text_encoder, processor)
+        self.pipeline = QwenImage21Pipeline(self)
+        self.print_and_status_update("Model Loaded")
+
+    # ------------------------------------------------------------------
+    # VAE. The latents are RGBA; toolkit images are RGB, so encode pads an
+    # opaque alpha channel and decode drops it again (unless model_kwargs.rgba).
+    # ------------------------------------------------------------------
+    @property
+    def output_rgba(self) -> bool:
+        return bool(self.model_config.model_kwargs.get("rgba", False))
+
+    def _latent_stats(self, device, dtype):
+        shape = (1, self.vae.config.z_dim, 1, 1, 1)
+        mean = torch.tensor(self.vae.config.latents_mean).view(shape).to(device, dtype)
+        std = torch.tensor(self.vae.config.latents_std).view(shape).to(device, dtype)
+        return mean, std
+
+    def encode_images(self, image_list, device=None, dtype=None):
+        device = device or self.vae_device_torch
+        dtype = dtype or self.vae_torch_dtype
+
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(device)
+        self.vae.eval()
+
+        images = torch.stack(
+            [image.to(device, dtype=dtype) for image in image_list]
+        ).to(device, dtype=dtype)
+        if images.shape[1] == 3:
+            # opaque alpha, in the [-1, 1] the VAE reads
+            images = torch.cat([images, torch.ones_like(images[:, :1])], dim=1)
+        images = images.unsqueeze(2)  # single-frame dim
+
+        latents = self.vae.encode(images).latent_dist.sample()
+        mean, std = self._latent_stats(latents.device, latents.dtype)
+        latents = (latents - mean) / std
+        return latents.squeeze(2).to(device, dtype=dtype)
+
+    def decode_latents(self, latents: torch.Tensor, device=None, dtype=None):
+        images = self._decode_rgba(latents, device=device, dtype=dtype)
+        if not self.output_rgba:
+            images = images[:, :3]
+        return images
+
+    def _decode_rgba(self, latents: torch.Tensor, device=None, dtype=None):
+        device = device or self.vae_device_torch
+        dtype = dtype or self.vae_torch_dtype
+
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(device)
+
+        latents = latents.to(device, dtype=dtype).unsqueeze(2)
+        mean, std = self._latent_stats(latents.device, latents.dtype)
+        latents = latents * std + mean
+
+        # A one-shot decode of this VAE's 16x upsample stack is heavy: 2048x2048
+        # -- the resolution Qwen recommends -- needs more than a 32 GB card has.
+        # Tile above 1 MP, and whenever low_vram is set.
+        pixels = latents.shape[-2] * latents.shape[-1] * self.vae_scale_factor**2
+        tiled = self.model_config.low_vram or pixels > TILE_DECODE_ABOVE_PIXELS
+        if tiled:
+            self.vae.enable_tiling(
+                tile_sample_min_height=1024,
+                tile_sample_min_width=1024,
+                tile_sample_stride_height=768,
+                tile_sample_stride_width=768,
+            )
+        try:
+            images = self.vae.decode(latents).sample
+        finally:
+            if tiled:
+                self.vae.disable_tiling()
+        return images.squeeze(2).to(device, dtype=dtype)
+
+    def decode_to_images(self, latents: torch.Tensor) -> List[Image.Image]:
+        """Decode to PIL, keeping the alpha channel when model_kwargs.rgba is set."""
+        return [
+            self.image_tensor_to_pil(image) for image in self.decode_latents(latents)
+        ]
+
+    @staticmethod
+    def image_tensor_to_pil(image: torch.Tensor) -> Image.Image:
+        """`(C, H, W)` in [-1, 1] -> PIL (RGB, or RGBA when C is 4)."""
+        array = (image.float().clamp(-1, 1) / 2 + 0.5).permute(1, 2, 0)
+        array = (array.cpu().numpy() * 255).round().astype(np.uint8)
+        return Image.fromarray(array, mode="RGBA" if array.shape[2] == 4 else "RGB")
+
+    # ------------------------------------------------------------------
+    # Reference images
+    # ------------------------------------------------------------------
+    @property
+    def control_image_max_pixels(self) -> int:
+        """Pixel budget a reference image is shrunk to fit. A reference that is
+        already smaller keeps its size (only snapped to the 32 px grid), so a
+        bucket-resized training reference stays at the target resolution."""
+        return int(
+            self.model_config.model_kwargs.get("control_image_max_pixels", 1024 * 1024)
+        )
+
+    def _normalize_control_images(self, control_images, batch_size: int) -> List[List]:
+        """Any of the shapes the toolkit hands over -> one list per batch item.
+
+        Control images arrive as a `(B, C, H, W)` batch tensor, a
+        `(B, N, C, H, W)` multi-reference tensor, a per-sample list of lists, or
+        a flat list for a single prompt (sampling / blank-embed caching).
+        """
+        if control_images is None:
+            return [[] for _ in range(batch_size)]
+        if isinstance(control_images, torch.Tensor):
+            if control_images.dim() == 5:
+                control_images = [list(sample) for sample in control_images]
+            else:
+                control_images = [[sample] for sample in control_images]
+        elif len(control_images) > 0 and not isinstance(control_images[0], list):
+            control_images = [list(control_images)]
+        control_images = [_drop_repeats(sample) for sample in control_images]
+        if len(control_images) == 1 and batch_size > 1:
+            control_images = control_images * batch_size
+        if len(control_images) != batch_size:
+            raise ValueError(
+                f"got {len(control_images)} control image sets for {batch_size} prompts"
+            )
+        return control_images
+
+    def _prepare_control_images(
+        self, control_images: List[List[torch.Tensor]]
+    ) -> List[List[torch.Tensor]]:
+        """Put every reference on the 32 px grid, as `(1, C, H, W)` in [0, 1]."""
+        budget = self.control_image_max_pixels
+        prepared = []
+        for sample in control_images:
+            images = []
+            for image in sample:
+                if image.dim() == 3:
+                    image = image.unsqueeze(0)
+                images.append(
+                    prepare_condition_image(image.to(self.device_torch), budget)
+                )
+            prepared.append(images)
+        return prepared
+
+    def encode_condition_images(self, control_images):
+        """Reference images -> `(B, N, C)` packed latents plus their latent grids.
+
+        `control_images` is a per-sample list of `(1, C, H, W)` tensors in
+        [0, 1], already at their final size.
+        """
+        if not control_images or not any(len(sample) for sample in control_images):
+            return None, []
+
+        sample_latents, shapes = [], []
+        for index, sample in enumerate(control_images):
+            packed = []
+            for image in sample:
+                latent = self.encode_images(
+                    [image[0].to(self.device_torch) * 2 - 1],
+                    device=self.device_torch,
+                    dtype=self.torch_dtype,
+                )
+                if index == 0:
+                    shapes.append((latent.shape[2], latent.shape[3]))
+                packed.append(pack_latents(latent))
+            sample_latents.append(torch.cat(packed, dim=1))
+
+        lengths = {latents.shape[1] for latents in sample_latents}
+        if len(lengths) > 1:
+            raise ValueError(
+                "every sample in a batch must contribute the same reference image "
+                f"token count, got {sorted(lengths)}. The dataloader bucket-resizes "
+                "reference images, so this means the samples disagree on how many."
+            )
+        return torch.cat(sample_latents, dim=0), shapes
+
+    # ------------------------------------------------------------------
+    # Prompts
+    # ------------------------------------------------------------------
+    def get_prompt_embeds(self, prompt, control_images=None) -> AdvancedPromptEmbeds:
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        if self.text_encoder[0].device != self.device_torch:
+            self.text_encoder[0].to(self.device_torch)
+
+        images = None
+        if control_images is not None:
+            samples = self._prepare_control_images(
+                self._normalize_control_images(control_images, len(prompt))
+            )
+            images = [[tensor_to_pil(image) for image in sample] for sample in samples]
+
+        embeds, masks, slot_masks = self.prompt_encoder.encode(
+            prompt, images=images, device=self.device_torch
+        )
+        pe = AdvancedPromptEmbeds(
+            text_embeds=[embed.to(self.torch_dtype) for embed in embeds],
+            attention_mask=masks,
+            image_slot_mask=[slots.to(torch.bool) for slots in slot_masks],
+        )
+        # the masks are bookkeeping, not activations: a .to(dtype) must not
+        # turn them into bf16
+        pe.frozen_dtype_keys = ["attention_mask", "image_slot_mask"]
+        return pe
+
+    def pad_prompt_embeds(self, prompt_embeds: AdvancedPromptEmbeds):
+        """`AdvancedPromptEmbeds` -> the padded `(embeds, mask, slot_mask)` batch."""
+        return pad_prompt_batch(
+            prompt_embeds.text_embeds,
+            prompt_embeds.attention_mask,
+            prompt_embeds.image_slot_mask,
+            self.device_torch,
+            self.torch_dtype,
+        )
+
+    # ------------------------------------------------------------------
+    # Training
+    # ------------------------------------------------------------------
+    def condition_noisy_latents(
+        self, latents: torch.Tensor, batch: "DataLoaderBatchDTO"
+    ):
+        # reference latents join the sequence in get_noise_prediction, clean
+        return latents.detach()
+
+    def get_noise_prediction(
+        self,
+        latent_model_input: torch.Tensor,  # (B, 64, h, w)
+        timestep: torch.Tensor,  # 0..1000 scale
+        text_embeddings: AdvancedPromptEmbeds,
+        batch: "DataLoaderBatchDTO" = None,
+        **kwargs,
+    ):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+        batch_size = latent_model_input.shape[0]
+
+        prompt_embeds, prompt_mask, slot_mask = self.pad_prompt_embeds(text_embeddings)
+
+        # The prompt is what decides: it reserved the slots the references go
+        # into, so a prompt encoded without them (a plain T2I dataset, a fully
+        # dropped caption) takes no references here either.
+        condition_latents, condition_shapes = None, []
+        if batch is not None and bool(slot_mask.any()):
+            with torch.no_grad():
+                control = batch.control_tensor_list
+                if control is None:
+                    control = batch.control_tensor
+                samples = self._prepare_control_images(
+                    self._normalize_control_images(control, batch_size)
+                )
+                condition_latents, condition_shapes = self.encode_condition_images(
+                    samples
+                )
+
+        # toolkit timestep (0..1000, 1000 = pure noise) -> the model's t in [0, 1];
+        # same direction, so a plain divide
+        t = timestep.to(self.device_torch, dtype=self.torch_dtype) / 1000
+        if t.dim() == 0:
+            t = t.unsqueeze(0)
+        if t.shape[0] != batch_size:
+            t = t.expand(batch_size)
+
+        return run_transformer(
+            self.transformer,
+            latent_model_input.to(self.device_torch, self.torch_dtype),
+            t,
+            prompt_embeds,
+            prompt_mask,
+            slot_mask,
+            condition_latents=condition_latents,
+            condition_shapes=condition_shapes,
+            **kwargs,
+        )
+
+    def get_loss_target(self, *args, **kwargs):
+        # flow-matching velocity target: noise - clean
+        noise = kwargs.get("noise")
+        batch = kwargs.get("batch")
+        return (noise - batch.latents).detach()
+
+    def get_model_has_grad(self):
+        return False
+
+    def get_te_has_grad(self):
+        return False
+
+    # ------------------------------------------------------------------
+    # Sampling
+    # ------------------------------------------------------------------
+    def get_generation_pipeline(self):
+        return QwenImage21Pipeline(self)
+
+    def generate_single_image(
+        self,
+        pipeline: QwenImage21Pipeline,
+        gen_config: GenerateImageConfig,
+        conditional_embeds: AdvancedPromptEmbeds,
+        unconditional_embeds: AdvancedPromptEmbeds,
+        generator: torch.Generator,
+        extra: dict,
+    ):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+
+        divisor = self.get_bucket_divisibility()
+        gen_config.width = int(gen_config.width // divisor * divisor)
+        gen_config.height = int(gen_config.height // divisor * divisor)
+
+        # the same list the sampler built for the prompt embeddings, so the two
+        # agree on how many references there are
+        paths = [
+            path
+            for path in (
+                gen_config.ctrl_img,
+                gen_config.ctrl_img_1,
+                gen_config.ctrl_img_2,
+                gen_config.ctrl_img_3,
+            )
+            if path is not None
+        ]
+        condition_images = None
+        if paths:
+            tensors = [
+                torch.from_numpy(
+                    np.array(Image.open(path).convert("RGBA"), dtype=np.float32) / 255.0
+                )
+                .permute(2, 0, 1)
+                .unsqueeze(0)
+                for path in paths
+            ]
+            condition_images = self._prepare_control_images(
+                self._normalize_control_images([tensors], 1)
+            )
+
+        return 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,
+            condition_images=condition_images,
+        )[0]
+
+    # ------------------------------------------------------------------
+    # Saving / bookkeeping
+    # ------------------------------------------------------------------
+    def save_model(self, output_path, meta, save_dtype):
+        # comfy-format single-file save; prequantized layers keep their storage
+        transformer: QwenImage21Transformer2DModel = unwrap_model(self.model)
+        if not output_path.endswith(".safetensors"):
+            output_path += ".safetensors"
+        transformer.save_model(
+            output_path,
+            dtype=save_dtype,
+            metadata=get_meta_for_safetensors(meta, name=self.arch),
+        )
+
+    def get_base_model_version(self):
+        return self.arch
+
+    def get_transformer_block_names(self) -> Optional[List[str]]:
+        return ["transformer_blocks"]
+
+    def get_quantization_exclude_modules(self):
+        return QwenImage21Transformer2DModel.get_quantization_exclude_modules()
+
+    lora_keys_use_comfy_prefix = True

extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py

@@ -0,0 +1,412 @@
+"""Prompt encoding, sequence packing and a preview sampler for Qwen-Image 2.1.
+
+Qwen-Image 2.1 puts text and images in ONE sequence. The Qwen3-VL encoder is
+handed the prompt with a `<|image_pad|>` slot per condition image; the
+transformer then throws the vision embeddings away and drops the VAE latents
+into those slots, four latent tokens per slot (a vision token covers 32 px, a
+latent token 16 px). Everything here exists to keep those two views lined up:
+`encode_prompt` returns the per-slot mask alongside the embeddings, and
+`run_transformer` appends one slot per 2x2 group of target latents before the
+call.
+
+Condition images are therefore resized ONCE, to a multiple of 32, and that same
+resize feeds both the Qwen3-VL processor and the VAE -- a different size on
+either side changes the slot/latent-token counts and the transformer rejects
+the sequence.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import List, Optional, Sequence
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from diffusers.utils.torch_utils import randn_tensor
+
+SYSTEM_PROMPT = "Comprehend and analyze the provided prompt."
+VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>"
+PROMPT_TEMPLATE = (
+    f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
+    "<|im_start|>user\n{}<|im_end|>\n"
+    "<|im_start|>assistant\n"
+)
+
+# a Qwen3-VL vision token covers 32 px (patch 16 x merge 2), a latent token 16 px
+VISION_TOKEN_PIXELS = 32
+VAE_SCALE_FACTOR = 16
+LATENT_TOKENS_PER_SLOT = (VISION_TOKEN_PIXELS // VAE_SCALE_FACTOR) ** 2
+
+
+def _snap(size: float) -> int:
+    """Round to the 32 px grid the vision tower and the VAE share."""
+    return max(
+        VISION_TOKEN_PIXELS, round(size / VISION_TOKEN_PIXELS) * VISION_TOKEN_PIXELS
+    )
+
+
+def calculate_dimensions(target_area: int, ratio: float) -> tuple[int, int]:
+    """Width/height with the given aspect ratio covering `target_area`, both on
+    the 32 px grid."""
+    width = math.sqrt(target_area * ratio)
+    return _snap(width), _snap(width / ratio)
+
+
+def prepare_condition_image(image: torch.Tensor, max_pixels: int) -> torch.Tensor:
+    """Put one `(1, C, H, W)` condition image on the 32 px grid, shrinking it
+    first if it is over the pixel budget.
+
+    Deterministic in the image alone, on purpose: the Qwen3-VL pass and the VAE
+    pass happen in different calls (`get_prompt_embeds` and
+    `get_noise_prediction`) and must agree on the size to the pixel, or the slot
+    count and the latent-token count disagree and the sequence is rejected.
+    """
+    height, width = image.shape[2], image.shape[3]
+    if height * width > max_pixels:
+        new_width, new_height = calculate_dimensions(max_pixels, width / height)
+    else:
+        new_width, new_height = _snap(width), _snap(height)
+    if (new_height, new_width) == (height, width):
+        return image
+    return F.interpolate(
+        image.float(), size=(new_height, new_width), mode="bicubic", antialias=True
+    ).clamp(0, 1)
+
+
+def tensor_to_pil(image: torch.Tensor) -> Image.Image:
+    """`(1, C, H, W)` or `(C, H, W)` in [0, 1] -> RGB PIL for the vision tower.
+
+    An alpha channel is composited over white, which is what the checkpoint was
+    trained with; the VAE still reads all four channels.
+    """
+    if image.dim() == 4:
+        image = image[0]
+    image = image.detach().float().clamp(0, 1).cpu()
+    if image.shape[0] == 4:
+        alpha = image[3:4]
+        image = image[:3] * alpha + (1.0 - alpha)
+    array = (image.permute(1, 2, 0).numpy() * 255).round().astype(np.uint8)
+    return Image.fromarray(array, mode="RGB")
+
+
+class QwenImage21PromptEncoder:
+    """Turns prompts (+ optional condition images) into the transformer's text
+    stream, its validity mask, and the image-slot mask."""
+
+    def __init__(self, text_encoder, processor):
+        self.text_encoder = text_encoder
+        self.processor = processor
+        # How many leading tokens the system turn takes, so it can be dropped
+        # from the hidden states. Derived from the processor, not hardcoded.
+        system_message = [
+            {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}
+        ]
+        system_tokens = processor.apply_chat_template(
+            system_message, tokenize=True, return_dict=False
+        )
+        self.drop_index = len(system_tokens[0])
+        self.image_token_id = processor.tokenizer.encode("<|image_pad|>")[0]
+
+    def _template(self, num_images: int) -> str:
+        if num_images == 0:
+            return PROMPT_TEMPLATE
+        refs = " ".join(f"<image{i + 1}>{VISION_BLOCK}" for i in range(num_images))
+        return PROMPT_TEMPLATE.replace("{}", refs + "{}", 1)
+
+    @torch.no_grad()
+    def encode(
+        self,
+        prompts: Sequence[str],
+        images: Optional[Sequence[Sequence[Image.Image]]] = None,
+        device: Optional[torch.device] = None,
+    ) -> tuple[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor]]:
+        """Encode a batch of prompts.
+
+        `images[i]` are the condition images for `prompts[i]`, already resized.
+        Returns three per-sample lists, each entry at its own natural length:
+        the `(L, dim)` embeddings, an `(L,)` validity mask, and an `(L,)` bool
+        mask marking the image slots.
+        """
+        device = device or self.text_encoder.device
+        # Qwen has no bos token, so an empty prompt leaves the encoder nothing to read
+        prompts = [p if p else " " for p in prompts]
+        images = list(images) if images is not None else [[] for _ in prompts]
+
+        texts = [
+            self._template(len(sample_images)).format(prompt)
+            for prompt, sample_images in zip(prompts, images)
+        ]
+        flat_images = [image for sample_images in images for image in sample_images]
+
+        processor_kwargs = {
+            # left padding, as the checkpoint was trained with: the padding is
+            # dropped either way, but the side decides the positions the
+            # encoder sees across prompts of different lengths
+            "text": texts,
+            "padding": True,
+            "padding_side": "left",
+            "return_tensors": "pt",
+        }
+        if flat_images:
+            processor_kwargs["images"] = flat_images
+        model_inputs = self.processor(**processor_kwargs).to(device)
+
+        forward_kwargs = {
+            "input_ids": model_inputs.input_ids,
+            "attention_mask": model_inputs.attention_mask,
+            "output_hidden_states": True,
+        }
+        if flat_images and hasattr(model_inputs, "pixel_values"):
+            forward_kwargs["pixel_values"] = model_inputs.pixel_values
+            forward_kwargs["image_grid_thw"] = model_inputs.image_grid_thw
+        if hasattr(model_inputs, "mm_token_type_ids"):
+            forward_kwargs["mm_token_type_ids"] = model_inputs.mm_token_type_ids
+
+        # The transformer was trained on the last decoder layer's output BEFORE
+        # the encoder's final RMSNorm. transformers 4.x returned that as
+        # hidden_states[-1]; 5.x ties that entry to last_hidden_state, so it
+        # comes back normalized instead. A forward hook returning the norm's
+        # input neutralizes it for this call on either version.
+        language_model = getattr(
+            self.text_encoder.model, "language_model", self.text_encoder.model
+        )
+        handle = language_model.norm.register_forward_hook(
+            lambda module, args, output: args[0]
+        )
+        try:
+            outputs = self.text_encoder(**forward_kwargs)
+        finally:
+            handle.remove()
+        hidden_states = outputs.hidden_states[-1]
+
+        embeds, masks, slot_masks = [], [], []
+        for sample_ids, sample_mask, sample_hidden in zip(
+            model_inputs.input_ids, model_inputs.attention_mask, hidden_states
+        ):
+            valid = sample_mask.bool()
+            sample_hidden = sample_hidden[valid][self.drop_index :]
+            slots = (sample_ids[valid] == self.image_token_id)[self.drop_index :]
+            embeds.append(sample_hidden)
+            masks.append(
+                torch.ones(sample_hidden.shape[0], dtype=torch.long, device=device)
+            )
+            slot_masks.append(slots)
+        return embeds, masks, slot_masks
+
+
+def pad_prompt_batch(
+    embeds: Sequence[torch.Tensor],
+    masks: Sequence[torch.Tensor],
+    slot_masks: Sequence[torch.Tensor],
+    device: torch.device,
+    dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Right-pad per-sample prompt tensors into one batch.
+
+    The transformer reads the image-slot layout from row 0 and applies it to
+    every row, which is sound because the template puts the slots before the
+    prompt text: the prefix is identical across samples and only the padded
+    tail differs.
+    """
+    max_len = max(e.shape[0] for e in embeds)
+    batch_size, dim = len(embeds), embeds[0].shape[-1]
+
+    out_embeds = torch.zeros(batch_size, max_len, dim, device=device, dtype=dtype)
+    out_mask = torch.zeros(batch_size, max_len, dtype=torch.bool, device=device)
+    out_slots = torch.zeros(batch_size, max_len, dtype=torch.bool, device=device)
+    for i, (embed, mask, slots) in enumerate(zip(embeds, masks, slot_masks)):
+        length = embed.shape[0]
+        out_embeds[i, :length] = embed.to(device, dtype)
+        out_mask[i, :length] = mask.to(device).bool()
+        out_slots[i, :length] = slots.to(device).bool()
+    return out_embeds, out_mask, out_slots
+
+
+def pack_latents(latents: torch.Tensor) -> torch.Tensor:
+    """`(B, C, H, W)` -> `(B, H * W, C)`. 2.1 consumes latents unpatched."""
+    batch_size, channels, height, width = latents.shape
+    return latents.view(batch_size, channels, height * width).transpose(1, 2)
+
+
+def unpack_latents(latents: torch.Tensor, height: int, width: int) -> torch.Tensor:
+    """`(B, H * W, C)` -> `(B, C, H, W)`."""
+    batch_size, _, channels = latents.shape
+    return latents.transpose(1, 2).reshape(batch_size, channels, height, width)
+
+
+def run_transformer(
+    transformer,
+    latents: torch.Tensor,  # (B, C, h, w) noisy target latents
+    timestep: torch.Tensor,  # (B,) in [0, 1]
+    prompt_embeds: torch.Tensor,  # (B, L, dim)
+    prompt_mask: torch.Tensor,  # (B, L) bool
+    image_slot_mask: torch.Tensor,  # (B, L) bool
+    condition_latents: Optional[torch.Tensor] = None,  # (B, N, C) packed, in slot order
+    condition_shapes: Sequence[tuple[int, int]] = (),  # per condition image (h, w)
+    **kwargs,
+) -> torch.Tensor:
+    """Run the DiT on the joint sequence and return the target velocity as
+    `(B, C, h, w)`."""
+    batch_size, channels, height, width = latents.shape
+    target_tokens = height * width
+
+    hidden_states = pack_latents(latents)
+    if condition_latents is not None:
+        hidden_states = torch.cat([condition_latents, hidden_states], dim=1)
+
+    # The prompt reserved one slot per 2x2 group of each reference's latent
+    # tokens. Embeddings encoded against different references than the ones
+    # passed here (a shared blank/unconditional embedding, say) would silently
+    # mis-slot the sequence, so check it before the DiT does.
+    expected_slots = sum(h * w for h, w in condition_shapes) // LATENT_TOKENS_PER_SLOT
+    if int(image_slot_mask[0].sum()) != expected_slots:
+        raise ValueError(
+            f"the prompt embeddings hold {int(image_slot_mask[0].sum())} reference "
+            f"image slots but {len(condition_shapes)} reference image(s) were passed, "
+            f"worth {expected_slots} slots. They must be encoded together -- an "
+            "embedding cached against other references (caption dropout, a blank "
+            "unconditional embedding) cannot be reused here."
+        )
+
+    img_shapes = [
+        [*[(1, h, w) for h, w in condition_shapes], (1, height, width)]
+    ] * batch_size
+
+    # one image slot per 2x2 group of target latent tokens
+    img_mask = torch.cat(
+        [
+            image_slot_mask,
+            image_slot_mask.new_ones(
+                batch_size, target_tokens // LATENT_TOKENS_PER_SLOT
+            ),
+        ],
+        dim=1,
+    )
+    # a mask with nothing masked costs the attention backends that reject one
+    encoder_mask = None if bool(prompt_mask.all()) else prompt_mask
+
+    output = transformer(
+        hidden_states=hidden_states,
+        encoder_hidden_states=prompt_embeds,
+        encoder_hidden_states_mask=encoder_mask,
+        timestep=timestep,
+        img_shapes=img_shapes,
+        img_mask=img_mask,
+        return_dict=False,
+        **kwargs,
+    )[0]
+    return unpack_latents(output[:, -target_tokens:], height, width)
+
+
+def calculate_shift(
+    image_seq_len: int,
+    base_seq_len: int = 256,
+    max_seq_len: int = 8192,
+    base_shift: float = 0.5,
+    max_shift: float = 0.9,
+) -> float:
+    slope = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    return image_seq_len * slope + base_shift - slope * base_seq_len
+
+
+class QwenImage21Pipeline:
+    """Minimal flow-matching sampler for ai-toolkit's preview generation."""
+
+    def __init__(self, model):
+        # `model` is the QwenImage2Model, so its encode/decode and config are reused
+        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
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        conditional_embeds,
+        unconditional_embeds=None,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = 40,
+        guidance_scale: float = 1.0,
+        latents: Optional[torch.Tensor] = None,
+        generator: Optional[torch.Generator] = None,
+        condition_images: Optional[List[torch.Tensor]] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        model = self.model
+        device, dtype = model.device_torch, model.torch_dtype
+        transformer = model.transformer
+
+        latent_height = height // VAE_SCALE_FACTOR
+        latent_width = width // VAE_SCALE_FACTOR
+        channels = transformer.config.in_channels
+
+        condition_latents, condition_shapes = model.encode_condition_images(
+            condition_images
+        )
+
+        if latents is None:
+            # randn_tensor, not torch.randn: the caller's generator may be a cpu one
+            latents = randn_tensor(
+                (1, channels, latent_height, latent_width),
+                generator=generator,
+                device=torch.device(device),
+                dtype=torch.float32,
+            )
+        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.set_begin_index(0)
+
+        # 2.1 is meant to be sampled without guidance; a scale of 1 skips it
+        do_cfg = guidance_scale > 1.0 and unconditional_embeds is not None
+        cond = model.pad_prompt_embeds(conditional_embeds)
+        uncond = model.pad_prompt_embeds(unconditional_embeds) if do_cfg else None
+
+        for timestep in scheduler.timesteps:
+            t = timestep.expand(latents.shape[0]).to(device, dtype=dtype) / 1000
+            noise_pred = run_transformer(
+                transformer,
+                latents,
+                t,
+                *cond,
+                condition_latents=condition_latents,
+                condition_shapes=condition_shapes,
+            )
+            if do_cfg:
+                uncond_pred = run_transformer(
+                    transformer,
+                    latents,
+                    t,
+                    *uncond,
+                    condition_latents=condition_latents,
+                    condition_shapes=condition_shapes,
+                )
+                noise_pred = uncond_pred + guidance_scale * (noise_pred - uncond_pred)
+
+            latents = scheduler.step(
+                noise_pred.to(torch.float32),
+                timestep,
+                latents.to(torch.float32),
+                return_dict=False,
+            )[0].to(dtype)
+
+        return model.decode_to_images(latents)

extensions_built_in/diffusion_models/qwen_image_2/src/text_encoder.py

@@ -0,0 +1,60 @@
+"""Qwen3-VL-8B conditioning stack for Qwen-Image 2.1.
+
+The Comfy-Org repack keeps the transformers key layout except that the language
+tower sits directly under `model.` instead of `model.language_model.`, so the
+single-file load only needs that prefix rewrite before the toolkit's comfy
+quantization importer attaches the int8/convrot layers.
+"""
+
+from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder
+
+# comfy prefix -> transformers prefix for the language tower
+_COMFY_PREFIXES = (
+    ("model.layers.", "model.language_model.layers."),
+    ("model.embed_tokens.", "model.language_model.embed_tokens."),
+    ("model.norm.", "model.language_model.norm."),
+)
+
+
+class QwenImage21TextEncoder(Qwen3VLTextEncoder):
+    aitk_config_repo = "Qwen/Qwen-Image-2.1"
+    aitk_tokenizer_subfolder = "processor"
+    aitk_processor_repo = "Qwen/Qwen-Image-2.1"
+    aitk_processor_subfolder = "processor"
+
+    aitk_comfy_repo = "Comfy-Org/Qwen-Image-2.1"
+    # convrot8 first: it is the toolkit's default qtype, so a matching
+    # pre-quantized file attaches as-is. The w4a8 file is a different backend
+    # and would have to be dequantized, so it is not a candidate.
+    _COMFY_FILES = [
+        "text_encoders/qwen3vl_8b_int8_convrot.safetensors",
+        "text_encoders/qwen3vl_8b_bf16.safetensors",
+    ]
+    aitk_comfy_weight_names = {
+        ("Comfy-Org/Qwen-Image-2.1", "text_encoder"): _COMFY_FILES,
+        ("Qwen/Qwen-Image-2.1", "text_encoder"): _COMFY_FILES,
+    }
+
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        if not any(k.startswith("model.layers.") for k in state_dict):
+            return state_dict
+        converted = {}
+        for key, value in state_dict.items():
+            for comfy_prefix, hf_prefix in _COMFY_PREFIXES:
+                if key.startswith(comfy_prefix):
+                    key = hf_prefix + key[len(comfy_prefix) :]
+                    break
+            converted[key] = value
+        return converted
+
+    @classmethod
+    def convert_state_dict_on_save(cls, state_dict):
+        converted = {}
+        for key, value in state_dict.items():
+            for comfy_prefix, hf_prefix in _COMFY_PREFIXES:
+                if key.startswith(hf_prefix):
+                    key = comfy_prefix + key[len(hf_prefix) :]
+                    break
+            converted[key] = value
+        return converted

extensions_built_in/diffusion_models/ui.tsx

@@ -533,6 +533,41 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
         "uint3|ostris/accuracy_recovery_adapters/qwen_image_edit_2511_torchao_uint3.safetensors",
     },
   },
+  {
+    name: "qwen_image_2",
+    label: "Qwen-Image-2.1",
+    group: "image",
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      "config.process[0].model.name_or_path": [
+        "Comfy-Org/Qwen-Image-2.1",
+        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"],
+      // the Comfy-Org weights are pre-quantized int8 convrot; these qtypes
+      // match the checkpoints exactly, so the load is unchanged. Picking a
+      // different qtype re-quantizes layer by layer into that format.
+      "config.process[0].model.qtype": ["convrot8", "qfloat8"],
+      "config.process[0].model.qtype_te": ["convrot8", "qfloat8"],
+      // 2.1 is meant to be sampled without guidance
+      "config.process[0].sample.guidance_scale": [1.0, 4.0],
+    },
+    disableSections: ["network.conv", "train.unload_text_encoder"],
+    // one model: it edits when the dataset has control paths, and is plain
+    // text to image when it does not
+    additionalSections: [
+      "datasets.multi_control_paths",
+      "sample.multi_ctrl_imgs",
+      "model.low_vram",
+      "model.layer_offloading",
+    ],
+  },
   {
     name: "hidream",
     label: "HiDream",