Initialize QMI pipeline implementation (#2)

* init

Signed-off-by: Isotr0py <mozf@inferact.ai>

* fix

Signed-off-by: Isotr0py <mozf@inferact.ai>

---------

Signed-off-by: Isotr0py <mozf@inferact.ai>

docs/models/supported_models.md

@@ -32,6 +32,7 @@ th {
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `QwenImage21Pipeline` | Qwen-Image-2.1 (T2I + image-conditioned generation, up to 4 condition images) | `Qwen/Qwen-Image-2.1` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2.1.md) |
 | `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — |
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | — |

docs/user_guide/diffusion_features.md

@@ -139,6 +139,7 @@ The following tables show which models support each feature:
 | **Qwen-Image-Edit**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Qwen-Image-Edit-2509** |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |        ✅ (decode)         |          ✅           |       ❌        |        ❌         |
 | **Qwen-Image-Layered**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **Qwen-Image-2.1**       |     ❌     |     ❌      |   ✅ (Ulysses only)    |       ✅        |         ❓         |          ❌          |   ❓    |             ❓             |          ❓           |       ❓        |        ✅         |
 | **SenseNova-U1 / U1.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Stable-Diffusion-XL**  |     ❌     |     ❌      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Stable-Diffusion3.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
@@ -154,6 +155,7 @@ The following tables show which models support each feature:
 > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
 > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
 > 6. BAGEL step execution supports image generation with `bagel.yaml`, `bagel_think.yaml`, and `bagel_single_stage.yaml`; two-stage Thinker execution and explicit single-stage text output remain on their existing complete-request paths. Image requests require `num_inference_steps >= 2`. BAGEL step execution cannot currently be combined with sequence parallelism or a diffusion cache backend; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
+> 7. Qwen-Image-2.1 (`QwenImage21Pipeline`) handles both text-to-image and image-conditioned generation (up to 4 condition images) and supports LoRA through `QwenImageLoraLoaderMixin`. Sequence parallelism is Ulysses-only (ring attention is not supported), and cache acceleration (TeaCache / Cache-DiT) is unavailable; its built-in KV cache for the text/condition prefix is a separate feature from those cache backends.
 
 ### VideoGen
 

examples/offline_inference/image_to_image/qwen_image_21_edit.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Example script for image-conditioned generation with Qwen-Image 2.1.
+
+Qwen-Image 2.1 uses a single pipeline (QwenImage21Pipeline) for both
+text-to-image and image-conditioned generation. Pass one or more condition
+images (at most 4) via --image to run the image-conditioned path.
+
+Usage (single image):
+    python qwen_image_21_edit.py \
+        --image input.png \
+        --prompt "Let this mascot dance under the moon, surrounded by floating stars" \
+        --negative-prompt "blurry, low quality, text, watermark" \
+        --output qwen_image_21_edit.png \
+        --num-inference-steps 50 \
+        --cfg-scale 4.0
+
+Usage (multiple images, up to 4):
+    python qwen_image_21_edit.py \
+        --image input1.png input2.png \
+        --prompt "Combine these images into a single scene" \
+        --negative-prompt "blurry, low quality" \
+        --output qwen_image_21_edit.png \
+        --num-inference-steps 50 \
+        --cfg-scale 4.0
+
+Notes:
+    - Qwen-Image 2.1 only supports true classifier-free guidance
+      (--cfg-scale); there is no guidance-distilled guidance_scale.
+      --cfg-scale > 1 only takes effect when --negative-prompt is provided.
+    - --height/--width must be multiples of 32. When omitted, the output size
+      is derived from the last condition image's aspect ratio at ~1024x1024.
+    - Cache acceleration backends (cache_dit / tea_cache) are not supported.
+
+For more options, run:
+    python qwen_image_21_edit.py --help
+"""
+
+import argparse
+import os
+import time
+from pathlib import Path
+
+import torch
+from PIL import Image
+
+from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.model_extras import (
+    build_image_to_image_prompt,
+    get_model_class_name,
+)
+from vllm_omni.platforms import current_omni_platform
+
+MAX_INPUT_IMAGES = 4
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Image-conditioned generation with Qwen-Image 2.1.")
+    parser.add_argument(
+        "--model",
+        default="Qwen/Qwen-Image-2.1",
+        help="Diffusion model name or local path.",
+    )
+    parser.add_argument(
+        "--image",
+        type=str,
+        nargs="+",
+        required=True,
+        help="Path(s) to input condition image file(s) (PNG, JPG, etc.). At most 4 images.",
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        required=True,
+        help="Text prompt describing the desired generation.",
+    )
+    parser.add_argument(
+        "--negative-prompt",
+        type=str,
+        default=None,
+        help="Negative prompt. Required for true classifier-free guidance (--cfg-scale > 1) to take effect.",
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=None,
+        metavar="W",
+        help="Output image width in pixels (multiple of 32). Default: None (derived from the condition image).",
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=None,
+        metavar="H",
+        help="Output image height in pixels (multiple of 32). Default: None (derived from the condition image).",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed for deterministic results.",
+    )
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=4.0,
+        help=(
+            "True classifier-free guidance scale (default: 4.0). Enabled by setting cfg_scale > 1 and providing "
+            "a negative prompt. Qwen-Image 2.1 supports true CFG only; there is no guidance_scale."
+        ),
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="qwen_image_21_edit.png",
+        help="Path to save the generated image (PNG).",
+    )
+    parser.add_argument(
+        "--num-outputs-per-prompt",
+        type=int,
+        default=1,
+        help="Number of images to generate for the given prompt.",
+    )
+    parser.add_argument(
+        "--num-inference-steps",
+        type=int,
+        default=50,
+        help="Number of denoising steps for the diffusion sampler.",
+    )
+    parser.add_argument(
+        "--ulysses-degree",
+        type=int,
+        default=1,
+        help="Number of GPUs used for ulysses sequence parallelism (ring attention is not supported).",
+    )
+    parser.add_argument(
+        "--cfg-parallel-size",
+        type=int,
+        default=1,
+        choices=[1, 2],
+        help="Number of GPUs used for classifier free guidance parallel size.",
+    )
+    parser.add_argument(
+        "--enforce-eager",
+        action="store_true",
+        default=None,
+        help="Disable torch.compile and force eager execution.",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    if len(args.image) > MAX_INPUT_IMAGES:
+        raise ValueError(f"Qwen-Image 2.1 supports at most {MAX_INPUT_IMAGES} input images, got {len(args.image)}")
+
+    input_images = []
+    for image_path in args.image:
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Input image not found: {image_path}")
+        input_images.append(Image.open(image_path).convert("RGB"))
+
+    # Use single image or list based on number of inputs
+    input_image = input_images[0] if len(input_images) == 1 else input_images
+
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+
+    omni_kwargs = {
+        "model": args.model,
+        "ulysses_degree": args.ulysses_degree,
+        "cfg_parallel_size": args.cfg_parallel_size,
+    }
+    if args.enforce_eager is not None:
+        omni_kwargs["enforce_eager"] = args.enforce_eager
+    omni = Omni(**omni_kwargs)
+    model_class_name = get_model_class_name(omni)
+    print("Pipeline loaded")
+
+    print(f"\n{'=' * 60}")
+    print("Generation Configuration:")
+    print(f"  Model: {args.model}")
+    print(f"  Inference steps: {args.num_inference_steps}")
+    print(f"  CFG scale (true CFG): {args.cfg_scale}")
+    if args.height is not None or args.width is not None:
+        print(f"  Output size: {args.width or 'auto'}x{args.height or 'auto'}")
+    if isinstance(input_image, list):
+        print(f"  Number of input images: {len(input_image)}")
+    else:
+        print(f"  Input image size: {input_image.size}")
+    print(
+        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, "
+        f"cfg_parallel_size={args.cfg_parallel_size}"
+    )
+    print(f"{'=' * 60}\n")
+
+    generation_start = time.perf_counter()
+
+    prompt_dict = build_image_to_image_prompt(
+        model_class_name=model_class_name,
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        input_image=input_image,
+        height=args.height,
+        width=args.width,
+    )
+
+    diffusion_params = OmniDiffusionSamplingParams(
+        generator=generator,
+        seed=args.seed,
+        true_cfg_scale=args.cfg_scale,
+        num_inference_steps=args.num_inference_steps,
+        num_outputs_per_prompt=args.num_outputs_per_prompt,
+        height=args.height,
+        width=args.width,
+    )
+
+    outputs = omni.generate(prompt_dict, sampling_params_list=[diffusion_params])
+    generation_end = time.perf_counter()
+    generation_time = generation_end - generation_start
+
+    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
+
+    if not outputs:
+        raise ValueError("No output generated from omni.generate()")
+
+    images = None
+    for output in outputs:
+        images = getattr(output, "images", None)
+        if images:
+            break
+
+    if not images:
+        images = extract_images_from_outputs(outputs)
+
+    if not images:
+        raise ValueError("No images found in request_output")
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    suffix = output_path.suffix or ".png"
+    stem = output_path.stem or "qwen_image_21_edit"
+
+    if len(images) <= 1:
+        images[0].save(output_path)
+        print(f"Saved generated image to {os.path.abspath(output_path)}")
+    else:
+        for idx, img in enumerate(images):
+            save_path = output_path.parent / f"{stem}_{idx}{suffix}"
+            img.save(save_path)
+            print(f"Saved generated image to {os.path.abspath(save_path)}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/image_to_image/run_qwen_image_21_edit.sh

@@ -0,0 +1,8 @@
+python qwen_image_21_edit.py \
+    --model Qwen/Qwen-Image-2.1 \
+    --image qwen_bear.png \
+    --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
+    --negative-prompt "blurry, low quality, text, watermark" \
+    --output qwen_image_21_edit.png \
+    --num-inference-steps 50 \
+    --cfg-scale 4.0

examples/offline_inference/text_to_image/text_to_image.py

@@ -86,7 +86,8 @@ def parse_args() -> argparse.Namespace:
         "--model",
         default="Qwen/Qwen-Image",
         help="Diffusion model name or local path. Supported models: "
-        "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1, "
+        "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, Qwen/Qwen-Image-2.1, "
+        "stepfun-ai/NextStep-1.1, "
         "black-forest-labs/FLUX.1-dev, black-forest-labs/FLUX.2-klein-9B, "
         "black-forest-labs/FLUX.2-dev, tencent/HunyuanImage-3.0-Instruct, "
         "HiDream-ai/HiDream-O1-Image, meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "

recipes/Qwen/Qwen-Image-2.1.md

@@ -0,0 +1,148 @@
+# Qwen-Image-2.1
+
+> Text-to-image and image-conditioned generation with Qwen-Image 2.1
+
+## Summary
+
+- Vendor: Qwen
+- Model: `Qwen/Qwen-Image-2.1`
+- Task: Text-to-image generation and image-conditioned generation (editing)
+- Mode: Offline inference and online serving with optional step-wise execution
+- Maintainer: Community
+
+> **Model id:** `Qwen/Qwen-Image-2.1` follows the reference diffusers
+> documentation and may change before the official release — the final
+> checkpoint name is subject to the official release.
+
+## When to use this recipe
+
+Use this recipe as a starting point for running `Qwen/Qwen-Image-2.1`. A single
+pipeline class, `QwenImage21Pipeline`, handles both pure text-to-image requests
+(no condition image) and image-conditioned requests with up to 4 condition
+images: prompt and condition images are encoded together by a Qwen3-VL text
+encoder, so editing-style requests go through the same serving path as
+text-to-image.
+
+## References
+
+- Related offline examples:
+  [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py),
+  [`examples/offline_inference/image_to_image/qwen_image_21_edit.py`](../../examples/offline_inference/image_to_image/qwen_image_21_edit.py)
+- Supported models table:
+  [`docs/models/supported_models.md`](../../docs/models/supported_models.md)
+- Feature compatibility matrix:
+  [`docs/user_guide/diffusion_features.md`](../../docs/user_guide/diffusion_features.md)
+- Sibling recipes: [`Qwen-Image.md`](./Qwen-Image.md),
+  [`Qwen-Image-Edit.md`](./Qwen-Image-Edit.md)
+
+## Offline Inference
+
+### Text-to-image
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image-2.1 \
+  --prompt "A ceramic teapot on a wooden table" \
+  --negative-prompt "blurry, low quality, text, watermark" \
+  --output qwen_image_21_t2i.png \
+  --num-inference-steps 50 \
+  --cfg-scale 4.0
+```
+
+### Image-conditioned generation (editing)
+
+```bash
+python examples/offline_inference/image_to_image/qwen_image_21_edit.py \
+  --image qwen_bear.png \
+  --prompt "Let this mascot dance under the moon, surrounded by floating stars" \
+  --negative-prompt "blurry, low quality, text, watermark" \
+  --output qwen_image_21_edit.png \
+  --num-inference-steps 50 \
+  --cfg-scale 4.0
+```
+
+or via the bundled launcher:
+
+```bash
+bash examples/offline_inference/image_to_image/run_qwen_image_21_edit.sh
+```
+
+Up to 4 condition images can be passed via `--image`:
+
+```bash
+python examples/offline_inference/image_to_image/qwen_image_21_edit.py \
+  --image input1.png input2.png \
+  --prompt "Combine these images into a single scene" \
+  --negative-prompt "blurry, low quality" \
+  --output qwen_image_21_multi.png \
+  --num-inference-steps 50 \
+  --cfg-scale 4.0
+```
+
+## Online Serving
+
+Start the server:
+
+```bash
+vllm serve Qwen/Qwen-Image-2.1 --omni --port 8091
+```
+
+To enable the step-wise runtime:
+
+```bash
+vllm serve Qwen/Qwen-Image-2.1 --omni \
+  --port 8091 \
+  --step-execution \
+  --max-num-seqs 8
+```
+
+### Verification
+
+For a direct API smoke test:
+
+```bash
+curl http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen-Image-2.1",
+    "prompt": "A ceramic teapot on a wooden table",
+    "negative_prompt": "blurry, low quality, text, watermark",
+    "size": "1024x1024",
+    "num_inference_steps": 50,
+    "true_cfg_scale": 4.0,
+    "seed": 42
+  }'
+```
+
+For image-conditioned generation, pass condition images through the
+OpenAI-compatible multimodal endpoint (`/v1/chat/completions` with
+`image_url` content parts, or `/v1/images/edits`), same as
+[`Qwen-Image-Edit.md`](./Qwen-Image-Edit.md).
+
+## Parameters
+
+| Parameter | Default | Notes |
+| --- | --- | --- |
+| `true_cfg_scale` | 4.0 | Only takes effect when `negative_prompt` is provided. Qwen-Image 2.1 supports true CFG only — there is no `guidance_scale` parameter. |
+| `num_inference_steps` | 50 | Number of denoising steps. |
+| `height` / `width` | 1024×1024 | Must be multiples of 32. For image-conditioned requests, leaving them unset derives the output size from the last condition image's aspect ratio at ~1024×1024. |
+| `negative_prompt` | None | Required to activate true classifier-free guidance (`true_cfg_scale > 1`). |
+| condition images | — | Up to 4 input images per request; more than 4 raises an error. |
+| `seed` | — | Fix for reproducible outputs. |
+
+### Prefix KV cache
+
+The transformer caches the text/condition-image prefix across denoising steps
+(a prefill/decode split keyed by CFG branch). The cache is enabled
+automatically when the checkpoint declares `causal_condition: true` and is
+re-created per generation, so no state leaks between requests. It changes the
+compute path — cached prefix vs. full recompute at every step — so exact
+numerical reproduction against a no-cache reference (e.g. the diffusers
+pipeline) requires matching the cache setting on both sides.
+
+## Known Limitations
+
+- Cache acceleration backends (`cache_dit`, `tea_cache`) are not supported;
+  `QwenImage21Pipeline` is registered in `_NO_CACHE_ACCELERATION`.
+- Sequence parallelism supports Ulysses only; ring attention is not supported.
+- Only true CFG is exposed; there is no `guidance_scale` knob.

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+from diffusers.models.autoencoders.vae import DecoderOutput
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+from vllm_omni.diffusion.models.qwen_image_21.autoencoder_kl_qwenimage21 import (
+    AutoencoderKLQwenImage21,
+    _unpatchify,
+)
+
+logger = init_logger(__name__)
+
+
+class DistributedAutoencoderKLQwenImage21(AutoencoderKLQwenImage21, DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+
+    def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        # mostly copy from AutoencoderKLQwenImage21.tiled_decode
+        _, _, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        tile_sample_stride_height = self.tile_sample_stride_height
+        tile_sample_stride_width = self.tile_sample_stride_width
+        if self.config.patch_size is not None:
+            sample_height = sample_height // self.config.patch_size
+            sample_width = sample_width // self.config.patch_size
+            tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
+            tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
+            blend_height = self.tile_sample_min_height // self.config.patch_size - tile_sample_stride_height
+            blend_width = self.tile_sample_min_width // self.config.patch_size - tile_sample_stride_width
+        else:
+            blend_height = self.tile_sample_min_height - tile_sample_stride_height
+            blend_width = self.tile_sample_min_width - tile_sample_stride_width
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        tiletask_list = []
+        for i in range(0, height, tile_latent_stride_height):
+            for j in range(0, width, tile_latent_stride_width):
+                time_list = []
+                for k in range(num_frames):
+                    tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                    time_list.append(tile)
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // tile_latent_stride_height, j // tile_latent_stride_width),
+                        time_list,
+                        workload=time_list[0].shape[3] * time_list[0].shape[4],
+                    )
+                )
+        tile_spec = {
+            "sample_height": sample_height,
+            "sample_width": sample_width,
+            "blend_height": blend_height,
+            "blend_width": blend_width,
+            "tile_sample_stride_height": tile_sample_stride_height,
+            "tile_sample_stride_width": tile_sample_stride_width,
+        }
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec=tile_spec,
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def tile_exec(self, task: TileTask) -> torch.Tensor:
+        """Decode a single latent tile into RGB space."""
+        self.clear_cache()
+        time = []
+        for k in range(len(task.tensor)):
+            self._conv_idx = [0]
+            tile = self.post_quant_conv(task.tensor[k])
+            decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
+            time.append(decoded)
+        result = torch.cat(time, dim=2)
+        return result
+
+    def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+        """Merge decoded tiles into a full image."""
+        grid_h, grid_w = grid_spec.grid_shape
+        self.clear_cache()
+
+        result_rows = []
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        :,
+                        : grid_spec.tile_spec["tile_sample_stride_height"],
+                        : grid_spec.tile_spec["tile_sample_stride_width"],
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+        dec = torch.cat(result_rows, dim=3)[
+            :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
+        ]
+
+        if self.config.patch_size is not None:
+            dec = _unpatchify(dec, patch_size=self.config.patch_size)
+
+        dec = torch.clamp(dec, min=-1.0, max=1.0)
+        return dec
+
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
+        if not self.is_distributed_enabled():
+            return super().tiled_decode(z, return_dict=return_dict)
+
+        logger.debug("Decode running with distributed executor")
+        result = self.distributed_executor.execute(
+            z,
+            DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
+            broadcast_result=True,
+        )
+        if not return_dict:
+            return (result,)
+
+        return DecoderOutput(sample=result)

