diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 6b0952eac..756f4be6a 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -32,6 +32,7 @@ th { | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | +| `QwenImage21Pipeline` | Qwen-Image-2.1 (T2I + image-conditioned generation, up to 4 condition images) | `Qwen/Qwen-Image-2.1` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2.1.md) | | `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — | | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | — | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index a8fb6d972..7eaecc0f0 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -139,6 +139,7 @@ The following tables show which models support each feature: | **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ | | **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | +| **Qwen-Image-2.1** | ❌ | ❌ | ✅ (Ulysses only) | ✅ | ❓ | ❌ | ❓ | ❓ | ❓ | ❓ | ✅ | | **SenseNova-U1 / U1.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **Stable-Diffusion-XL** | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | @@ -154,6 +155,7 @@ The following tables show which models support each feature: > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`. > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution). > 6. BAGEL step execution supports image generation with `bagel.yaml`, `bagel_think.yaml`, and `bagel_single_stage.yaml`; two-stage Thinker execution and explicit single-stage text output remain on their existing complete-request paths. Image requests require `num_inference_steps >= 2`. BAGEL step execution cannot currently be combined with sequence parallelism or a diffusion cache backend; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution). +> 7. Qwen-Image-2.1 (`QwenImage21Pipeline`) handles both text-to-image and image-conditioned generation (up to 4 condition images) and supports LoRA through `QwenImageLoraLoaderMixin`. Sequence parallelism is Ulysses-only (ring attention is not supported), and cache acceleration (TeaCache / Cache-DiT) is unavailable; its built-in KV cache for the text/condition prefix is a separate feature from those cache backends. ### VideoGen diff --git a/examples/offline_inference/image_to_image/qwen_image_21_edit.py b/examples/offline_inference/image_to_image/qwen_image_21_edit.py new file mode 100644 index 000000000..cdc03d39d --- /dev/null +++ b/examples/offline_inference/image_to_image/qwen_image_21_edit.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Example script for image-conditioned generation with Qwen-Image 2.1. + +Qwen-Image 2.1 uses a single pipeline (QwenImage21Pipeline) for both +text-to-image and image-conditioned generation. Pass one or more condition +images (at most 4) via --image to run the image-conditioned path. + +Usage (single image): + python qwen_image_21_edit.py \ + --image input.png \ + --prompt "Let this mascot dance under the moon, surrounded by floating stars" \ + --negative-prompt "blurry, low quality, text, watermark" \ + --output qwen_image_21_edit.png \ + --num-inference-steps 50 \ + --cfg-scale 4.0 + +Usage (multiple images, up to 4): + python qwen_image_21_edit.py \ + --image input1.png input2.png \ + --prompt "Combine these images into a single scene" \ + --negative-prompt "blurry, low quality" \ + --output qwen_image_21_edit.png \ + --num-inference-steps 50 \ + --cfg-scale 4.0 + +Notes: + - Qwen-Image 2.1 only supports true classifier-free guidance + (--cfg-scale); there is no guidance-distilled guidance_scale. + --cfg-scale > 1 only takes effect when --negative-prompt is provided. + - --height/--width must be multiples of 32. When omitted, the output size + is derived from the last condition image's aspect ratio at ~1024x1024. + - Cache acceleration backends (cache_dit / tea_cache) are not supported. + +For more options, run: + python qwen_image_21_edit.py --help +""" + +import argparse +import os +import time +from pathlib import Path + +import torch +from PIL import Image + +from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_extras import ( + build_image_to_image_prompt, + get_model_class_name, +) +from vllm_omni.platforms import current_omni_platform + +MAX_INPUT_IMAGES = 4 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Image-conditioned generation with Qwen-Image 2.1.") + parser.add_argument( + "--model", + default="Qwen/Qwen-Image-2.1", + help="Diffusion model name or local path.", + ) + parser.add_argument( + "--image", + type=str, + nargs="+", + required=True, + help="Path(s) to input condition image file(s) (PNG, JPG, etc.). At most 4 images.", + ) + parser.add_argument( + "--prompt", + type=str, + required=True, + help="Text prompt describing the desired generation.", + ) + parser.add_argument( + "--negative-prompt", + type=str, + default=None, + help="Negative prompt. Required for true classifier-free guidance (--cfg-scale > 1) to take effect.", + ) + parser.add_argument( + "--width", + type=int, + default=None, + metavar="W", + help="Output image width in pixels (multiple of 32). Default: None (derived from the condition image).", + ) + parser.add_argument( + "--height", + type=int, + default=None, + metavar="H", + help="Output image height in pixels (multiple of 32). Default: None (derived from the condition image).", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Random seed for deterministic results.", + ) + parser.add_argument( + "--cfg-scale", + type=float, + default=4.0, + help=( + "True classifier-free guidance scale (default: 4.0). Enabled by setting cfg_scale > 1 and providing " + "a negative prompt. Qwen-Image 2.1 supports true CFG only; there is no guidance_scale." + ), + ) + parser.add_argument( + "--output", + type=str, + default="qwen_image_21_edit.png", + help="Path to save the generated image (PNG).", + ) + parser.add_argument( + "--num-outputs-per-prompt", + type=int, + default=1, + help="Number of images to generate for the given prompt.", + ) + parser.add_argument( + "--num-inference-steps", + type=int, + default=50, + help="Number of denoising steps for the diffusion sampler.", + ) + parser.add_argument( + "--ulysses-degree", + type=int, + default=1, + help="Number of GPUs used for ulysses sequence parallelism (ring attention is not supported).", + ) + parser.add_argument( + "--cfg-parallel-size", + type=int, + default=1, + choices=[1, 2], + help="Number of GPUs used for classifier free guidance parallel size.", + ) + parser.add_argument( + "--enforce-eager", + action="store_true", + default=None, + help="Disable torch.compile and force eager execution.", + ) + return parser.parse_args() + + +def main(): + args = parse_args() + + if len(args.image) > MAX_INPUT_IMAGES: + raise ValueError(f"Qwen-Image 2.1 supports at most {MAX_INPUT_IMAGES} input images, got {len(args.image)}") + + input_images = [] + for image_path in args.image: + if not os.path.exists(image_path): + raise FileNotFoundError(f"Input image not found: {image_path}") + input_images.append(Image.open(image_path).convert("RGB")) + + # Use single image or list based on number of inputs + input_image = input_images[0] if len(input_images) == 1 else input_images + + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) + + omni_kwargs = { + "model": args.model, + "ulysses_degree": args.ulysses_degree, + "cfg_parallel_size": args.cfg_parallel_size, + } + if args.enforce_eager is not None: + omni_kwargs["enforce_eager"] = args.enforce_eager + omni = Omni(**omni_kwargs) + model_class_name = get_model_class_name(omni) + print("Pipeline loaded") + + print(f"\n{'=' * 60}") + print("Generation Configuration:") + print(f" Model: {args.model}") + print(f" Inference steps: {args.num_inference_steps}") + print(f" CFG scale (true CFG): {args.cfg_scale}") + if args.height is not None or args.width is not None: + print(f" Output size: {args.width or 'auto'}x{args.height or 'auto'}") + if isinstance(input_image, list): + print(f" Number of input images: {len(input_image)}") + else: + print(f" Input image size: {input_image.size}") + print( + f" Parallel configuration: ulysses_degree={args.ulysses_degree}, " + f"cfg_parallel_size={args.cfg_parallel_size}" + ) + print(f"{'=' * 60}\n") + + generation_start = time.perf_counter() + + prompt_dict = build_image_to_image_prompt( + model_class_name=model_class_name, + prompt=args.prompt, + negative_prompt=args.negative_prompt, + input_image=input_image, + height=args.height, + width=args.width, + ) + + diffusion_params = OmniDiffusionSamplingParams( + generator=generator, + seed=args.seed, + true_cfg_scale=args.cfg_scale, + num_inference_steps=args.num_inference_steps, + num_outputs_per_prompt=args.num_outputs_per_prompt, + height=args.height, + width=args.width, + ) + + outputs = omni.generate(prompt_dict, sampling_params_list=[diffusion_params]) + generation_end = time.perf_counter() + generation_time = generation_end - generation_start + + print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)") + + if not outputs: + raise ValueError("No output generated from omni.generate()") + + images = None + for output in outputs: + images = getattr(output, "images", None) + if images: + break + + if not images: + images = extract_images_from_outputs(outputs) + + if not images: + raise ValueError("No images found in request_output") + + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + suffix = output_path.suffix or ".png" + stem = output_path.stem or "qwen_image_21_edit" + + if len(images) <= 1: + images[0].save(output_path) + print(f"Saved generated image to {os.path.abspath(output_path)}") + else: + for idx, img in enumerate(images): + save_path = output_path.parent / f"{stem}_{idx}{suffix}" + img.save(save_path) + print(f"Saved generated image to {os.path.abspath(save_path)}") + + +if __name__ == "__main__": + main() diff --git a/examples/offline_inference/image_to_image/run_qwen_image_21_edit.sh b/examples/offline_inference/image_to_image/run_qwen_image_21_edit.sh new file mode 100644 index 000000000..6ed48ab2a --- /dev/null +++ b/examples/offline_inference/image_to_image/run_qwen_image_21_edit.sh @@ -0,0 +1,8 @@ +python qwen_image_21_edit.py \ + --model Qwen/Qwen-Image-2.1 \ + --image qwen_bear.png \ + --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \ + --negative-prompt "blurry, low quality, text, watermark" \ + --output qwen_image_21_edit.png \ + --num-inference-steps 50 \ + --cfg-scale 4.0 diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 2ee09e178..bc7ccbf31 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -86,7 +86,8 @@ def parse_args() -> argparse.Namespace: "--model", default="Qwen/Qwen-Image", help="Diffusion model name or local path. Supported models: " - "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1, " + "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, Qwen/Qwen-Image-2.1, " + "stepfun-ai/NextStep-1.1, " "black-forest-labs/FLUX.1-dev, black-forest-labs/FLUX.2-klein-9B, " "black-forest-labs/FLUX.2-dev, tencent/HunyuanImage-3.0-Instruct, " "HiDream-ai/HiDream-O1-Image, meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, " diff --git a/recipes/Qwen/Qwen-Image-2.1.md b/recipes/Qwen/Qwen-Image-2.1.md new file mode 100644 index 000000000..4c7aa441b --- /dev/null +++ b/recipes/Qwen/Qwen-Image-2.1.md @@ -0,0 +1,148 @@ +# Qwen-Image-2.1 + +> Text-to-image and image-conditioned generation with Qwen-Image 2.1 + +## Summary + +- Vendor: Qwen +- Model: `Qwen/Qwen-Image-2.1` +- Task: Text-to-image generation and image-conditioned generation (editing) +- Mode: Offline inference and online serving with optional step-wise execution +- Maintainer: Community + +> **Model id:** `Qwen/Qwen-Image-2.1` follows the reference diffusers +> documentation and may change before the official release — the final +> checkpoint name is subject to the official release. + +## When to use this recipe + +Use this recipe as a starting point for running `Qwen/Qwen-Image-2.1`. A single +pipeline class, `QwenImage21Pipeline`, handles both pure text-to-image requests +(no condition image) and image-conditioned requests with up to 4 condition +images: prompt and condition images are encoded together by a Qwen3-VL text +encoder, so editing-style requests go through the same serving path as +text-to-image. + +## References + +- Related offline examples: + [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py), + [`examples/offline_inference/image_to_image/qwen_image_21_edit.py`](../../examples/offline_inference/image_to_image/qwen_image_21_edit.py) +- Supported models table: + [`docs/models/supported_models.md`](../../docs/models/supported_models.md) +- Feature compatibility matrix: + [`docs/user_guide/diffusion_features.md`](../../docs/user_guide/diffusion_features.md) +- Sibling recipes: [`Qwen-Image.md`](./Qwen-Image.md), + [`Qwen-Image-Edit.md`](./Qwen-Image-Edit.md) + +## Offline Inference + +### Text-to-image + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model Qwen/Qwen-Image-2.1 \ + --prompt "A ceramic teapot on a wooden table" \ + --negative-prompt "blurry, low quality, text, watermark" \ + --output qwen_image_21_t2i.png \ + --num-inference-steps 50 \ + --cfg-scale 4.0 +``` + +### Image-conditioned generation (editing) + +```bash +python examples/offline_inference/image_to_image/qwen_image_21_edit.py \ + --image qwen_bear.png \ + --prompt "Let this mascot dance under the moon, surrounded by floating stars" \ + --negative-prompt "blurry, low quality, text, watermark" \ + --output qwen_image_21_edit.png \ + --num-inference-steps 50 \ + --cfg-scale 4.0 +``` + +or via the bundled launcher: + +```bash +bash examples/offline_inference/image_to_image/run_qwen_image_21_edit.sh +``` + +Up to 4 condition images can be passed via `--image`: + +```bash +python examples/offline_inference/image_to_image/qwen_image_21_edit.py \ + --image input1.png input2.png \ + --prompt "Combine these images into a single scene" \ + --negative-prompt "blurry, low quality" \ + --output qwen_image_21_multi.png \ + --num-inference-steps 50 \ + --cfg-scale 4.0 +``` + +## Online Serving + +Start the server: + +```bash +vllm serve Qwen/Qwen-Image-2.1 --omni --port 8091 +``` + +To enable the step-wise runtime: + +```bash +vllm serve Qwen/Qwen-Image-2.1 --omni \ + --port 8091 \ + --step-execution \ + --max-num-seqs 8 +``` + +### Verification + +For a direct API smoke test: + +```bash +curl http://localhost:8091/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Qwen/Qwen-Image-2.1", + "prompt": "A ceramic teapot on a wooden table", + "negative_prompt": "blurry, low quality, text, watermark", + "size": "1024x1024", + "num_inference_steps": 50, + "true_cfg_scale": 4.0, + "seed": 42 + }' +``` + +For image-conditioned generation, pass condition images through the +OpenAI-compatible multimodal endpoint (`/v1/chat/completions` with +`image_url` content parts, or `/v1/images/edits`), same as +[`Qwen-Image-Edit.md`](./Qwen-Image-Edit.md). + +## Parameters + +| Parameter | Default | Notes | +| --- | --- | --- | +| `true_cfg_scale` | 4.0 | Only takes effect when `negative_prompt` is provided. Qwen-Image 2.1 supports true CFG only — there is no `guidance_scale` parameter. | +| `num_inference_steps` | 50 | Number of denoising steps. | +| `height` / `width` | 1024×1024 | Must be multiples of 32. For image-conditioned requests, leaving them unset derives the output size from the last condition image's aspect ratio at ~1024×1024. | +| `negative_prompt` | None | Required to activate true classifier-free guidance (`true_cfg_scale > 1`). | +| condition images | — | Up to 4 input images per request; more than 4 raises an error. | +| `seed` | — | Fix for reproducible outputs. | + +### Prefix KV cache + +The transformer caches the text/condition-image prefix across denoising steps +(a prefill/decode split keyed by CFG branch). The cache is enabled +automatically when the checkpoint declares `causal_condition: true` and is +re-created per generation, so no state leaks between requests. It changes the +compute path — cached prefix vs. full recompute at every step — so exact +numerical reproduction against a no-cache reference (e.g. the diffusers +pipeline) requires matching the cache setting on both sides. + +## Known Limitations + +- Cache acceleration backends (`cache_dit`, `tea_cache`) are not supported; + `QwenImage21Pipeline` is registered in `_NO_CACHE_ACCELERATION`. +- Sequence parallelism supports Ulysses only; ring attention is not supported. +- Only true CFG is exposed; there is no `guidance_scale` knob. diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py new file mode 100644 index 000000000..f0037186c --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import torch +from diffusers.models.autoencoders.vae import DecoderOutput +from vllm.logger import init_logger + +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedOperator, + DistributedVaeMixin, + GridSpec, + TileTask, +) +from vllm_omni.diffusion.models.qwen_image_21.autoencoder_kl_qwenimage21 import ( + AutoencoderKLQwenImage21, + _unpatchify, +) + +logger = init_logger(__name__) + + +class DistributedAutoencoderKLQwenImage21(AutoencoderKLQwenImage21, DistributedVaeMixin): + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model + + def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + # mostly copy from AutoencoderKLQwenImage21.tiled_decode + _, _, 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. + tiletask_list = [] + for i in range(0, height, tile_latent_stride_height): + for j in range(0, width, tile_latent_stride_width): + time_list = [] + for k in range(num_frames): + tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width] + time_list.append(tile) + tiletask_list.append( + TileTask( + len(tiletask_list), + (i // tile_latent_stride_height, j // tile_latent_stride_width), + time_list, + workload=time_list[0].shape[3] * time_list[0].shape[4], + ) + ) + tile_spec = { + "sample_height": sample_height, + "sample_width": sample_width, + "blend_height": blend_height, + "blend_width": blend_width, + "tile_sample_stride_height": tile_sample_stride_height, + "tile_sample_stride_width": tile_sample_stride_width, + } + grid_spec = GridSpec( + split_dims=(3, 4), + grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + tile_spec=tile_spec, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + + def tile_exec(self, task: TileTask) -> torch.Tensor: + """Decode a single latent tile into RGB space.""" + self.clear_cache() + time = [] + for k in range(len(task.tensor)): + self._conv_idx = [0] + tile = self.post_quant_conv(task.tensor[k]) + decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0)) + time.append(decoded) + result = torch.cat(time, dim=2) + return result + + def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor: + """Merge decoded tiles into a full image.""" + grid_h, grid_w = grid_spec.grid_shape + self.clear_cache() + + result_rows = [] + for i in range(grid_h): + result_row = [] + for j in range(grid_w): + tile = coord_tensor_map[(i, j)] + if i > 0: + tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"]) + if j > 0: + tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"]) + result_row.append( + tile[ + :, + :, + :, + : grid_spec.tile_spec["tile_sample_stride_height"], + : grid_spec.tile_spec["tile_sample_stride_width"], + ] + ) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=3)[ + :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["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) + return dec + + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True): + if not self.is_distributed_enabled(): + return super().tiled_decode(z, return_dict=return_dict) + + logger.debug("Decode running with distributed executor") + result = self.distributed_executor.execute( + z, + DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge), + broadcast_result=True, + ) + if not return_dict: + return (result,) + + return DecoderOutput(sample=result) diff --git a/vllm_omni/diffusion/lora/loader.py b/vllm_omni/diffusion/lora/loader.py index a13363c7a..8c6713d58 100644 --- a/vllm_omni/diffusion/lora/loader.py +++ b/vllm_omni/diffusion/lora/loader.py @@ -22,6 +22,11 @@ lora_convert_mapping: dict[str, Callable] = { "QwenImagePipeline": _convert_non_diffusers_qwen_lora_to_diffusers, "QwenImageEditPipeline": _convert_non_diffusers_qwen_lora_to_diffusers, "QwenImageEditPlusPipeline": _convert_non_diffusers_qwen_lora_to_diffusers, + # 2.1's single-stream blocks keep the same leaf names the converter's + # protected n-grams cover (`attn.to_q/to_k/to_v`, `img_mlp`, ...), so the + # 2.0 converter applies. Non-diffusers LoRAs that target the fused `to_qkv` + # projection are not supported (the converter would split it into `to.qkv`). + "QwenImage21Pipeline": _convert_non_diffusers_qwen_lora_to_diffusers, "Wan22Pipeline": _convert_non_diffusers_wan_lora_to_diffusers, "Wan22I2VPipeline": _convert_non_diffusers_wan_lora_to_diffusers, } diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index c91d2ead5..329ce44dd 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -20,6 +20,8 @@ class DiffusionModelMetadata: QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4 +# Qwen-Image 2.1 image-conditioned generation caps condition images at 4. +QWEN_IMAGE_21_MAX_INPUT_IMAGES = 4 # Upstream HunyuanImage-3.0 "Multi-Image Fusion" caps reference images at 3. HUNYUAN_IMAGE3_MAX_INPUT_IMAGES = 3 # Boogu-Image editing (TI2I) supports a single reference image for now. @@ -31,6 +33,12 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { supports_multimodal_inputs=True, max_multimodal_image_inputs=QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES, ), + # Qwen-Image 2.1 handles both text-to-image (no image) and + # image-conditioned requests through the same pipeline class. + "QwenImage21Pipeline": DiffusionModelMetadata( + supports_multimodal_inputs=True, + max_multimodal_image_inputs=QWEN_IMAGE_21_MAX_INPUT_IMAGES, + ), "HunyuanImage3Pipeline": DiffusionModelMetadata( supports_multimodal_inputs=True, max_multimodal_image_inputs=HUNYUAN_IMAGE3_MAX_INPUT_IMAGES, diff --git a/vllm_omni/diffusion/models/qwen_image_21/__init__.py b/vllm_omni/diffusion/models/qwen_image_21/__init__.py new file mode 100644 index 000000000..ec77cf534 --- /dev/null +++ b/vllm_omni/diffusion/models/qwen_image_21/__init__.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen Image 2.1 diffusion model components.""" + +from vllm_omni.diffusion.models.qwen_image_21.autoencoder_kl_qwenimage21 import ( + AutoencoderKLQwenImage21, +) +from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import ( + QwenImage21Pipeline, + get_qwen_image_21_post_process_func, + get_qwen_image_21_pre_process_func, +) +from vllm_omni.diffusion.models.qwen_image_21.qwen_image_21_transformer import ( + QwenImage21Transformer2DModel, +) + +__all__ = [ + "AutoencoderKLQwenImage21", + "QwenImage21Pipeline", + "QwenImage21Transformer2DModel", + "get_qwen_image_21_post_process_func", + "get_qwen_image_21_pre_process_func", +] diff --git a/vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py b/vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py new file mode 100644 index 000000000..0feba0fc9 --- /dev/null +++ b/vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py @@ -0,0 +1,1519 @@ +# 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. + +# Copied from diffusers to avoid version coupling. + + +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.models.activations import get_activation +from diffusers.models.autoencoders.vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils import logging +from diffusers.utils.accelerate_utils import apply_forward_hook + +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__() + 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 in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.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 + + +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) + assert cache_x is None + 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 + + +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): + return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias + + +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): + 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=[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. + non_linearity (str, optional): Type of non-linearity to use. Default is "silu". + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + dropout: float = 0.0, + non_linearity: str = "silu", + ) -> None: + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.nonlinearity = get_activation(non_linearity) + + # 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=[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 + + +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. + non_linearity (str): Type of non-linearity to use. + """ + + def __init__(self, dim: int, dropout: float = 0.0, non_linearity: str = "silu", num_layers: int = 1): + super().__init__() + self.dim = dim + + # Create the components + resnets = [QwenImage21ResidualBlock(dim, dim, dropout, non_linearity)] + attentions = [] + for _ in range(num_layers): + attentions.append(QwenImage21AttentionBlock(dim)) + resnets.append(QwenImage21ResidualBlock(dim, dim, dropout, non_linearity)) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=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=[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. + non_linearity (str): Type of non-linearity to use. + """ + + 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, + non_linearity: str = "silu", + 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(non_linearity) + + # 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, non_linearity, 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=[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 + non_linearity (str): Type of non-linearity to use + """ + + 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, + non_linearity: str = "silu", + ): + 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, non_linearity)) + 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=[0], first_chunk=False): + """ + 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') + non_linearity (str): Type of non-linearity to use + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + dropout: float = 0.0, + upsample_mode: str | None = None, + non_linearity: str = "silu", + ): + 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, non_linearity)) + 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=[0], first_chunk=None): + """ + 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. + non_linearity (str): Type of non-linearity to use. + """ + + 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, + non_linearity: str = "silu", + 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(non_linearity) + + # 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, non_linearity, 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, + non_linearity=non_linearity, + ) + else: + up_block = QwenImage21UpBlock( + in_dim=in_dim, + out_dim=out_dim, + num_res_blocks=num_res_blocks, + dropout=dropout, + upsample_mode=upsample_mode, + non_linearity=non_linearity, + ) + 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=[0], first_chunk=False): + ## 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 + + +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 + + +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 + + +class AutoencoderKLQwenImage21(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): + 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"] + + @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 = 3, + out_channels: int = 3, + patch_size: int | None = None, + scale_factor_temporal: int | None = 8, + scale_factor_spatial: int | None = 8, + ) -> 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, + } + + 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 + + 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 + + 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 + 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) + + 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 + 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) + + 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 + + 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 + + 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 + + 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. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of 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/vllm_omni/diffusion/models/qwen_image_21/cfg_parallel.py b/vllm_omni/diffusion/models/qwen_image_21/cfg_parallel.py new file mode 100644 index 000000000..9635eaf30 --- /dev/null +++ b/vllm_omni/diffusion/models/qwen_image_21/cfg_parallel.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CFG Parallel Mixin for Qwen-Image 2.1. + +Differences from the 2.0 mixin (``qwen_image/cfg_parallel.py``): +- The transformer owns a KV cache for the text/condition-image prefix, so the + denoise loop passes a fresh per-generation ``kv_cache`` and names the CFG + branch (``cache_branch="cond"/"uncond"``) on every forward. +- The joint sequence is ``[condition images, ..., target image]`` — the target + noise prediction is the *tail* of the output, not the head. +""" + +import logging +from typing import Any + +import torch + +from vllm_omni.diffusion.distributed.cfg_parallel import ( + CFGParallelMixin, + _get_cfg_world_size_or_one, + _unwrap, + _wrap, +) +from vllm_omni.diffusion.distributed.parallel_state import ( + get_cfg_group, + get_classifier_free_guidance_rank, +) +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin + +logger = logging.getLogger(__name__) + + +class QwenImage21CFGParallelMixin(CFGParallelMixin, ProgressBarMixin): + """Base Mixin class for the Qwen-Image 2.1 pipeline providing shared CFG methods.""" + + def predict_noise_maybe_with_cfg( + self, + do_true_cfg: bool, + true_cfg_scale: float, + positive_kwargs: dict[str, Any], + negative_kwargs: dict[str, Any] | None, + cfg_normalize: bool = False, + output_slice: int | None = None, + kwargs: dict[str, Any] | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: + # The base implementation slices `[:, :output_slice]` per branch, which + # matches the 2.0 condition-last layout. 2.1 concatenates condition + # latents *before* the target, so the target prediction is the tail of + # the joint output. Each branch must be tail-sliced *before* the CFG + # combine: positive and negative prompts tokenize to different lengths, + # so the full joint outputs do not align positionally. + def tail_slice(pred: torch.Tensor | tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: + if output_slice is None: + return _wrap(pred) + return tuple(p[:, -output_slice:] for p in _wrap(pred)) + + if do_true_cfg: + if _get_cfg_world_size_or_one() > 1: + cfg_group = get_cfg_group() + cfg_rank = get_classifier_free_guidance_rank() + branch_kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs + local_pred = tail_slice(self.predict_noise(**branch_kwargs)) + gathered = [cfg_group.all_gather(p, separate_tensors=True) for p in local_pred] + positive_noise_pred = tuple(g[0] for g in gathered) + negative_noise_pred = tuple(g[1] for g in gathered) + else: + positive_noise_pred = tail_slice(self.predict_noise(**positive_kwargs)) + negative_noise_pred = tail_slice(self.predict_noise(**negative_kwargs)) + return self.combine_cfg_noise( + positive_noise_pred, + negative_noise_pred, + true_cfg_scale, + cfg_normalize, + **({} if kwargs is None else {"kwargs": kwargs}), + ) + + pred = self.predict_noise(**positive_kwargs) + if output_slice is not None: + pred = _unwrap(tail_slice(pred)) + return pred + + def diffuse( + self, + prompt_embeds: torch.Tensor, + prompt_embeds_mask: torch.Tensor, + negative_prompt_embeds: torch.Tensor | None, + negative_prompt_embeds_mask: torch.Tensor | None, + img_mask: torch.Tensor, + negative_img_mask: torch.Tensor | None, + latents: torch.Tensor, + image_latents: torch.Tensor | None, + img_shapes: list, + timesteps: torch.Tensor, + do_true_cfg: bool, + true_cfg_scale: float, + attention_kwargs: dict[str, Any] | None = None, + additional_transformer_kwargs: dict[str, Any] | None = None, + ) -> torch.Tensor: + """ + Diffusion loop with optional classifier-free guidance. + + Args: + prompt_embeds: Positive prompt embeddings + prompt_embeds_mask: Mask for positive prompt + negative_prompt_embeds: Negative prompt embeddings + negative_prompt_embeds_mask: Mask for negative prompt + img_mask: Joint-sequence image-slot mask for the positive branch + (VLM image slots plus one slot per 2x2 group of target tokens) + negative_img_mask: Same for the negative branch + latents: Noise latents to denoise (target image only, packed) + image_latents: Packed condition-image latents, prepended to the + noise latents (default: None for pure text-to-image) + img_shapes: Per-sample list of (frame, height, width) latent-token + shapes, condition images first and the target last + timesteps: Diffusion timesteps + do_true_cfg: Whether to apply CFG + true_cfg_scale: CFG scale factor + attention_kwargs: Passed through to the transformer (attn_path hint) + additional_transformer_kwargs: Extra kwargs for the transformer + + Returns: + Denoised latents + """ + # Fresh cache for every generation: the pipeline is a singleton shared + # across requests, so cache state must never leak between them. + cache_enabled = getattr(self.transformer, "causal_condition", False) + kv_cache = [{} for _ in range(len(self.transformer.transformer_blocks))] if cache_enabled else None + self.scheduler.set_begin_index(0) + self.transformer.do_true_cfg = do_true_cfg + additional_transformer_kwargs = additional_transformer_kwargs or {} + + with self.progress_bar(total=len(timesteps)) as pbar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + self._current_timestep = t + + # Broadcast timestep to match batch size + timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype) + + # Condition image latents go in front of the noise latents. + latent_model_input = latents + if image_latents is not None: + latent_model_input = torch.cat([image_latents, latents], dim=1) + + positive_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep / 1000, + "encoder_hidden_states_mask": prompt_embeds_mask, + "encoder_hidden_states": prompt_embeds, + "img_shapes": img_shapes, + "img_mask": img_mask, + "attention_kwargs": attention_kwargs, + "kv_cache": kv_cache, + "cache_branch": "cond", + **additional_transformer_kwargs, + } + if do_true_cfg: + negative_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep / 1000, + "encoder_hidden_states_mask": negative_prompt_embeds_mask, + "encoder_hidden_states": negative_prompt_embeds, + "img_shapes": img_shapes, + "img_mask": negative_img_mask, + "attention_kwargs": attention_kwargs, + "kv_cache": kv_cache, + "cache_branch": "uncond", + **additional_transformer_kwargs, + } + else: + negative_kwargs = None + + # Prefill (first step) returns the full joint sequence and + # caches the prefix; decode steps return target tokens only. + # Slicing the target tail is correct for both. + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg, + true_cfg_scale, + positive_kwargs, + negative_kwargs, + cfg_normalize=False, + output_slice=latents.size(1), + ) + + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + + pbar.update() + + return latents diff --git a/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py b/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py new file mode 100644 index 000000000..8119da80b --- /dev/null +++ b/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py @@ -0,0 +1,1296 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# Adapted from diffusers' pipeline_qwenimage21.py (Qwen-Image 2.1 T2I + image-conditioned generation). + +import copy +import json +import logging +import os +import random +from collections.abc import Iterable +from typing import TYPE_CHECKING, Any, ClassVar + +import numpy as np +import PIL.Image +import torch +import torch.nn.functional as F +from diffusers.image_processor import VaeImageProcessor +from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage21 import ( + DistributedAutoencoderKLQwenImage21, +) +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.lora.loader import QwenImageLoraLoaderMixin +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders +from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_21_MAX_INPUT_IMAGES +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import ( + calculate_shift, + retrieve_timesteps, +) +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import ( + calculate_dimensions, + retrieve_latents, +) +from vllm_omni.diffusion.models.qwen_image_21.cfg_parallel import QwenImage21CFGParallelMixin +from vllm_omni.diffusion.models.qwen_image_21.qwen_image_21_transformer import ( + QwenImage21Transformer2DModel, +) +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.prompt_utils import ( + validate_prompt_sequence_lengths, +) +from vllm_omni.diffusion.utils.size_utils import ( + normalize_min_aligned_size, +) +from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request +from vllm_omni.inputs.data import OmniTextPrompt + +if TYPE_CHECKING: + from vllm_omni.diffusion.worker.input_batch import InputBatch + from vllm_omni.diffusion.worker.utils import StepRequestState + +from vllm_omni.model_executor.model_loader.weight_utils import ( + download_weights_from_hf_specific, +) + +logger = logging.getLogger(__name__) + +# Target area (in pixels) used to derive the output size from a condition +# image's aspect ratio and to resize condition images before encoding. +OUTPUT_RESOLUTION = 1024 + +# Re-export the shared metadata value locally so this pipeline keeps a nearby, +# descriptive constant for validation without becoming the source of truth. +MAX_QWEN_IMAGE_21_INPUT_IMAGES = QWEN_IMAGE_21_MAX_INPUT_IMAGES + + +def _read_vae_scale_factor(od_config: OmniDiffusionConfig) -> int: + """Read the VAE spatial compression ratio from the checkpoint's vae/config.json.""" + model_name = od_config.model + if os.path.exists(model_name): + model_path = model_name + else: + model_path = download_weights_from_hf_specific(model_name, None, ["*"]) + vae_config_path = os.path.join(model_path, "vae/config.json") + with open(vae_config_path) as f: + vae_config = json.load(f) + if vae_config.get("scale_factor_spatial"): + return int(vae_config["scale_factor_spatial"]) + # 2.1 compresses 16x spatially (four downsampling stages). + if "temperal_downsample" in vae_config: + return 2 ** len(vae_config["temperal_downsample"]) + return 16 + + +def get_qwen_image_21_pre_process_func( + od_config: OmniDiffusionConfig, +): + """Pre-processing function for QwenImage21Pipeline. + + Reads condition images from ``request.prompt["multi_modal_data"]["image"]`` + (single image or list), resizes them once for both the vision-language + prompt encoder and the VAE, and stores the results under + ``prompt["additional_information"]``. Pure text-to-image requests (no + image) pass through unchanged. + """ + vae_scale_factor = _read_vae_scale_factor(od_config) + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor, do_convert_rgb=True) + # The Qwen-Image 2.1 VAE is RGBA (in/out_channels=4): condition images get an + # opaque alpha channel unless the input already carries one. + vae_image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor) + + def _to_pil(image: Any) -> PIL.Image.Image: + if isinstance(image, PIL.Image.Image): + return image + if isinstance(image, str): + return PIL.Image.open(image) + if isinstance(image, torch.Tensor): + arr = image.detach().cpu() + if arr.dim() == 3 and arr.shape[0] in (1, 3, 4): + arr = arr.permute(1, 2, 0) + arr = (arr * 255).round().to(torch.uint8) if arr.dtype.is_floating_point else arr + arr = arr.numpy() + if arr.shape[-1] == 1: + arr = arr.squeeze(-1) + return PIL.Image.fromarray(arr) + if isinstance(image, np.ndarray): + return PIL.Image.fromarray(image) + raise TypeError(f"Unsupported condition image type: {type(image)}") + + def pre_process_func( + request: OmniDiffusionRequest, + ): + """Pre-process requests for QwenImage21Pipeline.""" + prompt = request.prompt + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt) + if "additional_information" not in prompt: + prompt["additional_information"] = {} + + if not raw_image: # None or empty list: pure text-to-image request + request.prompt = prompt + return request + + if not isinstance(raw_image, list): + raw_image = [raw_image] + if len(raw_image) > MAX_QWEN_IMAGE_21_INPUT_IMAGES: + raise ValueError( + f"Received {len(raw_image)} input images. " + f"At most {MAX_QWEN_IMAGE_21_INPUT_IMAGES} images are supported by this model." + ) + images = [_to_pil(im) for im in raw_image] + + # One resize feeds both the text encoder and the VAE. + input_image_sizes = [] + prompt_images = [] + vae_images = [] + for img in images: + image_width, image_height = img.size + input_width, input_height = calculate_dimensions( + OUTPUT_RESOLUTION * OUTPUT_RESOLUTION, image_width / image_height + ) + input_image_sizes.append((input_width, input_height)) + prompt_images.append(image_processor.resize(img, height=input_height, width=input_width)) + vae_images.append( + vae_image_processor.preprocess(img.convert("RGBA"), height=input_height, width=input_width).unsqueeze(2) + ) + + # The generated image derives its aspect ratio from the last condition image. + last_width, last_height = images[-1].size + calculated_width, calculated_height = calculate_dimensions( + OUTPUT_RESOLUTION * OUTPUT_RESOLUTION, last_width / last_height + ) + height = request.sampling_params.height or calculated_height + width = request.sampling_params.width or calculated_width + height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2) + request.sampling_params.height = height + request.sampling_params.width = width + + prompt["additional_information"]["prompt_image"] = prompt_images + prompt["additional_information"]["vae_images"] = vae_images + prompt["additional_information"]["input_image_sizes"] = input_image_sizes + prompt["additional_information"]["calculated_height"] = calculated_height + prompt["additional_information"]["calculated_width"] = calculated_width + request.prompt = prompt + return request + + return pre_process_func + + +def get_qwen_image_21_post_process_func( + od_config: OmniDiffusionConfig, +): + """Post-processing function for QwenImage21Pipeline.""" + vae_scale_factor = _read_vae_scale_factor(od_config) + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor) + + def post_process_func( + images: torch.Tensor | dict[str, Any], + ): + if isinstance(images, dict) and isinstance(images.get("payload"), dict): + payload = dict(images["payload"]) + image_payload = payload.get("image") + if image_payload is None: + raise ValueError("Qwen-Image 2.1 postprocess expected payload['image'] in output envelope.") + payload["image"] = image_processor.postprocess(image_payload) + metadata = images.get("metadata") or {} + return { + "payload": payload, + "metadata": metadata if isinstance(metadata, dict) else {}, + } + return image_processor.postprocess(images) + + return post_process_func + + +class QwenImage21Pipeline( + nn.Module, + QwenImage21CFGParallelMixin, + DiffusionPipelineProfilerMixin, + SupportsComponentDiscovery, + QwenImageLoraLoaderMixin, +): + """Text-to-image and image-conditioned generation with Qwen-Image 2.1. + + Prompt and condition images are encoded together by a Qwen3-VL model, so a + condition image occupies the vision slots the encoder reserved for it and + the transformer sees one interleaved text/image sequence. + """ + + supports_request_batch = True + supports_step_execution: ClassVar[bool] = True + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + self.parallel_config = od_config.parallel_config + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ) + ] + + self.device = get_local_device() + model = od_config.model + # Check if model is a local path + local_files_only = os.path.isdir(model) + + # Guard against transformers v5 multi-worker race on partial subfolder + # shard sets (see pipeline_qwen_image_edit_plus, Buildkite #1043). + qwen_subfolders = ["scheduler", "text_encoder", "vae", "processor"] + prefetch_subfolders( + model, + qwen_subfolders, + local_files_only=local_files_only, + ) + + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, subfolder="scheduler", local_files_only=local_files_only + ) + # Unlike the 2.0 T2I pipeline, the vision tower is required: condition + # images are encoded jointly with the prompt by the VLM. + self.text_encoder = from_pretrained_with_prefetch( + Qwen3VLForConditionalGeneration.from_pretrained, + model, + subfolder="text_encoder", + prefetch_list=qwen_subfolders, + local_files_only=local_files_only, + ).to(self.device) + self.vae = from_pretrained_with_prefetch( + DistributedAutoencoderKLQwenImage21.from_pretrained, + model, + subfolder="vae", + prefetch_list=qwen_subfolders, + local_files_only=local_files_only, + ).to(self.device) + transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImage21Transformer2DModel) + self.transformer = QwenImage21Transformer2DModel( + od_config=od_config, + quant_config=od_config.quantization_config, + **transformer_kwargs, + ) + self.processor = from_pretrained_with_prefetch( + Qwen3VLProcessor.from_pretrained, + model, + subfolder="processor", + prefetch_list=qwen_subfolders, + local_files_only=local_files_only, + ) + + self.stage = None + + # The VAE compresses 16x spatially and the transformer consumes latents + # unpatched, so one token covers a 16x16 pixel tile. + self.vae_scale_factor = 16 + self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 64 + + self.sys_prompt = "Comprehend and analyze the provided prompt." + # The prompt is built as a raw template string and passed straight to + # `self.processor(text=..., images=...)`, rather than going through + # `apply_chat_template`: the two tokenize differently and the checkpoint + # expects this one. The "Picture 1: ..." vision prefix only appears in + # the image-conditioned template. + self.prompt_template_t2i = ( + f"<|im_start|>system\n{self.sys_prompt}<|im_end|>\n" + f"<|im_start|>user\n{{}}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + self.prompt_template_ti2i = ( + f"<|im_start|>system\n{self.sys_prompt}<|im_end|>\n" + f"<|im_start|>user\nPicture 1: <|vision_start|><|image_pad|><|vision_end|>{{}}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + self.ref_token_list = ["Picture ", "Image ", "图 ", "图片 "] + # Number of leading system-role tokens to drop from the hidden states. + # Derived from the tokenized system message rather than hardcoded, so it + # tracks the processor's template. + sys_message = [{"role": "system", "content": [{"type": "text", "text": self.sys_prompt}]}] + sys_tokens = self.processor.apply_chat_template(sys_message, tokenize=True, return_dict=False) + self._drop_idx = len(sys_tokens[0]) + self._img_token_id = self.processor.tokenizer.encode("<|image_pad|>")[0] + self._max_length = 8192 + + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + + def check_inputs( + self, + prompt, + height, + width, + negative_prompt=None, + max_sequence_length=None, + ): + if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: + logger.warning( + f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} " + f"but are {height} and {width}. Dimensions will be resized accordingly" + ) + + if prompt is None: + raise ValueError("Provide `prompt`. Cannot leave `prompt` undefined.") + elif not isinstance(prompt, str) and not isinstance(prompt, list): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if max_sequence_length is not None and max_sequence_length > self._max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self._max_length} but is {max_sequence_length}" + ) + + def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): + bool_mask = mask.bool() + valid_lengths = bool_mask.sum(dim=1) + selected = hidden_states[bool_mask] + split_result = torch.split(selected, valid_lengths.tolist(), dim=0) + + return split_result + + @staticmethod + def _downsample_image_pad_tokens( + hidden_states_list, + image_pad_mask_list, + expected_slots_list, + image_sizes_list, + ): + """Reconcile each contiguous ``<|image_pad|>`` run with one slot per 2x2 latent group. + + The transformer's joint sequence reserves one vision-language token slot for every 2x2 group + of latent tokens and expands each slot 4x when substituting the latent embeddings, so a + condition image resized to ``W x H`` pixels needs exactly ``(H // 16) * (W // 16) // 4`` + slots in ``img_mask``. + + A standard Qwen3-VL processor (patch=16, merge=2) already emits exactly that many tokens per + image, so the run is kept unchanged (identity). A processor without spatial merging (merge=1) + emits 4x as many; those runs are folded by keeping the first token of every 4 consecutive + tokens. Any other count raises: if an official checkpoint's processor tokenizes differently, + this is the place to adjust. + + Args: + hidden_states_list: Per-sample hidden states (after `drop_idx`, before padding). + image_pad_mask_list: Per-sample bool masks marking the processor's image tokens. + expected_slots_list: Per-sample list of expected slot counts, one per condition image. + image_sizes_list: Per-sample list of `(width, height)` pixel sizes, for error messages. + + Returns: + `tuple` of `(out_hs, out_mask)`, same per-sample list structure as the inputs. + """ + out_hs, out_mask = [], [] + for hidden_state, pad_mask, expected_slots, image_sizes in zip( + hidden_states_list, image_pad_mask_list, expected_slots_list, image_sizes_list + ): + pad_indices = torch.where(pad_mask)[0] + runs: list[tuple[int, int]] = [] # contiguous image-token runs, inclusive (start, end) + if len(pad_indices) > 0: + run_start = pad_indices[0].item() + prev = run_start + for idx in pad_indices[1:].tolist(): + if idx != prev + 1: + runs.append((run_start, prev)) + run_start = idx + prev = idx + runs.append((run_start, prev)) + + if len(runs) != len(expected_slots): + raise ValueError( + f"Prompt encodes {len(runs)} condition-image token run(s) but " + f"{len(expected_slots)} condition image(s) were provided (sizes: {image_sizes})." + ) + if not runs: + out_hs.append(hidden_state) + out_mask.append(pad_mask) + continue + + keep = ~pad_mask + for (run_start, run_end), expected, size in zip(runs, expected_slots, image_sizes): + run_len = run_end - run_start + 1 + if run_len == expected: + # Standard Qwen3-VL processor (patch=16, merge=2): already one + # token per 2x2 latent group. + keep[run_start : run_end + 1] = True + elif run_len == expected * 4: + # Processor without spatial merging (merge=1): fold every 4 + # consecutive tokens into the first one. + keep[run_start : run_end + 1 : 4] = True + else: + raise ValueError( + f"Condition image resized to {size[0]}x{size[1]} produced {run_len} image tokens from the " + f"processor, but the transformer expects {expected} slots (one per 2x2 latent group) or " + f"{expected * 4} unmerged tokens. If the official checkpoint's processor tokenizes " + "differently, adjust `_downsample_image_pad_tokens` accordingly." + ) + out_hs.append(hidden_state[keep]) + out_mask.append(pad_mask[keep]) + return out_hs, out_mask + + def _get_qwen_prompt_embeds( + self, + prompt: str | list[str] = None, + images_per_prompt: list[list[PIL.Image.Image] | None] | None = None, + max_sequence_length: int | None = None, + prompt_name: str = "prompt", + ): + """Encode prompts (and optional per-prompt condition images) with the Qwen3-VL encoder. + + Returns: + `tuple` of `(prompt_embeds, encoder_attention_mask, image_pad_mask)`, all padded to the + batch's maximum (post-downsample) sequence length. `image_pad_mask` is `True` at the + vision-language image slots. + """ + dtype = self.text_encoder.dtype + prompt = [prompt] if isinstance(prompt, str) else prompt + + has_images = images_per_prompt is not None and any(images_per_prompt) + if not has_images: + txt = [self.prompt_template_t2i.format(e) for e in prompt] + else: + if images_per_prompt is None or len(images_per_prompt) != len(prompt): + raise ValueError("`images_per_prompt` must align with `prompt`.") + txt = [] + condition_pil_per_prompt: list[list[PIL.Image.Image]] = [] + for t, images in zip(prompt, images_per_prompt): + images = [ + img if isinstance(img, PIL.Image.Image) else PIL.Image.fromarray(img) for img in (images or []) + ] + condition_pil_per_prompt.append(images) + if not images: + txt.append(self.prompt_template_t2i.format(t)) + continue + n_imgs = len(images) + replace = "Picture 1: <|vision_start|><|image_pad|><|vision_end|>" + for i in range(2, n_imgs + 1): + replace += f" Picture {i}: <|vision_start|><|image_pad|><|vision_end|>" + template = self.prompt_template_ti2i.replace( + "Picture 1: <|vision_start|><|image_pad|><|vision_end|>", + replace.replace("Picture ", random.choice(self.ref_token_list)), + ) + txt.append(template.format(t)) + condition_pil_list = [img for images in condition_pil_per_prompt for img in images] + + # Validate only the user prompt contribution against the text budget; + # image placeholder expansion happens later inside the processor. + template_tokens = self.processor.tokenizer( + [self.prompt_template_t2i.format("")], + padding=True, + truncation=False, + return_tensors="pt", + ) + txt_tokens = self.processor.tokenizer( + txt, + padding=True, + truncation=False, + return_tensors="pt", + ) + validate_prompt_sequence_lengths( + txt_tokens.attention_mask, + max_sequence_length=max_sequence_length or self._max_length, + supported_max_sequence_length=self._max_length, + prompt_name=prompt_name, + baseline_attention_mask=template_tokens.attention_mask, + error_context="after applying the Qwen-Image 2.1 prompt template", + ) + + processor_kwargs: dict[str, Any] = {"text": txt, "padding": True, "return_tensors": "pt"} + if has_images: + processor_kwargs["images"] = condition_pil_list + model_inputs = self.processor(**processor_kwargs).to(self.device) + + forward_kwargs: dict[str, Any] = { + "input_ids": model_inputs.input_ids, + "attention_mask": model_inputs.attention_mask, + "output_hidden_states": True, + } + if has_images and hasattr(model_inputs, "pixel_values"): + forward_kwargs.update( + pixel_values=model_inputs.pixel_values, + image_grid_thw=model_inputs.image_grid_thw, + ) + # transformers>=5 Qwen3-VL computes multimodal RoPE from the token type + # ids the processor emits alongside input_ids. + if getattr(model_inputs, "mm_token_type_ids", None) is not None: + forward_kwargs["mm_token_type_ids"] = model_inputs.mm_token_type_ids + + outputs = self.text_encoder(**forward_kwargs) + hidden_states = outputs.hidden_states[-1] + + split_hidden_states = list(self._extract_masked_hidden(hidden_states, model_inputs.attention_mask)) + split_hidden_states = [e[self._drop_idx :] for e in split_hidden_states] + + image_pad_mask = [ + (sample_ids[sample_mask.bool()] == self._img_token_id) + for sample_ids, sample_mask in zip(model_inputs.input_ids, model_inputs.attention_mask) + ] + image_pad_mask = [e[self._drop_idx :] for e in image_pad_mask] + + if has_images: + # Expected slot counts: one VLM token slot per 2x2 group of latent + # tokens, i.e. (H // vae_scale_factor) * (W // vae_scale_factor) // 4 + # for a condition image resized to W x H pixels. + expected_slots = [ + [(img.height // self.vae_scale_factor) * (img.width // self.vae_scale_factor) // 4 for img in images] + for images in condition_pil_per_prompt + ] + image_sizes = [[img.size for img in images] for images in condition_pil_per_prompt] + split_hidden_states, image_pad_mask = self._downsample_image_pad_tokens( + split_hidden_states, image_pad_mask, expected_slots, image_sizes + ) + + attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] + max_seq_len = max(e.size(0) for e in split_hidden_states) + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states] + ) + encoder_attention_mask = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list] + ) + image_pad_mask = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), dtype=torch.bool)]) for u in image_pad_mask] + ) + + prompt_embeds = prompt_embeds.to(dtype=dtype) + + return prompt_embeds, encoder_attention_mask, image_pad_mask + + def encode_prompt( + self, + prompt: str | list[str], + images_per_prompt: list[list[PIL.Image.Image] | None] | None = None, + num_images_per_prompt: int = 1, + max_sequence_length: int | None = None, + prompt_name: str = "prompt", + ): + r""" + + Args: + prompt (`str` or `list[str]`): + prompt to be encoded + images_per_prompt (`list[list[PIL.Image.Image]]`, *optional*): + Per-prompt condition images, encoded jointly with the prompt. + num_images_per_prompt (`int`): + number of images that should be generated per prompt + """ + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + prompt_embeds, prompt_embeds_mask, image_pad_mask = self._get_qwen_prompt_embeds( + prompt, + images_per_prompt, + max_sequence_length=max_sequence_length, + prompt_name=prompt_name, + ) + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) + prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + image_pad_mask = image_pad_mask.repeat(1, num_images_per_prompt) + image_pad_mask = image_pad_mask.view(batch_size * num_images_per_prompt, seq_len) + + return prompt_embeds, prompt_embeds_mask, image_pad_mask + + @staticmethod + def _pack_latents(latents, batch_size, num_channels_latents, height, width): + # 2.1 consumes latents unpatched, so packing is a plain spatial flatten. + return latents.view(batch_size, num_channels_latents, height * width).transpose(1, 2) + + @staticmethod + def _unpack_latents(latents, height, width, vae_scale_factor): + batch_size, _, channels = latents.shape + height = 2 * (int(height) // (vae_scale_factor * 2)) + width = 2 * (int(width) // (vae_scale_factor * 2)) + latents = latents.transpose(1, 2).reshape(batch_size, channels, 1, height, width) + return latents + + @staticmethod + def _append_target_slots(image_pad_mask: torch.Tensor, num_target_tokens: int) -> torch.Tensor: + """Append one True slot per 2x2 group of target-image tokens to the VLM image-slot mask.""" + target_slots = torch.ones( + [image_pad_mask.shape[0], num_target_tokens // 4], + dtype=image_pad_mask.dtype, + device=image_pad_mask.device, + ) + return torch.cat([image_pad_mask, target_slots], dim=1) + + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if isinstance(generator, list): + image_latents = [ + retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax") + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax") + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, self.latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + image_latents = (image_latents - latents_mean) / latents_std + + return image_latents + + def prepare_latents( + self, + vae_images: list[torch.Tensor] | None, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + + image_latents = None + if vae_images is not None: + all_image_latents = [] + for image in vae_images: + image = image.to(device=device, dtype=dtype) + encoded = image if image.shape[1] == self.latent_channels else self._encode_vae_image(image, generator) + if batch_size > encoded.shape[0]: + if batch_size % encoded.shape[0] != 0: + raise ValueError( + f"Cannot duplicate `image` of batch size {encoded.shape[0]} to {batch_size} text prompts." + ) + encoded = torch.cat([encoded] * (batch_size // encoded.shape[0]), dim=0) + image_latent_height, image_latent_width = encoded.shape[3:] + all_image_latents.append( + self._pack_latents( + encoded, batch_size, num_channels_latents, image_latent_height, image_latent_width + ) + ) + image_latents = torch.cat(all_image_latents, dim=1) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + if latents is None: + shape = (batch_size, 1, num_channels_latents, height, width) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) + else: + latents = latents.to(device=device, dtype=dtype) + + return latents, image_latents + + def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len): + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + sigmas=sigmas, + mu=mu, + ) + return timesteps, num_inference_steps + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + def _extract_prompts(self, prompts): + """Extract prompt and negative_prompt from OmniPromptType list.""" + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None + if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts): + negative_prompt = None + elif prompts: + negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts] + else: + negative_prompt = None + return prompt, negative_prompt + + @staticmethod + def _extract_request_images(prompts) -> list[dict[str, Any]]: + """Pull per-request condition-image data staged by the pre-process func.""" + per_request = [] + for p in prompts: + info = p.get("additional_information", {}) if isinstance(p, dict) else {} + per_request.append( + { + "prompt_images": info.get("prompt_image"), + "vae_images": info.get("vae_images"), + "input_image_sizes": info.get("input_image_sizes"), + } + ) + return per_request + + def _prepare_generation_context( + self, + *, + prompt, + negative_prompt, + height, + width, + num_inference_steps, + sigmas, + num_images_per_prompt, + generator, + true_cfg_scale, + max_sequence_length, + per_request_images: list[dict[str, Any]], + latents=None, + attention_kwargs=None, + ): + """Shared preparation logic for forward() and prepare_encode(). + + Validates inputs, encodes prompts (jointly with condition images), + VAE-encodes condition images, prepares latents and timesteps, and + returns all intermediate values as a dict. + """ + self.check_inputs( + prompt, + height, + width, + negative_prompt, + max_sequence_length, + ) + + self._attention_kwargs = attention_kwargs or {} + self._current_timestep = None + self._interrupt = False + + batch_size = len(prompt) + + has_neg_prompt = negative_prompt is not None + do_true_cfg = true_cfg_scale > 1 and has_neg_prompt + self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt) + if true_cfg_scale > 1 and not has_neg_prompt: + logger.warning( + f"true_cfg_scale is passed as {true_cfg_scale}, but classifier-free guidance is not enabled since no " + f"negative_prompt is provided." + ) + elif true_cfg_scale <= 1 and has_neg_prompt: + logger.warning( + "negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1" + ) + + images_per_prompt = [entry["prompt_images"] for entry in per_request_images] + has_images = any(images_per_prompt) + if has_images and not all(images_per_prompt): + raise ValueError( + "Cannot batch a mix of image-conditioned and text-to-image requests; " + "the transformer's joint sequence requires a shared layout." + ) + + prompt_embeds, prompt_embeds_mask, image_pad_mask = self.encode_prompt( + prompt=prompt, + images_per_prompt=images_per_prompt if has_images else None, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + if do_true_cfg: + negative_prompt_embeds, negative_prompt_embeds_mask, negative_image_pad_mask = self.encode_prompt( + prompt=negative_prompt, + images_per_prompt=images_per_prompt if has_images else None, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + prompt_name="negative_prompt", + ) + else: + negative_prompt_embeds = None + negative_prompt_embeds_mask = None + negative_image_pad_mask = None + + # Collate condition images per input slot across requests. Batched VAE + # encoding requires every request to share the same per-slot sizes. + vae_images_batched = None + input_image_sizes: list[tuple[int, int]] = [] + if has_images: + num_slots = len(per_request_images[0]["vae_images"]) + for entry in per_request_images: + if entry["input_image_sizes"] != per_request_images[0]["input_image_sizes"]: + raise ValueError( + "All image-conditioned requests in a batch must share condition image sizes " + "(the transformer's joint sequence requires a shared layout)." + ) + input_image_sizes = per_request_images[0]["input_image_sizes"] or [] + vae_images_batched = [ + torch.cat([entry["vae_images"][slot] for entry in per_request_images], dim=0) + for slot in range(num_slots) + ] + + num_channels_latents = self.transformer.in_channels + latents, image_latents = self.prepare_latents( + vae_images_batched, + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + self.device, + generator, + latents, + ) + + img_shapes = [ + [ + *[(1, h // self.vae_scale_factor, w // self.vae_scale_factor) for w, h in input_image_sizes], + (1, height // self.vae_scale_factor, width // self.vae_scale_factor), + ] + ] * batch_size + + timesteps, num_inference_steps = self.prepare_timesteps( + num_inference_steps, + sigmas, + latents.shape[1], + ) + self._num_timesteps = len(timesteps) + + # The transformer's `img_mask` spans the joint sequence, so append one + # slot per 2x2 group of target latents. + img_mask = self._append_target_slots(image_pad_mask, latents.shape[1]) + if do_true_cfg: + negative_img_mask = self._append_target_slots(negative_image_pad_mask, latents.shape[1]) + else: + negative_img_mask = None + + return { + "prompt_embeds": prompt_embeds, + "prompt_embeds_mask": prompt_embeds_mask, + "negative_prompt_embeds": negative_prompt_embeds, + "negative_prompt_embeds_mask": negative_prompt_embeds_mask, + "image_pad_mask": image_pad_mask, + "negative_image_pad_mask": negative_image_pad_mask, + "img_mask": img_mask, + "negative_img_mask": negative_img_mask, + "latents": latents, + "image_latents": image_latents, + "img_shapes": img_shapes, + "timesteps": timesteps, + "do_true_cfg": do_true_cfg, + } + + def _decode_latents( + self, + latents: torch.Tensor, + height: int, + width: int, + output_type: str = "pil", + ) -> DiffusionOutput: + """Unpack, denormalize, and VAE-decode latents into a DiffusionOutput.""" + if output_type == "latent": + return DiffusionOutput( + output=latents, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + + latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) + latents = latents.to(self.vae.dtype) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents = latents * latents_std + latents_mean + image = self.vae.decode(latents, return_dict=False)[0][:, :, 0] + return DiffusionOutput( + output=image, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: + sampling_params_list = req.sampling_params_list + common_sampling_params = sampling_params_list[0] + prompt, negative_prompt = self._extract_prompts(req.prompts) + per_request_images = self._extract_request_images(req.prompts) + + height = common_sampling_params.height or OUTPUT_RESOLUTION + width = common_sampling_params.width or OUTPUT_RESOLUTION + height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2) + num_inference_steps = common_sampling_params.num_inference_steps or 50 + sigmas = common_sampling_params.sigmas + max_sequence_length = common_sampling_params.max_sequence_length or self._max_length + num_images_per_prompt = ( + common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1 + ) + generator = req.collate_request_generators(num_images_per_prompt, None) + latents = req.collate_request_tensors("latents", None) + true_cfg_scale = common_sampling_params.true_cfg_scale or 4.0 + output_type = common_sampling_params.output_type or "pil" + + # Batch homogeneity decides the attention path: a homogeneous batch + # (same condition-image count and sizes per request) can use the + # piecewise-span backend; heterogeneous batches take the dense 4D mask. + signatures = [ + ( + len(entry["vae_images"] or []), + tuple(entry["input_image_sizes"] or ()), + ) + for entry in per_request_images + ] + homogeneous = all(sig == signatures[0] for sig in signatures) + attention_kwargs = {"attn_path": "auto" if homogeneous else "mask"} + + ctx = self._prepare_generation_context( + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + sigmas=sigmas, + num_images_per_prompt=num_images_per_prompt, + generator=generator, + true_cfg_scale=true_cfg_scale, + max_sequence_length=max_sequence_length, + per_request_images=per_request_images, + latents=latents, + attention_kwargs=attention_kwargs, + ) + + latents = self.diffuse( + ctx["prompt_embeds"], + ctx["prompt_embeds_mask"], + ctx["negative_prompt_embeds"], + ctx["negative_prompt_embeds_mask"], + ctx["img_mask"], + ctx["negative_img_mask"], + ctx["latents"], + ctx["image_latents"], + ctx["img_shapes"], + ctx["timesteps"], + ctx["do_true_cfg"], + true_cfg_scale, + attention_kwargs=self.attention_kwargs, + additional_transformer_kwargs={ + "return_dict": False, + }, + ) + + self._current_timestep = None + + result = self._decode_latents(latents, height, width, output_type) + return split_diffusion_output_by_request( + result, + req, + num_outputs_per_prompt=num_images_per_prompt, + ) + + # ── Step-execution protocol ── + + def prepare_encode( + self, + state: "StepRequestState", + **kwargs: Any, + ) -> "StepRequestState": + """Populate *state* with encoded prompts, latents, timesteps, and CFG config.""" + sampling = state.sampling + prompt, negative_prompt = self._extract_prompts([state.prompt] if state.prompt is not None else []) + per_request_images = self._extract_request_images([state.prompt] if state.prompt is not None else []) + + height = sampling.height or OUTPUT_RESOLUTION + width = sampling.width or OUTPUT_RESOLUTION + height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2) + num_images_per_prompt = sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else 1 + + ctx = self._prepare_generation_context( + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_inference_steps=sampling.num_inference_steps or 50, + sigmas=sampling.sigmas, + num_images_per_prompt=num_images_per_prompt, + generator=sampling.generator, + true_cfg_scale=sampling.true_cfg_scale or 4.0, + max_sequence_length=sampling.max_sequence_length or self._max_length, + per_request_images=per_request_images, + attention_kwargs=kwargs.get("attention_kwargs"), + ) + + # prepare_timesteps() has already materialized request-specific timestep + # state on self.scheduler, so deepcopy preserves dynamic-shifting state + # without replaying set_timesteps() on the per-request scheduler. + # Per-request scheduler (must not share state with self.scheduler) + req_scheduler = copy.deepcopy(self.scheduler) + req_scheduler.set_begin_index(0) + + state.prompt_embeds = ctx["prompt_embeds"] + state.prompt_embeds_mask = ctx["prompt_embeds_mask"] + state.negative_prompt_embeds = ctx["negative_prompt_embeds"] + state.negative_prompt_embeds_mask = ctx["negative_prompt_embeds_mask"] + state.latents = ctx["latents"] + state.timesteps = ctx["timesteps"] + state.step_index = 0 + state.scheduler = req_scheduler + state.do_true_cfg = ctx["do_true_cfg"] + state.guidance = None + state.img_shapes = ctx["img_shapes"] + # 2.1 applies plain CFG (no output normalization, unlike 2.0). + state.sampling.cfg_normalize = False + + # Pipeline-specific state: the image-slot mask is stored *before* the + # target slots are appended so a step batch can right-pad the prompt + # region to a common length first. The KV cache is per request — + # the pipeline is a singleton shared across requests. + state.extra["image_pad_mask"] = ctx["image_pad_mask"] + state.extra["negative_image_pad_mask"] = ctx["negative_image_pad_mask"] + state.extra["image_latents"] = ctx["image_latents"] + state.extra["height"] = height + state.extra["width"] = width + cache_enabled = getattr(self.transformer, "causal_condition", False) + state.extra["kv_cache"] = ( + [{} for _ in range(len(self.transformer.transformer_blocks))] if cache_enabled else None + ) + + return state + + @staticmethod + def _kv_cache_phase(kv_cache: list[dict] | None) -> str: + if kv_cache is None: + return "none" + first = kv_cache[0] + return "decode" if any("key" in branch for branch in first.values()) else "prefill" + + @staticmethod + def _assemble_kv_cache( + states: list["StepRequestState"], + ) -> tuple[list[dict] | None, bool]: + """Merge per-request KV caches for a batched forward. + + Returns `(kv_cache, take_ownership)`. When `take_ownership` is True the + returned cache is the freshly written prefill cache and its entries must + be scattered back into the per-request caches by `_scatter_kv_cache`. + """ + caches = [state.extra.get("kv_cache") for state in states] + if all(cache is None for cache in caches): + return None, False + phases = {QwenImage21Pipeline._kv_cache_phase(cache) for cache in caches} + if phases == {"prefill"}: + # Batched prefill: a fresh merged cache collects all requests' prefix + # K/V; written entries are split back per request after the step. + num_blocks = len(caches[0]) + return [{} for _ in range(num_blocks)], True + if phases == {"decode"}: + # Decode only reads the cache, so a batched view is enough and the + # per-request caches stay untouched. + num_blocks = len(caches[0]) + merged = [] + for block_idx in range(num_blocks): + branches: dict[str, dict[str, torch.Tensor]] = {} + branch_names = {branch for cache in caches for branch in cache[block_idx]} + for branch in branch_names: + branches[branch] = { + part: torch.cat([cache[block_idx][branch][part] for cache in caches], dim=0) + for part in ("key", "value") + } + merged.append(branches) + return merged, False + raise ValueError( + "Cannot batch requests at mixed KV-cache phases (a request starting its first denoise step " + "joined a batch of in-flight requests). Schedule them separately." + ) + + @staticmethod + def _scatter_kv_cache( + states: list["StepRequestState"], + row_counts: list[int], + kv_cache: list[dict], + ) -> None: + """Split a batched prefill cache back into the per-request caches.""" + start = 0 + for state, rows in zip(states, row_counts): + cache = state.extra["kv_cache"] + for block_idx, block_cache in enumerate(kv_cache): + for branch, parts in block_cache.items(): + cache[block_idx][branch] = { + part: tensor[start : start + rows].contiguous() for part, tensor in parts.items() + } + start += rows + + def denoise_step( + self, + input_batch: "InputBatch", + **kwargs: Any, + ) -> torch.Tensor | None: + """One denoise step: read from *input_batch*, delegate to CFGParallelMixin.""" + del kwargs + if self.interrupt: + return None + + t = input_batch.timesteps + self._current_timestep = t + self.transformer.do_true_cfg = input_batch.do_true_cfg + + states = list(input_batch.states) + if not states: + raise ValueError("QwenImage21Pipeline.denoise_step requires per-request states on the InputBatch.") + + latents = input_batch.latents + timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype) + + # Assemble the joint image-slot mask: right-pad each request's VLM mask + # to the batch's padded prompt length, then append the target slots. + max_seq_len = input_batch.prompt_embeds.shape[1] + mask_rows = [] + for state in states: + row = state.extra["image_pad_mask"] + if row.shape[1] < max_seq_len: + row = F.pad(row, (0, max_seq_len - row.shape[1]), value=False) + mask_rows.append(row) + img_mask = self._append_target_slots(torch.cat(mask_rows, dim=0), latents.shape[1]) + negative_img_mask = None + if input_batch.do_true_cfg: + neg_rows = [] + for state in states: + row = state.extra["negative_image_pad_mask"] + if row.shape[1] < max_seq_len: + row = F.pad(row, (0, max_seq_len - row.shape[1]), value=False) + neg_rows.append(row) + negative_img_mask = self._append_target_slots(torch.cat(neg_rows, dim=0), latents.shape[1]) + + image_latents_list = [state.extra["image_latents"] for state in states] + if all(image_latent is None for image_latent in image_latents_list): + image_latents = None + elif all(image_latent is not None for image_latent in image_latents_list): + image_latents = torch.cat(image_latents_list, dim=0) + else: + raise ValueError("Cannot batch a mix of image-conditioned and text-to-image requests in step mode.") + + latent_model_input = latents + if image_latents is not None: + latent_model_input = torch.cat([image_latents, latents], dim=1) + + row_counts = [state.latents.shape[0] for state in states] + kv_cache, take_ownership = self._assemble_kv_cache(states) + + positive_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep / 1000, + "encoder_hidden_states_mask": input_batch.prompt_embeds_mask, + "encoder_hidden_states": input_batch.prompt_embeds, + "img_shapes": input_batch.img_shapes, + "img_mask": img_mask, + "attention_kwargs": self.attention_kwargs, + "kv_cache": kv_cache, + "cache_branch": "cond", + "return_dict": False, + } + if input_batch.do_true_cfg: + negative_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep / 1000, + "encoder_hidden_states_mask": input_batch.negative_prompt_embeds_mask, + "encoder_hidden_states": input_batch.negative_prompt_embeds, + "img_shapes": input_batch.img_shapes, + "img_mask": negative_img_mask, + "attention_kwargs": self.attention_kwargs, + "kv_cache": kv_cache, + "cache_branch": "uncond", + "return_dict": False, + } + else: + negative_kwargs = None + + noise_pred = self.predict_noise_maybe_with_cfg( + input_batch.do_true_cfg, + input_batch.true_cfg_scale, + positive_kwargs, + negative_kwargs, + input_batch.cfg_normalize, + output_slice=latents.size(1), + ) + + if take_ownership: + self._scatter_kv_cache(states, row_counts, kv_cache) + + return noise_pred + + def step_scheduler( + self, + state: "StepRequestState", + noise_pred: torch.Tensor, + **kwargs: Any, + ) -> None: + """One scheduler step: update ``state.latents`` and advance ``step_index``.""" + if self.interrupt: + return + + t = state.current_timestep + state.latents = self.scheduler_step_maybe_with_cfg( + noise_pred, + t, + state.latents, + state.do_true_cfg, + per_request_scheduler=state.scheduler, + ) + + state.step_index += 1 + + def post_decode( + self, + state: "StepRequestState", + **kwargs: Any, + ) -> DiffusionOutput: + """Decode final latents from *state*.""" + self._current_timestep = None + + height = state.extra.get("height") or state.sampling.height or OUTPUT_RESOLUTION + width = state.extra.get("width") or state.sampling.width or OUTPUT_RESOLUTION + output_type = kwargs.get("output_type") or state.sampling.output_type or "pil" + + return self._decode_latents(state.latents, height, width, output_type) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py b/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py new file mode 100644 index 000000000..57b98aad7 --- /dev/null +++ b/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py @@ -0,0 +1,1106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# Adapted from diffusers' transformer_qwenimage21.py (Qwen-Image 2.1 single-stream DiT). + +from __future__ import annotations + +import copy +import math +from collections.abc import Iterable +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.base import CachedTransformer +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, +) +from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + +logger = init_logger(__name__) + + +def _apply_qwen_image21_rotary_emb(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Rotate interleaved pairs with complex FP32 multiplication. + + Same contract as the reference `apply_rotary_emb_qwen(..., use_real=False)`: + `x` is [B, S, H, D], `freqs` is the complex frequency tensor [S, D // 2]. + """ + paired = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + return torch.view_as_real(paired * freqs.unsqueeze(1)).flatten(3).to(x.dtype) + + +def _apply_qwen_image21_rotary_emb_native(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Platform fallback for `_apply_qwen_image21_rotary_emb` without complex tensors. + + Algebraically identical to the complex multiplication, computed in fp32: + (x1 + i x2)(cos + i sin) interleaved back as [x1 cos - x2 sin, x1 sin + x2 cos, ...]. + """ + cos = torch.real(freqs).float().unsqueeze(1) + sin = torch.imag(freqs).float().unsqueeze(1) + x_pairs = x.float().unflatten(-1, (-1, 2)) + x1, x2 = x_pairs.unbind(-1) + out = torch.stack((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1) + return out.flatten(-2).to(x.dtype) + + +def _normalize_qwen_image21_weight_name(name: str) -> str: + name = name.removeprefix("transformer.") + if ".to_out.0." in name: + name = name.replace(".to_out.0.", ".to_out.") + return name + + +def _resolve_qwen_image21_lookup_name( + name: str, + stacked_params_mapping: list[tuple[str, str, str]], +) -> tuple[str, str | None]: + lookup_name = _normalize_qwen_image21_weight_name(name) + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in lookup_name or param_name in lookup_name: + continue + return lookup_name.replace(weight_name, param_name), shard_id + return lookup_name, None + + +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. + + Ported from diffusers' transformer_qwenimage21._select_modulation_rows. + """ + 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) + + +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 QwenImage21TimestepEmbedder(nn.Module): + """MLP over the sinusoidal timestep projection. No biases, matching the 2.1 checkpoint.""" + + def __init__(self, in_dim: int, embedding_dim: int, prefix: str = ""): + super().__init__() + self.linear_1 = ReplicatedLinear( + in_dim, + embedding_dim, + bias=False, + return_bias=False, + quant_config=None, + prefix=f"{prefix}.linear_1", + ) + self.act = nn.SiLU() + self.linear_2 = ReplicatedLinear( + embedding_dim, + embedding_dim, + bias=False, + return_bias=False, + quant_config=None, + prefix=f"{prefix}.linear_2", + ) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + return self.linear_2(self.act(self.linear_1(sample))) + + +class QwenImage21TimestepProjEmbeddings(nn.Module): + def __init__(self, embedding_dim: int, prefix: str = ""): + super().__init__() + self.time_proj = QwenImage21TemporalTimesteps(timestep_dim=256) + self.timestep_embedder = QwenImage21TimestepEmbedder( + in_dim=256, embedding_dim=embedding_dim, prefix=f"{prefix}.timestep_embedder" + ) + + 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`. The weight is loaded as-is (no re-centering) by `load_weights`. + """ + + 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, prefix: str = ""): + super().__init__() + self.text_norm = QwenImage21ZeroCenterRMSNorm(context_in_dim, eps=eps) + self.in_layer = ReplicatedLinear( + context_in_dim, + hidden_size, + bias=False, + return_bias=False, + quant_config=None, + prefix=f"{prefix}.in_layer", + ) + self.act = nn.GELU(approximate="tanh") + self.out_layer = ReplicatedLinear( + hidden_size, + hidden_size, + bias=False, + return_bias=False, + quant_config=None, + prefix=f"{prefix}.out_layer", + ) + + 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): + def __init__( + self, + hidden_size: int, + mlp_hidden_size: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.proj = ColumnParallelLinear( + hidden_size, + mlp_hidden_size, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.proj", + ) + self.gate_layer = ColumnParallelLinear( + hidden_size, + mlp_hidden_size, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_layer", + ) + self.out = RowParallelLinear( + mlp_hidden_size, + hidden_size, + bias=False, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.out", + ) + self.activation_fn = nn.SiLU() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.out(self.activation_fn(self.gate_layer(hidden_states)) * self.proj(hidden_states)) + + +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, prefix: str = ""): + super().__init__() + self.silu = nn.SiLU() + self.linear = ReplicatedLinear( + conditioning_embedding_dim, + embedding_dim, + bias=False, + return_bias=False, + quant_config=None, + prefix=f"{prefix}.linear", + ) + 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) + + +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. + + Ported from diffusers' transformer_qwenimage21.QwenImage21Rope. + """ + + 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 QwenImage21Attention(nn.Module): + r"""Single-stream attention for Qwen-Image 2.1. + + Q/K/V projections are fused into one `QKVParallelLinear` (checkpoint shards load via + `packed_modules_mapping`); the output projection is a `RowParallelLinear`. The actual + attention runs through vllm-omni's `Attention` layer with Q/K/V in (B, S, H, D) layout, + so structural masking arrives via `AttentionMetadata` (piecewise `full_attn_spans` or a + 4D `attn_mask`) instead of a custom processor. + """ + + def __init__( + self, + dim: int, + heads: int, + dim_head: int, + eps: float = 1e-6, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.head_dim = dim_head + + self.to_qkv = QKVParallelLinear( + hidden_size=dim, + head_size=dim_head, + total_num_heads=heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_qkv", + ) + self.num_heads = self.to_qkv.num_heads + self.num_kv_heads = self.to_qkv.num_kv_heads + + # Per-head RMSNorm over the head dim; TP shards whole heads, so the norm is TP-safe. + self.norm_q = nn.RMSNorm(dim_head, eps=eps) + self.norm_k = nn.RMSNorm(dim_head, eps=eps) + + self.to_out = RowParallelLinear( + heads * dim_head, + dim, + bias=False, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_out.0", + ) + + self.attn = Attention( + num_heads=self.num_heads, + head_size=dim_head, + softmax_scale=1.0 / (dim_head**0.5), + causal=False, + num_kv_heads=self.num_kv_heads, + ) + + def _apply_rotary_emb(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + if x.device.type == "cuda": + return _apply_qwen_image21_rotary_emb(x, freqs) + # Platform fallback without complex-tensor kernels (the CustomOp rope layer + # dispatches by platform, not tensor device, so it cannot serve CPU tensors). + return _apply_qwen_image21_rotary_emb_native(x, freqs) + + def forward( + self, + hidden_states: torch.Tensor, + freqs: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + sp_prefix_len: int = 0, + sp_decode: bool = False, + kv_cache: dict[str, dict[str, torch.Tensor]] | None = None, + cache_branch: str = "cond", + cache_write_len: int | None = None, + ) -> torch.Tensor: + r""" + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, dim)`): + Modulated tokens. The layout is the full joint sequence normally; under sequence parallelism it is + `[replicated prefix, local target shard]`, and in KV-cache decode it is target tokens only. + freqs (`torch.Tensor`): Complex RoPE frequencies matching the local `hidden_states` layout. + attn_metadata (`AttentionMetadata`, *optional*): Mask/span metadata for the attention backend. This module + takes a shallow copy before attaching SP joint tensors, so callers may share one metadata object + across blocks. + sp_prefix_len (`int`): Number of leading tokens that are replicated across SP ranks. When > 0 they are + handed to the `Attention` layer as joint (front) tensors and only the trailing shard is all-to-all'd. + sp_decode (`bool`): KV-cache decode under sequence parallelism: the cached prefix K/V are replicated + and join through the Ulysses joint mechanism with an empty joint query. + kv_cache (dict, *optional*): Per-block cache keyed by CFG branch (`"cond"` / `"uncond"`); each branch + stores `{"key", "value"}` of the timestep-independent prefix (text + condition images). + cache_branch (`str`): CFG branch whose cache entry is read/written this forward. + cache_write_len (`int`, *optional*): Prefill mode — cache the first `cache_write_len` K/V positions. + `None` with a populated branch means decode — cached prefix K/V are prepended instead. + """ + batch_size, seq_len, _ = hidden_states.shape + + qkv, _ = self.to_qkv(hidden_states) + q_size = self.num_heads * self.head_dim + kv_size = self.num_kv_heads * self.head_dim + query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1) + + query = query.unflatten(-1, (self.num_heads, self.head_dim)) + key = key.unflatten(-1, (self.num_kv_heads, self.head_dim)) + value = value.unflatten(-1, (self.num_kv_heads, self.head_dim)) + + query = self.norm_q(query).to(value.dtype) + key = self.norm_k(key).to(value.dtype) + + query = self._apply_rotary_emb(query, freqs) + key = self._apply_rotary_emb(key, freqs) + + cached_key = cached_value = None + if kv_cache is not None: + branch_cache = kv_cache.setdefault(cache_branch, {}) + if cache_write_len is not None: + # Prefill: cache the timestep-independent prefix K/V for later denoising steps. + branch_cache["key"] = key[:, :cache_write_len].contiguous() + branch_cache["value"] = value[:, :cache_write_len].contiguous() + else: + cached_key = branch_cache["key"] + cached_value = branch_cache["value"] + + metadata = copy.copy(attn_metadata) if attn_metadata is not None else None + if sp_prefix_len > 0: + # Ulysses joint attention: the prefix is replicated across SP ranks, so it is + # concatenated (front) after the all-to-all instead of being sharded. + if metadata is None: + metadata = AttentionMetadata() + metadata.joint_query = query[:, :sp_prefix_len] + metadata.joint_key = key[:, :sp_prefix_len] + metadata.joint_value = value[:, :sp_prefix_len] + metadata.joint_strategy = "front" + attn_output = self.attn( + query[:, sp_prefix_len:], key[:, sp_prefix_len:], value[:, sp_prefix_len:], metadata + ) + elif cached_key is not None: + if sp_decode: + # Decode under SP: cached prefix K/V are replicated; pass them as joint K/V with an + # empty joint query so only the target shard goes through the all-to-all. + if metadata is None: + metadata = AttentionMetadata() + metadata.joint_query = query.new_zeros(batch_size, 0, self.num_heads, self.head_dim) + metadata.joint_key = cached_key + metadata.joint_value = cached_value + metadata.joint_strategy = "front" + attn_output = self.attn(query, key, value, metadata) + else: + # Decode: the block-causal mask degenerates to full attention for target rows. + key = torch.cat([cached_key, key], dim=1) + value = torch.cat([cached_value, value], dim=1) + attn_output = self.attn(query, key, value, metadata) + else: + attn_output = self.attn(query, key, value, metadata) + + attn_output = attn_output.flatten(2, 3).type_as(hidden_states) + return self.to_out(attn_output) + + +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, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + 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, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.img_mlp = QwenImage21SwiGLUFeedForward( + hidden_size=dim, + mlp_hidden_size=dim * mlp_ratio, + quant_config=quant_config, + prefix=f"{prefix}.img_mlp", + ) + + 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, + freqs: torch.Tensor, + target_token_mask: torch.Tensor | None = None, + attn_metadata: AttentionMetadata | None = None, + sp_prefix_len: int = 0, + sp_decode: bool = False, + kv_cache: dict[str, dict[str, torch.Tensor]] | None = None, + cache_branch: str = "cond", + cache_write_len: int | 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( + img_modulated, + freqs, + attn_metadata=attn_metadata, + sp_prefix_len=sp_prefix_len, + sp_decode=sp_decode, + kv_cache=kv_cache, + cache_branch=cache_branch, + cache_write_len=cache_write_len, + ) + 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 QwenImage21SequencePrepare(nn.Module): + r"""Assemble the joint text/image sequence and split it into prefix and target parts. + + This module boundary lets `_sp_plan` shard the target-image tokens (and their RoPE + frequencies) via `split_output=True` while the prefix (text + condition images) stays + replicated across SP ranks for the Ulysses joint-attention mechanism. + """ + + def __init__(self, img_in: nn.Module, pos_embed: QwenImage21Rope): + super().__init__() + self.img_in = img_in + self.pos_embed = pos_embed + + @staticmethod + def expand_image_pad_mask(img_mask_row: torch.Tensor) -> torch.Tensor: + """One VLM image slot stands for a 2x2 group of latent tokens; expand those positions four-fold.""" + repeats = torch.where(img_mask_row, 4, 1) + return torch.repeat_interleave(img_mask_row, repeats) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None, + img_shapes: list[tuple[int, int, int]], + img_mask: torch.Tensor, + is_decode: bool, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + r""" + Returns: + `tuple` of `(prefix_hidden_states, target_hidden_states, prefix_freqs, target_freqs)`. + `prefix_*` cover text + condition-image tokens (replicated under SP); `target_*` cover the + target image (sharded under SP via `_sp_plan`). In decode mode the prefix parts are empty. + """ + batch_size = hidden_states.shape[0] + image_pad_mask = self.expand_image_pad_mask(img_mask[0]) + rotary_emb = self.pos_embed(img_shapes, image_pad_mask, device=hidden_states.device) + + target_tokens = math.prod(img_shapes[-1]) + prefix_len = image_pad_mask.shape[0] - target_tokens + + if is_decode: + # Only the target image's tokens are recomputed; the prefix comes from the KV cache. + target_hidden_states = self.img_in(hidden_states[:, -target_tokens:]) + dim = target_hidden_states.shape[-1] + prefix_hidden_states = target_hidden_states.new_zeros(batch_size, 0, dim) + return prefix_hidden_states, target_hidden_states, rotary_emb[:0], rotary_emb[prefix_len:] + + hidden_states = self.img_in(hidden_states) + + 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, + ) + repeats = torch.where(img_mask[0], 4, 1) + joint_hidden_states = joint_hidden_states.repeat_interleave(repeats, dim=1) + joint_hidden_states[:, image_pad_mask] = hidden_states + + return ( + joint_hidden_states[:, :prefix_len], + joint_hidden_states[:, prefix_len:], + rotary_emb[:prefix_len], + rotary_emb[prefix_len:], + ) + + +class QwenImage21Transformer2DModel(CachedTransformer): + r""" + The single-stream Transformer used by Qwen-Image 2.1, ported to vllm-omni parallel layers. + + 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 from 2.0, both switched on by config and neither adding parameters: + + - `causal_block` — attention follows `(q_idx >= kv_idx) or same_image_block`, so the sequence is causal while + every image block stays internally bidirectional. Implemented through `AttentionMetadata`: piecewise + `full_attn_spans` (one span per image block; requires a piecewise-capable backend such as FLASH_ATTN) when the + batch is homogeneous and unpadded, or a dense 4D boolean `attn_mask` otherwise. + - `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 + (`kv_cache`). + """ + + # the small and frequently-repeated block(s) of a model + # -- typically a transformer layer + # used for torch compile optimizations + _repeated_blocks = ["QwenImage21TransformerBlock"] + _layerwise_offload_blocks_attrs = ["transformer_blocks"] + packed_modules_mapping = { + "to_qkv": ["to_q", "to_k", "to_v"], + } + + _hsdp_shard_conditions = [is_transformer_block_module] + + # Sequence Parallelism plan (corresponds to diffusers' _cp_plan). + # + # The joint sequence is `[prefix (text + condition images) | target image]`. The prefix is + # replicated across SP ranks and reaches the attention backend through the Ulysses joint + # mechanism (`joint_strategy="front"`); only the target tokens and their RoPE frequencies + # are sharded, via the `sequence_prepare` module boundary with auto-padding. SP padding + # positions sit at the very end of the joint sequence (the target is the tail block) and + # are masked out through the 4D-mask path built in `forward`. + # + # Limitation: prefix tokens are recomputed redundantly on every rank through all blocks + # (their K/V are needed by the joint attention); only the target image is actually + # sequence-parallel. + _sp_plan = { + # Shard sequence_prepare outputs: target hidden_states and target RoPE freqs. + # Indices 0 (prefix hidden_states) and 2 (prefix freqs) stay replicated. + "sequence_prepare": { + 1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True), + 3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), + }, + # Gather the (target-only) output after proj_out. + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + + def __init__( + self, + od_config: OmniDiffusionConfig, + 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, + causal_block: bool = True, + quant_config: QuantizationConfig | None = None, + prefix: str = "transformer", + ): + super().__init__() + self.parallel_config = od_config.parallel_config + self.in_channels = in_channels + self.out_channels = out_channels or in_channels + self.patch_size = patch_size + self.inner_dim = num_attention_heads * attention_head_dim + self.causal_condition = causal_condition + self.causal_block = causal_block + self.quant_config = quant_config + + self.pos_embed = QwenImage21Rope(theta=10000, axes_dim=list(axes_dims_rope)) + self.time_text_embed = QwenImage21TimestepProjEmbeddings( + embedding_dim=self.inner_dim, prefix=f"{prefix}.time_text_embed" + ) + self.txt_in = QwenImage21TextProjection(context_in_dim, self.inner_dim, eps=eps, prefix=f"{prefix}.txt_in") + + # Entry projections are kept full precision — small sensitive layers at the network boundary. + self.img_in = ReplicatedLinear( + in_channels * patch_size * patch_size, + self.inner_dim, + bias=False, + return_bias=False, + quant_config=None, + prefix=f"{prefix}.img_in", + ) + + # One shared modulation for every block: [mod1.scale, mod1.gate, mod2.scale, mod2.gate]. + self.modulation = nn.Sequential( + nn.SiLU(), + ReplicatedLinear( + self.inner_dim, + 4 * self.inner_dim, + bias=False, + return_bias=False, + quant_config=None, + prefix=f"{prefix}.modulation.1", + ), + ) + + 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, + quant_config=quant_config, + prefix=f"{prefix}.transformer_blocks.{i}", + ) + for i in range(num_layers) + ] + ) + + self.norm_out = QwenImage21AdaLayerNormContinuous( + self.inner_dim, self.inner_dim, eps=eps, prefix=f"{prefix}.norm_out" + ) + self.proj_out = ReplicatedLinear( + self.inner_dim, + patch_size * patch_size * self.out_channels, + bias=False, + return_bias=False, + quant_config=None, + prefix=f"{prefix}.proj_out", + ) + + # Module boundary where _sp_plan shards the target tokens + their RoPE freqs. + self.sequence_prepare = QwenImage21SequencePrepare(self.img_in, self.pos_embed) + + @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. + + 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 + + @staticmethod + def _build_full_attn_spans(image_ids: torch.Tensor) -> list[tuple[int, int]]: + """One `[start, end)` full-attention span per image block, in joint-sequence coordinates.""" + spans = [] + for block_id in range(int(image_ids.max()) + 1): + positions = (image_ids == block_id).nonzero(as_tuple=True)[0] + spans.append((int(positions[0]), int(positions[-1]) + 1)) + return spans + + def _build_block_causal_mask( + self, + image_ids: torch.Tensor, + joint_key_valid: torch.Tensor | None, + batch_size: int, + padded_seq_len: int, + device: torch.device, + ) -> torch.Tensor: + r"""Dense 4D block-causal mask `(q_idx >= kv_idx) or same_image_block`, excluding padding keys. + + Fallback path for heterogeneous/padded batches and for backends without piecewise-span support. + """ + seq_len = image_ids.shape[0] + ids = image_ids + if padded_seq_len > seq_len: + ids = F.pad(ids, (0, padded_seq_len - seq_len), value=-1) + q_idx = torch.arange(padded_seq_len, device=device).view(-1, 1) + kv_idx = torch.arange(padded_seq_len, device=device).view(1, -1) + same_image_block = (ids.view(-1, 1) == ids.view(1, -1)) & (ids.view(-1, 1) >= 0) + mask = (q_idx >= kv_idx) | same_image_block + mask = mask.view(1, 1, padded_seq_len, padded_seq_len).expand(batch_size, 1, -1, -1) + if joint_key_valid is not None: + mask = mask & joint_key_valid.view(batch_size, 1, 1, -1) + return mask.contiguous() + + def _resolve_attn_path( + self, + attn_path_hint: str, + has_padding: bool, + sp_padding: bool, + ) -> str: + r"""Choose between the piecewise-span and dense-mask attention paths. + + Returns one of `"piecewise"` (block-causal via `full_attn_spans`), `"mask"` (dense 4D mask) + or `"key_padding"` (no block-causal structure — `causal_block=False` or KV-cache decode). + """ + if not self.causal_block: + return "key_padding" + if attn_path_hint in ("piecewise", "mask"): + return attn_path_hint + if attn_path_hint != "auto": + raise ValueError(f"Unknown attn_path hint {attn_path_hint!r}; expected 'auto', 'piecewise' or 'mask'.") + if has_padding or sp_padding: + return "mask" + backend = self.transformer_blocks[0].attn.attn.attn_backend + if backend is not None and backend.supports_piecewise_spans: + return "piecewise" + return "mask" + + 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: list[dict[str, dict[str, torch.Tensor]]] | None = None, + cache_branch: str = "cond", + 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. + attention_kwargs (dict, *optional*): Extra options. Recognized key: `attn_path` — `"auto"` (default), + `"piecewise"` or `"mask"` to force the block-causal attention implementation. + kv_cache (list of dict, *optional*): One dict per block, keyed by CFG branch (`"cond"`/`"uncond"`). + Empty dicts prefill the text and condition-image keys and values; populated dicts switch to decode, + where only the target image's tokens are recomputed. Requires `causal_condition`. + cache_branch (`str`): CFG branch of the KV cache to read/write this forward. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a `Transformer2DModelOutput` instead of a plain tuple. + """ + attention_kwargs = attention_kwargs or {} + attn_path_hint = attention_kwargs.get("attn_path", "auto") + + batch_size = hidden_states.shape[0] + + if kv_cache is not None and not self.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." + ) + + # All samples must share a token layout (same interleaving of text and image blocks); + # per-sample variation enters only through `encoder_hidden_states_mask` padding. + layout = img_shapes[0] + for sample_shapes in img_shapes[1:]: + if list(sample_shapes) != list(layout): + raise ValueError("All samples in a batch must share the same img_shapes layout.") + if batch_size > 1 and not bool((img_mask == img_mask[:1]).all()): + raise ValueError("All samples in a batch must share the same img_mask layout.") + + sp_active = self.parallel_config is not None and self.parallel_config.sequence_parallel_size > 1 + ctx = get_forward_context() if is_forward_context_available() else None + sp_padding = 0 + if sp_active: + # The prefix is replicated across SP ranks for the Ulysses joint mechanism. + if ctx is not None: + ctx.split_text_embed_in_sp = False + if ( + getattr(self.parallel_config, "mask_sp_padding", False) + and ctx.sp_original_seq_len is not None + and ctx.sp_padding_size > 0 + ): + sp_padding = int(ctx.sp_padding_size) + + is_decode = kv_cache is not None and len(kv_cache) > 0 and "key" in kv_cache[0].get(cache_branch, {}) + + if is_decode: + txt_hidden_states = None + else: + txt_hidden_states = self.txt_in(encoder_hidden_states) + + prefix_hidden_states, target_hidden_states, prefix_freqs, target_freqs = self.sequence_prepare( + hidden_states, txt_hidden_states, layout, img_mask, is_decode + ) + # Local layout is always [prefix, target (shard)]: full sequence when SP is off. + joint_hidden_states = torch.cat([prefix_hidden_states, target_hidden_states], dim=1) + freqs = torch.cat([prefix_freqs, target_freqs], dim=0) + prefix_len = prefix_hidden_states.shape[1] + sp_prefix_len = prefix_len if sp_active and not is_decode else 0 + + image_pad_mask = self.sequence_prepare.expand_image_pad_mask(img_mask[0]) + _, target_token_mask = self.build_token_metadata(image_pad_mask, layout) + + timestep = timestep.to(hidden_states.dtype) + if self.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, joint_hidden_states) + modulation = self.modulation(temb) + + # Align the modulation mask with the local [prefix, target shard] layout. + if modulation_mask is not None: + local_mask = torch.cat( + [ + modulation_mask[:prefix_len], + torch.ones(target_hidden_states.shape[1], dtype=torch.bool, device=modulation_mask.device), + ] + ) + else: + local_mask = None + + # 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] + + image_ids, _ = self.build_token_metadata(image_pad_mask, layout) + full_seq_len = image_pad_mask.shape[0] + padded_seq_len = full_seq_len + sp_padding + if sp_padding and joint_key_valid is None: + joint_key_valid = torch.ones(batch_size, full_seq_len, dtype=torch.bool, device=hidden_states.device) + if sp_padding: + joint_key_valid = F.pad(joint_key_valid, (0, sp_padding), value=False) + + has_padding = joint_key_valid is not None and not bool(joint_key_valid.all()) + attn_metadata: AttentionMetadata | None = None + cache_write_len: int | None = None + if is_decode: + # Decode: block-causal degenerates to full attention for target rows; only key padding applies. + if joint_key_valid is not None: + attn_metadata = AttentionMetadata(attn_mask=joint_key_valid) + else: + attn_path = self._resolve_attn_path(attn_path_hint, has_padding, sp_padding > 0) + if kv_cache is not None: + cache_write_len = prefix_len + if attn_path == "piecewise": + spans = self._build_full_attn_spans(image_ids) + attn_metadata = AttentionMetadata(full_attn_spans=[list(spans) for _ in range(batch_size)]) + elif attn_path == "mask": + attn_metadata = AttentionMetadata( + attn_mask=self._build_block_causal_mask( + image_ids, joint_key_valid, batch_size, padded_seq_len, hidden_states.device + ) + ) + elif joint_key_valid is not None: + attn_metadata = AttentionMetadata(attn_mask=joint_key_valid) + + for index_block, block in enumerate(self.transformer_blocks): + block_kv_cache = kv_cache[index_block] if kv_cache is not None else None + joint_hidden_states = block( + hidden_states=joint_hidden_states, + modulation=modulation, + freqs=freqs, + target_token_mask=local_mask, + attn_metadata=attn_metadata, + sp_prefix_len=sp_prefix_len, + sp_decode=sp_active and is_decode, + kv_cache=block_kv_cache, + cache_branch=cache_branch, + cache_write_len=cache_write_len, + ) + + if sp_active and not is_decode: + # Only target tokens feed proj_out under SP: the prefix is replicated on every + # rank and gathering it would duplicate it. Output is the (gathered) target image. + joint_hidden_states = joint_hidden_states[:, prefix_len:] + if local_mask is not None: + local_mask = local_mask[prefix_len:] + + joint_hidden_states = self.norm_out(joint_hidden_states, temb, local_mask) + output = self.proj_out(joint_hidden_states) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".to_qkv", ".to_q", "q"), + (".to_qkv", ".to_k", "k"), + (".to_qkv", ".to_v", "v"), + ] + # Expose packed shard mappings for LoRA handling of fused projections. + self.stacked_params_mapping = stacked_params_mapping + + params_dict = dict(self.named_parameters()) + + loaded_params: set[str] = set() + for name, loaded_weight in weights: + original_name = name.removeprefix("transformer.") + lookup_name, shard_id = _resolve_qwen_image21_lookup_name( + original_name, + stacked_params_mapping, + ) + + if lookup_name.endswith(".bias") and lookup_name not in params_dict: + continue + + param = params_dict.get(lookup_name) + if param is None: + logger.warning("Skipping unexpected Qwen-Image 2.1 transformer weight %s", original_name) + continue + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if shard_id is None: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + + loaded_params.add(original_name) + loaded_params.add(lookup_name) + + unloaded = sorted(set(params_dict) - set(loaded_params)) + if unloaded: + logger.warning("Model parameters not loaded from checkpoint: %s", unloaded) + return loaded_params diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 39fff7df2..c4f923444 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -41,6 +41,11 @@ _DIFFUSION_MODELS = { "pipeline_qwen_image_layered", "QwenImageLayeredPipeline", ), + "QwenImage21Pipeline": ( + "qwen_image_21", + "pipeline_qwen_image_21", + "QwenImage21Pipeline", + ), "GlmImagePipeline": ( "glm_image", "pipeline_glm_image", @@ -378,6 +383,10 @@ _NO_CACHE_ACCELERATION = { # instead of erroring. "Pi0Pipeline", "LingBotWorldCausalDMDPipeline", + # Qwen-Image 2.1's transformer carries its own prefix KV cache (keyed by CFG + # branch) across denoising steps, which conflicts with cache_dit / tea_cache + # step-skipping hooks. + "QwenImage21Pipeline", } @@ -550,6 +559,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "QwenImagePipeline": "get_qwen_image_post_process_func", "QwenImageEditPipeline": "get_qwen_image_edit_post_process_func", "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_post_process_func", + "QwenImage21Pipeline": "get_qwen_image_21_post_process_func", "GlmImagePipeline": "get_glm_image_post_process_func", "ZImagePipeline": "get_post_process_func", "OvisImagePipeline": "get_ovis_image_post_process_func", @@ -629,6 +639,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "BooguImageTurboPipeline": "get_boogu_image_pre_process_func", "QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func", "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_pre_process_func", + "QwenImage21Pipeline": "get_qwen_image_21_pre_process_func", "LongCatImageEditPipeline": "get_longcat_image_edit_pre_process_func", "LongCatVideoAvatarPipeline": "get_longcat_video_avatar_pre_process_func", "QwenImageLayeredPipeline": "get_qwen_image_layered_pre_process_func",