From 3dc35694b3d458fc1461a8fe71832ac12e8c8cea Mon Sep 17 00:00:00 2001 From: NumberWan Date: Wed, 7 Oct 2026 21:23:59 +0800 Subject: [PATCH] [New model] Qwen-Image-2.1 support (rebase + review follow-ups) (#8099) Signed-off-by: NumberWan Signed-off-by: SamitHuang <285365963@qq.com> Signed-off-by: Samit <285365963@qq.com> Co-authored-by: Isotr0py Co-authored-by: Tianyu Guo Co-authored-by: Cursor Co-authored-by: SamitHuang <285365963@qq.com> Co-authored-by: Gao Han --- docs/models/supported_models.md | 1 + docs/user_guide/diffusion_features.md | 4 +- docs/user_guide/quantization/fp8.md | 25 +- .../image_to_image/image_edit.py | 11 +- .../image_to_image/run_qwen_image_21.sh | 16 + .../text_to_image/text_to_image.py | 3 +- recipes/Qwen/Qwen-Image-2.1.md | 418 +++++ recipes/README.md | 1 + tests/diffusion/attention/test_cudnn_attn.py | 92 +- .../test_distributed_vae_executor.py | 1 + .../distributed/test_vae_error_recovery.py | 71 + .../qwen_image_21/test_batch_compatibility.py | 192 ++ .../models/qwen_image_21/test_decode_graph.py | 318 ++++ .../models/qwen_image_21/test_pipeline.py | 402 ++++ .../test_qwen_image_21_cfg_parallel.py | 101 + .../models/qwen_image_21/test_registry.py | 64 + .../qwen_image_21/test_token_metadata.py | 46 + .../test_transformer_numerics.py | 82 + .../models/qwen_image_21/test_vae.py | 266 +++ .../offloader/test_sequential_backend.py | 98 +- .../test_diffusion_config_propagation.py | 15 + .../diffusion/test_diffusion_model_runner.py | 29 + .../e2e/online_serving/test_qwen_image_21.py | 66 + .../openai_api/test_image_server.py | 18 +- tests/model_extras/test_model_extras.py | 25 +- tests/model_tests/diffusion/test_alignment.py | 3 + vllm_omni/config/omni_config.py | 4 + vllm_omni/config/stage_config.py | 4 + .../attention/backends/cudnn_attn.py | 36 +- vllm_omni/diffusion/data.py | 17 +- .../autoencoder_kl_qwenimage21.py | 149 ++ .../autoencoders/distributed_vae_executor.py | 112 +- vllm_omni/diffusion/lora/loader.py | 5 + vllm_omni/diffusion/model_metadata.py | 8 + .../models/qwen_image_21/__init__.py | 50 + .../autoencoder_kl_qwenimage21.py | 1628 +++++++++++++++++ .../models/qwen_image_21/cfg_parallel.py | 190 ++ .../models/qwen_image_21/decode_graph.py | 457 +++++ .../qwen_image_21/pipeline_qwen_image_21.py | 1551 ++++++++++++++++ .../qwen_image_21_transformer.py | 1320 +++++++++++++ .../diffusion/offloader/module_residency.py | 17 +- .../diffusion/offloader/sequential_backend.py | 60 + vllm_omni/diffusion/registry.py | 11 + vllm_omni/diffusion/request.py | 2 + vllm_omni/diffusion/sched/step_scheduler.py | 14 + .../worker/diffusion_model_runner.py | 23 +- vllm_omni/engine/arg_utils.py | 4 + vllm_omni/entrypoints/cli/serve.py | 10 + vllm_omni/entrypoints/openai/api_server.py | 2 + vllm_omni/entrypoints/openai/serving_chat.py | 98 +- vllm_omni/model_extras/qwen_image_21.py | 28 + vllm_omni/model_extras/registry.py | 26 +- 52 files changed, 8058 insertions(+), 136 deletions(-) create mode 100644 examples/offline_inference/image_to_image/run_qwen_image_21.sh create mode 100644 recipes/Qwen/Qwen-Image-2.1.md create mode 100644 tests/diffusion/distributed/test_vae_error_recovery.py create mode 100644 tests/diffusion/models/qwen_image_21/test_batch_compatibility.py create mode 100644 tests/diffusion/models/qwen_image_21/test_decode_graph.py create mode 100644 tests/diffusion/models/qwen_image_21/test_pipeline.py create mode 100644 tests/diffusion/models/qwen_image_21/test_qwen_image_21_cfg_parallel.py create mode 100644 tests/diffusion/models/qwen_image_21/test_registry.py create mode 100644 tests/diffusion/models/qwen_image_21/test_token_metadata.py create mode 100644 tests/diffusion/models/qwen_image_21/test_transformer_numerics.py create mode 100644 tests/diffusion/models/qwen_image_21/test_vae.py create mode 100644 tests/e2e/online_serving/test_qwen_image_21.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py create mode 100644 vllm_omni/diffusion/models/qwen_image_21/__init__.py create mode 100644 vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py create mode 100644 vllm_omni/diffusion/models/qwen_image_21/cfg_parallel.py create mode 100644 vllm_omni/diffusion/models/qwen_image_21/decode_graph.py create mode 100644 vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py create mode 100644 vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py create mode 100644 vllm_omni/model_extras/qwen_image_21.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index f5d3034a5..0b9bdb21a 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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tongyi-MAI/Z-Image-Turbo.md) | | `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/krea/Krea-2.md) | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 2a2f54aa9..830b20597 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -163,6 +163,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** | ❌ | ❌ | ❓ | ✅ | ✅ | ❌ | ❓ | ❓ | ✅ (decode) | ✅ | ✅ | | **SenseNova-U1** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **SenseNova-U1.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | | **SenseNova-U1-A3B** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | @@ -181,7 +182,8 @@ The following tables show which models support each feature: > 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. MammothModa2 runs its DiT stage on the diffusion runner (`StageExecutionType.DIFFUSION`); Cache-DiT is enabled through the standard diffusion-stage knobs on the stage entry of the deploy YAML (`cache_backend: cache_dit`, plus optional `cache_config` / `enable_cache_dit_summary`). The runner installs the backend at startup and the pipeline adopts it per request. Only the repeated main-layer stack is cached; requests with `text_guidance_scale = 1.0` bypass the cache hooks. -> 8. SenseNova-U1-A3B: `--cfg-parallel-size 2` and TP=2 both complete. CFG-P covers t2i 2-branch (`cond`+`uncond`), img2img 2-branch (`cond`+`img_cond` when `img_cfg_scale=1`), and img2img 3-branch (`cond`+`img_cond`+`uncond` when `img_cfg_scale` is neither 1 nor equal to `cfg_scale`; 3 branches on 2 ranks dispatch as `[[0, 2], [1]]`). Cache-DiT does not: the adapter is declared on `SenseNovaU1Model`, but discovery walks `_dit_modules` (`language_model` / `SenseNovaU1ForCausalLM`) and raises. Layerwise CPU offload is a no-op for the same reason — that class has no `_layerwise_offload_blocks_attrs`, so the 72.2 GiB expert weights stay resident. +> 8. Qwen-Image-2.1 (`QwenImage21Pipeline`) handles both text-to-image and image-conditioned generation (up to 4 condition images) and supports LoRA through `QwenImageLoraLoaderMixin`. Ulysses sequence parallelism is wired and ring attention is not supported, but multi-GPU padding e2e and a parity number are not in this PR's H200 results, so SP stays ❓. 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. Online FP8 is documented in the quantization guide with H200 numbers. +> 9. SenseNova-U1-A3B: `--cfg-parallel-size 2` and TP=2 both complete. CFG-P covers t2i 2-branch (`cond`+`uncond`), img2img 2-branch (`cond`+`img_cond` when `img_cfg_scale=1`), and img2img 3-branch (`cond`+`img_cond`+`uncond` when `img_cfg_scale` is neither 1 nor equal to `cfg_scale`; 3 branches on 2 ranks dispatch as `[[0, 2], [1]]`). Cache-DiT does not: the adapter is declared on `SenseNovaU1Model`, but discovery walks `_dit_modules` (`language_model` / `SenseNovaU1ForCausalLM`) and raises. Layerwise CPU offload is a no-op for the same reason — that class has no `_layerwise_offload_blocks_attrs`, so the 72.2 GiB expert weights stay resident. ### VideoGen diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index 9dfe1310d..a6d40c3ef 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/docs/user_guide/quantization/fp8.md @@ -1,4 +1,4 @@ -# FP8 Quantization +# FP8 Quantization ## Overview @@ -17,7 +17,7 @@ in deep DiT blocks. ## Hardware Support | Device | Support | -|--------|---------| +| -------- | --------- | | NVIDIA Blackwell GPU (SM 100+) | ✅ | | NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | | NVIDIA Ampere GPU (SM 80+) | ✅ | @@ -93,8 +93,9 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)]) ### Diffusion Models | Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization | -|-------|-----------|:-------:|:------:|----------------|------------------|------------------| +| ------- | ----------- | :-------: | :------: | ---------------- | ------------------ | ------------------ | | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | | +| Qwen-Image-2.1 | `Qwen/Qwen-Image-2.1` | Yes | Not validated | Skip image-stream MLPs for best fidelity; all-layer FP8 is loadable. On H200 (1024×1024, teapot prompt) DiT FP8-all is ~30.4 GB / 6.06 s vs BF16 ~36.9 GB / 7.34 s (~33 dB vs BF16). Text encoder FP8 covers the language model only — the vision tower and LM head stay BF16 | `img_mlp` | ✅︎ (language model only) | | Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | | | LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | | | LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | | @@ -110,20 +111,20 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)]) ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) -| Model | Scope | Format | Status | -|-------|-------|--------|--------| +| Model | Scope | Format | Status | +| ---------- | ---------------------------- | ---------------------------------------- | ----------------------------------- | | Qwen3-Omni | Thinker language-model stage | [ModelOpt](modelopt.md) `quant_algo=FP8` | Tested for thinker memory reduction | -| Qwen3-TTS | TTS language-model stage | Checkpoint config | Not validated | +| Qwen3-TTS | TTS language-model stage | Checkpoint config | Not validated | Audio encoder, vision encoder, talker, and code2wav stay in BF16 unless a model-specific guide says otherwise. ### Multi-Stage Diffusion Model (BAGEL, GLM-Image) -| Model | Scope | Status | Notes | -|-------|-------|--------|-------| -| BAGEL | Stage-specific transformer or DiT module | Not validated | Route FP8 to the intended stage before enabling | -| GLM-Image | Stage-specific transformer or DiT module | Not validated | Validate quality against BF16 baseline | +| Model | Scope | Status | Notes | +| --------- | ---------------------------------------- | ------------- | ----------------------------------------------- | +| BAGEL | Stage-specific transformer or DiT module | Not validated | Route FP8 to the intended stage before enabling | +| GLM-Image | Stage-specific transformer or DiT module | Not validated | Validate quality against BF16 baseline | ## Configuration @@ -162,7 +163,7 @@ For a pipeline that exposes both a transformer and a quantization-aware text encoder, the scope is: | Configuration | Transformer | Text encoder | Components without supported quantizable layers | -|---------------|-------------|--------------|-------------------------------------------------| +| --------------- | ------------- | -------------- | ------------------------------------------------- | | `quantization="fp8"` | FP8 | FP8 | checkpoint precision | | `{"transformer": {"method": "fp8"}}` | FP8 | checkpoint precision | checkpoint precision | | `{"text_encoder": {"method": "fp8"}}` | checkpoint precision | FP8 | checkpoint precision | @@ -211,7 +212,7 @@ does not fit the workload. ## Parameters | Parameter | Type | Default | Description | -|-----------|------|---------|-------------| +| ----------- | ------ | --------- | ------------- | | `method` | str | - | Quantization method (`"fp8"`) | | `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16 | | `activation_scheme` | str | `"dynamic"` | `"dynamic"` selects online activation scaling, or `"static"` when scales are available | diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index cec7014de..4985117f6 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Example script for image editing with OmniGen2. @@ -209,14 +209,14 @@ def parse_profiler_config(value: str) -> dict[str, Any]: def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Edit an image with Qwen-Image-Edit.") + parser = argparse.ArgumentParser(description="Edit images with vLLM-Omni.") parser.add_argument( "--model", default="Qwen/Qwen-Image-Edit", help=( "Diffusion model name or local path. " - "For multiple image inputs, use Qwen/Qwen-Image-Edit-2509 or Qwen/Qwen-Image-Edit-2511" - "which supports QwenImageEditPlusPipeline." + "For multiple image inputs, use Qwen/Qwen-Image-Edit-2509, " + "Qwen/Qwen-Image-Edit-2511, or Qwen/Qwen-Image-2.1." ), ) parser.add_argument( @@ -368,7 +368,7 @@ def parse_args() -> argparse.Namespace: "--color-format", type=str, default="RGB", - help="For Qwen-Image-Layered, set to RGBA.", + help="For Qwen-Image-Layered or Qwen-Image-2.1, set to RGBA to preserve transparency.", ) # Cache-DiT specific parameters @@ -647,6 +647,7 @@ def main(): diffusion_params = OmniDiffusionSamplingParams( generator=generator, + seed=args.seed, true_cfg_scale=args.cfg_scale, guidance_scale=args.guidance_scale, guidance_scale_2=args.guidance_scale_2, diff --git a/examples/offline_inference/image_to_image/run_qwen_image_21.sh b/examples/offline_inference/image_to_image/run_qwen_image_21.sh new file mode 100644 index 000000000..07f1eb67e --- /dev/null +++ b/examples/offline_inference/image_to_image/run_qwen_image_21.sh @@ -0,0 +1,16 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +set -euo pipefail + +python "$(dirname "${BASH_SOURCE[0]}")/image_edit.py" \ + --model Qwen/Qwen-Image-2.1 \ + --color-format RGBA \ + --seed 42 \ + --image qwen_bear.png \ + --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \ + --output qwen_image_21_edit.png \ + --num-inference-steps 50 \ + --cfg-scale 1.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 3d7e3e868..a2f45c2ed 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -108,7 +108,8 @@ def build_parser() -> argparse.ArgumentParser: "--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..1a8f6495b --- /dev/null +++ b/recipes/Qwen/Qwen-Image-2.1.md @@ -0,0 +1,418 @@ +# 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/image_edit.py`](../../examples/offline_inference/image_to_image/image_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" \ + --output qwen_image_21_t2i.png \ + --num-inference-steps 50 \ + --cfg-scale 1.0 +``` + +### Image-conditioned generation (editing) + +```bash +python examples/offline_inference/image_to_image/image_edit.py \ + --model Qwen/Qwen-Image-2.1 \ + --color-format RGBA \ + --seed 42 \ + --image qwen_bear.png \ + --prompt "Let this mascot dance under the moon, surrounded by floating stars" \ + --output qwen_image_21_edit.png \ + --num-inference-steps 50 \ + --cfg-scale 1.0 +``` + +or via the bundled launcher: + +```bash +bash examples/offline_inference/image_to_image/run_qwen_image_21.sh +``` + +Up to 4 condition images can be passed via `--image`: + +```bash +python examples/offline_inference/image_to_image/image_edit.py \ + --model Qwen/Qwen-Image-2.1 \ + --color-format RGBA \ + --seed 42 \ + --image input1.png input2.png \ + --prompt "Combine these images into a single scene" \ + --output qwen_image_21_multi.png \ + --num-inference-steps 50 \ + --cfg-scale 1.0 +``` + +Use `--color-format RGBA` to preserve transparency in condition images. The +shared example otherwise loads inputs as RGB. Add `--width` and `--height` +(multiples of 32) to set the output size; when omitted, Qwen-Image 2.1 derives +it from the last condition image's aspect ratio at approximately 1024×1024. +Use `--num-outputs-per-prompt` to generate multiple images and `--vae-use-tiling` +to reduce VAE memory usage. The launcher accepts these options as well. + +## 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 +``` + +Step execution batches compatible requests at the same KV-cache phase. The +pipeline tags every request with a `batch_compatibility_key` (conditioning +mode plus condition-image layout), so a newly arriving request that would +otherwise join a batch already past its first denoising step is isolated by +the scheduler and deferred until the in-flight batch drains — it no longer +fails the whole batch. `--max-num-seqs` controls batch capacity. + +### CUDA Graph decode + +Qwen-Image-2.1 automatically uses CUDA Graph for supported fixed-shape denoising +decode steps when `enable_cuda_graph_decode=True` (the default) and +`enforce_eager=False`. Prefill remains eager. This works with both request +execution and `--step-execution`. + +To disable graph capture while keeping the rest of the compilation setup: + +```bash +vllm serve Qwen/Qwen-Image-2.1 --omni --no-enable-cuda-graph-decode +``` + +The same control is available as `Omni(..., enable_cuda_graph_decode=False)` +and as `enable_cuda_graph_decode: false` on a stage in the deployment YAML. +`--enforce-eager` / `enforce_eager: true` remains the broader switch: it +disables both graph capture and automatic `torch.compile`. The two +optimizations stack: the DiT blocks are regionally `torch.compile`'d and +graph capture records the compiled (fused) kernels, while inductor's own +cudagraphs stay disabled so the two graph layers never nest. Measured on a +single H200 (1024×1024, 50 steps, seed 42, BF16, `true_cfg_scale=1.0`, +warmup 1 + median of 3): compile+graph **5.93 s**, compile-only **6.88 s**, +graph-only **9.52 s**, eager **8.28 s** end-to-end per image; combo and +compile-only vs eager are **38.10 dB** PSNR (compile fusion drift), and +graph-only vs eager is bit-identical. + +The autoregressive engine's `compilation_config.cudagraph_mode` does not +control this diffusion path; use `enable_cuda_graph_decode` (or +`enforce_eager` for everything) to force eager decode execution. + +TP/SP/ring parallelism, HSDP, cache hooks, quantized KV caches, +dynamic LoRA, padded text masks, and a second in-flight request whose cache +aliases an already-owned graph key fall back to eager decode. +Model-level CPU offload does not disable decode graphs while +`enable_cuda_graph_decode` stays on (the default). DiT weights stay on fixed +device storage so captured pointers remain valid, and that storage is not +released while the text encoder runs. Pass `--no-enable-cuda-graph-decode` +(or `enable_cuda_graph_decode: false`) for the eager offload path, which +moves DiT weights off the device and frees that VRAM. Peak memory for the +default graph-plus-offload combination was not re-measured in this PR. +Graphs keep separate entries for different image layouts and copy +request-owned prefix K/V into static buffers before replay; each graph key +has a single live owner at a time, so concurrent same-shape requests never +overwrite each other's prefix. Those buffers require additional memory; graph +entries are bounded by the transformer's cache limit. + +### 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", + "size": "1024x1024", + "num_inference_steps": 50, + "true_cfg_scale": 1.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` | 1.0 | Disables the negative-prompt branch by default. Set above 1 with a `negative_prompt` to enable true CFG. 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 | Ignored at the default CFG scale of 1.0; used only when `true_cfg_scale > 1`. | +| condition images | — | Up to 4 input images per request; more than 4 raises an error. | +| `seed` | — | Fix for reproducible outputs. | + +The default CFG scale of 1.0 runs only the positive-prompt branch at each +denoising step. To enable negative-prompt guidance, explicitly set +`--cfg-scale` (offline) or `true_cfg_scale` (API) above 1 and provide a +negative prompt. This changes the generated output as well as compute cost. + +### 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. The prefix KV +cache is orthogonal to weight quantization and works unchanged under FP8. + +### FP8 quantization + +Online FP8 quantization of the transformer is supported; the fused +`to_qkv` projections load correctly, and the prefix KV cache path is +unaffected. Only block-internal linears are eligible — the boundary +projections (`img_in`/`txt_in`/`modulation`/`proj_out`/`time_text_embed`) stay +in BF16 by construction: + +```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" \ + --output qwen_image_21_fp8.png \ + --num-inference-steps 50 --cfg-scale 1.0 \ + --quantization fp8 --ignored-layers "img_mlp" +``` + +`ignored_layers` entries are matched as name patterns for this model (e.g. +`img_mlp` matches every block's `img_mlp.proj`/`img_mlp.gate_layer`/ +`img_mlp.out`); exact full prefixes like +`transformer.transformer_blocks.3.attn.to_qkv` also work. + +A per-module sensitivity sweep (fixed prompt embeds / latent / timestep, 5 +timesteps across the trajectory, relative L2 of the noise prediction vs BF16) +ranks the groups: `img_mlp` is the most sensitive (3.8% max error on its own, +driven by `img_mlp.proj` 2.8% and `img_mlp.out` 2.5%), then `attn.to_out` +(2.7%), then `attn.to_qkv` (1.6%); errors are largest at the late (low-sigma) +steps. End-to-end on H200, single GPU, same prompt as Quality Comparison +(`"A ceramic teapot on a wooden table"`, 1024×1024, 50 steps, seed 42, +eager): + +| Config | Quantized block linears | PSNR vs BF16 | Steady-state time/image | Peak memory | +| --- | --- | --- | --- | --- | +| BF16 | 0 / 160 | — | **7.34 s** | **36.9 GB** | +| FP8 all layers | 160 / 160 | **33.0 dB** | **6.06 s** | **30.4 GB** | +| FP8, `ignored_layers=["img_mlp"]` | 64 / 160 | **34.3 dB** | **6.98 s** | **34.9 GB** | + +Keeping `img_mlp` in BF16 is still the higher-fidelity DiT setting vs +all-layer FP8. On H200, all-layer FP8 is both leaner and faster than BF16 +at this shape. + +## Quantization + +Online FP8 is supported per component via `quantization_config`: + +```python +from vllm_omni import Omni + +omni = Omni( + model="Qwen/Qwen-Image-2.1", + quantization_config={ + # DiT: skip the sensitive image-stream MLPs if quality regresses. + "transformer": {"method": "fp8", "ignored_layers": ["img_mlp"]}, + # Qwen3-VL text encoder: only the language-model linear layers are + # quantized; the vision tower and the (unused) LM head stay BF16. + "text_encoder": {"method": "fp8"}, + }, +) +``` + +Either component can be quantized on its own. Measured on H200 at 1024×1024 +(seed 42, 50 steps, same teapot prompt, eager): text-encoder FP8 lowers peak +GPU memory from **36.9 GB** (BF16) to **30.4 GB** (−6.5 GB); T2I PSNR vs BF16 +**32.1 dB**. DiT FP8-all is **30.4 GB** / **6.06 s**. Combined DiT+text-encoder +FP8 was not re-measured on H200. The edit path routes condition images through +the BF16 vision tower, so it is unaffected by text-encoder FP8. See +[`docs/user_guide/quantization/fp8.md`](../../docs/user_guide/quantization/fp8.md) +for the scope rules. + +### FP8 prefix KV storage + +The cached prefix K/V can be stored in FP8 E4M3 (with per-token-per-head fp32 +scales), halving the prefix-cache memory — relevant for long prompts (up to +8192 text tokens) and up to 4 condition images (~4096 latent tokens each). +Prefill quantizes once; decode dequantizes per step, so attention compute is +unchanged. Enable it through the model-specific `extras` config (e.g. in the +deploy YAML's stage `extras` or as an `Omni(...)` engine kwarg): + +```yaml +extras: + prefix_kv_cache_dtype: "fp8" +``` + +Accepted values: `"fp8"` / `"fp8_e4m3"` (K and V quantized), `"fp8_v"` (V +only — K stays in the native dtype), and `None` / `"auto"` (default, native +dtype — behavior unchanged). This is independent of +`diffusion_kv_cache_dtype`, which quantizes attention Q/K/V *compute* per +forward pass on supported backends. + +Measured on one NVIDIA H200, T2I 1024x1024 (seed 42, 50 steps, true CFG 4.0, +eager, PSNR vs. the bf16 baseline on the same GPU): **`"fp8"` 34.2 dB**, +**`"fp8_v"` 44.0 dB**. The error is dominated by +K quantization: post-RoPE keys are the precision-sensitive half of the +cache, so `"fp8_v"` buys back about 9.8 dB at 75% (instead of 50%) of the original +cache size. The residual error is inherent e4m3 precision accumulated +coherently over the denoising trajectory; finer scale granularity +(per-tensor/per-head/per-token were compared) or Hadamard-rotated V did not +improve it. The saving scales with prefix length. These sizes are calculated at the +limit (8192 text tokens + 4 condition images, ~24.6k prefix tokens), not +measured on the 1024 T2I run above: ~12.9 GB per CFG branch in bf16, ~6.6 GB +in `"fp8"` and ~9.7 GB in `"fp8_v"`. That is where this option matters. Treat it as an opt-in for +memory-bound long-prompt / multi-image workloads, not a free lunch; +`"fp8_v"` is the better default trade-off when quality matters. + +Note: a quantized prefix cache is not CUDA-graph capturable — requests using +this option stay on the eager decode path (a warning is logged once). + +### Why not the scheduler-managed paged KV (`DiffusionKVCacheMode.PAGED_SCHEDULER`)? + +Evaluated and deliberately not adopted (see the `qwen21-p0-kv` commit message +for the full analysis): the paged scheduler targets *cross-request* KV +management — a statically sized page pool, block sharing/dedup, and eviction — +whereas this prefix cache lives exactly as long as one generation and is +sized once per request. Migrating would require routing decode attention +through the worker paged adapter (writing fresh target K/V into pages every +step), per-request page tables, and re-validating the block-causal +piecewise-span masking and Ulysses/CFG constraints, for no memory benefit over +the per-request dense tensors. The FP8 storage above captures the actual +memory win at ~1% of the integration cost. + +## Quality Comparison (head vs diffusers) + +The reference is the diffusers Qwen-Image 2.1 pipeline +([huggingface/diffusers#14804](https://github.com/huggingface/diffusers/pull/14804)). +Both sides must pin the same checkpoint, inputs, seed, precision (BF16), and +hardware. + +vLLM-Omni (this recipe): + +```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" \ + --seed 42 \ + --output qwen_image_21_vllm_omni.png \ + --num-inference-steps 50 \ + --cfg-scale 1.0 +``` + +diffusers reference: + +```python +import torch +from diffusers import QwenImage21Pipeline # diffusers PR #14804 + +pipe = QwenImage21Pipeline.from_pretrained( + "Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16 +).to("cuda") +image = pipe( + prompt="A ceramic teapot on a wooden table", + num_inference_steps=50, + true_cfg_scale=1.0, + generator=torch.Generator(device="cuda").manual_seed(42), + height=1024, + width=1024, +).images[0] +image.save("qwen_image_21_diffusers.png") +``` + +| Metric | vLLM-Omni | diffusers (#14804) | Status | +| --- | --- | --- | --- | +| Output image (1024×1024, 50 steps, seed 42, BF16, H200) | generated (eager + default CUDA-graph) | generated | measured | +| PSNR / LPIPS vs reference | eager vs Diffusers **47.17 dB**; **default compile+graph** vs Diffusers **38.02 dB**; **default compile+graph** vs eager **38.10 dB** | — | measured (H200) | +| Latency (time/image, same GPU, H200) | **default compile+graph** **6.04 s** (120.8 ms/step); eager **7.31 s** (146.2 ms/step) | **7.74 s** (154.9 ms/step) | measured | + +The prefix KV cache changes the compute path (cached prefix vs. full recompute +at every step), so exact numerical parity is not expected; match the cache +setting on both sides or compare with PSNR/LPIPS rather than exact pixels. +Eager stays numerically close to Diffusers (~47 dB). The **default** serving +path stacks regional `torch.compile` with CUDA-graph decode; on H200 that +combo lands at **~38.1 dB** vs eager (~1.21× on the Quality Comparison timed +run). Graph-only vs eager is bit-identical; the ~38 dB drift is the compile +fusion, not CUDA-graph replay. See "CUDA Graph decode" for the four-mode +H200 breakdown (warmup 1 + median of 3). + +Parity / feature measurements (same prompt/seed/steps/CFG, BF16, eager decode +unless noted; 4× NVIDIA H200): + +- [x] TP=2/4 output parity vs single-card — measured (1024×1024): TP1 vs TP2 + **39.32 dB**, TP1 vs TP4 **38.37 dB**, TP2 vs TP4 **45.77 dB**. Latency: + TP1 **7.30 s**, TP2 **4.89 s**, TP4 **5.12 s** per image. +- [x] VAE tiling vs non-tiled decode — measured: 1024×1024 nontile vs + `--vae-use-tiling` **43.55 dB** (peak mem ~37.7 GB → ~32.7 GB); 2048×2048 + nontile vs tiled **44.14 dB** (peak mem ~57.9 GB → ~33.3 GB; gen ~36.8 s vs + ~37.1 s). +- [x] CUDA-graph four-mode vs eager — measured on H200 (1024×1024, 50 steps, + seed 42, BF16, warmup 1 + median of 3): compile+graph **5.93 s**, + compile-only **6.88 s**, graph-only **9.52 s**, eager **8.28 s**. Combo / + compile-only vs eager **38.10 dB**; graph-only vs eager bit-identical. + Quality Comparison table uses a separate single-run A/B vs Diffusers + (compile+graph **6.04 s**, eager **7.31 s**). + +## Known Limitations + +- Cache acceleration backends (`cache_dit`, `tea_cache`) are not supported; + `QwenImage21Pipeline` is registered in `_NO_CACHE_ACCELERATION`. +- Ulysses sequence parallelism is wired; ring attention is not a supported path. + A short H200 smoke (1024×1024, 8 steps, seed 42, true CFG 1.0, eager) finished: + Ulysses SP=2 vs SP=1 is **46.5 dB**. That size's target length is divisible by 4, + so SP padding was not exercised, and this is not the 50-step parity next to the + TP tables. The features matrix therefore stays ❓. Ring SP=2 also ran, at + **29.9 dB** vs SP=1, which is why it stays unsupported. +- Only true CFG is exposed; there is no `guidance_scale` knob. +- VAE tiling (`--vae-use-tiling`, also implied by `--vae-patch-parallel-size > 1`) + decodes in 512px tiles with 384px stride (32 latent pixels per tile at 16x + spatial compression). On CUDA OOM the tile size is automatically halved + (keeping 25% overlap) and decoding retried, down to 128px tiles. Tile sizes + are not yet exposed as engine arguments. diff --git a/recipes/README.md b/recipes/README.md index e7fe44236..389d22024 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -87,6 +87,7 @@ recipes/ | [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB | | [`NVIDIA/SANA-WM.md`](./NVIDIA/SANA-WM.md) | First-frame image-to-video serving with camera control | 1x 24GB+ CUDA GPU; ~40GB disk | | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4; offline text-to-image on XPU | 1x A100 80GB / 2x B200 / 1x Intel Arc Pro B70 32GB | +| [`Qwen/Qwen-Image-2.1.md`](./Qwen/Qwen-Image-2.1.md) | Text-to-image and image-conditioned generation (editing) with prefix KV cache and CUDA graph decode | 1x GPU with CPU offload / 1x H200 | | [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 | | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB | | [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat; offline end-to-end text-to-speech on XPU | 1x A100 80GB / 5x Intel Arc Pro B70 32GB | diff --git a/tests/diffusion/attention/test_cudnn_attn.py b/tests/diffusion/attention/test_cudnn_attn.py index 5c787a91a..9aeb04bf3 100644 --- a/tests/diffusion/attention/test_cudnn_attn.py +++ b/tests/diffusion/attention/test_cudnn_attn.py @@ -12,9 +12,10 @@ import vllm_omni.diffusion.attention.backends.cudnn_attn as cudnn_backend from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend, CuDNNAttentionImpl -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] +@pytest.mark.cpu def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch): """Automatic CUDNN_ATTN (platform default) may use MATH for singleton K/V.""" selected_backends = [] @@ -40,6 +41,7 @@ def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch): assert selected_backends == [(SDPBackend.MATH,)] +@pytest.mark.cpu def test_explicit_cudnn_rejects_kv_seq_len_one(): impl = CuDNNAttentionImpl( num_heads=2, @@ -54,6 +56,7 @@ def test_explicit_cudnn_rejects_kv_seq_len_one(): impl.forward_cuda(query, singleton_kv, singleton_kv) +@pytest.mark.cpu def test_cudnn_backend_pins_cudnn_only_when_kv_seq_len_gt_one(monkeypatch): selected_backends = [] @@ -77,6 +80,7 @@ def test_cudnn_backend_pins_cudnn_only_when_kv_seq_len_gt_one(monkeypatch): assert selected_backends == [(SDPBackend.CUDNN_ATTENTION,)] +@pytest.mark.cpu def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch): observed: dict[str, Any] = {} @@ -112,6 +116,7 @@ def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch): assert observed["kwargs"]["attn_mask"] is None +@pytest.mark.cpu def test_cudnn_rejects_invalid_valid_kv_length(): impl = CuDNNAttentionImpl( num_heads=2, @@ -130,12 +135,97 @@ def test_cudnn_rejects_invalid_valid_kv_length(): @pytest.mark.parametrize("head_size", [8, 64, 128, 256]) +@pytest.mark.cpu def test_cudnn_backend_accepts_blackwell_fmha_head_sizes(head_size): assert CuDNNAttentionBackend.supports_head_size(head_size) assert head_size in CuDNNAttentionBackend.get_supported_head_sizes() @pytest.mark.parametrize("head_size", [0, 7, 12, 320]) +@pytest.mark.cpu def test_cudnn_backend_rejects_incompatible_head_sizes(head_size): assert not CuDNNAttentionBackend.supports_head_size(head_size) assert head_size not in CuDNNAttentionBackend.get_supported_head_sizes() + + +@pytest.mark.parametrize("batch_size", [1, 2]) +@pytest.mark.parametrize("kv_heads", [1, 2]) +@pytest.mark.parametrize("query_offset", [0, 4]) +@pytest.mark.cpu +def test_piecewise_matches_dense_block_causal_mask(monkeypatch, batch_size, kv_heads, query_offset): + monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext()) + torch.manual_seed(42) + query = torch.randn(batch_size, 12, 2, 8) + key = torch.randn(batch_size, 12, kv_heads, 8) + value = torch.randn_like(key) + spans = [(2, 5), (7, 10)] + positions = torch.arange(12) + mask = positions[:, None] >= positions[None, :] + for start, end in spans: + mask[start:end, start:end] = True + expected = torch.nn.functional.scaled_dot_product_attention( + query[:, query_offset:].transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=mask[None, None, query_offset:], + scale=0.5, + enable_gqa=kv_heads == 1, + ).transpose(1, 2) + impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5) + actual = impl.forward_cuda( + query[:, query_offset:], key, value, AttentionMetadata(full_attn_spans=[spans] * batch_size) + ) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.cpu +def test_piecewise_preserves_discontinuous_query_ranges(monkeypatch): + from vllm_omni.diffusion.attention.backends.abstract import QueryRange + + monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext()) + query, key, value = [torch.randn(2, 12, 2, 8) for _ in range(3)] + impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5) + metadata = AttentionMetadata(full_attn_spans=[[(2, 5), (7, 12)]] * 2) + expected = impl.forward_cuda(query, key, value, metadata) + metadata.query_ranges = (QueryRange(0, 3, 0), QueryRange(3, 6, 8)) + actual = impl.forward_cuda(torch.cat([query[:, :3], query[:, 8:11]], dim=1), key, value, metadata) + torch.testing.assert_close(actual, torch.cat([expected[:, :3], expected[:, 8:11]], dim=1)) + + +@pytest.mark.cpu +def test_piecewise_rejects_padding_mask(monkeypatch): + impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5) + query = torch.randn(1, 8, 2, 8) + metadata = AttentionMetadata(full_attn_spans=[[(2, 8)]], attn_mask=torch.ones(1, 8, dtype=torch.bool)) + with pytest.raises(ValueError, match="padding mask"): + impl.forward_cuda(query, query, query, metadata) + + +@pytest.mark.cpu +def test_dense_mask_takes_precedence_over_spans(monkeypatch): + monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext()) + query, key, value = [torch.randn(2, 8, 2, 8) for _ in range(3)] + mask = torch.ones(2, 1, 8, 8, dtype=torch.bool).tril() + mask[1, :, :, 3] = False + impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5) + expected = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask)) + actual = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask, full_attn_spans=[[(2, 8)]] * 2)) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.cuda +@pytest.mark.gpu +@pytest.mark.parametrize("batch_size", [1, 2]) +def test_cudnn_piecewise_matches_dense_on_device(batch_size): + generator = torch.Generator(device="cuda").manual_seed(42) + query, key, value = [ + torch.randn(batch_size, 96, 2, 128, generator=generator, device="cuda", dtype=torch.bfloat16) for _ in range(3) + ] + spans = [(8, 40), (48, 96)] + mask = torch.ones(96, 96, device="cuda", dtype=torch.bool).tril() + for start, end in spans: + mask[start:end, start:end] = True + impl = CuDNNAttentionImpl(num_heads=2, head_size=128, softmax_scale=128**-0.5) + actual = impl.forward_cuda(query, key, value, AttentionMetadata(full_attn_spans=[spans] * batch_size)) + expected = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask[None, None])) + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.004) diff --git a/tests/diffusion/distributed/test_distributed_vae_executor.py b/tests/diffusion/distributed/test_distributed_vae_executor.py index c61a680cc..29fdeb58f 100644 --- a/tests/diffusion/distributed/test_distributed_vae_executor.py +++ b/tests/diffusion/distributed/test_distributed_vae_executor.py @@ -64,6 +64,7 @@ class E2EOperator: @dataclass class FakeWorldGroup: device_group: object + cpu_group: object = None class DummyMixin(DistributedVaeMixin): diff --git a/tests/diffusion/distributed/test_vae_error_recovery.py b/tests/diffusion/distributed/test_vae_error_recovery.py new file mode 100644 index 000000000..1fcb8e009 --- /dev/null +++ b/tests/diffusion/distributed/test_vae_error_recovery.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist + +from vllm_omni.diffusion.distributed.autoencoders import distributed_vae_executor as vae_executor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class TileOperator: + def split(self, tensor): + return [vae_executor.TileTask(i, (i,), row[None]) for i, row in enumerate(tensor)], vae_executor.GridSpec( + split_dims=(0,), grid_shape=(tensor.shape[0],) + ) + + def exec(self, task): + return task.tensor * 2 + + def merge(self, tiles, grid): + return torch.cat([tiles[(i,)] for i in range(grid.grid_shape[0])]) + + +def run_recovery_worker(rank, init_file): + torch.set_num_threads(1) + dist.init_process_group( + "gloo", init_method=f"file://{init_file}", rank=rank, world_size=2, timeout=timedelta(seconds=30) + ) + try: + group = SimpleNamespace(device_group=dist.group.WORLD, cpu_group=dist.group.WORLD) + with pytest.MonkeyPatch.context() as patch: + patch.setattr(vae_executor, "get_world_group", lambda: group) + executor = vae_executor.DistributedVaeExecutor() + executor.set_parallel_size(2) + tensor = torch.arange(12, dtype=torch.float32).reshape(4, 3) + operator = TileOperator() + for stage, failing_rank in [("decode", 1), ("packing", 1), ("merge", 0), ("gather", 1)]: + with pytest.MonkeyPatch.context() as patch: + if rank == failing_rank: + + def fail(*args, **kwargs): + raise torch.OutOfMemoryError("injected OOM") + + target, attribute = { + "decode": (operator, "exec"), + "packing": (executor, "_pack_local_tiles"), + "merge": (operator, "merge"), + "gather": (torch, "empty_like"), + }[stage] + patch.setattr(target, attribute, fail) + with pytest.raises(torch.OutOfMemoryError, match="Distributed VAE"): + executor.execute(tensor, operator) + torch.testing.assert_close(executor.execute(tensor, operator), tensor * 2) + + with pytest.raises(RuntimeError, match="Distributed VAE mixed failures failed"): + with executor._sync_errors("mixed failures"): + if rank == 0: + raise torch.OutOfMemoryError("retryable") + raise ValueError("not retryable") + torch.testing.assert_close(executor.execute(tensor, operator), tensor * 2) + finally: + dist.destroy_process_group() + + +def test_all_ranks_observe_failures_and_can_retry(tmp_path): + torch.multiprocessing.spawn(run_recovery_worker, args=(str(tmp_path / "rendezvous"),), nprocs=2, join=True) diff --git a/tests/diffusion/models/qwen_image_21/test_batch_compatibility.py b/tests/diffusion/models/qwen_image_21/test_batch_compatibility.py new file mode 100644 index 000000000..c66b27e14 --- /dev/null +++ b/tests/diffusion/models/qwen_image_21/test_batch_compatibility.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import json +from types import SimpleNamespace + +import pytest +from PIL import Image + +from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import ( + get_qwen_image_21_pre_process_func, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler, StepScheduler +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.fixture +def preprocess(tmp_path): + vae_dir = tmp_path / "vae" + vae_dir.mkdir() + (vae_dir / "config.json").write_text(json.dumps({"temperal_downsample": [False] * 4})) + return get_qwen_image_21_pre_process_func(SimpleNamespace(model=str(tmp_path))) + + +def make_request(request_id, sizes, color="red"): + return OmniDiffusionRequest( + request_id=request_id, + prompt={ + "prompt": "A ceramic teapot", + "multi_modal_data": {"image": [Image.new("RGB", size, color) for size in sizes]}, + }, + sampling_params=OmniDiffusionSamplingParams( + height=1024, width=1024, num_inference_steps=2, true_cfg_scale=1.0, seed=42 + ), + ) + + +@pytest.mark.parametrize("scheduler_class", [RequestScheduler, StepScheduler]) +@pytest.mark.parametrize( + "first_sizes,second_sizes,compatible", + [ + pytest.param([], [(64, 64)], False, id="text-and-edit"), + pytest.param([(64, 64)], [], False, id="edit-and-text"), + pytest.param([(64, 64)], [(64, 64), (64, 64)], False, id="image-count"), + pytest.param([(64, 64)], [(128, 64)], False, id="image-aspect-ratio"), + pytest.param([(64, 64), (128, 64)], [(128, 64), (64, 64)], False, id="slot-order"), + pytest.param([], [], True, id="two-text-requests"), + pytest.param([(64, 64)], [(64, 64)], True, id="different-image-content"), + pytest.param([(64, 64)], [(128, 128)], True, id="same-resized-layout"), + pytest.param([(64, 64), (128, 64)], [(128, 128), (256, 128)], True, id="same-multi-image-layout"), + ], +) +def test_scheduler_groups_by_condition_layout(preprocess, scheduler_class, first_sizes, second_sizes, compatible): + first = preprocess(make_request("first", first_sizes)) + second = preprocess(make_request("second", second_sizes, color="blue")) + scheduler = scheduler_class() + scheduler.initialize(SimpleNamespace(max_num_seqs=2)) + scheduler.add_request(first) + scheduler.add_request(second) + + scheduled = scheduler.schedule() + expected_ids = ["first", "second"] if compatible else ["first"] + assert scheduled.scheduled_request_ids == expected_ids + + scheduler.finish_requests(expected_ids, DiffusionRequestStatus.FINISHED_COMPLETED) + if not compatible: + assert scheduler.schedule().scheduled_request_ids == ["second"] + scheduler.finish_requests("second", DiffusionRequestStatus.FINISHED_COMPLETED) + assert not scheduler.has_requests() + + +@pytest.mark.parametrize( + "prompt", ["A teapot", {"prompt": "A teapot"}, {"prompt": "A teapot", "multi_modal_data": {"image": None}}] +) +def test_text_input_forms_share_compatibility_key(preprocess, prompt): + request = make_request("text", []) + request.prompt = prompt + processed = preprocess(request) + empty_list_request = preprocess(make_request("empty-list", [])) + assert processed.batch_compatibility_key is not None + assert processed.batch_compatibility_key == empty_list_request.batch_compatibility_key + + +@pytest.mark.parametrize("image_type", ["numpy", "tensor"]) +def test_preprocess_accepts_array_images_without_truth_testing(preprocess, image_type): + import numpy as np + import torch + + pixels = np.full((32, 32, 4), 255, dtype=np.uint8) + image = torch.from_numpy(pixels).permute(2, 0, 1) if image_type == "tensor" else pixels + request = make_request("array", []) + request.prompt["multi_modal_data"]["image"] = image + processed = preprocess(request) + assert processed.batch_compatibility_key is not None + assert processed.allow_mixed_step_phases is False + + +def test_late_prefill_waits_until_all_decode_requests_finish(preprocess): + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput + + scheduler = StepScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=3)) + for request_id in ("a", "b"): + scheduler.add_request(preprocess(make_request(request_id, []))) + wave = scheduler.schedule() + assert wave.scheduled_request_ids == ["a", "b"] + scheduler.update_from_output( + wave, + BatchRunnerOutput.from_list( + [RunnerOutput(request_id=request_id, step_index=1, finished=False) for request_id in ("a", "b")] + ), + ) + scheduler.add_request(preprocess(make_request("late", []))) + assert scheduler.schedule().scheduled_request_ids == ["a", "b"] + scheduler.finish_requests("a", DiffusionRequestStatus.FINISHED_COMPLETED) + assert scheduler.schedule().scheduled_request_ids == ["b"] + scheduler.finish_requests("b", DiffusionRequestStatus.FINISHED_COMPLETED) + assert scheduler.schedule().scheduled_request_ids == ["late"] + + +def test_preprocess_normalizes_tuple_images_like_a_list(preprocess): + image = Image.new("RGB", (64, 64), "red") + tuple_request = make_request("tuple", []) + tuple_request.prompt["multi_modal_data"]["image"] = (image,) + + processed = preprocess(tuple_request) + reference = preprocess(make_request("list", [(64, 64)])) + + assert processed.batch_compatibility_key == reference.batch_compatibility_key + assert processed.batch_compatibility_key[0] == "qwen_image_21" + assert len(processed.batch_compatibility_key[1]) == 1 + info = processed.prompt["additional_information"] + assert len(info["prompt_image"]) == 1 + assert len(info["vae_images"]) == 1 + assert len(info["input_image_sizes"]) == 1 + + +def test_preprocess_rejects_more_images_than_the_model_supports(preprocess): + from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import ( + MAX_QWEN_IMAGE_21_INPUT_IMAGES, + ) + + request = make_request("too-many", [(64, 64)] * (MAX_QWEN_IMAGE_21_INPUT_IMAGES + 1)) + with pytest.raises(ValueError, match="At most"): + preprocess(request) + + +def test_preprocess_defaults_true_cfg_scale(preprocess): + request = make_request("cfg", []) + request.sampling_params.true_cfg_scale = None + processed = preprocess(request) + assert processed.sampling_params.true_cfg_scale == 1.0 + + request = make_request("cfg-set", []) + request.sampling_params.true_cfg_scale = 4.0 + assert preprocess(request).sampling_params.true_cfg_scale == 4.0 + + +@pytest.mark.parametrize("size,expected_wider", [((128, 64), True), ((64, 128), False)]) +def test_preprocess_derives_output_size_from_condition_aspect(preprocess, size, expected_wider): + request = make_request("aspect", [size]) + request.sampling_params.height = None + request.sampling_params.width = None + + processed = preprocess(request) + + height = processed.sampling_params.height + width = processed.sampling_params.width + assert height % 32 == 0 and width % 32 == 0 + assert (width > height) is expected_wider + info = processed.prompt["additional_information"] + assert info["calculated_height"] == height + assert info["calculated_width"] == width + + +def test_preprocess_stages_rgba_vae_images_and_resized_prompt_images(preprocess): + processed = preprocess(make_request("staged", [(96, 64)])) + info = processed.prompt["additional_information"] + + (vae_image,) = info["vae_images"] + # RGB input gains an opaque alpha channel for the RGBA VAE; the frame dim is + # inserted for the 3D (video-capable) VAE. + assert vae_image.shape[1] == 4 + assert vae_image.shape[2] == 1 + (prompt_image,) = info["prompt_image"] + assert isinstance(prompt_image, Image.Image) + (width, height) = info["input_image_sizes"][0] + assert prompt_image.size == (width, height) + assert processed.batch_compatibility_key == ("qwen_image_21", ((width, height),)) diff --git a/tests/diffusion/models/qwen_image_21/test_decode_graph.py b/tests/diffusion/models/qwen_image_21/test_decode_graph.py new file mode 100644 index 000000000..ea1141597 --- /dev/null +++ b/tests/diffusion/models/qwen_image_21/test_decode_graph.py @@ -0,0 +1,318 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import weakref + +import pytest +import torch + +from vllm_omni.diffusion.hooks import HookRegistry, ModelHook +from vllm_omni.diffusion.models.qwen_image_21 import decode_graph +from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +class DecodeModel(torch.nn.Module): + in_channels = 4 + + def _decode_graph_forward(self, entry): + return ( + entry.hidden.sin() + + entry.k[0].mean(dim=(1, 2, 3))[:, None, None] + + entry.v[0].mean(dim=(1, 2, 3))[:, None, None] + + entry.freqs[None, :, :] + ) + + +@pytest.fixture +def graph_case(request, monkeypatch): + device = "cuda" if request.node.get_closest_marker("cuda") else "cpu" + manager = decode_graph.QwenImage21DecodeGraphManager(DecodeModel(), max_entries=2) + monkeypatch.setattr(manager, "eligible", lambda: True) + monkeypatch.setattr(manager, "_backend_name", lambda: "test") + cleanup = [] + monkeypatch.setattr(decode_graph.current_omni_platform, "empty_cache", lambda: cleanup.append(True)) + return manager, device, cleanup + + +def make_cache(device, value): + return [ + { + "cond": { + "key": torch.full((1, 2, 1, 4), value, device=device), + "value": torch.full((1, 2, 1, 4), value + 1, device=device), + } + } + ] + + +def register(manager, cache, layout=(1, 2, 2)): + device = cache[0]["cond"]["key"].device + manager.register_prefill( + kv_cache=cache, + cache_branch="cond", + prefix_len=2, + img_shapes=[layout], + target_freqs=torch.full((4, 4), float(layout[1]), device=device), + joint_key_valid=None, + dtype=torch.float32, + ) + + +def decode(manager, cache, layout=(1, 2, 2)): + device = cache[0]["cond"]["key"].device + return manager.try_decode( + hidden_states=torch.ones(1, 4, 4, device=device), + timestep=torch.ones(1, device=device), + kv_cache=cache, + cache_branch="cond", + img_shapes=[layout], + img_mask=torch.tensor([[False, False, True]], device=device), + encoder_hidden_states_mask=None, + ) + + +@pytest.mark.cpu +def test_model_level_offload_without_staging_disables_decode_graphs(): + manager = decode_graph.QwenImage21DecodeGraphManager(DecodeModel(), model_level_offload=True) + assert manager._offload_reason() is not None + assert manager.eligible() is False + + +@pytest.mark.cpu +def test_top_level_offload_hook_without_staging_disables_decode_graphs(): + model = DecodeModel() + manager = decode_graph.QwenImage21DecodeGraphManager(model) + registry = HookRegistry.get_or_create(model) + # A present-but-empty registry (e.g. after hooks were removed) must not disqualify. + assert manager._offload_reason() is None + registry.register_hook("sequential_offload", ModelHook()) + assert manager._offload_reason() is not None + assert manager.eligible() is False + + +@pytest.mark.cpu +def test_persistent_staging_makes_model_level_offload_graph_eligible(): + model = DecodeModel() + manager = decode_graph.QwenImage21DecodeGraphManager(model, model_level_offload=True) + hook = SequentialOffloadHook(offload_targets=[], device=torch.device("cpu")) + HookRegistry.get_or_create(model).register_hook(SequentialOffloadHook._HOOK_NAME, hook) + # No staging storage yet: a swap would allocate fresh storage per generation. + assert manager._offload_reason() is not None + # Staging established: fixed device storage keeps captured pointers valid. + hook._stager = object() + assert manager._offload_reason() is None + # A second, non-offload hook on the top-level module still disqualifies. + model._hook_registry.register_hook("teacache", ModelHook()) + assert manager._offload_reason() is not None + + +@pytest.mark.cpu +def test_allocation_oom_releases_partial_buffers_and_remembers_failure(graph_case, monkeypatch): + manager, device, cleanup = graph_case + cache = make_cache(device, 1.0) + constructor = decode_graph.QwenImage21DecodeGraphEntry + allocated = [] + + def fail_allocation(**kwargs): + entry = constructor(**kwargs) + allocated.append(weakref.ref(entry.k[0])) + raise torch.OutOfMemoryError("injected allocation OOM") + + monkeypatch.setattr(decode_graph, "QwenImage21DecodeGraphEntry", fail_allocation) + register(manager, cache) + register(manager, cache) + assert len(allocated) == 1 + assert allocated[0]() is None + assert list(manager.entries.values()) == [None] + assert len(cleanup) == 1 + assert torch.equal(cache[0]["cond"]["key"], torch.ones(1, 2, 1, 4)) + + +@pytest.mark.cpu +def test_evict_before_allocating_replacement(graph_case, monkeypatch): + manager, device, _ = graph_case + manager.max_entries = 1 + cache = make_cache(device, 1.0) + register(manager, cache) + old_buffer = weakref.ref(next(iter(manager.entries.values())).k[0]) + constructor = decode_graph.QwenImage21DecodeGraphEntry + + def allocate_after_release(**kwargs): + assert old_buffer() is None + return constructor(**kwargs) + + monkeypatch.setattr(decode_graph, "QwenImage21DecodeGraphEntry", allocate_after_release) + register(manager, cache, (1, 1, 4)) + assert len(manager.entries) == 1 + + +@pytest.mark.cpu +def test_registration_preserves_request_cache_and_separates_layouts(graph_case): + manager, device, _ = graph_case + first, second = make_cache(device, 1.0), make_cache(device, 5.0) + original = first[0]["cond"]["key"] + register(manager, first) + register(manager, second) + assert first[0]["cond"]["key"] is original + assert torch.equal(original, torch.ones_like(original)) + register(manager, second, (1, 1, 4)) + assert len(manager.entries) == 2 + entries = list(manager.entries.values()) + assert not torch.equal(entries[0].freqs, entries[1].freqs) + + +@pytest.mark.cuda +@pytest.mark.gpu +@pytest.mark.parametrize("failure_phase", ["warmup", "capture"]) +def test_capture_failure_releases_buffers_and_keeps_request_cache(graph_case, monkeypatch, failure_phase): + manager, device, cleanup = graph_case + cache = make_cache(device, 1.0) + register(manager, cache) + entry = next(iter(manager.entries.values())) + entry_ref, buffer_ref = weakref.ref(entry), weakref.ref(entry.k[0]) + del entry + + original = manager.model._decode_graph_forward + + def fail_capture(entry): + output = original(entry) + if failure_phase == "warmup" or torch.cuda.is_current_stream_capturing(): + raise torch.OutOfMemoryError("injected capture OOM") + return output + + monkeypatch.setattr(manager.model, "_decode_graph_forward", fail_capture) + assert decode(manager, cache) is None + assert entry_ref() is None and buffer_ref() is None + assert list(manager.entries.values()) == [None] + register(manager, cache) + assert decode(manager, cache) is None + assert len(cleanup) == 1 + assert bool((cache[0]["cond"]["key"] == 1).all()) + + +@pytest.mark.cuda +@pytest.mark.gpu +def test_replay_uses_active_request_and_layout(graph_case): + manager, device, _ = graph_case + first, second = make_cache(device, 1.0), make_cache(device, 5.0) + first_key = weakref.ref(first[0]["cond"]["key"]) + with torch.inference_mode(): + register(manager, first) + # A second live request aliasing the same key is refused: it keeps its + # own cache and decodes eagerly while the owner is in flight. + register(manager, second) + assert decode(manager, second) is None + output = decode(manager, first) + expected = torch.ones_like(output).sin() + first[0]["cond"]["key"].mean() + first[0]["cond"]["value"].mean() + 2 + torch.testing.assert_close(output, expected) + # Once the owner's cache is released, the key is free to adopt. + del first + assert first_key() is None + register(manager, second) + output = decode(manager, second) + expected = ( + torch.ones_like(output).sin() + second[0]["cond"]["key"].mean() + second[0]["cond"]["value"].mean() + 2 + ) + torch.testing.assert_close(output, expected) + original = decode(manager, second).clone() + register(manager, second, (1, 1, 4)) + other = decode(manager, second, (1, 1, 4)) + torch.testing.assert_close(other, torch.ones_like(other).sin() + 5 + 6 + 1) + torch.testing.assert_close(decode(manager, second), original) + assert all(entry.captures == 1 for entry in manager.entries.values()) + + +@pytest.mark.cpu +@pytest.mark.parametrize("inference", [False, True]) +def test_prefix_reuse_refreshes_every_layer_without_retaining_requests(graph_case, inference): + manager, device, _ = graph_case + with torch.inference_mode(inference): + cache = make_cache(device, 1.0) + make_cache(device, 2.0) + register(manager, cache) + entry = next(iter(manager.entries.values())) + entry.refresh_prefix(cache) + versions = [tensor._version for tensor in entry.k + entry.v] + entry.refresh_prefix(cache) + assert versions == [tensor._version for tensor in entry.k + entry.v] + cache[1]["cond"]["value"] = torch.full_like(cache[1]["cond"]["value"], 7.0) + entry.refresh_prefix(cache) + assert bool((entry.v[1] == 7).all()) + source = weakref.ref(cache[0]["cond"]["key"]) + del cache + assert source() is None + + +@pytest.mark.cpu +def test_prefix_reuse_detects_inplace_changes_and_prefill_registration(graph_case): + manager, device, _ = graph_case + cache = make_cache(device, 1.0) + register(manager, cache) + entry = next(iter(manager.entries.values())) + entry.refresh_prefix(cache) + cache[0]["cond"]["value"].add_(3) + entry.refresh_prefix(cache) + assert bool((entry.v[0] == 5).all()) + register(manager, cache) + assert not entry.prefix_sources + + +@pytest.mark.cpu +def test_all_valid_mask_shares_graph_entry_and_padding_falls_back(graph_case): + manager, device, _ = graph_case + cache = make_cache(device, 1.0) + register(manager, cache) + entry = next(iter(manager.entries.values())) + for valid in (True, False): + manager.register_prefill( + kv_cache=cache, + cache_branch="cond", + prefix_len=2, + img_shapes=[(1, 2, 2)], + target_freqs=torch.ones(4, 4), + joint_key_valid=torch.full((1, 6), valid), + dtype=torch.float32, + ) + assert list(manager.entries.values()) == [entry] + assert entry.attn_metadata is None + + +@pytest.mark.cuda +@pytest.mark.gpu +def test_batch_changes_and_padding_do_not_reuse_stale_prefix(graph_case): + manager, device, _ = graph_case + first = make_cache(device, 1.0) + second = make_cache(device, 4.0) + merged = [ + {"cond": {part: torch.cat([first[0]["cond"][part], second[0]["cond"][part]]) for part in ("key", "value")}} + ] + with torch.inference_mode(): + register(manager, first) + single = decode(manager, first).clone() + register(manager, merged) + kwargs = dict( + hidden_states=torch.ones(2, 4, 4, device=device), + timestep=torch.ones(2, device=device), + kv_cache=merged, + cache_branch="cond", + img_shapes=[(1, 2, 2)], + img_mask=torch.tensor([[False, False, True]] * 2, device=device), + encoder_hidden_states_mask=torch.ones(2, 3, dtype=torch.bool, device=device), + ) + result = manager.try_decode(**kwargs) + assert result is not None and result.shape == (2, 4, 4) + expected = torch.ones_like(result).sin() + torch.tensor([5.0, 11.0], device=device)[:, None, None] + torch.testing.assert_close(result, expected) + # A rebuilt cache (e.g. reordered rows) is a different live cache + # aliasing the same key: decode refuses the shared entry until the + # previous owner's tensors are released. + flipped = [{"cond": {part: tensor.flip(0) for part, tensor in merged[0]["cond"].items()}}] + kwargs["kv_cache"] = flipped + assert manager.try_decode(**kwargs) is None + del merged + register(manager, flipped) + torch.testing.assert_close(manager.try_decode(**kwargs), expected.flip(0)) + kwargs["encoder_hidden_states_mask"][1, 0] = False + assert manager.try_decode(**kwargs) is None + torch.testing.assert_close(decode(manager, first), single, rtol=0, atol=0) diff --git a/tests/diffusion/models/qwen_image_21/test_pipeline.py b/tests/diffusion/models/qwen_image_21/test_pipeline.py new file mode 100644 index 000000000..0ff6d34dd --- /dev/null +++ b/tests/diffusion/models/qwen_image_21/test_pipeline.py @@ -0,0 +1,402 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU unit tests for QwenImage21Pipeline helpers: prepare_latents duplication +order, KV-cache assembly across step phases, and the post-process func.""" + +import json +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import ( + QwenImage21Pipeline, + get_qwen_image_21_post_process_func, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _toy_pipeline() -> QwenImage21Pipeline: + """A pipeline shell with the attributes prepare_latents reads, no weights.""" + pipeline = QwenImage21Pipeline.__new__(QwenImage21Pipeline) + nn.Module.__init__(pipeline) + pipeline.vae_scale_factor = 16 + pipeline.latent_channels = 4 + return pipeline + + +# --------------------------------------------------------------------------- +# prepare_latents +# --------------------------------------------------------------------------- + + +def test_prepare_latents_duplicates_condition_latents_request_major(): + """batch=2 prompts x n=2 outputs must order condition latents A,A,B,B. + + ``encode_prompt`` expands prompt embeddings request-major (A,A,B,B); the + condition latents must follow the same order or generations 2 and 3 swap + condition images. ``repeat`` would produce the wrong (A,B,A,B) tile-major + order; ``repeat_interleave`` is correct. + """ + pipeline = _toy_pipeline() + cond_a = torch.full((1, 4, 1, 8, 8), 1.0) + cond_b = torch.full((1, 4, 1, 8, 8), 2.0) + stacked = torch.cat([cond_a, cond_b], dim=0) # one row per request + + latents, image_latents = pipeline.prepare_latents( + [stacked], + batch_size=4, + num_channels_latents=4, + height=64, + width=64, + dtype=torch.float32, + device=torch.device("cpu"), + generator=torch.Generator().manual_seed(0), + ) + + assert latents.shape == (4, 16, 4) # packed: 4x4 latent tokens, 4 channels + assert image_latents.shape == (4, 64, 4) # packed: 8x8 condition tokens + row_means = image_latents.mean(dim=(1, 2)) + torch.testing.assert_close(row_means, torch.tensor([1.0, 1.0, 2.0, 2.0])) + + +def test_prepare_latents_without_duplication_when_rows_match_batch(): + pipeline = _toy_pipeline() + stacked = torch.arange(4, dtype=torch.float32).view(4, 1, 1, 1, 1).expand(4, 4, 1, 2, 2).contiguous() + + _, image_latents = pipeline.prepare_latents( + [stacked], + batch_size=4, + num_channels_latents=4, + height=32, + width=32, + dtype=torch.float32, + device=torch.device("cpu"), + generator=torch.Generator().manual_seed(0), + ) + + torch.testing.assert_close(image_latents.mean(dim=(1, 2)), torch.arange(4, dtype=torch.float32)) + + +def test_prepare_latents_rejects_non_divisible_duplication(): + pipeline = _toy_pipeline() + stacked = torch.zeros(2, 4, 1, 8, 8) + + with pytest.raises(ValueError, match="Cannot duplicate"): + pipeline.prepare_latents( + [stacked], + batch_size=3, + num_channels_latents=4, + height=64, + width=64, + dtype=torch.float32, + device=torch.device("cpu"), + generator=torch.Generator().manual_seed(0), + ) + + +def test_prepare_latents_rejects_generator_list_length_mismatch(): + pipeline = _toy_pipeline() + + with pytest.raises(ValueError, match="list of generators"): + pipeline.prepare_latents( + None, + batch_size=2, + num_channels_latents=4, + height=32, + width=32, + dtype=torch.float32, + device=torch.device("cpu"), + generator=[torch.Generator().manual_seed(0)], + ) + + +def test_prepare_latents_concatenates_multiple_condition_images(): + pipeline = _toy_pipeline() + first = torch.full((2, 4, 1, 4, 4), 1.0) + second = torch.full((2, 4, 1, 2, 2), 3.0) + + _, image_latents = pipeline.prepare_latents( + [first, second], + batch_size=2, + num_channels_latents=4, + height=32, + width=32, + dtype=torch.float32, + device=torch.device("cpu"), + generator=torch.Generator().manual_seed(0), + ) + + # Token dim concatenates the packed condition images: 16 + 4 per row. + assert image_latents.shape == (2, 20, 4) + torch.testing.assert_close(image_latents[:, :16].mean(), torch.tensor(1.0)) + torch.testing.assert_close(image_latents[:, 16:].mean(), torch.tensor(3.0)) + + +# --------------------------------------------------------------------------- +# KV-cache assembly across step phases +# --------------------------------------------------------------------------- + + +def _state(cache) -> SimpleNamespace: + return SimpleNamespace(extra={"kv_cache": cache}) + + +def _decode_cache(rows: int, value: float) -> list[dict]: + return [ + { + "cond": { + "key": torch.full((rows, 2), value), + "value": torch.full((rows, 2), value + 0.5), + } + } + ] + + +def test_kv_cache_phase_classification(): + assert QwenImage21Pipeline._kv_cache_phase(None) == "none" + assert QwenImage21Pipeline._kv_cache_phase([{}]) == "prefill" + assert QwenImage21Pipeline._kv_cache_phase(_decode_cache(1, 0.0)) == "decode" + + +def test_assemble_kv_cache_all_none(): + assert QwenImage21Pipeline._assemble_kv_cache([_state(None), _state(None)]) == (None, False) + + +def test_assemble_kv_cache_prefill_returns_fresh_owned_cache(): + states = [_state([{} for _ in range(3)]), _state([{} for _ in range(3)])] + + kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache(states) + + assert take_ownership is True + assert kv_cache == [{}, {}, {}] + # Fresh blocks, not aliases of either request's cache: the batched prefill + # writes into these and _scatter_kv_cache splits them back per request. + for block in kv_cache: + assert all(block is not request_block for state in states for request_block in state.extra["kv_cache"]) + + +def test_assemble_kv_cache_decode_concatenates_in_request_order(): + first, second = _decode_cache(2, 1.0), _decode_cache(1, 3.0) + states = [_state(first), _state(second)] + + kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache(states) + + assert take_ownership is False + assert len(kv_cache) == 1 + cond = kv_cache[0]["cond"] + torch.testing.assert_close(cond["key"], torch.tensor([[1.0, 1.0], [1.0, 1.0], [3.0, 3.0]])) + torch.testing.assert_close(cond["value"], torch.tensor([[1.5, 1.5], [1.5, 1.5], [3.5, 3.5]])) + # Decode only reads the cache: per-request tensors must stay untouched. + assert first[0]["cond"]["key"].shape == (2, 2) + assert second[0]["cond"]["key"].shape == (1, 2) + + +def test_assemble_kv_cache_decode_merges_every_part(): + first = [{"cond": {"key": torch.ones(1, 1), "value": torch.ones(1, 1), "key_scale": torch.ones(1)}}] + second = [{"cond": {"key": torch.zeros(1, 1), "value": torch.zeros(1, 1), "key_scale": torch.zeros(1)}}] + + kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache([_state(first), _state(second)]) + + assert take_ownership is False + assert set(kv_cache[0]["cond"]) == {"key", "value", "key_scale"} + torch.testing.assert_close(kv_cache[0]["cond"]["key_scale"], torch.tensor([1.0, 0.0])) + + +def test_assemble_kv_cache_mixed_phases_raises(): + """Defensive backstop: a late prefill must never join a decoding batch. + + The preprocessor sets ``allow_mixed_step_phases = False`` so the + StepScheduler defers late requests until the batch finishes decoding; this + raise is the last line of defense if a mixed batch is ever assembled. + """ + states = [_state([{}]), _state(_decode_cache(1, 0.0))] + + with pytest.raises(ValueError, match="mixed KV-cache phases"): + QwenImage21Pipeline._assemble_kv_cache(states) + + with pytest.raises(ValueError, match="mixed KV-cache phases"): + QwenImage21Pipeline._assemble_kv_cache([_state(None), _state(_decode_cache(1, 0.0))]) + + +def test_assemble_kv_cache_decode_rejects_mismatched_prefill_prompt_seq_len(): + """Homogeneous merge invariant: denoise_step must sub-batch before assembling.""" + short = _state(_decode_cache(1, 1.0)) + short.extra["prefill_prompt_seq_len"] = 4 + long = _state(_decode_cache(1, 2.0)) + long.extra["prefill_prompt_seq_len"] = 8 + + with pytest.raises(ValueError, match="sub-batch by prefill_prompt_seq_len"): + QwenImage21Pipeline._assemble_kv_cache([short, long]) + + # Same baked length remains batchable. + same_a = _state(_decode_cache(1, 1.0)) + same_a.extra["prefill_prompt_seq_len"] = 4 + same_b = _state(_decode_cache(1, 2.0)) + same_b.extra["prefill_prompt_seq_len"] = 4 + _, take_ownership = QwenImage21Pipeline._assemble_kv_cache([same_a, same_b]) + assert take_ownership is False + + +def test_split_decode_groups_by_prefill_prompt_seq_len(): + """Routine serving: separate prefills → joint decode must split, not fail.""" + short = _state(_decode_cache(1, 1.0)) + short.request_id = "short" + short.extra["prefill_prompt_seq_len"] = 4 + long = _state(_decode_cache(1, 2.0)) + long.request_id = "long" + long.extra["prefill_prompt_seq_len"] = 8 + short2 = _state(_decode_cache(1, 3.0)) + short2.request_id = "short2" + short2.extra["prefill_prompt_seq_len"] = 4 + + states = [short, long, short2] + assert QwenImage21Pipeline._needs_decode_seq_len_split(states) is True + groups = QwenImage21Pipeline._split_decode_groups(states) + assert [[s.request_id for s in g] for g in groups] == [["short", "short2"], ["long"]] + + same = [_state(_decode_cache(1, 1.0)), _state(_decode_cache(1, 2.0))] + for s in same: + s.extra["prefill_prompt_seq_len"] = 4 + assert QwenImage21Pipeline._needs_decode_seq_len_split(same) is False + assert QwenImage21Pipeline._split_decode_groups(same) == [same] + + +def test_scatter_kv_cache_splits_rows_back_per_request(): + states = [_state([{}]), _state([{}])] + merged = [ + { + "cond": { + "key": torch.tensor([[1.0], [2.0], [3.0]]), + "value": torch.tensor([[4.0], [5.0], [6.0]]), + } + } + ] + + QwenImage21Pipeline._scatter_kv_cache(states, [1, 2], merged) + + torch.testing.assert_close(states[0].extra["kv_cache"][0]["cond"]["key"], torch.tensor([[1.0]])) + torch.testing.assert_close(states[0].extra["kv_cache"][0]["cond"]["value"], torch.tensor([[4.0]])) + torch.testing.assert_close(states[1].extra["kv_cache"][0]["cond"]["key"], torch.tensor([[2.0], [3.0]])) + scattered = states[1].extra["kv_cache"][0]["cond"]["key"] + assert scattered.is_contiguous() + # Must own storage (clone), not a view into the merged prefill tensor — + # otherwise step-mode CUDA-graph decode never matches owner identity. + assert scattered.data_ptr() != merged[0]["cond"]["key"][1:3].data_ptr() + merged[0]["cond"]["key"].fill_(0.0) + torch.testing.assert_close(scattered, torch.tensor([[2.0], [3.0]])) + + +def test_encode_prompt_expands_masks_request_major(): + """Two prompts × n=2 must keep embeds and masks in A,A,B,B order. + + A 2D ``mask.repeat(1, n, 1).view(B*n, S)`` would wrongly produce A,B,A,B. + """ + pipeline = _toy_pipeline() + embeds = torch.tensor([[[1.0], [1.0]], [[2.0], [0.0]]]) # [B=2, S=2, D=1] + mask = torch.tensor([[True, True], [True, False]]) + image_pad = torch.tensor([[False, True], [False, False]]) + + def _fake_get(prompt, images_per_prompt, max_sequence_length=None, prompt_name="prompt"): + del prompt, images_per_prompt, max_sequence_length, prompt_name + return embeds.clone(), mask.clone(), image_pad.clone() + + pipeline._get_qwen_prompt_embeds = _fake_get # type: ignore[method-assign] + + out_embeds, out_mask, out_pad = pipeline.encode_prompt( + ["short", "loooooonger"], + num_images_per_prompt=2, + ) + + assert out_embeds.shape[0] == 4 + # embeds request-major: A,A,B,B + torch.testing.assert_close(out_embeds[:, 0, 0], torch.tensor([1.0, 1.0, 2.0, 2.0])) + # masks must follow the same order (not A,B,A,B) + assert out_mask.tolist() == [ + [True, True], + [True, True], + [True, False], + [True, False], + ] + assert out_pad.tolist() == [ + [False, True], + [False, True], + [False, False], + [False, False], + ] + + +def test_resolve_prompt_seq_len_keeps_prefill_padding_after_batch_shrink(): + """Longest-prompt cancel must not shrink the surviving request's KV layout.""" + short = _state([{}]) + short.extra["image_pad_mask"] = torch.tensor([[True, False]]) + short.extra["prefill_prompt_seq_len"] = 4 # baked in when batched with a longer prompt + + # Current InputBatch only sees the short prompt's native length. + assert QwenImage21Pipeline._resolve_prompt_seq_len([short], current_seq_len=2) == 4 + padded = QwenImage21Pipeline._pad_rows_to_seq_len([short.extra["image_pad_mask"]], 4) + assert padded.tolist() == [[True, False, False, False]] + + +# --------------------------------------------------------------------------- +# post-process func +# --------------------------------------------------------------------------- + + +@pytest.fixture +def postprocess(tmp_path): + vae_dir = tmp_path / "vae" + vae_dir.mkdir() + (vae_dir / "config.json").write_text(json.dumps({"temperal_downsample": [False] * 4})) + return get_qwen_image_21_post_process_func(SimpleNamespace(model=str(tmp_path))) + + +def test_postprocess_tensor_returns_pil_images(postprocess): + from PIL import Image + + images = torch.rand(2, 4, 16, 16) * 2 - 1 + + result = postprocess(images) + + assert len(result) == 2 + assert all(isinstance(image, Image.Image) for image in result) + + +def test_postprocess_envelope_unwraps_payload_and_keeps_metadata(postprocess): + from PIL import Image + + envelope = { + "payload": {"image": torch.rand(1, 4, 16, 16), "seed": 42}, + "metadata": {"request_id": "abc"}, + } + + result = postprocess(envelope) + + assert result["metadata"] == {"request_id": "abc"} + assert result["payload"]["seed"] == 42 + assert len(result["payload"]["image"]) == 1 + assert isinstance(result["payload"]["image"][0], Image.Image) + + +def test_postprocess_envelope_requires_image_payload(postprocess): + with pytest.raises(ValueError, match=r"payload\['image'\]"): + postprocess({"payload": {"seed": 42}, "metadata": {}}) + + +def test_release_captured_graphs_delegates_to_transformer() -> None: + pipeline = QwenImage21Pipeline.__new__(QwenImage21Pipeline) + released: list[bool] = [] + pipeline.transformer = SimpleNamespace(release_captured_graphs=lambda: released.append(True)) + + pipeline.release_captured_graphs() + + assert released == [True] + + +def test_postprocess_envelope_normalizes_non_dict_metadata(postprocess): + result = postprocess({"payload": {"image": torch.rand(1, 4, 16, 16)}, "metadata": "bogus"}) + + assert result["metadata"] == {} diff --git a/tests/diffusion/models/qwen_image_21/test_qwen_image_21_cfg_parallel.py b/tests/diffusion/models/qwen_image_21/test_qwen_image_21_cfg_parallel.py new file mode 100644 index 000000000..125cfd75f --- /dev/null +++ b/tests/diffusion/models/qwen_image_21/test_qwen_image_21_cfg_parallel.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU lock for the Qwen-Image-2.1 CFG tail slice. + +2.1 prepends condition latents, so the target noise is the tail. The shared +mixin slices the head. These tests stub ``predict_noise`` and check the +returned rows are the tail, including rank-0 / rank-1 gather order. +""" + +from typing import Any + +import pytest +import torch + +from vllm_omni.diffusion.models.qwen_image_21 import cfg_parallel as cfg_mod +from vllm_omni.diffusion.models.qwen_image_21.cfg_parallel import QwenImage21CFGParallelMixin + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +PREFIX = 3 +TARGET = 2 +SCALE = 4.0 + + +class _Stub(QwenImage21CFGParallelMixin): + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def predict_noise(self, **kwargs: Any) -> torch.Tensor: + self.calls.append(kwargs) + return kwargs["joint"] + + +class _Gather: + def __init__(self, rank_tails: list[torch.Tensor]) -> None: + self.rank_tails = rank_tails + self.seen: list[torch.Tensor] = [] + + def all_gather(self, input_: torch.Tensor, dim: int = 0, separate_tensors: bool = False) -> list[torch.Tensor]: + del dim + assert separate_tensors + self.seen.append(input_.clone()) + return [tensor.clone() for tensor in self.rank_tails] + + +def _joint(fill: float) -> torch.Tensor: + prefix = torch.full((1, PREFIX, 1), fill) + target = torch.full((1, TARGET, 1), fill + 100.0) + return torch.cat([prefix, target], dim=1) + + +def _cfg(positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor: + return negative + SCALE * (positive - negative) + + +def test_single_rank_cfg_uses_target_tail_not_condition_head() -> None: + pipe = _Stub() + positive = _joint(1.0) + negative = _joint(2.0) + actual = pipe.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=SCALE, + positive_kwargs={"joint": positive, "branch": "pos"}, + negative_kwargs={"joint": negative, "branch": "neg"}, + output_slice=TARGET, + ) + tail = _cfg(positive[:, -TARGET:], negative[:, -TARGET:]) + head = _cfg(positive[:, :TARGET], negative[:, :TARGET]) + torch.testing.assert_close(actual, tail) + assert not torch.allclose(actual, head) + assert [call["branch"] for call in pipe.calls] == ["pos", "neg"] + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_cfg_parallel_gather_order_uses_target_tails(monkeypatch: pytest.MonkeyPatch, rank: int) -> None: + pipe = _Stub() + positive = _joint(1.0) + negative = _joint(2.0) + pos_tail = positive[:, -TARGET:] + neg_tail = negative[:, -TARGET:] + group = _Gather([pos_tail, neg_tail]) + monkeypatch.setattr(cfg_mod, "_get_cfg_world_size_or_one", lambda: 2) + monkeypatch.setattr(cfg_mod, "get_classifier_free_guidance_rank", lambda: rank) + monkeypatch.setattr(cfg_mod, "get_cfg_group", lambda: group) + + actual = pipe.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=SCALE, + positive_kwargs={"joint": positive, "branch": "pos"}, + negative_kwargs={"joint": negative, "branch": "neg"}, + output_slice=TARGET, + ) + + expected_branch = "pos" if rank == 0 else "neg" + expected_local = pos_tail if rank == 0 else neg_tail + assert [call["branch"] for call in pipe.calls] == [expected_branch] + torch.testing.assert_close(group.seen[0], expected_local) + torch.testing.assert_close(actual, _cfg(pos_tail, neg_tail)) + head = _cfg(positive[:, :TARGET], negative[:, :TARGET]) + assert not torch.allclose(actual, head) diff --git a/tests/diffusion/models/qwen_image_21/test_registry.py b/tests/diffusion/models/qwen_image_21/test_registry.py new file mode 100644 index 000000000..2a09ea79d --- /dev/null +++ b/tests/diffusion/models/qwen_image_21/test_registry.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""L1 tests for checkpoint-to-pipeline routing of Qwen-Image 2.1.""" + +import json + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_registry_entry_loads_the_pipeline_class(): + from vllm_omni.diffusion.models.qwen_image_21 import QwenImage21Pipeline + from vllm_omni.diffusion.registry import DiffusionModelRegistry + + assert DiffusionModelRegistry._try_load_model_cls("QwenImage21Pipeline") is QwenImage21Pipeline + + +def test_pipeline_declares_request_batching_and_step_execution(): + from vllm_omni.diffusion.models.qwen_image_21 import QwenImage21Pipeline + + assert QwenImage21Pipeline.supports_request_batch + assert QwenImage21Pipeline.supports_step_execution + + +def test_process_funcs_resolve_to_the_pipeline_module(): + from vllm_omni.diffusion.models.qwen_image_21 import pipeline_qwen_image_21 + from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS, _DIFFUSION_PRE_PROCESS_FUNCS + + pre_name = _DIFFUSION_PRE_PROCESS_FUNCS["QwenImage21Pipeline"] + post_name = _DIFFUSION_POST_PROCESS_FUNCS["QwenImage21Pipeline"] + assert pre_name == "get_qwen_image_21_pre_process_func" + assert post_name == "get_qwen_image_21_post_process_func" + assert callable(getattr(pipeline_qwen_image_21, pre_name)) + assert callable(getattr(pipeline_qwen_image_21, post_name)) + + +def test_multimodal_metadata_caps_condition_images_at_four(): + from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_21_MAX_INPUT_IMAGES, get_diffusion_model_metadata + + metadata = get_diffusion_model_metadata("QwenImage21Pipeline") + + assert metadata.supports_multimodal_inputs + assert metadata.max_multimodal_image_inputs == 4 + assert metadata.max_multimodal_image_inputs == QWEN_IMAGE_21_MAX_INPUT_IMAGES + + +def test_model_index_resolves_and_enriches(tmp_path): + from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name + + (tmp_path / "model_index.json").write_text( + json.dumps({"_class_name": "QwenImage21Pipeline"}), + encoding="utf-8", + ) + + assert resolve_model_class_name(str(tmp_path)) == "QwenImage21Pipeline" + + config = OmniDiffusionConfig(model=str(tmp_path)) + config.enrich_config() + + assert config.model_class_name == "QwenImage21Pipeline" + assert config.supports_multimodal_inputs + assert config.max_multimodal_image_inputs == 4 diff --git a/tests/diffusion/models/qwen_image_21/test_token_metadata.py b/tests/diffusion/models/qwen_image_21/test_token_metadata.py new file mode 100644 index 000000000..dc75163cc --- /dev/null +++ b/tests/diffusion/models/qwen_image_21/test_token_metadata.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU lock: adjacent condition images stay separate token blocks. + +Two condition images with no text between them are one run of True in +``image_pad_mask``. Block ids come from ``img_shapes`` token counts, so they +must not collapse into one bidirectional attention span. +""" + +import pytest +import torch + +from vllm_omni.diffusion.models.qwen_image_21.qwen_image_21_transformer import ( + QwenImage21Transformer2DModel, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_adjacent_condition_images_stay_separate_blocks() -> None: + # text(2) | cond0(2) | cond1(3) | text(1) | target(2) + # cond0 and cond1 are one contiguous True run. + image_pad_mask = torch.tensor([False, False, True, True, True, True, True, False, True, True]) + img_shapes = [(1, 1, 2), (1, 1, 3), (1, 1, 2)] + + image_ids, target_token_mask = QwenImage21Transformer2DModel.build_token_metadata(image_pad_mask, img_shapes) + + assert image_ids.tolist() == [-1, -1, 0, 0, 1, 1, 1, -1, 2, 2] + assert target_token_mask.tolist() == [ + False, + False, + False, + False, + False, + False, + False, + False, + True, + True, + ] + assert QwenImage21Transformer2DModel._build_full_attn_spans(image_ids) == [ + (2, 4), + (4, 7), + (8, 10), + ] diff --git a/tests/diffusion/models/qwen_image_21/test_transformer_numerics.py b/tests/diffusion/models/qwen_image_21/test_transformer_numerics.py new file mode 100644 index 000000000..61252d01a --- /dev/null +++ b/tests/diffusion/models/qwen_image_21/test_transformer_numerics.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.models.qwen_image_21 import qwen_image_21_transformer as transformer + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +@pytest.mark.parametrize( + "device", + [pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=[pytest.mark.cuda, pytest.mark.gpu])], +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_qk_norm_matches_diffusers_rounding(monkeypatch, device, dtype): + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_tp_group", lambda: SimpleNamespace(rank_in_group=0, world_size=1) + ) + monkeypatch.setattr(transformer, "Attention", lambda **kwargs: torch.nn.Identity()) + attention = transformer.QwenImage21Attention(dim=256, heads=2, dim_head=128).to(device=device, dtype=dtype) + values = torch.randn(2, 17, 2, 128, generator=torch.Generator().manual_seed(42)).to(device=device, dtype=dtype) + for norm in (attention.norm_q, attention.norm_k): + with torch.no_grad(): + norm.weight.copy_(torch.linspace(0.5, 1.5, 128, device=device, dtype=dtype)) + variance = values.float().square().mean(-1, keepdim=True) + expected = (values * torch.rsqrt(variance + norm.eps)).to(dtype) * norm.weight + torch.testing.assert_close(norm(values), expected, rtol=0, atol=0) + + +@pytest.mark.parametrize( + "device", + [pytest.param("meta", marks=pytest.mark.cpu), pytest.param("cuda", marks=[pytest.mark.cuda, pytest.mark.gpu])], +) +def test_rope_frequencies_do_not_depend_on_loader_device(device): + reference = transformer.QwenImage21Rope(theta=10000, axes_dim=[16, 56, 56]) + with torch.device(device): + loaded = transformer.QwenImage21Rope(theta=10000, axes_dim=[16, 56, 56]) + for expected, actual in zip(reference.freqs, loaded.freqs): + assert actual.device.type == "cpu" + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.cuda +@pytest.mark.gpu +def test_scheduler_matches_diffusers_on_latent_device(): + import numpy as np + from diffusers import FlowMatchEulerDiscreteScheduler + + from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import QwenImage21Pipeline + + pipeline = object.__new__(QwenImage21Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.device = torch.device("cuda", 0) + pipeline.scheduler = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True) + timesteps, _ = pipeline.prepare_timesteps(50, None, 4096) + reference = FlowMatchEulerDiscreteScheduler.from_config(pipeline.scheduler.config) + reference.set_timesteps(sigmas=np.linspace(1.0, 1 / 50, 50), device=pipeline.device, mu=1.15) + sample = torch.randn(1, 16, 64, generator=torch.Generator().manual_seed(42)).to( + device=pipeline.device, dtype=torch.bfloat16 + ) + noise = sample.cos() + actual = pipeline.scheduler.step(noise, timesteps[0], sample, return_dict=False)[0] + expected = reference.step(noise, reference.timesteps[0], sample, return_dict=False)[0] + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert pipeline.scheduler.sigmas.device == sample.device + assert timesteps.device == sample.device + + +@pytest.mark.cuda +@pytest.mark.gpu +def test_timestep_frequencies_match_cpu_initialization(): + reference = transformer.QwenImage21TemporalTimesteps(timestep_dim=256) + with torch.device("cuda"): + loaded = transformer.QwenImage21TemporalTimesteps(timestep_dim=256) + assert loaded.freqs.device.type == "cuda" + torch.testing.assert_close(loaded.freqs.cpu(), reference.freqs, rtol=0, atol=0) + timesteps = torch.tensor([0.0, 0.984375, 1.0], device="cuda") + torch.testing.assert_close(loaded(timesteps), reference.to("cuda")(timesteps), rtol=0, atol=0) diff --git a/tests/diffusion/models/qwen_image_21/test_vae.py b/tests/diffusion/models/qwen_image_21/test_vae.py new file mode 100644 index 000000000..00159b87b --- /dev/null +++ b/tests/diffusion/models/qwen_image_21/test_vae.py @@ -0,0 +1,266 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import math + +import pytest +import torch + +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage21 import ( + DistributedAutoencoderKLQwenImage21, +) +from vllm_omni.diffusion.models.qwen_image_21.autoencoder_kl_qwenimage21 import ( + AutoencoderKLQwenImage21, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.fixture(autouse=True) +def single_thread(): + previous = torch.get_num_threads() + torch.set_num_threads(1) + yield + torch.set_num_threads(previous) + + +def _make_vae(cls=AutoencoderKLQwenImage21, **overrides): + """Toy VAE with the 2.1 residual architecture; 4x spatial compression.""" + kwargs = dict( + base_dim=16, + decoder_base_dim=16, + z_dim=4, + dim_mult=[1, 2, 2], + num_res_blocks=1, + temperal_downsample=[False, False], + is_residual=True, + in_channels=4, + out_channels=4, + latents_mean=[0.0] * 4, + latents_std=[1.0] * 4, + scale_factor_spatial=None, + scale_factor_temporal=None, + ) + kwargs.update(overrides) + return cls(**kwargs) + + +def test_spatial_compression_ratio_derived_from_architecture(): + # One 2x downsampling stage per `temperal_downsample` entry. + vae = _make_vae() + assert vae.spatial_compression_ratio == 4 + + # The real Qwen-Image-2.1 checkpoint ships scale_factor_spatial=8, but its + # four downsampling stages compress 16x; the architecture wins. + with torch.device("meta"): + vae = AutoencoderKLQwenImage21( + base_dim=96, + decoder_base_dim=144, + z_dim=64, + dim_mult=[1, 2, 4, 8, 8], + num_res_blocks=2, + temperal_downsample=[False, True, True, True], + is_residual=True, + in_channels=4, + out_channels=4, + scale_factor_spatial=8, + ) + assert vae.spatial_compression_ratio == 16 + + +def test_tiled_decode_output_shape_and_closeness(): + torch.manual_seed(0) + vae = _make_vae() + z = torch.randn(1, 4, 1, 16, 16) + ratio = vae.spatial_compression_ratio + + with torch.no_grad(): + ref = vae._decode(z, return_dict=False)[0] + vae.enable_tiling( + tile_sample_min_height=32, + tile_sample_min_width=32, + tile_sample_stride_height=24, + tile_sample_stride_width=24, + ) + tiled = vae.decode(z, return_dict=False)[0] + + # Regression: with the wrong (config-derived) ratio the tiled output was + # cropped to a fraction of the sample size. + assert tiled.shape == (1, 4, 1, 16 * ratio, 16 * ratio) + assert tiled.shape == ref.shape + + # Tiled decode blends overlapping tiles, so it is not bitwise-identical to + # untiled decode; on this toy model the seam error stays small. + mse = ((ref - tiled) ** 2).mean().item() + psnr = 10 * math.log10(4.0 / mse) + assert psnr > 20.0 + + # Tiled output is clamped to [-1, 1] like the untiled path. + assert tiled.min() >= -1.0 and tiled.max() <= 1.0 + + +def test_distributed_tile_hooks_match_tiled_decode(): + """tile_split -> tile_exec -> tile_merge must reproduce tiled_decode exactly.""" + torch.manual_seed(0) + vae = _make_vae() + vae.enable_tiling( + tile_sample_min_height=32, + tile_sample_min_width=32, + tile_sample_stride_height=24, + tile_sample_stride_width=24, + ) + distributed = _make_vae(DistributedAutoencoderKLQwenImage21) + distributed.enable_tiling( + tile_sample_min_height=32, + tile_sample_min_width=32, + tile_sample_stride_height=24, + tile_sample_stride_width=24, + ) + distributed.load_state_dict(vae.state_dict()) + + z = torch.randn(1, 4, 1, 16, 16) + with torch.no_grad(): + tiled = vae.tiled_decode(z, return_dict=False)[0] + + tasks, grid_spec = DistributedAutoencoderKLQwenImage21.tile_split(distributed, z) + # 16x16 latent, 8x8 latent tiles, stride 6 -> range(0, 16, 6) = 3 stops/axis. + assert len(tasks) == 9 + assert grid_spec.grid_shape == (3, 3) + assert grid_spec.split_dims == (3, 4) + # Every tile carries all frames (T=1 for images). + assert all(len(task.tensor) == z.shape[2] for task in tasks) + + coord_tensor_map = {task.grid_coord: distributed.tile_exec(task) for task in tasks} + merged = distributed.tile_merge(coord_tensor_map, grid_spec) + + assert merged.shape == tiled.shape + torch.testing.assert_close(merged, tiled) + + +def test_tiled_encode_output_shape_and_closeness(): + """tiled_encode is engaged for image-conditioned generation when tiling is on.""" + torch.manual_seed(0) + vae = _make_vae() + x = torch.randn(1, 4, 1, 64, 64) + + with torch.no_grad(): + ref = vae._encode(x) + vae.enable_tiling( + tile_sample_min_height=32, + tile_sample_min_width=32, + tile_sample_stride_height=24, + tile_sample_stride_width=24, + ) + tiled = vae._encode(x) + + ratio = vae.spatial_compression_ratio + assert tiled.shape == (1, 8, 1, 64 // ratio, 64 // ratio) + assert tiled.shape == ref.shape + # Blended tile borders deviate from untiled encode but stay small. + assert (ref - tiled).abs().max().item() < 1.0 + + +def test_encode_decode_round_trip_shape(): + torch.manual_seed(0) + vae = _make_vae() + x = torch.randn(1, 4, 1, 32, 32) + with torch.no_grad(): + latents = vae.encode(x, return_dict=False)[0].mode() + decoded = vae.decode(latents, return_dict=False)[0] + ratio = vae.spatial_compression_ratio + assert latents.shape == (1, 4, 1, 32 // ratio, 32 // ratio) + assert decoded.shape == x.shape + + +@pytest.mark.parametrize("persistent", [False, True]) +def test_adaptive_oom_restores_tile_configuration(monkeypatch, persistent): + vae = _make_vae() + vae.enable_tiling( + tile_sample_min_height=512, + tile_sample_min_width=512, + tile_sample_stride_height=384, + tile_sample_stride_width=384, + ) + attempts = [] + z = torch.zeros(1, 4, 1, 4, 4) + + def decode_tile(z, return_dict): + attempts.append(vae.tile_sample_min_height) + if persistent or vae.tile_sample_min_height > 128: + raise torch.OutOfMemoryError("injected tile OOM") + return (z,) + + monkeypatch.setattr(vae, "tiled_decode", decode_tile) + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None) + if persistent: + with pytest.raises(torch.OutOfMemoryError): + vae._tiled_decode_adaptive(z, return_dict=False) + else: + assert vae._tiled_decode_adaptive(z, return_dict=False)[0] is z + assert attempts == [512, 256, 128] + assert ( + vae.tile_sample_min_height, + vae.tile_sample_min_width, + vae.tile_sample_stride_height, + vae.tile_sample_stride_width, + ) == (512, 512, 384, 384) + + +def test_adaptive_oom_clears_feat_map_when_retries_exhausted(monkeypatch): + """Final min-tile OOM must not leave decoder activations on the device.""" + vae = _make_vae() + vae.enable_tiling( + tile_sample_min_height=512, + tile_sample_min_width=512, + tile_sample_stride_height=384, + tile_sample_stride_width=384, + ) + vae.clear_cache() + z = torch.zeros(1, 4, 1, 4, 4) + leaked = torch.ones(2, 2, device="cpu") # stand-in for a cached activation + + def decode_tile(z, return_dict): + # Mimic a partial tiled decode that already wrote into _feat_map before OOM. + vae._feat_map[0] = leaked + raise torch.OutOfMemoryError("injected tile OOM") + + monkeypatch.setattr(vae, "tiled_decode", decode_tile) + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None) + + with pytest.raises(torch.OutOfMemoryError): + vae._tiled_decode_adaptive(z, return_dict=False) + + assert all(slot is None for slot in vae._feat_map) + assert all(slot is None for slot in vae._enc_feat_map) + assert ( + vae.tile_sample_min_height, + vae.tile_sample_min_width, + vae.tile_sample_stride_height, + vae.tile_sample_stride_width, + ) == (512, 512, 384, 384) + + +def test_tiled_encode_oom_clears_enc_feat_map(monkeypatch): + """Tiled _encode must clear encoder activations if tiled_encode OOMs.""" + vae = _make_vae() + vae.enable_tiling( + tile_sample_min_height=32, + tile_sample_min_width=32, + tile_sample_stride_height=24, + tile_sample_stride_width=24, + ) + vae.clear_cache() + leaked = torch.ones(2, 2, device="cpu") + x = torch.zeros(1, 4, 1, 64, 64) + + def encode_tile(x_in): + vae._enc_feat_map[0] = leaked + raise torch.OutOfMemoryError("injected tiled encode OOM") + + monkeypatch.setattr(vae, "tiled_encode", encode_tile) + + with pytest.raises(torch.OutOfMemoryError): + vae._encode(x) + + assert all(slot is None for slot in vae._enc_feat_map) + assert all(slot is None for slot in vae._feat_map) diff --git a/tests/diffusion/offloader/test_sequential_backend.py b/tests/diffusion/offloader/test_sequential_backend.py index 2d3e0db3b..0916df564 100644 --- a/tests/diffusion/offloader/test_sequential_backend.py +++ b/tests/diffusion/offloader/test_sequential_backend.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for SequentialOffloadBackend.""" @@ -284,6 +284,102 @@ def test_component_selective_model_offload_requires_swap_counterpart(component, backend.enable(pipeline) +def test_persistent_dit_staging_reuses_fixed_device_storage(accelerator_device) -> None: + dit = _create_simple_module().to(accelerator_device) + encoder = _create_simple_module().to(accelerator_device) + original_weight = dit.linear.weight.detach().clone() + + apply_sequential_offload( + dit_modules=[dit], + encoder_modules=[encoder], + device=accelerator_device, + pin_memory=False, + persistent_dit_staging=True, + ) + + dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) + encoder_hook = encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) + assert dit_hook._stager is not None + assert encoder_hook._stager is None # encoders keep plain move semantics + + # Registration snapshots the CPU master and rebinds to it. + assert dit.linear.weight.device.type == "cpu" + + dit_hook._to_gpu(dit) + first_ptr = dit.linear.weight.data_ptr() + assert dit.linear.weight.device == accelerator_device + torch.testing.assert_close(dit.linear.weight.data, original_weight) + retained_storage = dit_hook._stager._device_storages[0] + retained_ptr = retained_storage.data_ptr() + + # Offload via a *different* module's hook (the cross-module swap path). + # The module parameters move back to the CPU master, but the fixed device + # storage stays allocated so CUDA-graph-captured weight pointers stay valid. + encoder_hook._to_cpu(dit) + assert dit.linear.weight.device.type == "cpu" + assert dit_hook._stager._device_storages + assert dit_hook._stager._device_storages[0].data_ptr() == retained_ptr + dit_hook._to_gpu(dit) + assert dit.linear.weight.data_ptr() == first_ptr + torch.testing.assert_close(dit.linear.weight.data, original_weight) + + # A second offload/load cycle through the DiT hook itself stays put too. + dit_hook._to_cpu(dit) + dit_hook._to_gpu(dit) + assert dit.linear.weight.data_ptr() == first_ptr + + remove_sequential_offload([dit, encoder]) + + +def test_default_sequential_offload_releases_device_storage(accelerator_device) -> None: + """Without persistent staging, offloading keeps plain move semantics.""" + dit = _create_simple_module().to(accelerator_device) + encoder = _create_simple_module().to(accelerator_device) + + apply_sequential_offload( + dit_modules=[dit], + encoder_modules=[encoder], + device=accelerator_device, + pin_memory=False, + ) + + dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) + assert dit_hook._stager is None # no retained staging: GPU storage is freed on offload + + dit_hook._to_gpu(dit) + assert dit.linear.weight.device == accelerator_device + dit_hook._to_cpu(dit) + assert dit.linear.weight.device.type == "cpu" + assert all(parameter.device.type == "cpu" for parameter in dit.parameters()) + + remove_sequential_offload([dit, encoder]) + + +@pytest.mark.parametrize( + ("decode_graphs", "expect_stager"), + [(True, True), (False, False)], +) +def test_model_level_persistent_staging_gated_on_decode_graph( + accelerator_device, decode_graphs: bool, expect_stager: bool +) -> None: + """ModelLevelOffloadBackend only retains DiT staging for decode-graph models.""" + pipeline = nn.Module() + pipeline.transformer = _create_simple_module() + if decode_graphs: + pipeline.transformer.enable_cuda_graph_decode = True + pipeline.text_encoder = _create_simple_module() + backend = ModelLevelOffloadBackend( + OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False), + accelerator_device, + ) + backend.enable(pipeline) + + hook = pipeline.transformer._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) + assert (hook._stager is not None) is expect_stager + + backend.disable() + + def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None: dit = _create_simple_module() encoder = _create_simple_module() diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index 65f309dab..701ce6af3 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -46,6 +46,21 @@ def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig: return OmniDiffusionConfig(**{name: value for name, value in diffusion_kwargs.items() if value is not None}) +def test_direct_config_extras_are_independent_dicts(): + first = OmniDiffusionConfig(model="test-model") + second = OmniDiffusionConfig(model="test-model") + + assert first.extras.get("prefix_kv_cache_dtype") is None + first.extras["prefix_kv_cache_dtype"] = "fp8" + assert second.extras == {} + + +def test_direct_config_preserves_explicit_extras(): + config = OmniDiffusionConfig(model="test-model", extras={"prefix_kv_cache_dtype": "fp8"}) + + assert config.extras["prefix_kv_cache_dtype"] == "fp8" + + class TestParallelConfigPropagation: """Core regression tests: parallel_config must survive serialization.""" diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 8b9a8f385..b3dcf5aac 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -521,6 +521,35 @@ def test_compile_transformer_regionally_compiles_blocks(monkeypatch, use_hsdp): ] +@pytest.mark.core_model +@pytest.mark.cpu +def test_compile_transformer_scopes_inductor_cudagraphs_to_decode_graph(monkeypatch): + model = SimpleNamespace(enable_cuda_graph_decode=True) + runner = _make_compile_runner(model) + compile_calls = [] + + def _regionally_compile(target, *args, **kwargs): + compile_calls.append(kwargs) + return target + + monkeypatch.setattr(model_runner_module, "regionally_compile", _regionally_compile) + import torch._inductor.config as inductor_config + + before = inductor_config.triton.cudagraphs + DiffusionModelRunner._compile_transformer(runner, "transformer") + + assert inductor_config.triton.cudagraphs == before + assert compile_calls == [ + { + "dynamic": True, + "options": { + "triton.cudagraphs": False, + "triton.cudagraph_trees": False, + }, + } + ] + + @pytest.mark.core_model @pytest.mark.cpu def test_compile_transformer_uses_regional_dynamic_false_config(monkeypatch): diff --git a/tests/e2e/online_serving/test_qwen_image_21.py b/tests/e2e/online_serving/test_qwen_image_21.py new file mode 100644 index 000000000..0a0fc6cf3 --- /dev/null +++ b/tests/e2e/online_serving/test_qwen_image_21.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import os + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.media import generate_synthetic_image +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data + +MODEL = os.environ.get("QWEN_IMAGE_21_TEST_MODEL", "Qwen/Qwen-Image-2.1") +SERVER_CASES = [ + pytest.param( + OmniServerParams(model=MODEL, server_args=["--vae-use-tiling", "--max-num-seqs", "2"] + extra), + id=name, + marks=hardware_marks(res={"cuda": "H100"}), + ) + for name, extra in [("request", []), ("step", ["--step-execution"])] +] +pytestmark = [ + pytest.mark.core_model, + pytest.mark.advanced_model, + pytest.mark.diffusion, + pytest.mark.skipif( + not os.environ.get("QWEN_IMAGE_21_TEST_MODEL"), + reason="Set QWEN_IMAGE_21_TEST_MODEL to a provisioned checkpoint until public weights are available.", + ), +] + + +@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True) +def test_text_to_image(omni_server: OmniServer, online_client: OnlineOmniClient): + online_client.send_diffusion_request( + { + "model": omni_server.model, + "messages": dummy_messages_from_mix_data(content_text="A red ceramic teapot on a wooden table."), + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 4, + "seed": 42, + "negative_prompt": "blurry, low quality", + "true_cfg_scale": 4.0, + }, + } + ) + + +@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True) +def test_image_edit(omni_server: OmniServer, online_client: OnlineOmniClient): + image = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" + online_client.send_diffusion_request( + { + "model": omni_server.model, + "messages": dummy_messages_from_mix_data(content_text="Make the image blue.", image_data_url=image), + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 4, + "seed": 42, + "negative_prompt": "blurry, low quality", + "true_cfg_scale": 4.0, + }, + } + ) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index a5eeed06c..f1be5b522 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -933,7 +933,7 @@ def test_image_edits_streaming_returns_ar_delta_then_image(streaming_image_edit_ assert payloads[0]["index"] == 0 assert payloads[1]["delta"] == " done" assert payloads[2]["output_format"] == "png" - assert payloads[2]["size"] == "16x16" + assert payloads[2]["size"] == "32x24" image_payload = payloads[2]["data"][0] img = Image.open(io.BytesIO(base64.b64decode(image_payload["b64_json"]))) @@ -2290,8 +2290,8 @@ def test_extract_images_from_result(): class EmptyResult: pass - result = EmptyResult() - images = _extract_images_from_result(result) + empty_result = EmptyResult() + images = _extract_images_from_result(empty_result) assert images == [] # Test nested batch: [np.array(shape=(3, 64, 64, 3))] @@ -2301,8 +2301,8 @@ def test_extract_images_from_result(): def __init__(self): self.images = [batch] - result = BatchResult() - images = _extract_images_from_result(result) + batch_result = BatchResult() + images = _extract_images_from_result(batch_result) assert len(images) == 3 assert all(isinstance(img, Image.Image) for img in images) assert all(img.size == (64, 64) for img in images) @@ -2312,8 +2312,8 @@ def test_extract_images_from_result(): def __init__(self): self.images = [np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)] - result = DictRequestOutput() - images = _extract_images_from_result(result) + dict_result = DictRequestOutput() + images = _extract_images_from_result(dict_result) assert len(images) == 1 assert isinstance(images[0], Image.Image) @@ -2322,8 +2322,8 @@ def test_extract_images_from_result(): def __init__(self): self.images = [np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)] - result = AttrRequestOutput() - images = _extract_images_from_result(result) + attr_result = AttrRequestOutput() + images = _extract_images_from_result(attr_result) assert len(images) == 1 assert isinstance(images[0], Image.Image) assert images[0].size == (32, 32) diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index a187c570c..0966d6baa 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -271,7 +271,7 @@ def test_ltx_reference_image_size_policy(tmp_path, model_version: str, expected: def test_reference_image_size_policy_threads_revision(monkeypatch: pytest.MonkeyPatch) -> None: - captured = {} + captured: dict[str, Any] = {} def fake_policy(*, model, revision=None): captured.update(model=model, revision=revision) @@ -495,6 +495,27 @@ def test_bagel_image_to_image_prompt_builder() -> None: assert result["negative_prompt"] == "ugly" +@pytest.mark.core_model +@pytest.mark.cpu +def test_qwen_image_21_image_to_image_prompt_builder() -> None: + dummy_image = Image.new("RGBA", (64, 64)) + result = build_image_to_image_prompt( + "QwenImage21Pipeline", + prompt="make it dance", + negative_prompt="blurry", + input_image=dummy_image, + height=1024, + width=768, + ) + assert result == { + "prompt": "make it dance", + "multi_modal_data": {"image": dummy_image}, + "height": 1024, + "width": 768, + "negative_prompt": "blurry", + } + + @pytest.mark.core_model @pytest.mark.cpu def test_unknown_pipeline_preserves_canonical_text_to_image_prompt() -> None: @@ -785,7 +806,7 @@ def _patch_hunyuan_autotokenizer(monkeypatch: pytest.MonkeyPatch, tokenizer: obj return tokenizer fake_transformers = type(sys)("transformers") - fake_transformers.AutoTokenizer = _FakeAutoTokenizer + setattr(fake_transformers, "AutoTokenizer", _FakeAutoTokenizer) monkeypatch.setitem(sys.modules, "transformers", fake_transformers) return calls diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 3b2b96adc..5daeded11 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -25,6 +25,9 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] # for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS. EXCLUDED_MODELS = [ "QwenImageLayeredPipeline", + # Dedicated CPU/unit coverage in tests/diffusion/models/qwen_image_21/. + # Tiny-model offline suite is not wired yet (2.1 transformer/VAE configs). + "QwenImage21Pipeline", "GlmImagePipeline", "ZImagePipeline", "OvisImagePipeline", diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index ba8c3f634..85712ade6 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -845,6 +845,10 @@ class _DiffusionConfigProjection: pin_cpu_memory: bool = True diffusion_compile_granularity: Literal["regional", "full"] = "regional" diffusion_compile_dynamic: bool = Field(default=True, strict=True) + # Mirrors OmniDiffusionConfig.enable_cuda_graph_decode; enforced by the + # model (Qwen-Image-2.1) together with enforce_eager, which lives in the + # shared engine config. + enable_cuda_graph_decode: bool = True fa_deterministic: bool = False vae_use_slicing: bool = False vae_use_tiling: bool = False diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index ea73c339a..d3b184be6 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -480,6 +480,10 @@ class StageDeployConfig: # Diffusion execution, cache, and VAE behavior. diffusion_compile_granularity: str | None = None diffusion_compile_dynamic: bool | None = None + # CUDA graph capture of fixed-shape KV-cache decode steps (Qwen-Image-2.1 + # today). Independent of compilation_config.cudagraph_mode; + # enforce_eager=True also disables it. + enable_cuda_graph_decode: bool | None = None fa_deterministic: bool | None = None cache_backend: str | None = None cache_config: dict[str, Any] | None = None diff --git a/vllm_omni/diffusion/attention/backends/cudnn_attn.py b/vllm_omni/diffusion/attention/backends/cudnn_attn.py index ef239a82e..fad5e528c 100644 --- a/vllm_omni/diffusion/attention/backends/cudnn_attn.py +++ b/vllm_omni/diffusion/attention/backends/cudnn_attn.py @@ -10,11 +10,13 @@ from vllm_omni.diffusion.attention.backends.abstract import ( AttentionMetadata, ) from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask +from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import piecewise_attn class CuDNNAttentionBackend(AttentionBackend): accept_output_buffer: bool = True supports_prefix_kv_slicing: bool = True + supports_piecewise_spans: bool = True # cuDNN 9.5+ FMHA on Blackwell: head_dim divisible by 8 and at most 256 # for BF16/FP16. Used by automatic platform selection; explicit CUDNN_ATTN @@ -74,6 +76,23 @@ class CuDNNAttentionImpl(AttentionImpl): value: torch.Tensor, attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: + if ( + attn_metadata is not None + and attn_metadata.full_attn_spans is not None + and (attn_metadata.attn_mask is None or attn_metadata.attn_mask.ndim != 4) + ): + if attn_metadata.attn_mask is not None or "valid_kv_length" in attn_metadata.extra: + raise ValueError("Piecewise cuDNN attention cannot combine spans with a padding mask.") + return piecewise_attn( + query, + key, + value, + attn_metadata.full_attn_spans, + self.softmax_scale, + self._piecewise_attention, + query_ranges=attn_metadata.query_ranges, + ) + attention_mask = None if attn_metadata: valid_kv_length = attn_metadata.extra.get("valid_kv_length") @@ -97,6 +116,19 @@ class CuDNNAttentionImpl(AttentionImpl): mask_mode="broadcast_k", ) + return self._attention(query, key, value, attention_mask, self.causal, self.softmax_scale) + + def _piecewise_attention(self, query, key, value, *, causal, softmax_scale): + mask = None + if causal: + # SDPA's causal flag aligns at the top left. Text segments need + # bottom-right alignment so every query can also see its prefix. + q_len, kv_len = query.shape[1], key.shape[1] + mask = torch.ones(q_len, kv_len, dtype=torch.bool, device=query.device).tril(kv_len - q_len) + mask = mask[None, None] + return self._attention(query, key, value, mask, False, softmax_scale) + + def _attention(self, query, key, value, attention_mask, causal, softmax_scale): enable_gqa = query.shape[2] != key.shape[2] kv_seq_len = key.shape[1] query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) @@ -130,8 +162,8 @@ class CuDNNAttentionImpl(AttentionImpl): value, attn_mask=attention_mask, dropout_p=0.0, - is_causal=self.causal, - scale=self.softmax_scale, + is_causal=causal, + scale=softmax_scale, enable_gqa=enable_gqa, ) return output.permute(0, 2, 1, 3) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 477a62720..a258452bf 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -16,7 +16,7 @@ import diffusers import huggingface_hub import torch from PIL import Image -from pydantic import Field, model_validator +from pydantic import model_validator from typing_extensions import Self from vllm.config.utils import config from vllm.logger import init_logger @@ -981,6 +981,19 @@ class OmniDiffusionConfig: # Compilation enforce_eager: bool = False + # Capture fixed-shape KV-cache decode (denoising) steps into CUDA graphs. + # Currently only Qwen-Image-2.1's transformer implements this; other models + # ignore the flag. This is a per-model diffusion path, independent of the + # AR engine's ``compilation_config.cudagraph_mode`` (which does not apply + # to diffusion stages). It stacks with ``diffusion_compile_granularity``: + # the DiT blocks are still torch.compile'd and graph capture records the + # compiled (fused) kernels, while inductor's own cudagraphs stay off. + # ``enforce_eager=True`` forces eager decode and + # disables capture regardless of this flag; unsupported configurations + # (SP/TP, ring, HSDP, dynamic LoRA, padded masks, quantized prefix KV, or + # a second in-flight request aliasing the same graph key) log and fall + # back to eager decode. + enable_cuda_graph_decode: bool = True # Controls the generic compilation path used when a pipeline does not # provide its own setup_compile() implementation. diffusion_compile_granularity: str = "regional" @@ -1140,7 +1153,7 @@ class OmniDiffusionConfig: request_batch_max_wait_ms: float = 0.0 # Supplementary model specific parameters - extras: dict[str, Any] = Field(default_factory=dict) + extras: dict[str, Any] = field(default_factory=dict) @property def is_moe(self) -> bool: 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..ffda5c621 --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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, + _cudnn_deterministic, + _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: list[TileTask] = [] + 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 = [] + with _cudnn_deterministic(): + 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/distributed/autoencoders/distributed_vae_executor.py b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py index 552541d25..4880f410f 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py +++ b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py @@ -3,6 +3,8 @@ from __future__ import annotations +from collections.abc import Callable +from contextlib import contextmanager from dataclasses import dataclass, field import torch @@ -34,9 +36,9 @@ class TileTask: @dataclass class DistributedOperator: - split: callable - exec: callable - merge: callable + split: Callable[[torch.Tensor], tuple[list[TileTask], GridSpec]] + exec: Callable[[TileTask], torch.Tensor] + merge: Callable[[dict[tuple[int, ...], torch.Tensor], GridSpec], torch.Tensor] class DistributedVaeExecutor: @@ -46,7 +48,11 @@ class DistributedVaeExecutor: def __init__(self): # Use a dedicated process group spanning the complete worker WORLD. - self.group = get_world_group().device_group + world_group = get_world_group() + self.group = world_group.device_group + self._error_group = world_group.cpu_group + # Reporting an OOM must not require another device allocation. + self._error_status = torch.zeros((), dtype=torch.int32, device="cpu") self.world_size = dist.get_world_size(self.group) self.rank = dist.get_rank(self.group) self.parallel_size = 1 @@ -56,8 +62,34 @@ class DistributedVaeExecutor: self.parallel_size = parallel_size self.parallel_mode = mode + @contextmanager + def _sync_errors(self, stage: str): + """Agree on local-stage failures before entering the next collective.""" + status = 0 + message = "" + try: + yield + except torch.OutOfMemoryError as error: + status = 1 + message = str(error) + except Exception as error: + # A non-retryable error takes precedence over OOM on another rank. + status = 2 + message = f"{type(error).__name__}: {error}" + + self._error_status.fill_(status) + dist.all_reduce(self._error_status, op=dist.ReduceOp.MAX, group=self._error_group) + global_status = self._error_status.item() + if global_status == 2: + detail = message if status == 2 else "another rank failed" + raise RuntimeError(f"Distributed VAE {stage} failed: {detail}") + if global_status == 1: + detail = message if status == 1 else "another rank ran out of memory" + raise torch.OutOfMemoryError(f"Distributed VAE {stage} failed: {detail}") + def gather_tensors(self, tensor: torch.Tensor): - gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)] + with self._sync_errors("gather allocation"): + gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)] dist.all_gather(gather_list, tensor, group=self.group) return gather_list if self.rank == 0 else None @@ -66,11 +98,12 @@ class DistributedVaeExecutor: return tensor def _compute_global_padding_shape(self, local_results, output_ndim: int, device): - local_tile_max_dims = [0] * output_ndim - for _, tile_tensor in local_results: - for dim_idx, dim_size in enumerate(tile_tensor.shape): - local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size) - local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device) + with self._sync_errors("shape preparation"): + local_tile_max_dims = [0] * output_ndim + for _, tile_tensor in local_results: + for dim_idx, dim_size in enumerate(tile_tensor.shape): + local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size) + local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device) dist.all_reduce( local_shape_stat, op=dist.ReduceOp.MAX, @@ -100,17 +133,16 @@ class DistributedVaeExecutor: tid = int(meta_src[idx, 0]) if tid < 0: continue - slices = [slice(None)] * tiles_src[idx].ndim + slices_list = [slice(None)] * tiles_src[idx].ndim for i, dim in enumerate(grid_spec.split_dims): - slices[dim] = slice(0, int(meta_src[idx, i + 1])) - slices = tuple(slices) - coord_tensor_map[tid_coord_map[tid]] = tiles_src[idx][slices] + slices_list[dim] = slice(0, int(meta_src[idx, i + 1])) + coord_tensor_map[tid_coord_map[tid]] = tiles_src[idx][tuple(slices_list)] return coord_tensor_map def _balance_tasks(self, task_list, num_rank): workloads = [0] * num_rank - assigned = [[] for _ in range(num_rank)] + assigned: list[list[TileTask]] = [[] for _ in range(num_rank)] for task in sorted(task_list, key=lambda t: t.workload, reverse=True): r = workloads.index(min(workloads)) @@ -122,49 +154,51 @@ class DistributedVaeExecutor: def execute(self, z: torch.Tensor, operator: DistributedOperator, broadcast_result: bool = True): pp_size = min(self.parallel_size, self.world_size) - # 1. Split into tiles - tiletask_list, grid_spec = operator.split(z) - tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list} - - # 2. local decode - assigned = self._balance_tasks(tiletask_list, pp_size) - local_tasks = assigned[self.rank] if self.rank < pp_size else [] - local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks] + with self._sync_errors("tile decode"): + tiletask_list, grid_spec = operator.split(z) + tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list} + assigned = self._balance_tasks(tiletask_list, pp_size) + local_tasks = assigned[self.rank] if self.rank < pp_size else [] + local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks] # 3. compute shape per rank global_padding_shape = self._compute_global_padding_shape(local_results, z.ndim, z.device) # 5. prepare tile tensors output_dtype = grid_spec.output_dtype if grid_spec.output_dtype is not None else z.dtype - local_tile_tensor, local_meta_tensor = self._pack_local_tiles( - local_results, global_padding_shape, grid_spec, z.device, output_dtype - ) + with self._sync_errors("tile packing"): + local_tile_tensor, local_meta_tensor = self._pack_local_tiles( + local_results, global_padding_shape, grid_spec, z.device, output_dtype + ) # 6. gather tiles & meta meta_gather = self.gather_tensors(local_meta_tensor) tile_gather = self.gather_tensors(local_tile_tensor) - if self.rank != 0: - result = torch.empty(0, device=z.device) # Dummy return for non-zero ranks - else: - # 7. reconstruct full tensor (rank 0) - coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map) - result = operator.merge(coord_tensor_map, grid_spec) + with self._sync_errors("tile merge"): + if self.rank != 0: + result = torch.empty(0, device=z.device) # Dummy return for non-zero ranks + else: + # 7. reconstruct full tensor (rank 0) + coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map) + result = operator.merge(coord_tensor_map, grid_spec) if broadcast_result: result = self._sync_final_result(result, z.ndim, z.device, output_dtype) return result def _sync_final_result(self, rank0_result, output_ndim, output_device, output_dtype): - shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64) - if self.rank == 0: - shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64)) + with self._sync_errors("output shape preparation"): + shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64) + if self.rank == 0: + shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64)) dist.broadcast(shape_tensor, src=0, group=self.group) - if self.rank != 0: - sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype) - else: - sync_result = rank0_result + with self._sync_errors("output allocation"): + if self.rank != 0: + sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype) + else: + sync_result = rank0_result dist.broadcast(sync_result, src=0, group=self.group) return sync_result diff --git a/vllm_omni/diffusion/lora/loader.py b/vllm_omni/diffusion/lora/loader.py index 8ed60ea06..d4b3cdb7c 100644 --- a/vllm_omni/diffusion/lora/loader.py +++ b/vllm_omni/diffusion/lora/loader.py @@ -23,6 +23,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 9e924370d..ba91e2e21 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -26,6 +26,8 @@ class DiffusionModelMetadata: # FLUX.2 Klein supports up to four reference images. FLUX2_KLEIN_MAX_INPUT_IMAGES = 4 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. @@ -42,6 +44,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..8d1dcc168 --- /dev/null +++ b/vllm_omni/diffusion/models/qwen_image_21/__init__.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Qwen Image 2.1 diffusion model components. + +Lazy by necessity: the pipeline imports the distributed VAE +(`vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage21`), +which in turn imports the plain VAE module from this package. Eager imports +here would make that a circular import whenever the distributed module is +imported first. +""" + +from __future__ import annotations + +from typing import Any + +__all__ = [ + "AutoencoderKLQwenImage21", + "QwenImage21Pipeline", + "QwenImage21Transformer2DModel", + "get_qwen_image_21_post_process_func", + "get_qwen_image_21_pre_process_func", +] + + +def __getattr__(name: str) -> Any: + if name == "AutoencoderKLQwenImage21": + from .autoencoder_kl_qwenimage21 import AutoencoderKLQwenImage21 + + return AutoencoderKLQwenImage21 + if name == "QwenImage21Transformer2DModel": + from .qwen_image_21_transformer import QwenImage21Transformer2DModel + + return QwenImage21Transformer2DModel + if name in ( + "QwenImage21Pipeline", + "get_qwen_image_21_post_process_func", + "get_qwen_image_21_pre_process_func", + ): + from .pipeline_qwen_image_21 import ( + QwenImage21Pipeline, + get_qwen_image_21_post_process_func, + get_qwen_image_21_pre_process_func, + ) + + return { + "QwenImage21Pipeline": QwenImage21Pipeline, + "get_qwen_image_21_post_process_func": get_qwen_image_21_post_process_func, + "get_qwen_image_21_pre_process_func": get_qwen_image_21_pre_process_func, + }[name] + raise AttributeError(name) 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..57c51de7a --- /dev/null +++ b/vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py @@ -0,0 +1,1628 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +# 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 functools +from contextlib import contextmanager + +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 + + +@contextmanager +def _cudnn_deterministic(): + """Scope cuDNN to deterministic conv algorithms, restoring on exit. + + VAE decode/encode runs many odd-shaped conv3d shapes; under GPU memory + pressure cuDNN may fall back to workspace-limited or non-deterministic + algorithms, which makes tiled decode results vary run-to-run on shared + GPUs. Pinning deterministic algorithms removes that source of variance. + """ + previous = torch.backends.cudnn.deterministic + torch.backends.cudnn.deterministic = True + try: + yield + finally: + torch.backends.cudnn.deterministic = previous + + +def _with_cudnn_deterministic(fn): + @functools.wraps(fn) + def wrapper(*args, **kwargs): + with _cudnn_deterministic(): + return fn(*args, **kwargs) + + return wrapper + + +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): + needs_fp32_normalize = x.dtype in (torch.float16, torch.bfloat16) or any( + t in str(x.dtype) for t in ("float4_", "float8_") + ) + normalized = F.normalize(x.float() if needs_fp32_normalize else x, dim=(1 if self.channel_first else -1)).to( + x.dtype + ) + return normalized * self.scale * self.gamma + self.bias + + +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 + temporal_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, + temporal_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 temporal_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 temporal_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. + temporal_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=[], + temporal_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.temporal_upsample = temporal_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 temporal_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, + temporal_upsample=temporal_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.temporal_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, + temporal_upsample=self.temporal_upsample, + dropout=dropout, + out_channels=out_channels, + is_residual=is_residual, + ) + + # Derive the spatial compression ratio from the architecture: one + # downsampling stage per `temperal_downsample` entry. The checkpoint's + # `scale_factor_spatial` (8) under-reports it — with + # `dim_mult=[1, 2, 4, 8, 8]` there are four 2x stages, i.e. 16x — and a + # wrong ratio corrupts tiled decode/encode geometry. + self.spatial_compression_ratio = 2 ** len(temperal_downsample) + if scale_factor_spatial is not None and scale_factor_spatial != self.spatial_compression_ratio: + logger.warning( + f"VAE config scale_factor_spatial={scale_factor_spatial} does not match the architecture's " + f"spatial compression ratio {self.spatial_compression_ratio}; using the architecture-derived value." + ) + + # 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. + # Sized for the architecture-derived 16x compression ratio: the + # diffusers default of 256/192 targets 8x VAEs (32x32 latent tiles); + # at 16x it would yield only 16x16 latent tiles, which starves the + # decoder's receptive field and visibly degrades quality (probe: + # PSNR vs untiled 39.6 dB @ 256/192 -> 47.7 dB @ 512/384). + self.tile_sample_min_height = 512 + self.tile_sample_min_width = 512 + + # The minimal distance between two spatial tiles (75% of the tile + # size, i.e. 25% overlap) + self.tile_sample_stride_height = 384 + self.tile_sample_stride_width = 384 + + # 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 + + @_with_cudnn_deterministic + 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) + + try: + 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) + + return self.quant_conv(out) + finally: + # Drop cached activations even if tiled_encode or the loop OOMs. + self.clear_cache() + + @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) + + @_with_cudnn_deterministic + 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_adaptive(z, return_dict=return_dict) + + self.clear_cache() + try: + 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) + finally: + # Drop cached activations even if the loop OOMs mid-cat. + 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 + + @_with_cudnn_deterministic + 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_adaptive(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: + r"""Decode with the configured tile size, halving it on CUDA OOM. + + Tiled decode favors the largest tile that fits in memory, since larger + latent tiles give the decoder more spatial context and decode closer to + the untiled result. On OOM the tile size is halved (keeping the 25% + overlap ratio) and decoding is retried, down to a floor of 128px + sample-space tiles. + """ + saved = ( + self.tile_sample_min_height, + self.tile_sample_min_width, + self.tile_sample_stride_height, + self.tile_sample_stride_width, + ) + try: + while True: + try: + return self.tiled_decode(z, return_dict=return_dict) + except torch.OutOfMemoryError: + next_height = self.tile_sample_min_height // 2 + next_width = self.tile_sample_min_width // 2 + if min(next_height, next_width) < 128: + raise + logger.warning( + f"VAE tiled decode OOM with {self.tile_sample_min_height}x{self.tile_sample_min_width} " + f"tiles; retrying with {next_height}x{next_width}." + ) + self.tile_sample_min_height = next_height + self.tile_sample_min_width = next_width + self.tile_sample_stride_height = next_height * 3 // 4 + self.tile_sample_stride_width = next_width * 3 // 4 + self.clear_cache() + torch.accelerator.empty_cache() + finally: + # Always drop decoder activations: exhausting the min-tile OOM path + # raises before tiled_decode's trailing clear_cache, and _decode + # returns through this adaptive branch so it never hits its own + # cleanup finally either. + self.clear_cache() + ( + self.tile_sample_min_height, + self.tile_sample_min_width, + self.tile_sample_stride_height, + self.tile_sample_stride_width, + ) = saved + + @_with_cudnn_deterministic + 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/decode_graph.py b/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py new file mode 100644 index 000000000..655e4d6cc --- /dev/null +++ b/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py @@ -0,0 +1,457 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA Graph capture for Qwen-Image-2.1 KV-cache decode steps. + +Decode (denoising steps 1..N) recomputes only the target image's tokens: the +shape is identical every step and the timestep-independent prefix K/V come from +the cross-step KV cache. That is the textbook CUDA-graph case, except the +default cache stores prefix K/V in freshly allocated tensors, so a captured +graph would bind addresses that the next request (or a step-mode cache merge) +no longer owns. + +Requests retain ownership of their prefix K/V. Decode refreshes separate +fixed-address graph buffers when the active prefix tensors change. Graphs +are keyed by branch, batch, prefix length, complete image layout, +dtype, device and backend, so equal token counts with different RoPE layouts stay separate. + +Each key admits a single live owner, tracked by a weak reference to the +registering request's prefix K tensor. A second in-flight request whose cache +aliases the same key is refused graph decode and stays eager until the +owner's cache is released, so shared static buffers are never overwritten +under a concurrent owner. + +Memory note: graph entries own an additional copy of the prefix K/V. Decode +attention concatenates that prefix with target K/V inside the captured region; +transient tensors use the graph memory pool. + +The model runner regionally torch.compile's the transformer blocks even when +decode graphs are enabled, so capture records the compiled (fused) kernels; +inductor's own cudagraphs stay off so the two graph layers never stack. + +Fallbacks (all logged, all eager): CUDA unavailable, model-level CPU offload +without persistent weight staging, sequence/ring/tensor parallelism, +HSDP/offload hooks, KV-cache quantization, dynamic LoRA wrappers, padded text +masks (the masked attention path branches on mask contents, which cannot be +captured), unknown graph key at decode, a second live request aliasing an +owned graph key, and capture failure. +""" + +from __future__ import annotations + +import weakref +from collections import OrderedDict +from typing import TYPE_CHECKING + +import torch +from vllm.logger import init_logger +from vllm.platforms import current_platform + +from vllm_omni.platforms import current_omni_platform + +if TYPE_CHECKING: + from vllm_omni.diffusion.models.qwen_image_21.qwen_image_21_transformer import ( + QwenImage21Transformer2DModel, + ) + +logger = init_logger(__name__) + + +def _dynamic_lora_wrappers_present(module: torch.nn.Module) -> bool: + """True once ``DiffusionLoRAManager`` has wrapped layers under ``module``. + + A captured graph records the layers and branches that ran at capture time; + binding or rescaling an adapter afterwards changes the math without + changing the shapes the graph was keyed on, so graphs and dynamic LoRA + must not mix. Same reasoning as ``sensenova_u1.paged_decode``. + """ + try: + from vllm.lora.layers import BaseLayerWithLoRA + except ImportError: # pragma: no cover - depends on the wheel + return False + return any(isinstance(m, BaseLayerWithLoRA) for m in module.modules()) + + +class QwenImage21DecodeGraphEntry: + """Static state for one captured decode graph. + + Buffers are created outside inference mode so replay-side ``copy_`` is + legal regardless of the caller's grad context. ``graph`` is None until the + first decode step captures it. Failed entries are replaced by a ``None`` + marker in the manager so their buffers can be released. + """ + + def __init__( + self, + *, + branch: str, + batch_size: int, + prefix_len: int, + target_tokens: int, + in_channels: int, + dtype: torch.dtype, + device: torch.device, + kv_shapes: list[torch.Size], + kv_dtype: torch.dtype, + freqs: torch.Tensor, + ): + self.branch = branch + self.target_tokens = target_tokens + self.prefix_len = prefix_len + with torch.inference_mode(False): + self.hidden = torch.zeros(batch_size, target_tokens, in_channels, dtype=dtype, device=device) + self.timestep = torch.zeros(batch_size, dtype=dtype, device=device) + self.k = [torch.zeros(shape, dtype=kv_dtype, device=device) for shape in kv_shapes] + self.v = [torch.zeros(shape, dtype=kv_dtype, device=device) for shape in kv_shapes] + self.target_token_mask = torch.ones(target_tokens, dtype=torch.bool, device=device) + # RoPE frequencies are layout-dependent but step-independent. Clone + # into a normal (non-inference) tensor: the prefill runs under + # inference_mode, and the captured body takes views of this tensor + # outside it, which is forbidden for inference tensors. + self.freqs = freqs.clone() + self.attn_metadata = None + # Per-block cache dicts with the same shape as the legacy protocol. + self.block_caches = [{branch: {"key": self.k[i], "value": self.v[i]}} for i in range(len(kv_shapes))] + self.graph: torch.cuda.CUDAGraph | None = None + self.output: torch.Tensor | None = None + self.captures = 0 + self.prefix_sources: list[tuple[weakref.ReferenceType[torch.Tensor], int | None]] = [] + + def refresh_prefix(self, kv_cache: list[dict[str, dict[str, torch.Tensor]]]) -> None: + sources = [block[self.branch][name] for block in kv_cache for name in ("key", "value")] + versions = [None if tensor.is_inference() else tensor._version for tensor in sources] + if len(sources) == len(self.prefix_sources) and all( + previous() is tensor and version == current_version + for (previous, version), tensor, current_version in zip(self.prefix_sources, sources, versions) + ): + return + # Prefill replaces inference-mode tensors; decode treats them as immutable. + # Weak references avoid retaining completed requests or mistaking a reused + # allocation for the same prefix. Normal tensors also track in-place writes. + self.prefix_sources = [] + for i, block in enumerate(kv_cache): + self.k[i].copy_(block[self.branch]["key"]) + self.v[i].copy_(block[self.branch]["value"]) + self.prefix_sources = [(weakref.ref(tensor), version) for tensor, version in zip(sources, versions)] + + def capture(self, body) -> None: + """Warm up on a side stream, then capture one decode step. + + Warm-up runs the same computation it will record, allocating cuBLAS + workspaces and any lazy buffers outside the capture. Decode writes + nothing to the cache (prefix K/V are read-only), so warm-up is + idempotent. + """ + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + try: + with torch.cuda.stream(side), torch.inference_mode(False), torch.no_grad(): + for _ in range(2): + body() + finally: + torch.cuda.current_stream().wait_stream(side) + + graph = torch.cuda.CUDAGraph() + # The platform-wide pool is shared with every other captured path in + # the tree; entries only keep their *output* alive, whose tiny size + # makes reuse safe because the caller clones it before any other graph + # on the pool replays. + with torch.inference_mode(False), torch.no_grad(): + with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + self.output = body() + self.graph = graph + self.captures += 1 + + +class QwenImage21DecodeGraphManager: + """Owns the static KV buffers and captured graphs of one transformer.""" + + def __init__(self, model: QwenImage21Transformer2DModel, max_entries: int = 8, model_level_offload: bool = False): + self.model = model + self.max_entries = max_entries + self.model_level_offload = model_level_offload + self.entries: OrderedDict[tuple, QwenImage21DecodeGraphEntry | None] = OrderedDict() + # Key -> weak reference to the prefix K tensor of the request that owns + # the entry's static buffers. Dead references mean the owning request + # finished and its cache was released, so the key is free to adopt. + self._owners: dict[tuple, weakref.ReferenceType[torch.Tensor]] = {} + self._static_eligible: bool | None = None + + # ── eligibility ── + + def _offload_reason(self) -> str | None: + """Reason decode graphs are incompatible with model-level offload, if any. + + Model-level (sequential) offload registers its hook on the top-level + transformer module, not on individual blocks. With persistent staging + the hook keeps the DiT on fixed device storage and only rebinds + ``p.data`` around each generation, and the swap completes in + ``pre_forward`` before capture or replay begin, so captured weight + pointers stay valid. Without staging every swap allocates fresh + storage and replay would read stale pointers, so capture stays + disabled and decode falls back to eager. + """ + registry = getattr(self.model, "_hook_registry", None) + hooks = registry._hooks if registry is not None else {} + if hooks: + from vllm_omni.diffusion.offloader.sequential_backend import ( + SequentialOffloadHook, + sequential_offload_staging_active, + ) + + if set(hooks) <= {SequentialOffloadHook._HOOK_NAME} and sequential_offload_staging_active(self.model): + return None + return "model module carries offload/cache hooks" + if self.model_level_offload: + return "model-level CPU offload without persistent staging swaps weight storage after capture" + return None + + def _check_static_eligibility(self) -> bool: + """One-time checks that cannot change after weights are loaded.""" + model = self.model + reason = self._offload_reason() + param = next(model.parameters(), None) + if reason is None and (param is None or param.device.type != "cuda"): + reason = f"model is not on CUDA (device={None if param is None else param.device})" + elif reason is None: + parallel_config = getattr(model, "parallel_config", None) + if parallel_config is not None: + sp = getattr(parallel_config, "sequence_parallel_size", None) or 1 + if sp > 1 or getattr(parallel_config, "ring_degree", 1) > 1: + reason = "sequence/ring parallelism changes decode shapes per rank" + elif getattr(parallel_config, "use_hsdp", False): + reason = "HSDP wraps blocks with gather hooks outside the captured region" + if reason is None: + from vllm.distributed import get_tensor_model_parallel_world_size + + if get_tensor_model_parallel_world_size() > 1: + reason = "tensor parallelism runs collectives inside the captured region" + if reason is None: + for block in model.transformer_blocks: + if hasattr(block, "_hook_registry") or hasattr(block, "_omni_original_forward"): + reason = "transformer blocks carry offload/cache hooks" + break + if reason is None and model.transformer_blocks: + attn_layer = model.transformer_blocks[0].attn.attn + if getattr(attn_layer, "_kv_cache_dtype", None) is not None: + reason = "KV-cache quantization is enabled" + + if reason is not None: + logger.warning_once("Qwen-Image-2.1 CUDA graph decode disabled: %s. Falling back to eager decode.", reason) + return False + return True + + def eligible(self) -> bool: + if self._static_eligible is None: + self._static_eligible = self._check_static_eligibility() + return self._static_eligible + + def _backend_name(self) -> str: + backend = self.model.transformer_blocks[0].attn.attn.attn_backend + return backend.get_name() if backend is not None else "custom" + + # ── prefill registration ── + + def register_prefill( + self, + *, + kv_cache: list[dict[str, dict[str, torch.Tensor]]], + cache_branch: str, + prefix_len: int, + img_shapes: list[tuple[int, int, int]], + target_freqs: torch.Tensor, + joint_key_valid: torch.Tensor | None, + dtype: torch.dtype, + ) -> None: + """Prepare graph buffers without changing the request's prefix K/V. + + The first in-flight request to register a key owns its static buffers; + registrations for the same key from another live request are refused + so their decode stays eager until the owner completes. + """ + if not self.eligible() or torch.compiler.is_compiling(): + return + if _dynamic_lora_wrappers_present(self.model): + if self.entries: + logger.warning_once( + "Qwen-Image-2.1 CUDA graph decode: LoRA wrappers appeared; " + "dropping captured graphs and falling back to eager decode." + ) + self.entries.clear() + return + if joint_key_valid is not None and not bool(joint_key_valid.all()): + # The masked decode path (varlen-unpack or 4D mask) branches on mask + # contents and cannot be captured; this request stays eager. + logger.warning_once( + "Qwen-Image-2.1 CUDA graph decode: padded text mask present; this request falls back to eager decode." + ) + return + + first = kv_cache[0].get(cache_branch) + if first is None or "key" not in first: + return + if any(name.endswith("_scale") for name in first): + # Quantized (FP8) prefix storage: the static K/V buffers assume a single + # dtype and the captured body would not see the dequant scales (a bf16 K + + # fp8 V mix would even copy_ fp8 into bf16 buffers, silently dropping the + # scales). Keep quantized-cache requests on the eager decode path. + logger.warning_once( + "Qwen-Image-2.1 CUDA graph decode: quantized prefix KV cache " + "(prefix_kv_cache_dtype) is not graph-capturable; falling back to eager decode." + ) + return + batch_size = first["key"].shape[0] + key = ( + cache_branch, + batch_size, + prefix_len, + tuple(tuple(shape) for shape in img_shapes), + dtype, + first["key"].device.index, + self._backend_name(), + ) + + if key not in self.entries: + while len(self.entries) >= self.max_entries: + evicted_key, evicted_entry = self.entries.popitem(last=False) + self._owners.pop(evicted_key, None) + del evicted_entry + logger.debug("Evicting decode graph entry %s (max_entries=%d)", evicted_key, self.max_entries) + allocation_error = None + try: + self.entries[key] = QwenImage21DecodeGraphEntry( + branch=cache_branch, + batch_size=batch_size, + prefix_len=prefix_len, + target_tokens=target_freqs.shape[0], + in_channels=self.model.in_channels, + dtype=dtype, + device=first["key"].device, + kv_shapes=[block_cache[cache_branch]["key"].shape for block_cache in kv_cache], + kv_dtype=first["key"].dtype, + freqs=target_freqs, + ) + except torch.OutOfMemoryError as exc: + allocation_error = str(exc) + if allocation_error is not None: + # Keep a bounded failure marker without retaining tensors or the exception traceback. + self.entries[key] = None + current_omni_platform.empty_cache() + logger.warning( + "Qwen-Image-2.1 CUDA graph allocation failed for key=%s: %s. Falling back to eager decode.", + key, + allocation_error, + ) + return + else: + self.entries.move_to_end(key) + + entry = self.entries[key] + if entry is not None: + owner = self._owners.get(key) + owner_prefix = owner() if owner is not None else None + if owner_prefix is not None and owner_prefix is not first["key"]: + # Another in-flight request still owns this entry's static + # buffers; refreshing them would overwrite that request's + # prefix mid-generation. This request decodes eagerly until + # the owner's cache is released and the weak reference dies. + logger.warning_once( + "Qwen-Image-2.1 CUDA graph decode: entry key=%s is owned by another in-flight " + "request; this request falls back to eager decode.", + key, + ) + return + self._owners[key] = weakref.ref(first["key"]) + entry.prefix_sources = [] + + logger.debug( + "Registered decode graph entry key=%s (entries=%d, captured=%d)", + key, + len(self.entries), + sum(1 for e in self.entries.values() if e is not None and e.graph is not None), + ) + + # ── decode replay ── + + def try_decode( + self, + *, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + kv_cache: list[dict[str, dict[str, torch.Tensor]]], + cache_branch: str, + img_shapes: list[tuple[int, int, int]], + img_mask: torch.Tensor, + encoder_hidden_states_mask: torch.Tensor | None, + ) -> torch.Tensor | None: + """Replay the captured decode step, or return None to fall back to eager.""" + if not self.eligible() or torch.compiler.is_compiling(): + return None + if _dynamic_lora_wrappers_present(self.model): + self.clear() + return None + if encoder_hidden_states_mask is not None: + text_positions = ~img_mask[0, : encoder_hidden_states_mask.shape[1]] + if not bool(encoder_hidden_states_mask[:, text_positions].all()): + return None + first = kv_cache[0][cache_branch] + cached_key = first["key"] + if cached_key.device.type != "cuda": + # Registration is device-agnostic (unit tests exercise the buffer + # machinery on CPU), but capture/replay is CUDA-only: an empty + # graph over CPU ops would replay as a stale no-op. + return None + key = ( + cache_branch, + cached_key.shape[0], + cached_key.shape[1], + tuple(tuple(shape) for shape in img_shapes), + hidden_states.dtype, + cached_key.device.index, + self._backend_name(), + ) + entry = self.entries.get(key) + if entry is None: + return None + self.entries.move_to_end(key) + + owner = self._owners.get(key) + owner_prefix = owner() if owner is not None else None + if owner_prefix is not None and owner_prefix is not cached_key: + # The entry is owned by a different in-flight request (see + # register_prefill); this request stays on the eager decode path + # until the owner's cache is released. + return None + + entry.hidden.copy_(hidden_states[:, -entry.target_tokens :]) + entry.timestep.copy_(timestep) + # Graph buffers are scratch space; request caches must never alias them. + entry.refresh_prefix(kv_cache) + + if entry.graph is None: + capture_error = None + try: + entry.capture(lambda: self.model._decode_graph_forward(entry)) + except Exception as exc: + capture_error = str(exc) + if capture_error is not None: + self.entries[key] = None + del entry + current_omni_platform.empty_cache() + logger.warning( + "Qwen-Image-2.1 CUDA graph capture failed for key=%s: %s. " + "Falling back to eager decode for this shape.", + key, + capture_error, + ) + return None + logger.info("Captured Qwen-Image-2.1 decode graph for key=%s", key) + entry.graph.replay() + assert entry.output is not None + # Clone before returning: the output lives in the shared graph pool and + # the next replay (of this or another graph on the pool) overwrites it. + return entry.output.clone() + + def clear(self) -> None: + """Drop all entries, freeing buffers and captured graphs.""" + self.entries.clear() + self._owners.clear() 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..c38cf7111 --- /dev/null +++ b/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py @@ -0,0 +1,1551 @@ +# 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 dataclasses +import json +import logging +import os +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 AutoConfig, AutoModelForImageTextToText, 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.models.utils import create_transformers_model +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 +from vllm_omni.quantization import resolve_component_quant_config +from vllm_omni.quantization.component_config import resolve_encoder_quant_config + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + + 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) + # The architecture compresses 2x spatially per downsampling stage (one per + # `temperal_downsample` entry), e.g. 16x for Qwen-Image 2.1. The shipped + # `scale_factor_spatial` (8) under-reports it, so prefer the + # architecture-derived value. + if "temperal_downsample" in vae_config: + return 2 ** len(vae_config["temperal_downsample"]) + if vae_config.get("scale_factor_spatial"): + return int(vae_config["scale_factor_spatial"]) + return 16 + + +def _resolve_text_encoder_quant_config( + quant_config: "QuantizationConfig | None", +) -> "QuantizationConfig | None": + """Resolve the text-encoder slice of a global or per-component quant config. + + Pre-quantized formats (modelopt, svdquant, ...) require serialized scale + tensors the BF16 text-encoder checkpoint does not ship, so they are stripped + and the encoder stays in checkpoint precision. + """ + resolved = resolve_component_quant_config(quant_config, "text_encoder") + return resolve_encoder_quant_config(resolved) + + +# Subtrees of the HF Qwen3-VL text encoder that must stay in checkpoint +# precision: the vision tower (condition-image encoding is numerically +# sensitive) and the LM head (unused — the pipeline reads hidden states). +_TEXT_ENCODER_QUANT_EXCLUDED_PREFIXES = ("model.visual", "lm_head") + + +def _exclude_text_encoder_subtrees_from_quant( + quant_config: "QuantizationConfig | None", +) -> "QuantizationConfig | None": + """Keep the vision tower and LM head in checkpoint precision. + + Uses the config's standard ``ignored_layers`` mechanism, with substring + matching so the prefixes cover whole subtrees (same convention as the + DiT's ``_enable_pattern_ignored_layers``). + """ + if quant_config is None: + return None + if not hasattr(quant_config, "ignored_layers"): + logger.warning( + "Quantization config %s has no ignored_layers; the Qwen3-VL vision tower " + "and LM head cannot be excluded from quantization.", + type(quant_config).__name__, + ) + return quant_config + config = copy.copy(quant_config) + config.ignored_layers = [*config.ignored_layers, *_TEXT_ENCODER_QUANT_EXCLUDED_PREFIXES] + if hasattr(config, "ignored_layers_match_mode"): + config.ignored_layers_match_mode = "substring" + return config + + +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.""" + request.allow_mixed_step_phases = False + if request.sampling_params.true_cfg_scale is None: + request.sampling_params.true_cfg_scale = 1.0 + 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 isinstance(raw_image, tuple): + raw_image = list(raw_image) + if raw_image is None or (isinstance(raw_image, list) and len(raw_image) == 0): + request.batch_compatibility_key = ("qwen_image_21", ()) + 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.batch_compatibility_key = ("qwen_image_21", tuple(input_image_sizes)) + 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 release_captured_graphs(self) -> None: + """Drop decode graphs before sleep level 2 discards their capture storage. + + The runner looks this up on the pipeline, not on the transformer. + """ + transformer = getattr(self, "transformer", None) + release = getattr(transformer, "release_captured_graphs", None) + if callable(release): + release() + + 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. + text_encoder_quant_config = _resolve_text_encoder_quant_config(od_config.quantization_config) + if text_encoder_quant_config is None: + 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) + else: + # Quantized path: build the HF model on meta, swap its nn.Linear + # modules for vLLM quantizable linears, and stream the weights + # through the pipeline loader (same integration as Z-Image). The + # encoder-scoped quant config rides on a copied od_config so the + # shared create_transformers_model signature stays untouched. + text_encoder_config = AutoConfig.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only + ) + encoder_od_config = dataclasses.replace( + od_config, + quantization_config=_exclude_text_encoder_subtrees_from_quant(text_encoder_quant_config), + ) + self.text_encoder = create_transformers_model( + AutoModelForImageTextToText, + encoder_od_config, + hf_config=text_encoder_config, + ).to(self.device) + self.weights_sources.append( + DiffusersPipelineLoader.ComponentSource( + model_or_path=model, + subfolder="text_encoder", + revision=None, + prefix="text_encoder.", + ) + ) + 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=resolve_component_quant_config(od_config.quantization_config, "transformer"), + **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 "..." 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\n<|vision_start|><|image_pad|><|vision_end|>{{}}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + # 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 + prompt = [" " if not p else p for p in 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 = "<|vision_start|><|image_pad|><|vision_end|>" + for i in range(2, n_imgs + 1): + replace += f" <|vision_start|><|image_pad|><|vision_end|>" + template = self.prompt_template_ti2i.replace( + "<|vision_start|><|image_pad|><|vision_end|>", replace + ) + txt.append(template.format(t)) + condition_pil_list = [] + for images in condition_pil_per_prompt: + for img in images: + if img.mode == "RGBA": + white = PIL.Image.new("RGB", img.size, (255, 255, 255)) + white.paste(img, mask=img.getchannel("A")) + img = white + condition_pil_list.append(img) + + # 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, + "padding_side": "left", + "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 + + text_model = getattr(self.text_encoder.model, "language_model", self.text_encoder.model) + handle = text_model.norm.register_forward_hook(lambda module, args, output: args[0]) + try: + outputs = self.text_encoder(**forward_kwargs) + finally: + handle.remove() + hidden_states = outputs.hidden_states[-1] + + 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 + + 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, + ) + + # Request-major expansion (A,A,B,B). ``repeat`` on the 2D masks would + # prepend a dimension and reshape to A,B,A,B — wrong pairing with embeds. + prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) + prompt_embeds_mask = prompt_embeds_mask.repeat_interleave(num_images_per_prompt, dim=0) + image_pad_mask = image_pad_mask.repeat_interleave(num_images_per_prompt, dim=0) + + 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 = encoded.repeat_interleave(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, + device=self.device, + 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 if common_sampling_params.true_cfg_scale is not None else 1.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 if sampling.true_cfg_scale is not None else 1.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. Merge every stored part — FP8 + # cache entries carry per-request scale tensors alongside "key"/"value". + # Callers must pass a homogeneous baked-prefix-length group: denoise_step + # sub-batches by prefill_prompt_seq_len before assembling. + prefix_lens = {state.extra.get("prefill_prompt_seq_len") for state in states} + if len(prefix_lens) > 1: + raise ValueError( + "Cannot merge decode KV for requests baked at different padded prompt " + "lengths. denoise_step must sub-batch by prefill_prompt_seq_len first." + ) + 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: + part_names = {part for cache in caches for part in cache[block_idx][branch]} + branches[branch] = { + part: torch.cat([cache[block_idx][branch][part] for cache in caches], dim=0) + for part in sorted(part_names) + } + 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 _decode_group_key(state: "StepRequestState") -> tuple[Any, ...]: + """Group decode requests by the padded VLM length baked into their KV.""" + return ( + state.extra.get("prefill_prompt_seq_len"), + state.extra.get("prefill_negative_prompt_seq_len"), + ) + + @staticmethod + def _split_decode_groups( + states: list["StepRequestState"], + ) -> list[list["StepRequestState"]]: + """Split a decode batch into homogeneous baked-prefix-length groups. + + Preserves first-seen group order and relative order within each group. + Callers stitch per-group noise preds back into the input request order. + """ + groups: dict[tuple[Any, ...], list[StepRequestState]] = {} + order: list[tuple[Any, ...]] = [] + for state in states: + key = QwenImage21Pipeline._decode_group_key(state) + if key not in groups: + order.append(key) + groups[key] = [] + groups[key].append(state) + return [groups[key] for key in order] + + @staticmethod + def _needs_decode_seq_len_split(states: list["StepRequestState"]) -> bool: + if len(states) <= 1: + return False + phases = {QwenImage21Pipeline._kv_cache_phase(state.extra.get("kv_cache")) for state in states} + if phases != {"decode"}: + return False + return len({QwenImage21Pipeline._decode_group_key(state) for state in states}) > 1 + + @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. + + Clone (not ``.contiguous()``) so each request owns independent storage. + Dim-0 slices of a contiguous merged prefill tensor are already contiguous + views; keeping those views alive would pin the merged tensor and break + the ``owner_prefix is not cached_key`` identity check in step-mode + CUDA-graph decode (every decode would fall back to eager). + """ + 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].clone() for part, tensor in parts.items() + } + start += rows + + @staticmethod + def _resolve_prompt_seq_len( + states: list["StepRequestState"], + current_seq_len: int, + extra_key: str = "prefill_prompt_seq_len", + ) -> int: + """Prefer the padded VLM length baked into each request's KV at prefill.""" + seq_len = current_seq_len + for state in states: + stored = state.extra.get(extra_key) + if stored is not None: + seq_len = max(seq_len, int(stored)) + return seq_len + + @staticmethod + def _pad_rows_to_seq_len(rows: list[torch.Tensor], seq_len: int) -> torch.Tensor: + padded = [] + for row in rows: + if row.shape[1] < seq_len: + row = F.pad(row, (0, seq_len - row.shape[1]), value=False) + elif row.shape[1] > seq_len: + row = row[:, :seq_len] + padded.append(row) + return torch.cat(padded, dim=0) + + def denoise_step( + self, + input_batch: "InputBatch", + **kwargs: Any, + ) -> torch.Tensor | None: + """One denoise step: read from *input_batch*, delegate to CFGParallelMixin. + + Decode batches that mix different baked ``prefill_prompt_seq_len`` values + (common when the mixed-phase gate forces separate prefills) are + sub-batched by that length so both requests keep running. + """ + 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.") + + if self._needs_decode_seq_len_split(states): + # Local import: InputBatch is only TYPE_CHECKING at module scope. + from vllm_omni.diffusion.worker.input_batch import InputBatch as _InputBatch + + preds_by_id: dict[str, torch.Tensor] = {} + for group in self._split_decode_groups(states): + group_pred = self._denoise_step_homogeneous(_InputBatch.make_batch(group)) + if group_pred is None: + return None + offset = 0 + for state in group: + rows = state.latents.shape[0] + preds_by_id[state.request_id] = group_pred[offset : offset + rows] + offset += rows + return torch.cat([preds_by_id[state.request_id] for state in states], dim=0) + + return self._denoise_step_homogeneous(input_batch) + + def _denoise_step_homogeneous( + self, + input_batch: "InputBatch", + ) -> torch.Tensor | None: + """Run one denoise forward for a KV-compatible state group.""" + 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. + # After a mixed-length prefill, each request's KV prefix stays at that + # prefill's padded length — keep using it even if later decode batches + # shrink (longest prompt finished) so mask / RoPE layout still match. + max_seq_len = self._resolve_prompt_seq_len(states, input_batch.prompt_embeds.shape[1]) + img_mask = self._append_target_slots( + self._pad_rows_to_seq_len([state.extra["image_pad_mask"] for state in states], max_seq_len), + latents.shape[1], + ) + prompt_embeds = input_batch.prompt_embeds + prompt_embeds_mask = input_batch.prompt_embeds_mask + if prompt_embeds is not None and prompt_embeds.shape[1] < max_seq_len: + pad = max_seq_len - prompt_embeds.shape[1] + prompt_embeds = F.pad(prompt_embeds, (0, 0, 0, pad), value=0.0) + if prompt_embeds_mask is not None: + prompt_embeds_mask = F.pad(prompt_embeds_mask, (0, pad), value=False) + negative_img_mask = None + negative_prompt_embeds = input_batch.negative_prompt_embeds + negative_prompt_embeds_mask = input_batch.negative_prompt_embeds_mask + negative_seq_len = 0 + if input_batch.do_true_cfg: + negative_seq_len = self._resolve_prompt_seq_len( + states, + input_batch.negative_prompt_embeds.shape[1], + extra_key="prefill_negative_prompt_seq_len", + ) + negative_img_mask = self._append_target_slots( + self._pad_rows_to_seq_len( + [state.extra["negative_image_pad_mask"] for state in states], negative_seq_len + ), + latents.shape[1], + ) + if negative_prompt_embeds is not None and negative_prompt_embeds.shape[1] < negative_seq_len: + pad = negative_seq_len - negative_prompt_embeds.shape[1] + negative_prompt_embeds = F.pad(negative_prompt_embeds, (0, 0, 0, pad), value=0.0) + if negative_prompt_embeds_mask is not None: + negative_prompt_embeds_mask = F.pad(negative_prompt_embeds_mask, (0, pad), value=False) + + 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": prompt_embeds_mask, + "encoder_hidden_states": 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": negative_prompt_embeds_mask, + "encoder_hidden_states": 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: + # Remember the padded VLM length baked into the per-request KV so a + # later smaller decode batch can rebuild matching masks. + for state in states: + state.extra["prefill_prompt_seq_len"] = max_seq_len + if input_batch.do_true_cfg: + state.extra["prefill_negative_prompt_seq_len"] = negative_seq_len + 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..2e82ee4d8 --- /dev/null +++ b/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py @@ -0,0 +1,1320 @@ +# 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 diffusers.models.normalization import RMSNorm +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 +from vllm_omni.diffusion.models.qwen_image_21.decode_graph import QwenImage21DecodeGraphManager +from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload_strategy + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + +logger = init_logger(__name__) + +_FP8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max + +# Accepted values for the `prefix_kv_cache_dtype` switch (see QwenImage21Transformer2DModel). +_PREFIX_KV_FP8_ALIASES = {"fp8", "fp8_e4m3", "fp8_e4m3fn"} +_PREFIX_KV_FP8_V_ALIASES = {"fp8_v", "fp8_e4m3_v"} +_PREFIX_KV_FP8_V = "fp8_e4m3_v" + + +def _quantize_prefix_kv_fp8(t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Per-token-per-head symmetric FP8 E4M3 quantization of a (B, S, H, D) K/V tensor. + + One fp32 scale per (batch, token, head): it tracks the heavy-tailed per-token magnitude of + post-RoPE keys and outlier-channel values far better than a per-tensor scale, and the scale + tensor stays mergeable across the batch dim for batched decode (`_assemble_kv_cache`). + """ + scale = t.abs().amax(dim=-1, keepdim=True).float().clamp_min(1e-12) / _FP8_E4M3_MAX + quantized = (t / scale).clamp(-_FP8_E4M3_MAX, _FP8_E4M3_MAX).to(torch.float8_e4m3fn) + return quantized, scale + + +def _dequantize_prefix_kv_fp8(quantized: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + return (quantized.float() * scale).to(dtype) + + +def _normalize_prefix_kv_cache_dtype(value: Any) -> str | None: + """Normalize the `prefix_kv_cache_dtype` switch: None/"auto" = off, "fp8*" = FP8 E4M3 storage.""" + if value is None or value == "auto": + return None + if value in _PREFIX_KV_FP8_ALIASES: + return "fp8_e4m3" + if value in _PREFIX_KV_FP8_V_ALIASES: + # V-only FP8 storage: K stays in the native dtype. On H200 T2I 1024 + # (CFG 4, 50 steps) PSNR vs bf16 is 44.0 dB for V-only vs 34.2 dB for + # K+V fp8 — post-RoPE K is the precision-sensitive half of the cache. + return _PREFIX_KV_FP8_V + raise ValueError( + f"Unknown prefix_kv_cache_dtype {value!r}; expected None, 'auto' or one of " + f"{sorted(_PREFIX_KV_FP8_ALIASES | _PREFIX_KV_FP8_V_ALIASES)}." + ) + + +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 + # Compute the fixed table on CPU like Diffusers, then honor the loader's device. + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device="cpu") / half + ).to(torch.get_default_device()) + 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 + + # Match Diffusers' CPU initialization even when the loader sets a CUDA + # default device; CPU and CUDA trigonometric kernels round differently. + with torch.device("cpu"): + 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) + + height_index = torch.where(height_index < 0, height_index + self.freqs[1].shape[0], height_index) + width_index = torch.where(width_index < 0, width_index + self.freqs[2].shape[0], width_index) + return torch.cat( + [ + self.freqs[0].index_select(0, frame_index), + self.freqs[1].index_select(0, height_index), + self.freqs[2].index_select(0, width_index), + ], + dim=-1, + ) + + +def _enable_pattern_ignored_layers(quant_config: QuantizationConfig | None) -> QuantizationConfig | None: + """Switch vLLM FP8-style `ignored_layers` to pattern (substring) matching. + + vLLM's ``Fp8Config`` matches ``ignored_layers`` against full layer prefixes + exactly, but the diffusion quantization docs and CLI describe them as name + patterns (e.g. ``img_mlp``) — under exact matching such patterns silently + skip nothing. Substring matching is a strict superset for exact full + prefixes, so existing exact-prefix configs keep working. + """ + if quant_config is None: + return quant_config + component_configs = getattr(quant_config, "component_configs", None) + if component_configs is not None: + configs = [*component_configs.values(), getattr(quant_config, "default_config", None)] + else: + configs = [quant_config] + for config in configs: + if config is None: + continue + if getattr(config, "ignored_layers", None) and hasattr(config, "ignored_layers_match_mode"): + config.ignored_layers_match_mode = "substring" + return quant_config + + +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 = "", + prefix_kv_cache_dtype: str | None = None, + ): + super().__init__() + self.head_dim = dim_head + self.prefix_kv_cache_dtype = prefix_kv_cache_dtype + + 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 + + # Diffusers normalizes in FP32, then casts before the learned scale. + # Keep that rounding order for BF16 Q/K; TP shards whole heads. + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = 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). With + FP8 cache storage enabled, the branch additionally carries `{"key_scale", "value_scale"}`. + 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. + if self.prefix_kv_cache_dtype is not None: + # FP8 storage: quantize once at prefill; decode dequantizes per step. + # The fp32 scales ride along as extra branch entries so the pipeline's + # batched decode merge can concatenate them like the K/V tensors. The + # quantized tensors own their storage, so no clone is needed here. + if self.prefix_kv_cache_dtype == "fp8_e4m3": + prefix_key, branch_cache["key_scale"] = _quantize_prefix_kv_fp8(key[:, :cache_write_len]) + branch_cache["key"] = prefix_key + else: + # V-only FP8 ("fp8_e4m3_v"): K stays in the native dtype. + branch_cache["key"] = key[:, :cache_write_len].clone() + prefix_value, branch_cache["value_scale"] = _quantize_prefix_kv_fp8(value[:, :cache_write_len]) + branch_cache["value"] = prefix_value + else: + # Own the prefix storage instead of retaining the full prefill tensors. + branch_cache["key"] = key[:, :cache_write_len].clone() + branch_cache["value"] = value[:, :cache_write_len].clone() + else: + cached_key = branch_cache["key"] + cached_value = branch_cache["value"] + if self.prefix_kv_cache_dtype is not None: + # Dequantize only entries that were actually quantized (scale present). + if "key_scale" in branch_cache: + cached_key = _dequantize_prefix_kv_fp8(cached_key, branch_cache["key_scale"], key.dtype) + if "value_scale" in branch_cache: + cached_value = _dequantize_prefix_kv_fp8(cached_value, branch_cache["value_scale"], value.dtype) + + 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 = "", + prefix_kv_cache_dtype: str | None = None, + ): + 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", + prefix_kv_cache_dtype=prefix_kv_cache_dtype, + ) + 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 cached prefix K/V are stored in the running dtype by default, or in FP8 E4M3 (with + per-token-per-head fp32 scales) when `od_config.extras["prefix_kv_cache_dtype"]` is `"fp8"` (K and V) + or `"fp8_v"` (V only — K stays in the native dtype). + """ + + # 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", + cuda_graph_max_decode_graphs: int = 8, + ): + 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 = _enable_pattern_ignored_layers(quant_config) + quant_config = self.quant_config + + # Storage dtype for the cross-step prefix KV cache, from the supplementary + # model-specific config entry `extras["prefix_kv_cache_dtype"]` (None = bf16/fp32 + # native storage, unchanged behavior; "fp8" = FP8 E4M3 quantized storage, halving + # prefix K/V memory). This is independent of `diffusion_kv_cache_dtype`, which + # quantizes attention Q/K/V *compute* per forward on supported backends. + extras = getattr(od_config, "extras", None) or {} + self.prefix_kv_cache_dtype = _normalize_prefix_kv_cache_dtype(extras.get("prefix_kv_cache_dtype")) + + 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}", + prefix_kv_cache_dtype=self.prefix_kv_cache_dtype, + ) + 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) + + # Use decode graphs unless the caller disables them or requests eager + # execution. Model-level (sequential) offload registers its swap hook + # on this top-level module rather than on individual blocks; graphs + # stay eligible only when that hook keeps weights on persistent + # staging storage (see QwenImage21DecodeGraphManager._offload_reason). + self.enable_cuda_graph_decode = ( + getattr(od_config, "enable_cuda_graph_decode", True) and not od_config.enforce_eager + ) + self._decode_graph_manager = ( + QwenImage21DecodeGraphManager( + self, + max_entries=cuda_graph_max_decode_graphs, + model_level_offload=resolve_offload_strategy(od_config) + in {OffloadStrategy.MODEL_LEVEL, OffloadStrategy.DISTRIBUTED_LAYER_WISE}, + ) + if self.enable_cuda_graph_decode + else None + ) + + @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 + if sp_active and ctx is not None: + # The prefix is replicated across SP ranks for the Ulysses joint mechanism. + ctx.split_text_embed_in_sp = False + + is_decode = kv_cache is not None and len(kv_cache) > 0 and "key" in kv_cache[0].get(cache_branch, {}) + + if is_decode and self._decode_graph_manager is not None: + graph_output = self._decode_graph_manager.try_decode( + hidden_states=hidden_states, + timestep=timestep, + kv_cache=kv_cache, + cache_branch=cache_branch, + img_shapes=layout, + img_mask=img_mask, + encoder_hidden_states_mask=encoder_hidden_states_mask, + ) + if graph_output is not None: + if not return_dict: + return (graph_output,) + return Transformer2DModelOutput(sample=graph_output) + + 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 + ) + # Read SP pad size after sequence_prepare: the auto-pad hook only writes + # ctx.sp_padding_size during the split inside that call (and once per + # context). Reading earlier always sees 0 while the local shard is already + # padded — same ordering as Cosmos3. + sp_padding = 0 + if ( + sp_active + and ctx is not None + and 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) + # 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()) + if not has_padding: + # Match Diffusers and use the same mask-free dispatch in eager and graph decode. + joint_key_valid = None + 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 ( + not is_decode + and cache_write_len is not None + and kv_cache is not None + and self._decode_graph_manager is not None + ): + # Allocate graph scratch buffers; the request keeps its own prefix K/V. + self._decode_graph_manager.register_prefill( + kv_cache=kv_cache, + cache_branch=cache_branch, + prefix_len=cache_write_len, + img_shapes=layout, + target_freqs=target_freqs, + joint_key_valid=joint_key_valid, + dtype=hidden_states.dtype, + ) + + 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 _decode_graph_forward(self, entry) -> torch.Tensor: + """The exact decode-step computation, captured into a CUDA graph. + + Mirrors the eager decode path of ``forward``: target tokens only, + prefix K/V read from the entry's static buffers, the same mask-free + attention dispatch as eager decode, and an all-ones modulation mask. + All inputs live in the entry's fixed-address buffers. + """ + hidden_states = self.img_in(entry.hidden) + timestep = torch.cat([entry.timestep, entry.timestep.new_zeros(1)], dim=0) + temb = self.time_text_embed(timestep, hidden_states) + modulation = self.modulation(temb) + for index_block, block in enumerate(self.transformer_blocks): + hidden_states = block( + hidden_states=hidden_states, + modulation=modulation, + freqs=entry.freqs, + target_token_mask=entry.target_token_mask, + attn_metadata=entry.attn_metadata, + sp_prefix_len=0, + sp_decode=False, + kv_cache=entry.block_caches[index_block], + cache_branch=entry.branch, + cache_write_len=None, + ) + hidden_states = self.norm_out(hidden_states, temb, entry.target_token_mask) + return self.proj_out(hidden_states) + + def release_captured_graphs(self) -> None: + """Drop captured decode graphs and their static buffers (e.g. sleep mode).""" + if self._decode_graph_manager is not None: + self._decode_graph_manager.clear() + + 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/offloader/module_residency.py b/vllm_omni/diffusion/offloader/module_residency.py index ab584c790..0e579e931 100644 --- a/vllm_omni/diffusion/offloader/module_residency.py +++ b/vllm_omni/diffusion/offloader/module_residency.py @@ -98,6 +98,12 @@ class PinnedModuleStager: A module iterable is treated as one staging group. It uses one copy stream and one reusable completion event. Tensors sharing storage keep their shapes, strides, offsets, dtypes, and aliases across every transition. + + With ``retain_device_storage`` the device storages survive ``offload``: + ``load`` then only rebinds to the same fixed storages (weights are + immutable, so no H2D copy repeats). The stable storage addresses keep + CUDA-graph-captured weight pointers valid across offload/load swaps, at + the cost of keeping the device storage resident while offloaded. """ def __init__( @@ -108,6 +114,7 @@ class PinnedModuleStager: pin_memory: bool = True, copy_stream: Any | None = None, cache_retention: BoundedAllocatorCache | None = None, + retain_device_storage: bool = False, ) -> None: modules = (module,) if isinstance(module, nn.Module) else tuple(module) if not modules or not all(isinstance(item, nn.Module) for item in modules): @@ -117,6 +124,7 @@ class PinnedModuleStager: self.copy_stream = copy_stream if copy_stream is not None else current_omni_platform.Stream() self._ready_event = current_omni_platform.Event() self.cache_retention = cache_retention + self.retain_device_storage = retain_device_storage self.loaded = False self._groups = self._snapshot_groups(modules, pin_memory=pin_memory) self._device_storages: list[torch.Tensor] = [] @@ -234,6 +242,12 @@ class PinnedModuleStager: def load(self) -> None: if self.loaded: return + if self._device_storages: + # Storage retained from a previous load still holds the immutable + # weights; rebinding suffices, no H2D copy repeats. + self._bind(self._device_storages) + self.loaded = True + return try: self._load_once() except torch.OutOfMemoryError: @@ -263,7 +277,8 @@ class PinnedModuleStager: self.loaded = False self._release_cache(force=True) raise - self._device_storages.clear() + if not self.retain_device_storage: + self._device_storages.clear() self.loaded = False self._release_cache() diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index b03be3687..8023480ba 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -13,6 +13,7 @@ from vllm_omni.diffusion.hooks import HookRegistry, ModelHook from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload +from .module_residency import PinnedModuleStager from .plan_resolver import resolve_offload_plan logger = init_logger(__name__) @@ -60,6 +61,7 @@ class SequentialOffloadHook(ModelHook): pin_memory: bool = True, use_hsdp: bool = False, offload_after_context: bool = True, + persistent_staging: bool = False, ): # Modules to offload to CPU before this module runs self.offload_targets = offload_targets @@ -67,6 +69,21 @@ class SequentialOffloadHook(ModelHook): self.pin_memory = pin_memory self.use_hsdp = use_hsdp self.offload_after_context = offload_after_context + self.persistent_staging = persistent_staging + self._stager: PinnedModuleStager | None = None + + def initialize_hook(self, module: nn.Module) -> nn.Module: + if self.persistent_staging and not self.use_hsdp: + # Fixed staging storage keeps CUDA-graph-captured weight pointers + # valid across offload/load swaps. HSDP keeps the plain move + # semantics (decode graphs refuse HSDP anyway). + self._stager = PinnedModuleStager( + module, + self.device, + pin_memory=self.pin_memory, + retain_device_storage=True, + ) + return module @staticmethod def _move_params( @@ -101,7 +118,23 @@ class SequentialOffloadHook(ModelHook): moved = True return moved + @staticmethod + def _module_stager(module: nn.Module) -> PinnedModuleStager | None: + """The persistent stager owned by ``module``'s own sequential hook, if any. + + Cross-module swaps call ``_to_cpu(target)`` through the *running* + module's hook, so staging state must be resolved on the target's hook + to keep its ``loaded`` bookkeeping consistent. + """ + registry: HookRegistry | None = getattr(module, "_hook_registry", None) + hook = registry.get_hook(SequentialOffloadHook._HOOK_NAME) if registry is not None else None + return hook._stager if isinstance(hook, SequentialOffloadHook) else None + def _to_cpu(self, module: nn.Module) -> None: + stager = self._module_stager(module) + if stager is not None: + stager.offload() + return # XPU's allocator doesn't respect stream dependencies in empty_cache, # so non-blocking copies can race with cache eviction. Use blocking # copies on XPU to avoid NULL pointer errors during DMA. @@ -116,6 +149,10 @@ class SequentialOffloadHook(ModelHook): current_omni_platform.empty_cache() def _to_gpu(self, module: nn.Module) -> None: + stager = self._module_stager(module) + if stager is not None: + stager.load() + return self._move_params(module, self.device, non_blocking=False) def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple, dict]: @@ -147,6 +184,7 @@ def apply_sequential_offload( offload_initial_dits: bool = False, offload_dit_modules: Collection[nn.Module] | None = None, offload_encoder_modules: Collection[nn.Module] | None = None, + persistent_dit_staging: bool = False, ) -> None: """Apply sequential offloading hooks to DiT and encoder modules. @@ -165,6 +203,10 @@ def apply_sequential_offload( every DiT for backward compatibility. offload_encoder_modules: Encoder/stage modules allowed to move to CPU. None selects every supplied module for backward compatibility. + persistent_dit_staging: Keep DiT weights on fixed device staging + storage across swaps, so CUDA-graph-captured weight pointers stay + valid. Encoders keep plain move semantics so their memory is + actually freed while the DiT runs. Example: >>> apply_sequential_offload( @@ -193,6 +235,7 @@ def apply_sequential_offload( pin_memory=pin_memory, use_hsdp=use_hsdp, offload_after_context=id(dit_mod) in selected_dit_ids, + persistent_staging=persistent_dit_staging, ) registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook) logger.debug("Registered offload hook for %s", dit_mod.__class__.__name__) @@ -261,6 +304,15 @@ def _get_sequential_offload_hook(module: nn.Module) -> SequentialOffloadHook: return hook +def sequential_offload_staging_active(module: nn.Module) -> bool: + """True when the module's sequential offload hook serves fixed device storage. + + Fixed staging storage keeps CUDA-graph-captured weight pointers valid + across offload/load swaps. + """ + return SequentialOffloadHook._module_stager(module) is not None + + @contextmanager def sequential_offload_component(module: nn.Module) -> Iterator[None]: """Activate and release a hooked component called outside ``forward``.""" @@ -362,6 +414,14 @@ class ModelLevelOffloadBackend(OffloadBackend): use_hsdp=self.config.use_hsdp, offload_dit_modules=selected_dits, offload_encoder_modules=selected_encoders, + # Fixed DiT staging keeps decode-graph weight pointers valid; + # only models that capture weight pointers (Qwen-Image-2.1's + # CUDA-graph decode) pay for it — every other model keeps + # plain move semantics so offloading actually frees DiT VRAM + # while the encoders run. + persistent_dit_staging=( + self.device.type == "cuda" and any(getattr(dit, "enable_cuda_graph_decode", False) for dit in dits) + ), ) except BaseException: try: diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 9d05a8251..2c4cfb024 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -56,6 +56,11 @@ _DIFFUSION_MODELS = { "pipeline_qwen_image_layered", "QwenImageLayeredPipeline", ), + "QwenImage21Pipeline": ( + "qwen_image_21", + "pipeline_qwen_image_21", + "QwenImage21Pipeline", + ), "GlmImagePipeline": ( "glm_image", "pipeline_glm_image", @@ -426,6 +431,10 @@ _NO_CACHE_ACCELERATION = { "Pi0Pipeline", "Pi05Pipeline", "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", } @@ -651,6 +660,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "AnimaPipeline": "get_anima_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", @@ -740,6 +750,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", diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index b4125dcde..9af1fda59 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -76,6 +76,8 @@ class OmniDiffusionRequest: # A model preprocessor may keep selected requests on the legacy full-forward # path even when the engine is globally configured for step execution. use_step_execution: bool = True + # Whether first and later denoising steps may share a step batch. + allow_mixed_step_phases: bool = True # KV-recv wall-clock (ms), set by the runner's _prepare_request_for_forward # and carried to DiffusionOutput for the vllm_omni:diffusion_kv_load_s metric. kv_recv_ms: float = 0.0 diff --git a/vllm_omni/diffusion/sched/step_scheduler.py b/vllm_omni/diffusion/sched/step_scheduler.py index 2db6292fe..61a0aed36 100644 --- a/vllm_omni/diffusion/sched/step_scheduler.py +++ b/vllm_omni/diffusion/sched/step_scheduler.py @@ -13,6 +13,7 @@ from vllm_omni.diffusion.sched.base_scheduler import BaseScheduler from vllm_omni.diffusion.sched.interface import ( DiffusionRequestStatus, DiffusionSchedulerOutput, + SchedulerRequestState, ) if TYPE_CHECKING: @@ -62,6 +63,19 @@ class StepScheduler(BaseScheduler): ) return request_id + def _can_schedule_waiting(self, state: SchedulerRequestState) -> bool: + if not super()._can_schedule_waiting(state): + return False + + is_first_step = self._request_progress[state.request_id].current_step == 0 + for request_id in self._running: + running = self._request_states[request_id] + if state.req.allow_mixed_step_phases and running.req.allow_mixed_step_phases: + continue + if (self._request_progress[request_id].current_step == 0) != is_first_step: + return False + return True + def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: BaseRunnerOutput) -> set[str]: scheduled_request_ids = sched_output.scheduled_request_ids if not scheduled_request_ids and not sched_output.finished_req_ids: diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 19a59797c..07b3224c9 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -206,7 +206,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin): self.kv_transfer_manager = ( payload_transfer_manager if getattr(od_config, "kv_transfer_config", None) is None else None ) - self.init_omni_connectors(od_config, payload_transfer_manager, synchronous=True) + self.init_omni_connectors(od_config, payload_transfer_manager, synchronous=True) # type: ignore[arg-type] self._kv_connector = None from vllm_omni.diffusion.diffusion_kv.kv_connector import KVReceiveProgress, native_prefetch_enabled @@ -250,14 +250,25 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin): if model is None: return + compile_kwargs: dict[str, Any] = {"dynamic": self.od_config.diffusion_compile_dynamic} + if getattr(model, "enable_cuda_graph_decode", False): + # Decode-graph models still compile their blocks: graph capture + # records the compiled (fused) kernels. Scope inductor cudagraphs + # to this compile so the two graph layers never stack, without + # changing torch.compile for every other model in the process. + compile_kwargs["options"] = { + "triton.cudagraphs": False, + "triton.cudagraph_trees": False, + } + logger.info("Model runner: %s combines CUDA graph decode with torch.compile.", attr_name) + compile_granularity = self.od_config.diffusion_compile_granularity - compile_dynamic = self.od_config.diffusion_compile_dynamic try: if compile_granularity == "full": - model.compile(dynamic=compile_dynamic) + model.compile(**compile_kwargs) compiled_model = model else: - compiled_model = regionally_compile(model, dynamic=compile_dynamic) + compiled_model = regionally_compile(model, **compile_kwargs) setattr(self.pipeline, attr_name, compiled_model) except Exception as e: logger.warning( @@ -275,7 +286,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin): "compilation errors may surface on the first request.", attr_name, compile_granularity, - compile_dynamic, + compile_kwargs["dynamic"], ) def load_model( @@ -1454,7 +1465,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin): else req.denoise_completed ) if finished and result is not None: - self._maybe_send_stage_payload([req], [result]) + self._maybe_send_stage_payload([req], [result]) # type: ignore[list-item] runner_output_list.append( RunnerOutput( request_id=req.request_id, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 4bea0e05b..b3a23d8bd 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -585,6 +585,10 @@ class OrchestratorArgs: diffusion_attention_config: str | None = None diffusion_compile_granularity: str | None = None diffusion_compile_dynamic: bool | None = None + # CUDA graph capture of fixed-shape KV-cache decode steps (Qwen-Image-2.1 + # today). None defers to the OmniDiffusionConfig default (enabled); + # --enforce-eager also disables it. + enable_cuda_graph_decode: bool | None = None cache_backend: str = "none" cache_config: str | None = None video_output_transport: dict[str, object] | None = None diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index f551d5efe..765d540a3 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -614,6 +614,16 @@ class OmniServeCommand(CLISubcommand): "Disable for fixed-shape workloads with --no-diffusion-compile-dynamic." ), ) + omni_config_group.add_argument( + "--enable-cuda-graph-decode", + action=argparse.BooleanOptionalAction, + default=None, + help=( + "Capture fixed-shape KV-cache decode steps into CUDA graphs on models that implement them " + "(currently Qwen-Image-2.1). Independent of compilation_config.cudagraph_mode; " + "disable with --no-enable-cuda-graph-decode. --enforce-eager also disables it." + ), + ) omni_config_group.add_argument( "--fa-deterministic", dest="fa_deterministic", diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index cda13f31f..67d27dc51 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -2279,6 +2279,7 @@ async def generate_images( @router.post( "/v1/images/edits", + response_model=None, responses={ HTTPStatus.OK.value: {"model": ImageGenerationResponse}, HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, @@ -2286,6 +2287,7 @@ async def generate_images( HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, }, ) +@with_cancellation async def edit_images( raw_request: Request, image: list[UploadFile] | None = File(None), diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index bae1a3e63..421fd3337 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -473,7 +473,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): nested = getattr(request, "extra_body", None) explicit = getattr(request, "model_fields_set", None) if explicit is None: - explicit = getattr(request, "__fields_set__", ()) + explicit = getattr(request, "__fields_set__", None) or () explicit_root_args = { key: getattr(request, key) for key in explicit if key != "extra_body" and hasattr(request, key) } @@ -540,7 +540,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): def _truthy_extra_body_flag(request: Any, key: str) -> bool: if isinstance(request, dict): extra_body = request - model_extra = {} + model_extra: dict[str, Any] = {} else: extra_body = getattr(request, "extra_body", None) or {} model_extra = getattr(request, "model_extra", None) or {} @@ -778,7 +778,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Override the prompts produced by chat-template preprocessing. is_img2img = engine_prompt_image is not None - tprompt: OmniTextPrompt = {"prompt": extracted_prompt} + tprompt: OmniTextPrompt = OmniTextPrompt(prompt=extracted_prompt) if is_img2img: tprompt["modalities"] = ["img2img"] else: @@ -833,19 +833,18 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if not extracted_prompt: return self.create_error_response("No text prompt found in messages") - tprompt: OmniTextPrompt = {"prompt": extracted_prompt} - tprompt["modalities"] = ["text"] + text_prompt: OmniTextPrompt = OmniTextPrompt(prompt=extracted_prompt, modalities=["text"]) if reference_images: try: img_bytes = base64.b64decode(reference_images[0]) img = Image.open(BytesIO(img_bytes)) - tprompt["multi_modal_data"] = {"image": img} - tprompt["multi_modal_uuids"] = {"image": [f"{request_id}-image-0"]} + text_prompt["multi_modal_data"] = {"image": img} + text_prompt["multi_modal_uuids"] = {"image": [f"{request_id}-image-0"]} except Exception: pass - engine_prompts = [tprompt] + engine_prompts = [text_prompt] except Exception as e: logger.warning("Failed to build text-output prompt for single-stage diffusion: %s", e) _image_gen_height = None @@ -1422,7 +1421,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Pydantic v2 uses `model_fields_set`; keep v1 fallback for compatibility. explicit_fields = getattr(request, "model_fields_set", None) if explicit_fields is None: - explicit_fields = getattr(request, "__fields_set__", set()) + explicit_fields = getattr(request, "__fields_set__", None) or set() for field_name in self._OPENAI_SAMPLING_FIELDS: if field_name not in explicit_fields: @@ -2175,6 +2174,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # finish_reason is: # "tool_calls" for "auto" or "required" tool calls, # and "stop" for named tool calls. + finish_reason_: str | None if ( auto_tools_called or (tools_streamed[i] and not tool_choice_function_name) @@ -2626,7 +2626,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if output_text: # Get the corresponding output token IDs output_token_ids = None - if choice.index < len(final_res.outputs): + if final_res is not None and choice.index < len(final_res.outputs): output_token_ids = final_res.outputs[choice.index].token_ids self.request_logger.log_outputs( @@ -3234,7 +3234,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): resolve_stop_token_ids, ) - build_kwargs: dict[str, Any] = { + build_kwargs = { "task": "it2i" if reference_images else "t2i", "sys_type": use_system_prompt, "custom_system_prompt": custom_system_prompt, @@ -3282,7 +3282,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): **ar_stop_kwargs, ) - engine_prompt: OmniTextPrompt = {"prompt": prompt} + engine_prompt: OmniTextPrompt = OmniTextPrompt(prompt=prompt) if prompt_token_ids is not None: engine_prompt["prompt_token_ids"] = prompt_token_ids if system_prompt_type is not None: @@ -3327,7 +3327,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # AR stop tokens: use stage_type=="llm" instead of comprehension_idx # (None for DictConfig where is_comprehension is nested in engine_args). if stage_type == "llm" and ar_stop_token_ids is not None: - default_stage_params.stop_token_ids = ar_stop_token_ids + setattr(default_stage_params, "stop_token_ids", ar_stop_token_ids) if stage_type == "llm" and idx == 0 and ar_prompt_info is not None: if should_init_extra_args_for_non_diffusion_stages(model_class_name): @@ -3480,11 +3480,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): except Exception as e: # pragma: no cover - safeguard logger.warning("Failed to parse LoRA request: %s", e) - gen_prompt: OmniTextPrompt = { - "prompt": prompt, - "negative_prompt": negative_prompt, - "modalities": ["image"], - } + gen_prompt: OmniTextPrompt = OmniTextPrompt( + prompt=prompt, + modalities=["image"], + ) + if negative_prompt is not None: + gen_prompt["negative_prompt"] = negative_prompt if pil_images: if len(pil_images) == 1: gen_prompt["multi_modal_data"] = {"image": pil_images[0]} @@ -3612,20 +3613,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # ``images`` field, so fall back to the shared extractor used by the # offline example path before giving up. images = getattr(result, "images", None) or extract_images_from_outputs(result) - stage_durations = result.stage_durations - peak_memory_mb = result.peak_memory_mb + stage_durations = getattr(result, "stage_durations", None) + peak_memory_mb = getattr(result, "peak_memory_mb", None) response_metrics = getattr(result, "metrics", None) if return_stage_metrics else None cot_output = None - - req_out = result - if req_out: - prompt_obj = getattr(req_out, "prompt", None) - if isinstance(prompt_obj, dict): - extra = prompt_obj.get("extra", {}) - if isinstance(extra, dict): - ar_text = extra.get("ar_generated_text") - if isinstance(ar_text, str) and ar_text.strip(): - cot_output = ar_text + prompt_obj = getattr(result, "prompt", None) + if isinstance(prompt_obj, dict): + extra = prompt_obj.get("extra", {}) + if isinstance(extra, dict): + ar_text = extra.get("ar_generated_text") + if isinstance(ar_text, str) and ar_text.strip(): + cot_output = ar_text return self._flatten_diffusion_images(images), stage_durations, peak_memory_mb, cot_output, response_metrics @@ -3653,14 +3651,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): text = completion.text or "" if not text: continue - chunk = ImageEditARDeltaChunk( + delta_chunk = ImageEditARDeltaChunk( delta=text, index=completion.index, created=created, model=model, metrics=metrics, ) - yield f"data: {chunk.model_dump_json()}\n\n" + yield f"data: {delta_chunk.model_dump_json()}\n\n" elif final_output_type == "image": images = self._flatten_diffusion_images(getattr(output, "images", [])) if not images: @@ -3676,15 +3674,25 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ) for img in images ] - chunk = ImageEditImageChunk( + actual_size = f"{images[0].width}x{images[0].height}" + # Echo the request size only when it already matches the + # decoded pixels. Use actual WxH for size=auto/None or when + # the pipeline diverged (e.g. Qwen-Image-2.1 aspect-derived + # outputs). Edit models that honor an explicit WxH keep the + # same wire string because size == actual_size. + if size in (None, "auto") or size != actual_size: + chunk_size = actual_size + else: + chunk_size = size + img_chunk = ImageEditImageChunk( data=image_data, output_format=output_format, - size=size, + size=chunk_size, created=created, model=model, metrics=metrics, ) - yield f"data: {chunk.model_dump_json()}\n\n" + yield f"data: {img_chunk.model_dump_json()}\n\n" emitted_image = True if not emitted_image: raise RuntimeError("Streaming image edit completed without a final image output.") @@ -3704,7 +3712,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ) except OmniClientError as exc: logger.info("Client error during streaming image edit: %s", exc) - chunk = ImageEditStreamError( + client_error_chunk = ImageEditStreamError( created=created, model=model, error={ @@ -3713,10 +3721,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): "code": exc.status_code, }, ) - yield f"data: {chunk.model_dump_json()}\n\n" + yield f"data: {client_error_chunk.model_dump_json()}\n\n" except Exception as exc: logger.exception("Streaming image edit failed: %s", exc) - chunk = ImageEditStreamError( + server_error_chunk = ImageEditStreamError( created=created, model=model, error={ @@ -3725,7 +3733,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): "code": 500, }, ) - yield f"data: {chunk.model_dump_json()}\n\n" + yield f"data: {server_error_chunk.model_dump_json()}\n\n" yield "data: [DONE]\n\n" @staticmethod @@ -3832,11 +3840,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): logger.warning("Failed to decode reference image: %s", e) # Build generation kwargs - gen_prompt: OmniTextPrompt = { - "prompt": prompt, - "negative_prompt": negative_prompt, - "modalities": ["image"], - } + gen_prompt: OmniTextPrompt = OmniTextPrompt( + prompt=prompt, + modalities=["image"], + ) + if negative_prompt is not None: + gen_prompt["negative_prompt"] = negative_prompt gen_params = OmniDiffusionSamplingParams( height=height, width=width, @@ -4062,6 +4071,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ) # Build response + content: str | list[dict[str, Any]] if not image_contents: content = "Image generation completed but no images were produced." else: diff --git a/vllm_omni/model_extras/qwen_image_21.py b/vllm_omni/model_extras/qwen_image_21.py new file mode 100644 index 000000000..b45fa87bc --- /dev/null +++ b/vllm_omni/model_extras/qwen_image_21.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from typing import Any + +from PIL import Image + + +def build_image_to_image_prompt( + prompt: str, + negative_prompt: str | None, + input_image: Image.Image | list[Image.Image], + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + img_prompt: dict[str, Any] = { + "prompt": prompt, + "multi_modal_data": {"image": input_image}, + } + if height is not None: + img_prompt["height"] = height + if width is not None: + img_prompt["width"] = width + if negative_prompt is not None: + img_prompt["negative_prompt"] = negative_prompt + return img_prompt diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 233516db8..dbe5d797d 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -91,6 +91,9 @@ from vllm_omni.model_extras.ming_image import ( from vllm_omni.model_extras.ming_image import ( build_text_to_image_prompt as build_ming_image_text_to_image_prompt, ) +from vllm_omni.model_extras.qwen_image_21 import ( + build_image_to_image_prompt as build_qwen_image_21_image_to_image_prompt, +) from vllm_omni.model_extras.sana_video import SANA_VIDEO_EXTRA_BODY_PARAMS from vllm_omni.model_extras.sensenova_u1 import ( SENSENOVA_U1_EXTRA_BODY_PARAMS, @@ -329,6 +332,9 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "text_to_image_prompt_builder": build_ming_image_text_to_image_prompt, "image_to_image_prompt_builder": build_ming_image_image_to_image_prompt, }, + "QwenImage21Pipeline": { + "image_to_image_prompt_builder": build_qwen_image_21_image_to_image_prompt, + }, } _EXTRA_SPECS["MingImageLayeredDiffusionPipeline"] = _EXTRA_SPECS["MingImageDiffusionPipeline"] @@ -473,10 +479,10 @@ def build_text_to_image_prompt( if builder is None: return prompt return builder( - prompt=str(prompt["prompt"]), - negative_prompt=prompt.get("negative_prompt"), - height=height, - width=width, + str(prompt["prompt"]), + prompt.get("negative_prompt"), + height, + width, ) @@ -513,10 +519,10 @@ def build_image_to_video_prompt( if not isinstance(media_inputs, Mapping): raise TypeError("Canonical I2V prompt multi_modal_data must be a mapping.") return builder( - prompt=str(prompt["prompt"]), - negative_prompt=prompt.get("negative_prompt"), - media_inputs=media_inputs, - height=height, - width=width, - num_frames=num_frames, + str(prompt["prompt"]), + prompt.get("negative_prompt"), + media_inputs, + height, + width, + num_frames, )