vllm_omni/diffusion/lora/loader.py

@@ -22,6 +22,11 @@ lora_convert_mapping: dict[str, Callable] = {
     "QwenImagePipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
     "QwenImageEditPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
     "QwenImageEditPlusPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
+    # 2.1's single-stream blocks keep the same leaf names the converter's
+    # protected n-grams cover (`attn.to_q/to_k/to_v`, `img_mlp`, ...), so the
+    # 2.0 converter applies. Non-diffusers LoRAs that target the fused `to_qkv`
+    # projection are not supported (the converter would split it into `to.qkv`).
+    "QwenImage21Pipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
     "Wan22Pipeline": _convert_non_diffusers_wan_lora_to_diffusers,
     "Wan22I2VPipeline": _convert_non_diffusers_wan_lora_to_diffusers,
 }

vllm_omni/diffusion/model_metadata.py

@@ -20,6 +20,8 @@ class DiffusionModelMetadata:
 
 
 QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4
+# Qwen-Image 2.1 image-conditioned generation caps condition images at 4.
+QWEN_IMAGE_21_MAX_INPUT_IMAGES = 4
 # Upstream HunyuanImage-3.0 "Multi-Image Fusion" caps reference images at 3.
 HUNYUAN_IMAGE3_MAX_INPUT_IMAGES = 3
 # Boogu-Image editing (TI2I) supports a single reference image for now.
