From c2622ed250f709d37396097457e7dcaffd6039d0 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 20 Sep 2026 09:50:57 -0600 Subject: [PATCH] Add support for Qwen Image 2.1 --- README.md | 9 +- .../diffusion_models/__init__.py | 2 + .../diffusion_models/qwen_image_2/__init__.py | 3 + .../qwen_image_2/qwen_image_2.py | 570 +++++ .../qwen_image_2/src/__init__.py | 0 .../qwen_image_2/src/pipeline.py | 412 ++++ .../qwen_image_2/src/text_encoder.py | 60 + .../qwen_image_2/src/transformer.py | 1267 +++++++++++ .../diffusion_models/qwen_image_2/src/vae.py | 1872 +++++++++++++++++ extensions_built_in/diffusion_models/ui.tsx | 35 + testing/test_model_loading.py | 7 + toolkit/models/v2/_mixin.py | 11 +- version.py | 2 +- 13 files changed, 4247 insertions(+), 3 deletions(-) create mode 100644 extensions_built_in/diffusion_models/qwen_image_2/__init__.py create mode 100644 extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py create mode 100644 extensions_built_in/diffusion_models/qwen_image_2/src/__init__.py create mode 100644 extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py create mode 100644 extensions_built_in/diffusion_models/qwen_image_2/src/text_encoder.py create mode 100644 extensions_built_in/diffusion_models/qwen_image_2/src/transformer.py create mode 100644 extensions_built_in/diffusion_models/qwen_image_2/src/vae.py diff --git a/README.md b/README.md index 95c0760..1608536 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t - [Alpha-VLLM/Lumina-Image-2.0](https://huggingface.co/Alpha-VLLM/Lumina-Image-2.0) (Lumina2) - [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) - [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) @@ -41,6 +42,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t - [Qwen/Qwen-Image-Edit](https://huggingface.co/Qwen/Qwen-Image-Edit) (Qwen-Image-Edit) - [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 - [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) @@ -57,15 +59,20 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t - [Wan-AI/Wan2.2-TI2V-5B-Diffusers](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers) (Wan 2.2 TI2V 5B) - [Lightricks/LTX-2](https://huggingface.co/Lightricks/LTX-2) (LTX-2) - [Lightricks/LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3) (LTX-2.3) +- [Lightricks/LTX-2.5](https://huggingface.co/Lightricks/LTX-2.5) (LTX-2.5) - [MiniMaxAI/MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) (MiniMaxAI/MiniMax-H3) +- [MiniMaxAI/MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) (MiniMax-H3 Ref2V) - reference image to video ### Audio - [ACE-Step/Ace-Step1.5](https://huggingface.co/ACE-Step/Ace-Step1.5) (Ace Step 1.5) - [ACE-Step/acestep-v15-xl-base](https://huggingface.co/ACE-Step/acestep-v15-xl-base) (Ace Step 1.5 XL) +- [m-a-p/YuE2-3B](https://huggingface.co/m-a-p/YuE2-3B) (YuE2) - the official audio-to-token encoder is unreleased; training uses the community tokenizer by Kytra ([@sin_ceriously](https://x.com/sin_ceriously)), [Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4](https://huggingface.co/Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4). + +### LLM +- [Qwen/Qwen2.5-Omni-7B](https://huggingface.co/Qwen/Qwen2.5-Omni-7B) (Qwen2.5-Omni) ### Experimental - [lodestones/Zeta-Chroma](https://huggingface.co/lodestones/Zeta-Chroma) (Zeta Chroma) -- [m-a-p/YuE2-3B](https://huggingface.co/m-a-p/YuE2-3B) (YuE2) - audio. The official audio-to-token encoder is unreleased; training uses the community tokenizer by Kytra ([@sin_ceriously](https://x.com/sin_ceriously)), [Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4](https://huggingface.co/Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4). ## Installation diff --git a/extensions_built_in/diffusion_models/__init__.py b/extensions_built_in/diffusion_models/__init__.py index d0c1859..aeedcf0 100644 --- a/extensions_built_in/diffusion_models/__init__.py +++ b/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, diff --git a/extensions_built_in/diffusion_models/qwen_image_2/__init__.py b/extensions_built_in/diffusion_models/qwen_image_2/__init__.py new file mode 100644 index 0000000..7831437 --- /dev/null +++ b/extensions_built_in/diffusion_models/qwen_image_2/__init__.py @@ -0,0 +1,3 @@ +from .qwen_image_2 import QwenImage2Model + +__all__ = ["QwenImage2Model"] 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 new file mode 100644 index 0000000..843f512 --- /dev/null +++ b/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 diff --git a/extensions_built_in/diffusion_models/qwen_image_2/src/__init__.py b/extensions_built_in/diffusion_models/qwen_image_2/src/__init__.py new file mode 100644 index 0000000..e69de29 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 new file mode 100644 index 0000000..3a3fe67 --- /dev/null +++ b/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"{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) diff --git a/extensions_built_in/diffusion_models/qwen_image_2/src/text_encoder.py b/extensions_built_in/diffusion_models/qwen_image_2/src/text_encoder.py new file mode 100644 index 0000000..a3223dc --- /dev/null +++ b/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 diff --git a/extensions_built_in/diffusion_models/qwen_image_2/src/transformer.py b/extensions_built_in/diffusion_models/qwen_image_2/src/transformer.py new file mode 100644 index 0000000..6ecc845 --- /dev/null +++ b/extensions_built_in/diffusion_models/qwen_image_2/src/transformer.py @@ -0,0 +1,1267 @@ +# Copyright 2026 Qwen-Image Team, The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen-Image 2.1 transformer, vendored from diffusers (Apache 2.0). + +Vendored rather than imported because `QwenImage21Transformer2DModel` landed in +diffusers after the version this repo pins. Two deliberate changes from +upstream: + + - the SwiGLU feed forward keeps the ComfyUI **fused** `img_mlp.gate_up` + projection ([gate; up] row-concatenated) instead of the split + `gate_layer`/`proj` pair, so a Comfy-Org checkpoint (bf16 or int8_convrot) + loads with no key surgery and a LoRA trained here targets the same module + names ComfyUI has, + - `OstrisModelMixin` on the model class, for the toolkit's universal + load/quantize/offload path. +""" + +import math +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin +from diffusers.utils import logging +from diffusers.utils.peft_utils import apply_lora_scale +from diffusers.utils.torch_utils import maybe_allow_in_graph +from diffusers.models.attention import AttentionMixin, AttentionModuleMixin +from diffusers.models.attention_dispatch import dispatch_attention_fn +from diffusers.models.cache_utils import CacheMixin +from diffusers.models.embeddings import TimestepEmbedding +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.normalization import RMSNorm + +from toolkit.models.v2._mixin import OstrisModelMixin + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +# Each vision-language image slot represents a 2×2 group of latent tokens. +_IMG_TOKENS_PER_SLOT = 4 + +# `create_block_mask` quantizes the mask to 128-token blocks. +_FLEX_BLOCK_SIZE = 128 + +# flex_attention is optional: `QwenImage21FlexAttnProcessor` needs it, `QwenImage21AttnProcessor` does not. +_FLEX_AVAILABLE = False +try: + import torch.nn.attention.flex_attention as flex_attention_module + from torch.nn.attention.flex_attention import BlockMask, create_block_mask + + _FLEX_AVAILABLE = True +except ImportError: + BlockMask = None + flex_attention_module = None + + +class QwenImage21KVLayerCache: + """Per-layer KV cache for text and condition-image prefix tokens. + + Stores K and V projections (post-RoPE) for the prefix extracted during the first denoising step. Tensor format: + ``(batch_size, num_prefix_tokens, num_heads, head_dim)``. + """ + + def __init__(self): + self.k: torch.Tensor | None = None + self.v: torch.Tensor | None = None + + def store(self, k: torch.Tensor, v: torch.Tensor): + self.k = k + self.v = v + + def get(self) -> tuple[torch.Tensor, torch.Tensor]: + if self.k is None: + raise RuntimeError("KV cache has not been populated yet.") + return self.k, self.v + + +class QwenImage21KVCache: + """Container for all transformer blocks' prefix KV caches.""" + + def __init__(self, num_layers: int): + self.layer_caches = [QwenImage21KVLayerCache() for _ in range(num_layers)] + + def get_layer(self, layer_idx: int) -> QwenImage21KVLayerCache: + return self.layer_caches[layer_idx] + + +# Copied from diffusers.models.transformers.transformer_qwenimage.apply_rotary_emb_qwen +def apply_rotary_emb_qwen( + x: torch.Tensor, + freqs_cis: torch.Tensor | tuple[torch.Tensor], + use_real: bool = True, + use_real_unbind_dim: int = -1, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings + to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are + reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting + tensors contain rotary embeddings and are returned as real tensors. + + Args: + x (`torch.Tensor`): + Query or key tensor to apply rotary embeddings. [B, S, H, D] xk (torch.Tensor): Key tensor to apply + freqs_cis (`tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],) + + Returns: + tuple[torch.Tensor, torch.Tensor]: tuple of modified query tensor and key tensor with rotary embeddings. + """ + if use_real: + cos, sin = freqs_cis # [S, D] + cos = cos[None, None] + sin = sin[None, None] + cos, sin = cos.to(x.device), sin.to(x.device) + + if use_real_unbind_dim == -1: + # Used for flux, cogvideox, hunyuan-dit + x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind( + -1 + ) # [B, S, H, D//2] + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) + elif use_real_unbind_dim == -2: + # Used for Stable Audio, OmniGen, CogView4 and Cosmos + x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind( + -2 + ) # [B, S, H, D//2] + x_rotated = torch.cat([-x_imag, x_real], dim=-1) + else: + raise ValueError( + f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2." + ) + + out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) + + return out + else: + x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + freqs_cis = freqs_cis.unsqueeze(1) + x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) + + return x_out.type_as(x) + + +class QwenImage21TemporalTimesteps(nn.Module): + r"""Sinusoidal timestep embedding. `cos` occupies the first half of the channels and `sin` the second.""" + + def __init__( + self, timestep_dim: int, max_period: int = 10000, time_factor: float = 1000.0 + ): + super().__init__() + self.timestep_dim = timestep_dim + self.time_factor = time_factor + + half = timestep_dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32) + / half + ) + self.register_buffer("freqs", freqs, persistent=False) + + def forward(self, timestep: torch.Tensor) -> torch.Tensor: + timestep = self.time_factor * timestep.float() + args = timestep[:, None] * self.freqs[None].to(timestep.device) + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if self.timestep_dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 + ) + return embedding.to(timestep.dtype) + + +class QwenImage21TimestepProjEmbeddings(nn.Module): + def __init__(self, embedding_dim: int): + super().__init__() + self.time_proj = QwenImage21TemporalTimesteps(timestep_dim=256) + self.timestep_embedder = TimestepEmbedding( + in_channels=256, time_embed_dim=embedding_dim, sample_proj_bias=False + ) + + def forward( + self, timestep: torch.Tensor, hidden_states: torch.Tensor + ) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + return self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype)) + + +class QwenImage21ZeroCenterRMSNorm(nn.Module): + r""" + RMSNorm whose learnable weight is stored zero-centered: the effective scale is `weight + 1`, computed in fp32. + Checkpoints therefore store `scale - 1`. + """ + + def __init__(self, dim: int, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.zeros(dim)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.float() + rrms = torch.rsqrt( + torch.mean(hidden_states**2, dim=-1, keepdim=True) + self.eps + ) + return (hidden_states * rrms * (self.weight.float() + 1)).to(input_dtype) + + +class QwenImage21TextProjection(nn.Module): + def __init__(self, context_in_dim: int, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.text_norm = QwenImage21ZeroCenterRMSNorm(context_in_dim, eps=eps) + self.in_layer = nn.Linear(context_in_dim, hidden_size, bias=False) + self.act = nn.GELU(approximate="tanh") + self.out_layer = nn.Linear(hidden_size, hidden_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.text_norm(hidden_states) + hidden_states = self.in_layer(hidden_states) + hidden_states = self.act(hidden_states) + return self.out_layer(hidden_states) + + +class QwenImage21SwiGLUFeedForward(nn.Module): + r"""SwiGLU with the gate and up projections fused into one GEMM, the layout + the ComfyUI checkpoints ship: `gate_up.weight` is `[gate_layer; proj]` + concatenated on the out dim.""" + + def __init__(self, hidden_size: int, mlp_hidden_size: int): + super().__init__() + self.gate_up = nn.Linear(hidden_size, 2 * mlp_hidden_size, bias=False) + self.out = nn.Linear(mlp_hidden_size, hidden_size, bias=False) + self.activation_fn = nn.SiLU() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + gate, up = self.gate_up(hidden_states).chunk(2, dim=-1) + return self.out(self.activation_fn(gate) * up) + + +class QwenImage21AdaLayerNormContinuous(nn.Module): + r""" + Final adaptive norm. Scale only — this variant emits no shift, so `linear` maps to `embedding_dim` rather than `2 * + embedding_dim`. + """ + + def __init__( + self, embedding_dim: int, conditioning_embedding_dim: int, eps: float = 1e-6 + ): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=False) + self.norm = nn.LayerNorm( + embedding_dim, eps, elementwise_affine=False, bias=False + ) + + def forward( + self, + hidden_states: torch.Tensor, + conditioning_embedding: torch.Tensor, + target_token_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + scale = self.linear(self.silu(conditioning_embedding).to(hidden_states.dtype)) + scale = _select_modulation_rows(scale, target_token_mask) + return self.norm(hidden_states) * (1 + scale) + + +def _select_modulation_rows( + params: torch.Tensor, target_token_mask: torch.Tensor | None +) -> torch.Tensor: + r""" + Broadcast per-sample modulation `params` over the token axis. + + With `causal_condition`, `params` holds `batch_size + 1` rows: rows `[0, batch_size)` come from the real timestep + and the trailing row from `t = 0`. Text and condition-image tokens take the `t = 0` row, target-image tokens take + their own sample's row. + + Args: + params (`torch.Tensor`): `(batch_size, dim)` without `causal_condition`, else `(batch_size + 1, dim)`. + target_token_mask (`torch.Tensor`, *optional*): `(seq_len,)` bool, `True` at target-image positions. `None` + disables the split and every token uses its own sample's row. + """ + if target_token_mask is None: + return params.unsqueeze(1) + real, zero = params[:-1].unsqueeze(1), params[-1:].unsqueeze(0) + return torch.where(target_token_mask.view(1, -1, 1), real, zero) + + +def build_qwenimage21_block_causal_mask( + image_ids: torch.Tensor, + encoder_hidden_states_mask: torch.Tensor | None, + batch_size: int, + device: torch.device, +): + r""" + Build the block-causal [`~torch.nn.attention.flex_attention.BlockMask`] for Qwen-Image 2.1. + + The mask is `(q_idx >= kv_idx) or same_image_block`: the joint text/image sequence is causal, while every image + block — each condition image and the target image — is internally bidirectional. Text tokens are strictly causal. + Cross-sample isolation is implicit in diffusers because samples live on the batch axis. + + Positions masked out by `encoder_hidden_states_mask` are excluded as *keys* so right-padded prompts cannot be + attended to. They are kept as queries so their rows are never fully masked. + + Provided as a module-level function so callers can build the mask outside the transformer's compiled region. + + Args: + image_ids (`torch.Tensor`): `(seq_len,)` int, `-1` at text positions and a unique non-negative id per image + block. See [`~QwenImage21Transformer2DModel.build_token_metadata`]. + encoder_hidden_states_mask (`torch.Tensor`, *optional*): `(batch_size, seq_len)` bool over the joint sequence. + batch_size (`int`): Number of samples; the mask varies across the batch only through + `encoder_hidden_states_mask`. + """ + seq_len = image_ids.shape[0] + padded_seq_len = int(math.ceil(seq_len / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) + + image_ids = F.pad(image_ids, (0, padded_seq_len - seq_len), value=-1) + if encoder_hidden_states_mask is None: + key_valid = torch.ones( + batch_size, padded_seq_len, dtype=torch.bool, device=device + ) + else: + key_valid = F.pad( + encoder_hidden_states_mask.bool(), + (0, padded_seq_len - seq_len), + value=False, + ) + + def mask_mod(batch_idx, head_idx, q_idx, kv_idx): + is_padding = (q_idx >= seq_len) | (kv_idx >= seq_len) + q_image_id, kv_image_id = image_ids[q_idx], image_ids[kv_idx] + same_image_block = (q_image_id == kv_image_id) & (q_image_id >= 0) + allowed = ((q_idx >= kv_idx) | same_image_block) & key_valid[batch_idx, kv_idx] + return allowed & ~is_padding + + return create_block_mask( + mask_mod, + B=batch_size, + H=None, + Q_LEN=padded_seq_len, + KV_LEN=padded_seq_len, + device=device, + _compile=False, + ) + + +def _qwenimage21_prefix_segments( + image_ids: torch.Tensor, prefix_len: int +) -> list[tuple[int, int, bool]]: + """Split the prefix into `(start, end, is_text)` runs of equal `image_ids`. + + This is the block-causal structure in the form [`QwenImage21AttnProcessor`] consumes it, the way + [`~build_qwenimage21_block_causal_mask`] is the form [`QwenImage21FlexAttnProcessor`] consumes. It only depends on + `image_ids` and `prefix_len`, so the model derives it once per forward rather than in every processor call — + `tolist()` is a device sync, and there is one processor call per layer. + """ + prefix_ids = image_ids[:prefix_len].tolist() + segments = [] + start = 0 + for index in range(1, prefix_len + 1): + if index == prefix_len or prefix_ids[index] != prefix_ids[start]: + segments.append((start, index, prefix_ids[start] < 0)) + start = index + return segments + + +def _qwenimage21_prepare_qkv( + attn: "QwenImage21Attention", + hidden_states: torch.Tensor, + rotary_emb: torch.Tensor | None, + layer_cache: QwenImage21KVLayerCache | None, + kv_cache_mode: str | None, + cache_write_slice: slice | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Shared QKV projection, norm, RoPE and KV-cache bookkeeping for both processors.""" + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = query.unflatten(-1, (attn.heads, -1)) + key = key.unflatten(-1, (attn.heads, -1)) + value = value.unflatten(-1, (attn.heads, -1)) + + query = attn.norm_q(query).to(value.dtype) + key = attn.norm_k(key).to(value.dtype) + + if rotary_emb is not None: + query = apply_rotary_emb_qwen(query, rotary_emb, use_real=False) + key = apply_rotary_emb_qwen(key, rotary_emb, use_real=False) + + if layer_cache is not None: + if kv_cache_mode == "extract" and cache_write_slice is not None: + # `clone()`, not `contiguous()`: at batch size 1 the prefix slice already counts as contiguous + # (size-1 dims are ignored), so `contiguous()` returns the same view and the cache would pin the + # whole prefill K/V for every step of the denoising loop. + layer_cache.store( + key[:, cache_write_slice].clone(), + value[:, cache_write_slice].clone(), + ) + elif kv_cache_mode == "cached": + cached_k, cached_v = layer_cache.get() + key = torch.cat([cached_k, key], dim=1) + value = torch.cat([cached_v, value], dim=1) + + seq_len_q = query.shape[1] + return query, key, value, seq_len_q + + +class QwenImage21FlexAttnProcessor: + r""" + Attention processor for Qwen-Image 2.1 that runs the block-causal prefill as one `flex_attention` call driven by a + `BlockMask`, and the cached decode steps through the configured attention backend. + + Compile the model before using it, as the docs show. An uncompiled `flex_attention` falls back to a dense fp32 + score matrix, which is far slower and runs out of memory at high resolution. Use `QwenImage21AttnProcessor` when + you do not want to compile. + """ + + # Set by `set_attention_backend()` and only meaningful for the decode steps; the prefill needs the flex kernel + # for its `BlockMask` and is not configurable. + _attention_backend = None + _parallel_config = None + _warned_uncompiled = False + + def __init__(self): + if not _FLEX_AVAILABLE: + raise ImportError( + "`QwenImage21FlexAttnProcessor` requires `torch.nn.attention.flex_attention`, which needs " + "torch>=2.5. Use `QwenImage21AttnProcessor` instead." + ) + + def __call__( + self, + attn: "QwenImage21Attention", + hidden_states: torch.Tensor, + attention_mask: Any | None = None, + rotary_emb: torch.Tensor | None = None, + layer_cache: QwenImage21KVLayerCache | None = None, + kv_cache_mode: str | None = None, + cache_write_slice: slice | None = None, + segments: list[tuple[int, int, bool]] | None = None, + key_valid: torch.Tensor | None = None, + ) -> torch.Tensor: + query, key, value, seq_len_q = _qwenimage21_prepare_qkv( + attn, + hidden_states, + rotary_emb, + layer_cache, + kv_cache_mode, + cache_write_slice, + ) + + seq_len_kv = key.shape[1] + if isinstance(attention_mask, BlockMask): + # prefill: the BlockMask expresses the block-causal structure in one flex call. Query and key are + # padded up to the mask's block-quantized length. + # `dispatch_attention_fn` reaches flex_attention through its module, so a user who compiled it — directly + # or by compiling the model — is picked up here. Tracing means the model is compiled. + if ( + not self._warned_uncompiled + and not torch.compiler.is_compiling() + and not hasattr( + flex_attention_module.flex_attention, "_torchdynamo_orig_callable" + ) + ): + logger.warning( + "`QwenImage21FlexAttnProcessor` is running an uncompiled `flex_attention`, which materializes the " + "full attention score matrix in fp32 and will run out of memory at high resolution. Compile the " + "model with `transformer.compile()`, or switch to `QwenImage21AttnProcessor`." + ) + QwenImage21FlexAttnProcessor._warned_uncompiled = True + pad_q = ( + int(math.ceil(seq_len_q / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) + - seq_len_q + ) + pad_kv = ( + int(math.ceil(seq_len_kv / _FLEX_BLOCK_SIZE) * _FLEX_BLOCK_SIZE) + - seq_len_kv + ) + # Pad the sequence axis. `F.pad` counts from the last dimension, so the head and channel axes are + # padded by zero first. The result stays contiguous, which the compiled flex kernel requires. + if pad_q: + query = F.pad(query, (0, 0, 0, 0, 0, pad_q)) + if pad_kv: + key = F.pad(key, (0, 0, 0, 0, 0, pad_kv)) + value = F.pad(value, (0, 0, 0, 0, 0, pad_kv)) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + backend="flex", + parallel_config=self._parallel_config, + ) + else: + # decode: full attention over [cached prefix, target] + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states[:, :seq_len_q] + hidden_states = hidden_states.flatten(2, 3).type_as(query) + + hidden_states = attn.to_out[0](hidden_states) + return attn.to_out[1](hidden_states) + + +class QwenImage21AttnProcessor: + r""" + Attention processor for Qwen-Image 2.1 that needs neither `flex_attention` nor a compiled model. + + The prefill decomposes the block-causal mask into one attention call per prefix segment plus one for the target + image, which is exact but slower than [`QwenImage21FlexAttnProcessor`]. The segment boundaries are computed once + per forward by the model and passed in as `segments`. + """ + + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "QwenImage21Attention", + hidden_states: torch.Tensor, + attention_mask: Any | None = None, + rotary_emb: torch.Tensor | None = None, + layer_cache: QwenImage21KVLayerCache | None = None, + kv_cache_mode: str | None = None, + cache_write_slice: slice | None = None, + segments: list[tuple[int, int, bool]] | None = None, + key_valid: torch.Tensor | None = None, + ) -> torch.Tensor: + query, key, value, seq_len_q = _qwenimage21_prepare_qkv( + attn, + hidden_states, + rotary_emb, + layer_cache, + kv_cache_mode, + cache_write_slice, + ) + + if segments is None: + # decode: full attention over [cached prefix, target] + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + else: + # prefill: every segment attends to the keys `[0, end)` (everything before it plus its own block); text + # segments additionally get a causal triangle over their own keys; padded text keys are dropped. + # `attention_mask` may hold the flex `BlockMask` of the same structure, which is not used here. + prefix_len = segments[-1][1] if segments else 0 + outputs = [] + for start, end, is_text in segments: + seg_mask = None + if is_text: + seg_len = end - start + seg_mask = torch.cat( + [ + torch.ones( + seg_len, start, dtype=torch.bool, device=query.device + ), + torch.tril( + torch.ones( + seg_len, + seg_len, + dtype=torch.bool, + device=query.device, + ) + ), + ], + dim=1, + )[None, None] + if key_valid is not None: + seg_key_valid = key_valid[:, None, None, :end] + seg_mask = ( + seg_key_valid + if seg_mask is None + else (seg_mask & seg_key_valid) + ) + outputs.append( + dispatch_attention_fn( + query[:, start:end], + key[:, :end], + value[:, :end], + attn_mask=seg_mask, + dropout_p=0.0, + backend=None, + parallel_config=self._parallel_config, + ) + ) + outputs.append( + dispatch_attention_fn( + query[:, prefix_len:], + key, + value, + attn_mask=None + if key_valid is None + else key_valid[:, None, None, :], + dropout_p=0.0, + backend=None, + parallel_config=self._parallel_config, + ) + ) + hidden_states = torch.cat(outputs, dim=1) + hidden_states = hidden_states[:, :seq_len_q] + hidden_states = hidden_states.flatten(2, 3).type_as(query) + + hidden_states = attn.to_out[0](hidden_states) + return attn.to_out[1](hidden_states) + + +class QwenImage21Attention(torch.nn.Module, AttentionModuleMixin): + r""" + Attention module for [`QwenImage21TransformerBlock`]. Projection layout matches the legacy + [`~models.attention_processor.Attention`] so Qwen-Image 2.x checkpoints load into it unchanged. + """ + + # The default must not depend on the caller having compiled the model: an uncompiled `flex_attention` falls + # back to a dense fp32 score matrix and runs out of memory at high resolution. `QwenImage21FlexAttnProcessor` + # is the faster path once compiled, and the docs show how to opt into it. + _default_processor_cls = QwenImage21AttnProcessor + _available_processors = [QwenImage21AttnProcessor, QwenImage21FlexAttnProcessor] + + def __init__( + self, + dim: int, + heads: int, + dim_head: int, + eps: float = 1e-6, + processor: Any | None = None, + ): + super().__init__() + self.heads = heads + self.inner_dim = heads * dim_head + # Read by `AttentionModuleMixin.fuse_projections`; 2.1 has no biases anywhere. + self.use_bias = False + + self.to_q = nn.Linear(dim, self.inner_dim, bias=False) + self.to_k = nn.Linear(dim, self.inner_dim, bias=False) + self.to_v = nn.Linear(dim, self.inner_dim, bias=False) + self.to_out = nn.ModuleList( + [nn.Linear(self.inner_dim, dim, bias=False), nn.Dropout(0.0)] + ) + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + + self.set_processor( + processor if processor is not None else self._default_processor_cls() + ) + + def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return self.processor(self, hidden_states, **kwargs) + + +@maybe_allow_in_graph +class QwenImage21TransformerBlock(nn.Module): + r""" + Single-stream block. Modulation is not learned per block — the parent model computes one shared `modulation` tensor + and every block slices its own scales and gates out of it. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: int = 3, + eps: float = 1e-6, + ): + super().__init__() + self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.attn = QwenImage21Attention( + dim=dim, heads=num_attention_heads, dim_head=attention_head_dim, eps=eps + ) + self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.img_mlp = QwenImage21SwiGLUFeedForward( + hidden_size=dim, mlp_hidden_size=dim * mlp_ratio + ) + + def _modulate( + self, + hidden_states: torch.Tensor, + mod_params: torch.Tensor, + target_token_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + scale, gate = mod_params.chunk(2, dim=-1) + scale = _select_modulation_rows(scale, target_token_mask) + gate = _select_modulation_rows(gate, target_token_mask) + return hidden_states * (1 + scale), gate + + def forward( + self, + hidden_states: torch.Tensor, + modulation: torch.Tensor, + rotary_emb: torch.Tensor | None = None, + attention_mask: Any | None = None, + target_token_mask: torch.Tensor | None = None, + layer_cache: QwenImage21KVLayerCache | None = None, + kv_cache_mode: str | None = None, + cache_write_slice: slice | None = None, + segments: list[tuple[int, int, bool]] | None = None, + key_valid: torch.Tensor | None = None, + ) -> torch.Tensor: + mod1, mod2 = modulation.chunk(2, dim=-1) + + img_modulated, img_gate1 = self._modulate( + self.img_norm1(hidden_states), mod1, target_token_mask + ) + attn_output = self.attn( + hidden_states=img_modulated, + attention_mask=attention_mask, + rotary_emb=rotary_emb, + layer_cache=layer_cache, + kv_cache_mode=kv_cache_mode, + cache_write_slice=cache_write_slice, + segments=segments, + key_valid=key_valid, + ) + hidden_states = hidden_states + img_gate1.tanh() * attn_output + + img_modulated2, img_gate2 = self._modulate( + self.img_norm2(hidden_states), mod2, target_token_mask + ) + hidden_states = hidden_states + img_gate2.tanh() * self.img_mlp(img_modulated2) + + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + return hidden_states + + +class QwenImage21Rope(nn.Module): + r""" + 3-axis (frame, height, width) rotary embedding over the joint text/image sequence. + + Text tokens advance a shared position on all three axes. Every image block freezes the frame axis at the position + reached by the preceding text and lays its tokens out on a height/width grid centred on zero, so a block's spatial + positions do not depend on where it sits in the sequence. + """ + + def __init__(self, theta: int, axes_dim: list[int]): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + + pos_index = torch.arange(8192) + neg_index = torch.arange(1024).flip(0) * -1 - 1 + self.freqs = [ + torch.cat( + [ + self.rope_params(pos_index, dim, theta), + self.rope_params(neg_index, dim, theta), + ], + dim=0, + ) + for dim in axes_dim + ] + + def rope_params( + self, index: torch.Tensor, dim: int, theta: int = 10000 + ) -> torch.Tensor: + freqs = torch.outer( + index, + 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)), + ) + return torch.polar(torch.ones_like(freqs), freqs) + + def forward( + self, + img_shapes: list[tuple[int, int, int]], + image_pad_mask: torch.Tensor, + device: torch.device, + ) -> torch.Tensor: + self.freqs = [freq.to(device) for freq in self.freqs] + + frame_index, height_index, width_index = [], [], [] + image_height_index, image_width_index = [], [] + cursor, position = 0, 0 + total_len = image_pad_mask.shape[-1] + is_image_token = image_pad_mask.tolist() + + for _, height, width in img_shapes: + block_start = is_image_token.index(True, cursor) + text_len = block_start - cursor + frame_index.extend(range(position, position + text_len)) + position += text_len + + cursor = block_start + height * width + frame_index.extend([position] * (height * width)) + position += max(height, width) + + image_height_index.extend( + [ + h + for h in range(-(height - height // 2), height // 2) + for _ in range(width) + ] + ) + image_width_index.extend( + [ + w + for _ in range(height) + for w in range(-(width - width // 2), width // 2) + ] + ) + + if cursor < total_len: + frame_index.extend(range(position, position + total_len - cursor)) + + frame_index = torch.tensor(frame_index, dtype=torch.long, device=device) + height_index = frame_index.clone() + width_index = frame_index.clone() + height_index[image_pad_mask] = torch.tensor( + image_height_index, dtype=torch.long, device=device + ) + width_index[image_pad_mask] = torch.tensor( + image_width_index, dtype=torch.long, device=device + ) + + return torch.cat( + [ + self.freqs[0][frame_index], + self.freqs[1][height_index], + self.freqs[2][width_index], + ], + dim=-1, + ) + + +class QwenImage21Transformer2DModel( + ModelMixin, + ConfigMixin, + PeftAdapterMixin, + FromOriginalModelMixin, + CacheMixin, + AttentionMixin, + OstrisModelMixin, +): + r""" + The single-stream Transformer used by Qwen-Image 2.1. + + Text and image latents share one sequence: condition-image tokens are substituted into the text stream at the + positions the vision-language encoder reserved for them, and the target image's tokens are appended. A single + shared `modulation` projection feeds every block, so blocks hold no modulation parameters of their own. + + Two behaviours distinguish 2.1: + + - **Block-causal attention** — attention follows `(q_idx >= kv_idx) or same_image_block`, so the sequence is causal + while each image block stays internally bidirectional. The `flex` attention backend gives efficient single-pass + attention; without it the model uses an exact multi-pass SDPA prefill that processes each block separately. + - `causal_condition` — text and condition-image tokens are modulated from `t = 0` instead of the sampled timestep, + which also makes their activations timestep-independent and so cacheable across denoising steps. + + Args: + patch_size (`int`, defaults to `1`): + Side length of the latent patch folded into the channel dim. 2.1 consumes latents unpatched. + in_channels (`int`, defaults to `64`): + Latent channels of the input. + out_channels (`int`, *optional*, defaults to `64`): + Latent channels of the output. Falls back to `in_channels`. + num_layers (`int`, defaults to `32`): + Number of single-stream blocks. + attention_head_dim (`int`, defaults to `128`): + Channels per attention head. + num_attention_heads (`int`, defaults to `32`): + Number of attention heads. + context_in_dim (`int`, defaults to `4096`): + Channel dim of `encoder_hidden_states`. + mlp_ratio (`int`, defaults to `3`): + Feed-forward expansion factor. + axes_dims_rope (`tuple[int]`, defaults to `(16, 56, 56)`): + Rotary dims for the frame, height and width axes. + eps (`float`, defaults to `1e-6`): + Epsilon for the norm layers. + causal_condition (`bool`, defaults to `True`): + Modulate text and condition-image tokens from `t = 0`. Required for KV caching. + """ + + _supports_gradient_checkpointing = True + _no_split_modules = ["QwenImage21TransformerBlock"] + _skip_layerwise_casting_patterns = ["pos_embed", "norm"] + _repeated_blocks = ["QwenImage21TransformerBlock"] + _skip_keys = ["kv_cache"] + + # ---- toolkit loading (OstrisModelMixin) ---- + aitk_subfolder = "transformer" + aitk_config_repo = "Qwen/Qwen-Image-2.1" + + aitk_comfy_repo = "Comfy-Org/Qwen-Image-2.1" + # convrot8 first: it is the toolkit's default qtype, and a matching + # pre-quantized file attaches as-is instead of paying a quantize pass. + # Registered under both repo ids so name_or_path can be either the comfy + # repack (the default) or the original repo. + _COMFY_FILES = [ + "diffusion_models/qwen_image_2.1_int8_convrot.safetensors", + "diffusion_models/qwen_image_2.1_bf16.safetensors", + ] + aitk_comfy_weight_names = { + "Comfy-Org/Qwen-Image-2.1": _COMFY_FILES, + "Qwen/Qwen-Image-2.1": _COMFY_FILES, + } + + @classmethod + def get_transformer_block_names(cls): + return ["transformer_blocks"] + + @classmethod + def get_quantization_exclude_modules(cls): + # sensitive modules kept in full precision (fnmatch patterns on module + # names). These are exactly the ones the Comfy-Org int8_convrot file + # leaves in bf16, so the shipped quantization and a fresh one agree: + # img_in - latent input projection + # txt_in* - text feature -> model width projection + # time_text_embed* - timestep embedder, feeds every block's modulation + # modulation* - the one shared modulation projection + # norm_out* / proj_out - final modulated norm + output projection + return [ + "img_in", + "txt_in*", + "time_text_embed*", + "modulation*", + "norm_out*", + "proj_out", + ] + + @classmethod + def convert_state_dict_on_load(cls, state_dict): + """Accept the diffusers key layout (split `img_mlp.gate_layer`/`proj`) + as well as the comfy one this class uses natively (fused + `img_mlp.gate_up`). Comfy checkpoints pass straight through.""" + if not any(k.endswith(".img_mlp.gate_layer.weight") for k in state_dict): + return state_dict + + state_dict = dict(state_dict) + new_sd = {} + for key, value in state_dict.items(): + if key.endswith(".img_mlp.gate_layer.weight"): + prefix = key[: -len(".gate_layer.weight")] + # gate first, matching the comfy [gate; up] concat order + new_sd[f"{prefix}.gate_up.weight"] = torch.cat( + [value, state_dict[f"{prefix}.proj.weight"]], dim=0 + ) + continue + if key.endswith(".img_mlp.proj.weight"): + continue + new_sd[key] = value + return new_sd + + @register_to_config + def __init__( + self, + patch_size: int = 1, + in_channels: int = 64, + out_channels: int | None = 64, + num_layers: int = 32, + attention_head_dim: int = 128, + num_attention_heads: int = 32, + context_in_dim: int = 4096, + mlp_ratio: int = 3, + axes_dims_rope: tuple[int, int, int] = (16, 56, 56), + eps: float = 1e-6, + causal_condition: bool = True, + ): + super().__init__() + self.out_channels = out_channels or in_channels + self.inner_dim = num_attention_heads * attention_head_dim + + self.pos_embed = QwenImage21Rope(theta=10000, axes_dim=list(axes_dims_rope)) + self.time_text_embed = QwenImage21TimestepProjEmbeddings( + embedding_dim=self.inner_dim + ) + self.txt_in = QwenImage21TextProjection(context_in_dim, self.inner_dim, eps=eps) + self.img_in = nn.Linear( + in_channels * patch_size * patch_size, self.inner_dim, bias=False + ) + + # One shared modulation for every block: [mod1.scale, mod1.gate, mod2.scale, mod2.gate]. + self.modulation = nn.Sequential( + nn.SiLU(), nn.Linear(self.inner_dim, 4 * self.inner_dim, bias=False) + ) + + self.transformer_blocks = nn.ModuleList( + [ + QwenImage21TransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_ratio=mlp_ratio, + eps=eps, + ) + for _ in range(num_layers) + ] + ) + + self.norm_out = QwenImage21AdaLayerNormContinuous( + self.inner_dim, self.inner_dim, eps=eps + ) + self.proj_out = nn.Linear( + self.inner_dim, patch_size * patch_size * self.out_channels, bias=False + ) + + self.gradient_checkpointing = False + + @staticmethod + def build_token_metadata( + image_pad_mask: torch.Tensor, img_shapes: list[tuple[int, int, int]] + ) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Label every token of the joint sequence with the image block it belongs to. + + Block boundaries come from the token counts in `img_shapes`, not from runs of `True` in `image_pad_mask`: two + condition images that happen to sit next to each other with no text between them form one run but must stay + separate blocks, otherwise they would attend to each other bidirectionally. + + Args: + image_pad_mask (`torch.Tensor`): `(seq_len,)` bool, `True` at image-token positions. + img_shapes (`list[tuple[int, int, int]]`): Per-image `(frame, height, width)` in latent tokens, condition + images first and the target image last. + + Returns: + `tuple[torch.Tensor, torch.Tensor]`: `image_ids` `(seq_len,)` with `-1` at text positions and a unique id + per image block, and `target_token_mask` `(seq_len,)` marking the target image's tokens. + """ + image_positions = image_pad_mask.nonzero(as_tuple=True)[0] + block_lengths = [math.prod(shape) for shape in img_shapes] + if sum(block_lengths) != image_positions.numel(): + raise ValueError( + f"img_shapes accounts for {sum(block_lengths)} image tokens but image_pad_mask marks " + f"{image_positions.numel()}." + ) + + image_ids = torch.full_like(image_pad_mask, -1, dtype=torch.long) + block_ids = torch.repeat_interleave( + torch.arange(len(block_lengths), device=image_pad_mask.device), + torch.tensor(block_lengths, device=image_pad_mask.device), + ) + image_ids[image_positions] = block_ids + + target_token_mask = torch.zeros_like(image_pad_mask) + target_token_mask[image_positions[-block_lengths[-1] :]] = True + return image_ids, target_token_mask + + @apply_lora_scale("attention_kwargs") + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + img_shapes: list[list[tuple[int, int, int]]], + img_mask: torch.Tensor, + encoder_hidden_states_mask: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + kv_cache: QwenImage21KVCache | None = None, + kv_cache_mode: str | None = None, + return_dict: bool = True, + ) -> torch.Tensor | Transformer2DModelOutput: + r""" + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`): + Packed latents, condition images first and the target image last. + encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, context_in_dim)`): + Text embeddings from the vision-language encoder. + timestep (`torch.Tensor`): + Current denoising step, scaled to `[0, 1]`. + img_shapes (`list[list[tuple[int, int, int]]]`): + Per-sample list of `(frame, height, width)` in latent tokens, condition images first and the target + image last. All samples must share a layout. + img_mask (`torch.Tensor` of shape `(batch_size, vlm_sequence_length)`): + `True` at the vision-language encoder's image slots, each standing for a `2x2` group of latent tokens. + encoder_hidden_states_mask (`torch.Tensor`, *optional*): + `(batch_size, text_sequence_length)` bool marking valid text tokens. Padded positions are excluded from + attention. + kv_cache (`QwenImage21KVCache`, *optional*): + Cache container. Pass together with `kv_cache_mode` to enable prefix KV caching. + kv_cache_mode (`str`, *optional*): + `"extract"` to prefill the cache (first denoising step), `"cached"` to decode from it (later steps). + Requires `causal_condition=True`. + attention_kwargs (`dict`, *optional*): + Forwarded to the attention processors, and carries `scale` for the LoRA layers. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.modeling_outputs.Transformer2DModelOutput`] instead of a plain + tuple. + + Returns: + [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple`: + [`~models.modeling_outputs.Transformer2DModelOutput`] if `return_dict` is True, otherwise a plain + `tuple` whose first element is the denoised latents. + """ + + batch_size = hidden_states.shape[0] + if kv_cache is not None and not self.config.causal_condition: + raise ValueError( + "kv_cache requires `causal_condition=True`. The cache is only valid because text and condition-image " + "tokens modulate from t=0, which makes their activations independent of the denoising step." + ) + if kv_cache is not None and kv_cache_mode not in ("extract", "cached"): + raise ValueError( + f"kv_cache_mode must be 'extract' or 'cached' when kv_cache is provided, got {kv_cache_mode!r}." + ) + if kv_cache is None and kv_cache_mode is not None: + raise ValueError( + f"kv_cache_mode is {kv_cache_mode!r} but no kv_cache was passed to hold the prefix." + ) + + hidden_states = self.img_in(hidden_states) + encoder_hidden_states = self.txt_in(encoder_hidden_states) + + # Each vision-language image slot stands for 2x2 latent tokens, so expand those positions four-fold and drop + # the actual latents into them. Samples share a layout, hence the single row. + repeats = torch.where(img_mask, _IMG_TOKENS_PER_SLOT, 1)[0] + image_pad_mask = torch.repeat_interleave(img_mask[0], repeats) + + target_tokens = math.prod(img_shapes[0][-1]) + joint_hidden_states = torch.cat( + [ + encoder_hidden_states, + encoder_hidden_states.new_zeros( + batch_size, target_tokens // 4, encoder_hidden_states.shape[2] + ), + ], + dim=1, + ) + joint_hidden_states = joint_hidden_states.repeat_interleave(repeats, dim=1) + joint_hidden_states[:, image_pad_mask] = hidden_states + + rotary_emb = self.pos_embed( + img_shapes[0], image_pad_mask, device=hidden_states.device + ) + image_ids, target_token_mask = self.build_token_metadata( + image_pad_mask, img_shapes[0] + ) + + timestep = timestep.to(hidden_states.dtype) + if self.config.causal_condition: + # Extra t=0 row; text and condition-image tokens modulate from it. `modulation_mask` selects which row + # each token reads, and is `None` when every token shares the sampled timestep. + timestep = torch.cat([timestep, timestep.new_zeros(1)], dim=0) + modulation_mask = target_token_mask + else: + modulation_mask = None + temb = self.time_text_embed(timestep, hidden_states) + modulation = self.modulation(temb) + + # Right-padded prompt positions must never be attended to, on any path. Text positions of the joint sequence + # line up, in order, with the non-image positions of the vision-language sequence — the two are interleaved, + # so the mask cannot be sliced off as a prefix. + joint_key_valid = None + if encoder_hidden_states_mask is not None: + joint_key_valid = torch.ones( + batch_size, + image_pad_mask.shape[0], + dtype=torch.bool, + device=hidden_states.device, + ) + text_positions = (~image_pad_mask).nonzero(as_tuple=True)[0] + vlm_text_positions = ~img_mask[0][: encoder_hidden_states_mask.shape[1]] + joint_key_valid[:, text_positions] = encoder_hidden_states_mask.bool()[ + :, vlm_text_positions + ] + + prefix_len = int((~target_token_mask).sum()) + + if kv_cache_mode == "cached": + # decode: only the target image's queries are recomputed. The block-causal mask degenerates to full + # attention for target rows (they see the entire prefix + their own block), so only the padding mask is + # needed. + joint_hidden_states = joint_hidden_states[:, prefix_len:] + rotary_emb = rotary_emb[prefix_len:] + modulation_mask = modulation_mask[prefix_len:] + attention_mask = ( + None if joint_key_valid is None else joint_key_valid[:, None, None, :] + ) + cache_write_slice = None + block_segments, block_key_valid = None, None + else: + # prefill: the whole joint sequence. The block-causal structure goes down in whichever form the + # installed processors read it — a flex `BlockMask`, per-segment boundaries, or both for a mixed set — + # so neither path pays for building the other's metadata. + processors = [block.attn.processor for block in self.transformer_blocks] + needs_block_mask = any( + isinstance(processor, QwenImage21FlexAttnProcessor) + for processor in processors + ) + attention_mask = ( + build_qwenimage21_block_causal_mask( + image_ids, joint_key_valid, batch_size, hidden_states.device + ) + if needs_block_mask + else None + ) + block_segments = ( + None + if all( + isinstance(processor, QwenImage21FlexAttnProcessor) + for processor in processors + ) + else _qwenimage21_prefix_segments(image_ids, prefix_len) + ) + cache_write_slice = ( + slice(0, prefix_len) if kv_cache_mode == "extract" else None + ) + block_key_valid = joint_key_valid + + for index_block, block in enumerate(self.transformer_blocks): + layer_cache = ( + kv_cache.get_layer(index_block) if kv_cache is not None else None + ) + if torch.is_grad_enabled() and self.gradient_checkpointing: + joint_hidden_states = self._gradient_checkpointing_func( + block, + joint_hidden_states, + modulation, + rotary_emb, + attention_mask, + modulation_mask, + layer_cache, + kv_cache_mode, + cache_write_slice, + block_segments, + block_key_valid, + ) + else: + joint_hidden_states = block( + hidden_states=joint_hidden_states, + modulation=modulation, + rotary_emb=rotary_emb, + attention_mask=attention_mask, + target_token_mask=modulation_mask, + layer_cache=layer_cache, + kv_cache_mode=kv_cache_mode, + cache_write_slice=cache_write_slice, + segments=block_segments, + key_valid=block_key_valid, + ) + + joint_hidden_states = self.norm_out(joint_hidden_states, temb, modulation_mask) + output = self.proj_out(joint_hidden_states) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) diff --git a/extensions_built_in/diffusion_models/qwen_image_2/src/vae.py b/extensions_built_in/diffusion_models/qwen_image_2/src/vae.py new file mode 100644 index 0000000..ba37e80 --- /dev/null +++ b/extensions_built_in/diffusion_models/qwen_image_2/src/vae.py @@ -0,0 +1,1872 @@ +# Copyright 2026 The Qwen Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen-Image 2.1 VAE, vendored from diffusers (Apache 2.0). + +Vendored because `AutoencoderKLQwenImage21` landed in diffusers after the +version this repo pins. RGBA in/out (4 channels), 16x spatial compression, 64 +latent channels; it is a Wan-2.2-shaped stack whose convolutions are 2D, so a +still image rides in a single-frame dim. The only addition is +`OstrisModelMixin` plus the ComfyUI key conversion, since the Comfy-Org VAE +file is in the Wan-native layout with a size-1 temporal conv axis. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import FromOriginalModelMixin +from diffusers.utils import logging +from diffusers.utils.accelerate_utils import apply_forward_hook +from diffusers.models.activations import get_activation +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.autoencoders.vae import ( + AutoencoderMixin, + DecoderOutput, + DiagonalGaussianDistribution, +) + +from toolkit.models.v2._mixin import OstrisModelMixin + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +CACHE_T = 2 + + +class QwenImage21AvgDown3D(nn.Module): + def __init__( + self, + in_channels, + out_channels, + factor_t, + factor_s=1, + ): + super().__init__() + factor = factor_t * factor_s * factor_s + if in_channels * factor % out_channels != 0: + raise ValueError( + f"`in_channels` ({in_channels}) times the downsampling factor ({factor}) must be divisible by " + f"`out_channels` ({out_channels})." + ) + + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = factor + self.group_size = in_channels * factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + B, C, T, H, W = x.shape + x = x.view( + B, + C, + T // self.factor_t, + self.factor_t, + H // self.factor_s, + self.factor_s, + W // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + B, + C * self.factor, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.view( + B, + self.out_channels, + self.group_size, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.mean(dim=2) + return x + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.DupUp3D with DupUp3D->QwenImage21DupUp3D +class QwenImage21DupUp3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1 :, :, :] + return x + + +class QwenImage21CausalConv3d(nn.Conv2d): + r""" + A custom 3D causal convolution layer with feature caching support. + + This layer extends the standard Conv3D layer by ensuring causality in the time dimension and handling feature + caching for efficient inference. + + Args: + in_channels (int): Number of channels in the input image + out_channels (int): Number of channels produced by the convolution + kernel_size (int or tuple): Size of the convolving kernel + stride (int or tuple, optional): Stride of the convolution. Default: 1 + padding (int or tuple, optional): Zero-padding added to all three sides of the input. Default: 0 + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int | int | int], + stride: int | tuple[int | int | int] = 1, + padding: int | tuple[int | int | int] = 0, + ) -> None: + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + ) + + # Set up causal padding + self._padding = ( + self.padding[1], + self.padding[1], + self.padding[0], + self.padding[0], + ) + self.padding = (0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None: + raise ValueError( + "This convolution is the image specialization of Wan's causal 3D one: it folds the single frame away " + "and has no temporal context to prepend, so it cannot take a feature cache." + ) + x = x.squeeze(2) # Remove the temporal dimension + x = F.pad(x, padding) + x = super().forward(x) + x = x.unsqueeze(2) # Add the temporal dimension back + return x + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.WanRMS_norm with Wan->QwenImage21 +class QwenImage21RMS_norm(nn.Module): + r""" + A custom RMS normalization layer. + + Args: + dim (int): The number of dimensions to normalize over. + channel_first (bool, optional): Whether the input tensor has channels as the first dimension. + Default is True. + images (bool, optional): Whether the input represents image data. Default is True. + bias (bool, optional): Whether to include a learnable bias term. Default is False. + """ + + def __init__( + self, + dim: int, + channel_first: bool = True, + images: bool = True, + bias: bool = False, + ) -> None: + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + needs_fp32_normalize = x.dtype in (torch.float16, torch.bfloat16) or any( + t in str(x.dtype) for t in ("float4_", "float8_") + ) + normalized = F.normalize( + x.float() if needs_fp32_normalize else x, + dim=(1 if self.channel_first else -1), + ).to(x.dtype) + + return normalized * self.scale * self.gamma + self.bias + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.WanUpsample with Wan->QwenImage21 +class QwenImage21Upsample(nn.Upsample): + r""" + Perform upsampling while ensuring the output tensor has the same data type as the input. + + Args: + x (torch.Tensor): Input tensor to be upsampled. + + Returns: + torch.Tensor: Upsampled tensor with the same data type as the input. + """ + + def forward(self, x): + if self.mode in ("nearest", "nearest-exact"): + # nearest upsampling is a gather, not arithmetic, so upstream's + # fp32 round-trip is bit-identical to doing it in the input dtype + # -- and it doubles the largest tensor in the decoder, which is + # what makes a one-shot 2048px decode run out of memory + return super().forward(x) + return super().forward(x.float()).type_as(x) + + +class QwenImage21Resample(nn.Module): + r""" + A custom resampling module for 2D and 3D data. + + Args: + dim (int): The number of input/output channels. + mode (str): The resampling mode. Must be one of: + - 'none': No resampling (identity operation). + - 'upsample2d': 2D upsampling with nearest-exact interpolation and convolution. + - 'upsample3d': 3D upsampling with nearest-exact interpolation, convolution, and causal 3D convolution. + - 'downsample2d': 2D downsampling with zero-padding and convolution. + - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution. + """ + + def __init__(self, dim: int, mode: str, upsample_out_dim: int = None) -> None: + super().__init__() + self.dim = dim + self.mode = mode + + # default to dim //2 + if upsample_out_dim is None: + upsample_out_dim = dim // 2 + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + QwenImage21Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, upsample_out_dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + QwenImage21Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, upsample_out_dim, 3, padding=1), + ) + self.time_conv = QwenImage21CausalConv3d( + dim, dim * 2, (1, 1), padding=(0, 0) + ) + + elif mode == "downsample2d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + elif mode == "downsample3d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + self.time_conv = QwenImage21CausalConv3d( + dim, dim, (1, 1), stride=(1, 1), padding=(0, 0) + ) + + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = "Rep" + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if ( + cache_x.shape[2] < 2 + and feat_cache[idx] is not None + and feat_cache[idx] != "Rep" + ): + # cache last frame of last two chunk + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + if ( + cache_x.shape[2] < 2 + and feat_cache[idx] is not None + and feat_cache[idx] == "Rep" + ): + cache_x = torch.cat( + [torch.zeros_like(cache_x).to(cache_x.device), cache_x], + dim=2, + ) + if feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) + x = self.resample(x) + x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv( + torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2) + ) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + +class QwenImage21ResidualBlock(nn.Module): + r""" + A custom residual block module. + + Args: + in_dim (int): Number of input channels. + out_dim (int): Number of output channels. + dropout (float, optional): Dropout rate for the dropout layer. Default is 0.0. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.nonlinearity = get_activation("silu") + + # layers + self.norm1 = QwenImage21RMS_norm(in_dim, images=False) + self.conv1 = QwenImage21CausalConv3d(in_dim, out_dim, 3, padding=1) + self.norm2 = QwenImage21RMS_norm(out_dim, images=False) + self.dropout = nn.Dropout(dropout) + self.conv2 = QwenImage21CausalConv3d(out_dim, out_dim, 3, padding=1) + self.conv_shortcut = ( + QwenImage21CausalConv3d(in_dim, out_dim, 1) + if in_dim != out_dim + else nn.Identity() + ) + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] + # Apply shortcut connection + h = self.conv_shortcut(x) + + # First normalization and activation + x = self.norm1(x) + x = self.nonlinearity(x) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x, + ], + dim=2, + ) + + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + # Second normalization and activation + x = self.norm2(x) + x = self.nonlinearity(x) + + # Dropout + x = self.dropout(x) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x, + ], + dim=2, + ) + + x = self.conv2(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv2(x) + + # Add residual connection + return x + h + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.WanAttentionBlock with Wan->QwenImage21 +class QwenImage21AttentionBlock(nn.Module): + r""" + Causal self-attention with a single head. + + Args: + dim (int): The number of channels in the input tensor. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = QwenImage21RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + def forward(self, x): + identity = x + batch_size, channels, time, height, width = x.size() + + x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, height, width) + x = self.norm(x) + + # compute query, key, value + qkv = self.to_qkv(x) + qkv = qkv.reshape(batch_size * time, 1, channels * 3, -1) + qkv = qkv.permute(0, 1, 3, 2).contiguous() + q, k, v = qkv.chunk(3, dim=-1) + + # apply attention + x = F.scaled_dot_product_attention(q, k, v) + + x = ( + x.squeeze(1) + .permute(0, 2, 1) + .reshape(batch_size * time, channels, height, width) + ) + + # output projection + x = self.proj(x) + + # Reshape back: [(b*t), c, h, w] -> [b, c, t, h, w] + x = x.view(batch_size, time, channels, height, width) + x = x.permute(0, 2, 1, 3, 4) + + return x + identity + + +class QwenImage21MidBlock(nn.Module): + """ + Middle block for QwenVAE encoder and decoder. + + Args: + dim (int): Number of input/output channels. + dropout (float): Dropout rate. + """ + + def __init__(self, dim: int, dropout: float = 0.0, num_layers: int = 1): + super().__init__() + self.dim = dim + + # Create the components + resnets = [QwenImage21ResidualBlock(dim, dim, dropout)] + attentions = [] + for _ in range(num_layers): + attentions.append(QwenImage21AttentionBlock(dim)) + resnets.append(QwenImage21ResidualBlock(dim, dim, dropout)) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] + # First residual block + x = self.resnets[0](x, feat_cache=feat_cache, feat_idx=feat_idx) + + # Process through attention and residual blocks + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + x = attn(x) + + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) + + return x + + +class QwenImage21ResidualDownBlock(nn.Module): + def __init__( + self, + in_dim, + out_dim, + dropout, + num_res_blocks, + temperal_downsample=False, + down_flag=False, + ): + super().__init__() + + # Shortcut path with downsample + self.avg_shortcut = QwenImage21AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + + # Main path with residual blocks and downsample + resnets = [] + for _ in range(num_res_blocks): + resnets.append(QwenImage21ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + self.resnets = nn.ModuleList(resnets) + + # Add the final downsample block + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + self.downsampler = QwenImage21Resample(out_dim, mode=mode) + else: + self.downsampler = None + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] + x_copy = x.clone() + for resnet in self.resnets: + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) + if self.downsampler is not None: + x = self.downsampler(x, feat_cache=feat_cache, feat_idx=feat_idx) + + return x + self.avg_shortcut(x_copy) + + +class QwenImage21Encoder3d(nn.Module): + r""" + A 3D encoder module. + + Args: + dim (int): The base number of channels in the first layer. + z_dim (int): The dimensionality of the latent space. + dim_mult (list of int): Multipliers for the number of channels in each block. + num_res_blocks (int): Number of residual blocks in each block. + attn_scales (list of float): Scales at which to apply attention mechanisms. + temperal_downsample (list of bool): Whether to downsample temporally in each block. + dropout (float): Dropout rate for the dropout layers. + """ + + def __init__( + self, + in_channels: int = 3, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + is_residual: bool = False, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.nonlinearity = get_activation("silu") + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv_in = QwenImage21CausalConv3d(in_channels, dims[0], 3, padding=1) + + # downsample blocks + self.down_blocks = nn.ModuleList([]) + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if is_residual: + self.down_blocks.append( + QwenImage21ResidualDownBlock( + in_dim, + out_dim, + dropout, + num_res_blocks, + temperal_downsample=temperal_downsample[i] + if i != len(dim_mult) - 1 + else False, + down_flag=i != len(dim_mult) - 1, + ) + ) + else: + for _ in range(num_res_blocks): + self.down_blocks.append( + QwenImage21ResidualBlock(in_dim, out_dim, dropout) + ) + if scale in attn_scales: + self.down_blocks.append(QwenImage21AttentionBlock(out_dim)) + in_dim = out_dim + + # downsample block + if i != len(dim_mult) - 1: + mode = "downsample3d" if temperal_downsample[i] else "downsample2d" + self.down_blocks.append(QwenImage21Resample(out_dim, mode=mode)) + scale /= 2.0 + + # middle blocks + self.mid_block = QwenImage21MidBlock(out_dim, dropout, num_layers=1) + + # output blocks + self.norm_out = QwenImage21RMS_norm(out_dim, images=False) + self.conv_out = QwenImage21CausalConv3d(out_dim, z_dim, 3, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv_in(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_in(x) + + ## downsamples + for layer in self.down_blocks: + if feat_cache is not None: + x = layer(x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = layer(x) + + ## middle + x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx) + + ## head + x = self.norm_out(x) + x = self.nonlinearity(x) + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv_out(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_out(x) + + return x + + +class QwenImage21ResidualUpBlock(nn.Module): + """ + A block that handles upsampling for the QwenVAE decoder. + + Args: + in_dim (int): Input dimension + out_dim (int): Output dimension + num_res_blocks (int): Number of residual blocks + dropout (float): Dropout rate + temperal_upsample (bool): Whether to upsample on temporal dimension + up_flag (bool): Whether to upsample or not + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + dropout: float = 0.0, + temperal_upsample: bool = False, + up_flag: bool = False, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + if up_flag: + self.avg_shortcut = QwenImage21DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2, + ) + else: + self.avg_shortcut = None + + # create residual blocks + resnets = [] + current_dim = in_dim + for _ in range(num_res_blocks + 1): + resnets.append(QwenImage21ResidualBlock(current_dim, out_dim, dropout)) + current_dim = out_dim + + self.resnets = nn.ModuleList(resnets) + + # Add upsampling layer if needed + if up_flag: + upsample_mode = "upsample3d" if temperal_upsample else "upsample2d" + self.upsampler = QwenImage21Resample( + out_dim, mode=upsample_mode, upsample_out_dim=out_dim + ) + else: + self.upsampler = None + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False): + if feat_idx is None: + feat_idx = [0] + """ + Forward pass through the upsampling block. + + Args: + x (torch.Tensor): Input tensor + feat_cache (list, optional): Feature cache for causal convolutions + feat_idx (list, optional): Feature index for cache management + + Returns: + torch.Tensor: Output tensor + """ + x_copy = x.clone() + + for resnet in self.resnets: + if feat_cache is not None: + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = resnet(x) + + if self.upsampler is not None: + if feat_cache is not None: + x = self.upsampler(x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = self.upsampler(x) + + if self.avg_shortcut is not None: + x = x + self.avg_shortcut(x_copy, first_chunk=first_chunk) + + return x + + +class QwenImage21UpBlock(nn.Module): + """ + A block that handles upsampling for the QwenVAE decoder. + + Args: + in_dim (int): Input dimension + out_dim (int): Output dimension + num_res_blocks (int): Number of residual blocks + dropout (float): Dropout rate + upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d') + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + dropout: float = 0.0, + upsample_mode: str | None = None, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # Create layers list + resnets = [] + # Add residual blocks and attention if needed + current_dim = in_dim + for _ in range(num_res_blocks + 1): + resnets.append(QwenImage21ResidualBlock(current_dim, out_dim, dropout)) + current_dim = out_dim + + self.resnets = nn.ModuleList(resnets) + + # Add upsampling layer if needed + self.upsamplers = None + if upsample_mode is not None: + self.upsamplers = nn.ModuleList( + [QwenImage21Resample(out_dim, mode=upsample_mode)] + ) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=None): + if feat_idx is None: + feat_idx = [0] + """ + Forward pass through the upsampling block. + + Args: + x (torch.Tensor): Input tensor + feat_cache (list, optional): Feature cache for causal convolutions + feat_idx (list, optional): Feature index for cache management + + Returns: + torch.Tensor: Output tensor + """ + for resnet in self.resnets: + if feat_cache is not None: + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = resnet(x) + + if self.upsamplers is not None: + if feat_cache is not None: + x = self.upsamplers[0](x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = self.upsamplers[0](x) + return x + + +class QwenImage21Decoder3d(nn.Module): + r""" + A 3D decoder module. + + Args: + dim (int): The base number of channels in the first layer. + z_dim (int): The dimensionality of the latent space. + dim_mult (list of int): Multipliers for the number of channels in each block. + num_res_blocks (int): Number of residual blocks in each block. + attn_scales (list of float): Scales at which to apply attention mechanisms. + temperal_upsample (list of bool): Whether to upsample temporally in each block. + dropout (float): Dropout rate for the dropout layers. + """ + + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0, + out_channels: int = 3, + is_residual: bool = False, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + self.nonlinearity = get_activation("silu") + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + + # init block + self.conv_in = QwenImage21CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.mid_block = QwenImage21MidBlock(dims[0], dropout, num_layers=1) + + # upsample blocks + self.up_blocks = nn.ModuleList([]) + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if i > 0 and not is_residual: + in_dim = in_dim // 2 + + # determine if we need upsampling + up_flag = i != len(dim_mult) - 1 + # determine upsampling mode, if not upsampling, set to None + upsample_mode = None + if up_flag and temperal_upsample[i]: + upsample_mode = "upsample3d" + elif up_flag: + upsample_mode = "upsample2d" + # Create and add the upsampling block + if is_residual: + up_block = QwenImage21ResidualUpBlock( + in_dim=in_dim, + out_dim=out_dim, + num_res_blocks=num_res_blocks, + dropout=dropout, + temperal_upsample=temperal_upsample[i] if up_flag else False, + up_flag=up_flag, + ) + else: + up_block = QwenImage21UpBlock( + in_dim=in_dim, + out_dim=out_dim, + num_res_blocks=num_res_blocks, + dropout=dropout, + upsample_mode=upsample_mode, + ) + self.up_blocks.append(up_block) + + # output blocks + self.norm_out = QwenImage21RMS_norm(out_dim, images=False) + self.conv_out = QwenImage21CausalConv3d(out_dim, out_channels, 3, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False): + if feat_idx is None: + feat_idx = [0] + ## conv1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv_in(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_in(x) + + ## middle + x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx) + + ## upsamples + for up_block in self.up_blocks: + x = up_block( + x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk + ) + + ## head + x = self.norm_out(x) + x = self.nonlinearity(x) + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv_out(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_out(x) + return x + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.patchify with patchify->_patchify +def _patchify(x, patch_size): + if patch_size == 1: + return x + + if x.dim() != 5: + raise ValueError(f"Invalid input shape: {x.shape}") + # x shape: [batch_size, channels, frames, height, width] + batch_size, channels, frames, height, width = x.shape + + # Ensure height and width are divisible by patch_size + if height % patch_size != 0 or width % patch_size != 0: + raise ValueError( + f"Height ({height}) and width ({width}) must be divisible by patch_size ({patch_size})" + ) + + # Reshape to [batch_size, channels, frames, height//patch_size, patch_size, width//patch_size, patch_size] + x = x.view( + batch_size, + channels, + frames, + height // patch_size, + patch_size, + width // patch_size, + patch_size, + ) + + # Rearrange to [batch_size, channels * patch_size * patch_size, frames, height//patch_size, width//patch_size] + x = x.permute(0, 1, 6, 4, 2, 3, 5).contiguous() + x = x.view( + batch_size, + channels * patch_size * patch_size, + frames, + height // patch_size, + width // patch_size, + ) + + return x + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.unpatchify with unpatchify->_unpatchify +def _unpatchify(x, patch_size): + if patch_size == 1: + return x + + if x.dim() != 5: + raise ValueError(f"Invalid input shape: {x.shape}") + # x shape: [batch_size, (channels * patch_size * patch_size), frame, height, width] + batch_size, c_patches, frames, height, width = x.shape + channels = c_patches // (patch_size * patch_size) + + # Reshape to [b, c, patch_size, patch_size, f, h, w] + x = x.view(batch_size, channels, patch_size, patch_size, frames, height, width) + + # Rearrange to [b, c, f, h * patch_size, w * patch_size] + x = x.permute(0, 1, 4, 5, 3, 6, 2).contiguous() + x = x.view(batch_size, channels, frames, height * patch_size, width * patch_size) + + return x + + +# Comfy `.residual.` / `shortcut` -> the diffusers resnet submodule. +_COMFY_RESNET_PARTS = { + "residual.0": "norm1", + "residual.2": "conv1", + "residual.3": "norm2", + "residual.6": "conv2", + "shortcut": "conv_shortcut", +} +# Comfy's `middle` Sequential is resnet, attention, resnet. +_COMFY_MID_PARTS = {"0": "resnets.0", "1": "attentions.0", "2": "resnets.1"} + + +def _comfy_resnet_suffix(inner: list) -> list: + """`residual..*` / `shortcut.*` -> the diffusers resnet submodule path.""" + consumed = 2 if inner[0] == "residual" else 1 + return [_COMFY_RESNET_PARTS[".".join(inner[:consumed])]] + inner[consumed:] + + +def _comfy_vae_key(key: str) -> str: + """One ComfyUI VAE parameter name -> its diffusers name.""" + parts = key.split(".") + # the quant convs sit at the checkpoint root + if parts[0] == "conv1": + return ".".join(["quant_conv"] + parts[1:]) + if parts[0] == "conv2": + return ".".join(["post_quant_conv"] + parts[1:]) + + side, rest = parts[0], parts[1:] + if rest[0] == "conv1": + return ".".join([side, "conv_in"] + rest[1:]) + if rest[0] == "head": + # head.0 is the output norm and head.2 the output conv (head.1 is SiLU) + tail = "norm_out" if rest[1] == "0" else "conv_out" + return ".".join([side, tail] + rest[2:]) + if rest[0] == "middle": + part, inner = _COMFY_MID_PARTS[rest[1]], rest[2:] + if part.startswith("resnets"): + inner = _comfy_resnet_suffix(inner) + return ".".join([side, "mid_block", part] + inner) + + # encoder.downsamples..downsamples..* / decoder.upsamples..upsamples..* + block = "down_blocks" if side == "encoder" else "up_blocks" + sampler = "downsampler" if side == "encoder" else "upsampler" + stage, inner_index, inner = rest[1], rest[3], rest[4:] + if inner[0] in ("resample", "time_conv"): + # the stage's last entry is the resampler, not a resnet + return ".".join([side, block, stage, sampler] + inner) + return ".".join( + [side, block, stage, "resnets", inner_index] + _comfy_resnet_suffix(inner) + ) + + +class AutoencoderKLQwenImage21( + ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin, OstrisModelMixin +): + r""" + A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. + Introduced in [Qwen Image 2]. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + """ + + _supports_gradient_checkpointing = False + _group_offload_block_modules = [ + "quant_conv", + "post_quant_conv", + "encoder", + "decoder", + ] + # keys toignore when AlignDeviceHook moves inputs/outputs between devices + # these are shared mutable state modified in-place + _skip_keys = ["feat_cache", "feat_idx"] + + # ---- toolkit loading (OstrisModelMixin) ---- + aitk_subfolder = "vae" + aitk_config_repo = "Qwen/Qwen-Image-2.1" + + aitk_comfy_repo = "Comfy-Org/Qwen-Image-2.1" + _COMFY_FILES = ["vae/qwen_image_2.1_vae_bf16.safetensors"] + aitk_comfy_weight_names = { + "Comfy-Org/Qwen-Image-2.1": _COMFY_FILES, + "Qwen/Qwen-Image-2.1": _COMFY_FILES, + } + + @classmethod + def convert_state_dict_on_load(cls, state_dict): + """Convert the ComfyUI (Wan-native) VAE layout to the diffusers one. + + Comfy runs this VAE through its Wan 2.2 stack with a temporal kernel of + 1, so every conv weight carries a size-1 depth axis the 2D modules here + do not have; names differ throughout as well. Already-diffusers state + dicts pass through untouched. + """ + if "encoder.conv1.weight" not in state_dict: + return state_dict + return { + _comfy_vae_key(key): (value.squeeze(2) if value.ndim == 5 else value) + for key, value in state_dict.items() + } + + @register_to_config + def __init__( + self, + base_dim: int = 96, + decoder_base_dim: int | None = 144, + z_dim: int = 64, + dim_mult: list[int] = [1, 2, 4, 8, 8], + num_res_blocks: int = 2, + attn_scales: list[float] = [], + temperal_downsample: list[bool] = [False, True, True, True], + dropout: float = 0.0, + latents_mean: list[float] = [ + 0.5126, + 0.7721, + -0.0631, + 1.3506, + -0.7855, + -2.1025, + -0.3458, + 1.3722, + 1.8873, + -1.7177, + -0.6510, + 0.2732, + 0.7562, + -0.6163, + -1.0277, + 3.8363, + 2.0210, + 0.0472, + 0.9320, + 2.0087, + 2.4954, + -0.1391, + -1.4249, + 1.8464, + -0.5236, + 1.2826, + 3.7046, + -1.3035, + 2.7286, + -1.4518, + -1.9036, + -1.9955, + -0.0342, + -1.0265, + -0.7636, + 3.0555, + 0.0746, + -3.0751, + -0.1076, + 1.7376, + -1.0914, + -1.9435, + -0.2784, + -1.3680, + 0.4809, + -0.4433, + 0.3764, + 0.5729, + -2.0595, + 1.0960, + -1.3260, + -2.0211, + -5.0179, + 0.5275, + 4.0162, + 1.8505, + 0.3026, + 1.9373, + 1.4937, + 0.2632, + 0.5547, + -1.7121, + -0.1562, + 0.0304, + ], + latents_std: list[float] = [ + 3.2001, + 3.2936, + 3.4321, + 3.0091, + 3.1061, + 4.0379, + 4.0705, + 3.7910, + 3.0785, + 3.6500, + 3.9308, + 3.0904, + 2.8778, + 3.7675, + 3.7320, + 5.0756, + 3.2864, + 4.0397, + 3.1317, + 4.0443, + 2.9249, + 3.9454, + 3.0988, + 4.2489, + 3.4896, + 3.8513, + 3.9323, + 3.4719, + 3.7498, + 4.2830, + 3.5694, + 4.2467, + 3.9037, + 3.2947, + 5.0770, + 3.5075, + 3.2700, + 3.4767, + 2.8063, + 5.1125, + 3.5327, + 4.7833, + 3.1286, + 4.1819, + 3.8527, + 3.8312, + 3.5605, + 4.3875, + 3.9624, + 4.0168, + 3.5643, + 4.0550, + 5.5614, + 4.2963, + 4.4080, + 3.4959, + 3.8747, + 3.7608, + 3.5735, + 3.1490, + 3.7662, + 3.6746, + 3.4563, + 3.8161, + ], + is_residual: bool = True, + in_channels: int = 4, + out_channels: int = 4, + patch_size: int | None = None, + scale_factor_temporal: int | None = 8, + scale_factor_spatial: int | None = 16, + ) -> None: + super().__init__() + + self.z_dim = z_dim + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + if decoder_base_dim is None: + decoder_base_dim = base_dim + + self.encoder = QwenImage21Encoder3d( + in_channels=in_channels, + dim=base_dim, + z_dim=z_dim * 2, + dim_mult=dim_mult, + num_res_blocks=num_res_blocks, + attn_scales=attn_scales, + temperal_downsample=temperal_downsample, + dropout=dropout, + is_residual=is_residual, + ) + self.quant_conv = QwenImage21CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.post_quant_conv = QwenImage21CausalConv3d(z_dim, z_dim, 1) + + self.decoder = QwenImage21Decoder3d( + dim=decoder_base_dim, + z_dim=z_dim, + dim_mult=dim_mult, + num_res_blocks=num_res_blocks, + attn_scales=attn_scales, + temperal_upsample=self.temperal_upsample, + dropout=dropout, + out_channels=out_channels, + is_residual=is_residual, + ) + + self.spatial_compression_ratio = scale_factor_spatial + + # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension + # to perform decoding of a single video latent at a time. + self.use_slicing = False + + # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent + # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the + # intermediate tiles together, the memory requirement can be lowered. + self.use_tiling = False + + # The minimal tile height and width for spatial tiling to be used + self.tile_sample_min_height = 256 + self.tile_sample_min_width = 256 + + # The minimal distance between two spatial tiles + self.tile_sample_stride_height = 192 + self.tile_sample_stride_width = 192 + + # Precompute and cache conv counts for encoder and decoder for clear_cache speedup + self._cached_conv_counts = { + "decoder": sum( + isinstance(m, QwenImage21CausalConv3d) for m in self.decoder.modules() + ) + if self.decoder is not None + else 0, + "encoder": sum( + isinstance(m, QwenImage21CausalConv3d) for m in self.encoder.modules() + ) + if self.encoder is not None + else 0, + } + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.enable_tiling + def enable_tiling( + self, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_stride_height: float | None = None, + tile_sample_stride_width: float | None = None, + ) -> None: + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + + Args: + tile_sample_min_height (`int`, *optional*): + The minimum height required for a sample to be separated into tiles across the height dimension. + tile_sample_min_width (`int`, *optional*): + The minimum width required for a sample to be separated into tiles across the width dimension. + tile_sample_stride_height (`int`, *optional*): + The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are + no tiling artifacts produced across the height dimension. + tile_sample_stride_width (`int`, *optional*): + The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling + artifacts produced across the width dimension. + """ + self.use_tiling = True + self.tile_sample_min_height = ( + tile_sample_min_height or self.tile_sample_min_height + ) + self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width + self.tile_sample_stride_height = ( + tile_sample_stride_height or self.tile_sample_stride_height + ) + self.tile_sample_stride_width = ( + tile_sample_stride_width or self.tile_sample_stride_width + ) + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.clear_cache + def clear_cache(self): + # Use cached conv counts for decoder and encoder to avoid re-iterating modules each call + self._conv_num = self._cached_conv_counts["decoder"] + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + # cache encode + self._enc_conv_num = self._cached_conv_counts["encoder"] + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan._encode with patchify->_patchify + def _encode(self, x: torch.Tensor): + _, _, num_frame, height, width = x.shape + + self.clear_cache() + if self.config.patch_size is not None: + x = _patchify(x, patch_size=self.config.patch_size) + + if self.use_tiling and ( + width > self.tile_sample_min_width or height > self.tile_sample_min_height + ): + return self.tiled_encode(x) + + iter_ = 1 + (num_frame - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], 2) + + enc = self.quant_conv(out) + self.clear_cache() + return enc + + @apply_forward_hook + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.encode + def encode( + self, x: torch.Tensor, return_dict: bool = True + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: + r""" + Encode a batch of images into latents. + + Args: + x (`torch.Tensor`): Input batch of images. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded videos. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self._encode(x) + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan._decode with unpatchify->_unpatchify + def _decode(self, z: torch.Tensor, return_dict: bool = True): + _, _, num_frame, height, width = z.shape + tile_latent_min_height = ( + self.tile_sample_min_height // self.spatial_compression_ratio + ) + tile_latent_min_width = ( + self.tile_sample_min_width // self.spatial_compression_ratio + ) + + if self.use_tiling and ( + width > tile_latent_min_width or height > tile_latent_min_height + ): + return self.tiled_decode(z, return_dict=return_dict) + + self.clear_cache() + x = self.post_quant_conv(z) + for i in range(num_frame): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + x[:, :, i : i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=True, + ) + else: + out_ = self.decoder( + x[:, :, i : i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + ) + out = torch.cat([out, out_], 2) + + if self.config.patch_size is not None: + out = _unpatchify(out, patch_size=self.config.patch_size) + + out = torch.clamp(out, min=-1.0, max=1.0) + + self.clear_cache() + if not return_dict: + return (out,) + + return DecoderOutput(sample=out) + + @apply_forward_hook + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.decode + def decode( + self, z: torch.Tensor, return_dict: bool = True + ) -> DecoderOutput | torch.Tensor: + r""" + Decode a batch of images. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z).sample + + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.blend_v + def blend_v( + self, a: torch.Tensor, b: torch.Tensor, blend_extent: int + ) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * ( + 1 - y / blend_extent + ) + b[:, :, :, y, :] * (y / blend_extent) + return b + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.blend_h + def blend_h( + self, a: torch.Tensor, b: torch.Tensor, blend_extent: int + ) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * ( + 1 - x / blend_extent + ) + b[:, :, :, :, x] * (x / blend_extent) + return b + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.tiled_encode + def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput: + r"""Encode a batch of images using a tiled encoder. + + Args: + x (`torch.Tensor`): Input batch of videos. + + Returns: + `torch.Tensor`: + The latent representation of the encoded videos. + """ + + _, _, num_frames, height, width = x.shape + encode_spatial_compression_ratio = self.spatial_compression_ratio + if self.config.patch_size is not None: + assert encode_spatial_compression_ratio % self.config.patch_size == 0 + encode_spatial_compression_ratio = ( + self.spatial_compression_ratio // self.config.patch_size + ) + + latent_height = height // encode_spatial_compression_ratio + latent_width = width // encode_spatial_compression_ratio + + tile_latent_min_height = ( + self.tile_sample_min_height // encode_spatial_compression_ratio + ) + tile_latent_min_width = ( + self.tile_sample_min_width // encode_spatial_compression_ratio + ) + tile_latent_stride_height = ( + self.tile_sample_stride_height // encode_spatial_compression_ratio + ) + tile_latent_stride_width = ( + self.tile_sample_stride_width // encode_spatial_compression_ratio + ) + + blend_height = tile_latent_min_height - tile_latent_stride_height + blend_width = tile_latent_min_width - tile_latent_stride_width + + # Split x into overlapping tiles and encode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, self.tile_sample_stride_height): + row = [] + for j in range(0, width, self.tile_sample_stride_width): + self.clear_cache() + time = [] + frame_range = 1 + (num_frames - 1) // 4 + for k in range(frame_range): + self._enc_conv_idx = [0] + if k == 0: + tile = x[ + :, + :, + :1, + i : i + self.tile_sample_min_height, + j : j + self.tile_sample_min_width, + ] + else: + tile = x[ + :, + :, + 1 + 4 * (k - 1) : 1 + 4 * k, + i : i + self.tile_sample_min_height, + j : j + self.tile_sample_min_width, + ] + tile = self.encoder( + tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx + ) + tile = self.quant_conv(tile) + time.append(tile) + row.append(torch.cat(time, dim=2)) + rows.append(row) + self.clear_cache() + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + result_row.append( + tile[:, :, :, :tile_latent_stride_height, :tile_latent_stride_width] + ) + result_rows.append(torch.cat(result_row, dim=-1)) + + enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width] + return enc + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.tiled_decode with unpatchify->_unpatchify + def tiled_decode( + self, z: torch.Tensor, return_dict: bool = True + ) -> DecoderOutput | torch.Tensor: + r""" + Decode a batch of images using a tiled decoder. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + _, _, num_frames, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = ( + self.tile_sample_min_height // self.spatial_compression_ratio + ) + tile_latent_min_width = ( + self.tile_sample_min_width // self.spatial_compression_ratio + ) + tile_latent_stride_height = ( + self.tile_sample_stride_height // self.spatial_compression_ratio + ) + tile_latent_stride_width = ( + self.tile_sample_stride_width // self.spatial_compression_ratio + ) + tile_sample_stride_height = self.tile_sample_stride_height + tile_sample_stride_width = self.tile_sample_stride_width + if self.config.patch_size is not None: + sample_height = sample_height // self.config.patch_size + sample_width = sample_width // self.config.patch_size + tile_sample_stride_height = ( + tile_sample_stride_height // self.config.patch_size + ) + tile_sample_stride_width = ( + tile_sample_stride_width // self.config.patch_size + ) + blend_height = ( + self.tile_sample_min_height // self.config.patch_size + - tile_sample_stride_height + ) + blend_width = ( + self.tile_sample_min_width // self.config.patch_size + - tile_sample_stride_width + ) + else: + blend_height = self.tile_sample_min_height - tile_sample_stride_height + blend_width = self.tile_sample_min_width - tile_sample_stride_width + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, tile_latent_stride_height): + row = [] + for j in range(0, width, tile_latent_stride_width): + self.clear_cache() + time = [] + for k in range(num_frames): + self._conv_idx = [0] + tile = z[ + :, + :, + k : k + 1, + i : i + tile_latent_min_height, + j : j + tile_latent_min_width, + ] + tile = self.post_quant_conv(tile) + decoded = self.decoder( + tile, + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=(k == 0), + ) + time.append(decoded) + row.append(torch.cat(time, dim=2)) + rows.append(row) + self.clear_cache() + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + result_row.append( + tile[:, :, :, :tile_sample_stride_height, :tile_sample_stride_width] + ) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width] + + if self.config.patch_size is not None: + dec = _unpatchify(dec, patch_size=self.config.patch_size) + + dec = torch.clamp(dec, min=-1.0, max=1.0) + + if not return_dict: + return (dec,) + return DecoderOutput(sample=dec) + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: + """ + Args: + sample (`torch.Tensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior instead of taking its mode. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + generator (`torch.Generator`, *optional*): + Generator used when `sample_posterior` is `True`. + + Returns: + [`~models.autoencoders.vae.DecoderOutput`] or `tuple`: + [`~models.autoencoders.vae.DecoderOutput`] if `return_dict` is True, otherwise a plain `tuple`. + """ + x = sample + posterior = self.encode(x).latent_dist + + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z, return_dict=return_dict) + return dec diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index ef4c83f..3e7f074 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/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", diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index a928639..108a7c5 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -178,6 +178,13 @@ MODEL_TESTS = { "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 4.0}, "needs_control_image": True, }, + "qwen_image_2": { + "model": {"name_or_path": "Comfy-Org/Qwen-Image-2.1", "quantize": True, "quantize_te": True}, + # 2.1 is meant to be sampled without guidance. One arch: the control + # image exercises the reference path, which plain t2i is a subset of + "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 1.0}, + "needs_control_image": True, + }, # ---- legacy monolith archs (components adopted into v2 on load) ---- "sd1": { "model": {"name_or_path": "stable-diffusion-v1-5/stable-diffusion-v1-5"}, diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index dc16118..0b66ab7 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -672,6 +672,11 @@ class OstrisModelMixin: (and importer-assigned bias) legitimately report as missing keys.""" from toolkit.util.ostris_quant import OstrisLinear + # Anything the importer already materialized is not missing, it is + # done: a convrot embedding becomes a plain nn.Embedding holding the + # unrotated table, so its key is gone from the state dict on purpose. + # Taken before the load, while everything else is still on meta. + preloaded = {n for n, p in model.named_parameters() if not p.is_meta} result = model.load_state_dict(state_dict, assign=True, strict=False) quantized_param_keys = set() for name, m in model.named_modules(): @@ -679,7 +684,11 @@ class OstrisModelMixin: quantized_param_keys.add(f"{name}.weight") if m.bias is not None: quantized_param_keys.add(f"{name}.bias") - bad_missing = [k for k in result.missing_keys if k not in quantized_param_keys] + bad_missing = [ + k + for k in result.missing_keys + if k not in quantized_param_keys and k not in preloaded + ] if bad_missing or result.unexpected_keys: raise ValueError( f"{type(model).__name__} load mismatch: missing {bad_missing[:8]}, " diff --git a/version.py b/version.py index 5c22c0c..3fbbd16 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.13.16" +VERSION = "0.13.17"