@@ -31,6 +33,12 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
         supports_multimodal_inputs=True,
         max_multimodal_image_inputs=QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
     ),
+    # Qwen-Image 2.1 handles both text-to-image (no image) and
+    # image-conditioned requests through the same pipeline class.
+    "QwenImage21Pipeline": DiffusionModelMetadata(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=QWEN_IMAGE_21_MAX_INPUT_IMAGES,
+    ),
     "HunyuanImage3Pipeline": DiffusionModelMetadata(
         supports_multimodal_inputs=True,
         max_multimodal_image_inputs=HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,

vllm_omni/diffusion/models/qwen_image_21/__init__.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Qwen Image 2.1 diffusion model components."""
+
+from vllm_omni.diffusion.models.qwen_image_21.autoencoder_kl_qwenimage21 import (
+    AutoencoderKLQwenImage21,
+)
+from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import (
+    QwenImage21Pipeline,
+    get_qwen_image_21_post_process_func,
+    get_qwen_image_21_pre_process_func,
+)
+from vllm_omni.diffusion.models.qwen_image_21.qwen_image_21_transformer import (
+    QwenImage21Transformer2DModel,
+)
+
+__all__ = [
+    "AutoencoderKLQwenImage21",
+    "QwenImage21Pipeline",
+    "QwenImage21Transformer2DModel",
+    "get_qwen_image_21_post_process_func",
+    "get_qwen_image_21_pre_process_func",
+]

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

vllm_omni/diffusion/registry.py

@@ -41,6 +41,11 @@ _DIFFUSION_MODELS = {
         "pipeline_qwen_image_layered",
         "QwenImageLayeredPipeline",
     ),
+    "QwenImage21Pipeline": (
+        "qwen_image_21",
+        "pipeline_qwen_image_21",
+        "QwenImage21Pipeline",
+    ),
     "GlmImagePipeline": (
         "glm_image",
         "pipeline_glm_image",
@@ -378,6 +383,10 @@ _NO_CACHE_ACCELERATION = {
     # instead of erroring.
     "Pi0Pipeline",
     "LingBotWorldCausalDMDPipeline",
+    # Qwen-Image 2.1's transformer carries its own prefix KV cache (keyed by CFG
+    # branch) across denoising steps, which conflicts with cache_dit / tea_cache
+    # step-skipping hooks.
+    "QwenImage21Pipeline",
 }
 
 
@@ -550,6 +559,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "QwenImagePipeline": "get_qwen_image_post_process_func",
     "QwenImageEditPipeline": "get_qwen_image_edit_post_process_func",
     "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_post_process_func",
+    "QwenImage21Pipeline": "get_qwen_image_21_post_process_func",
     "GlmImagePipeline": "get_glm_image_post_process_func",
     "ZImagePipeline": "get_post_process_func",
     "OvisImagePipeline": "get_ovis_image_post_process_func",
@@ -629,6 +639,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "BooguImageTurboPipeline": "get_boogu_image_pre_process_func",
     "QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func",
     "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_pre_process_func",
+    "QwenImage21Pipeline": "get_qwen_image_21_pre_process_func",
     "LongCatImageEditPipeline": "get_longcat_image_edit_pre_process_func",
     "LongCatVideoAvatarPipeline": "get_longcat_video_avatar_pre_process_func",
     "QwenImageLayeredPipeline": "get_qwen_image_layered_pre_process_func",