Enable VAE tiling for Qwen-Image-2.1 (#5)

* init

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

* fix

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

* test(diffusion): add Qwen-Image-2.1 unit tests

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

* fix(diffusion): Qwen-Image-2.1 VAE spatial ratio, tiling e2e verification

- AutoencoderKLQwenImage21: derive spatial_compression_ratio from the
  architecture (2 ** len(temperal_downsample) = 16x) instead of trusting
  the checkpoint's scale_factor_spatial=8, which corrupted tiled
  decode/encode geometry (output cropped to half size) and tiling gating.
- pipeline: _read_vae_scale_factor prefers the architecture-derived value
  for the same reason.
- models/qwen_image_21/__init__: lazy exports to break the circular import
  distributed VAE -> package __init__ -> pipeline -> distributed VAE.
- Add VAE unit tests: ratio derivation, tiled vs untiled decode/encode
  consistency, distributed tile_split/exec/merge == tiled_decode.
- e2e matrix (GB200, seed 42, 20 steps) all green; see e2e_out/qwen21_vae/.
- examples: qwen_image_21_edit.py gains --vae-use-tiling.
- docs: mark VAE patch-parallel verified for Qwen-Image-2.1; recipe notes
  the default-tile banding limitation at 2048px+.

* clean tests

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

* fix cudnn determination

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

---------

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

docs/user_guide/diffusion_features.md

@@ -139,7 +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)    |       ✅        |         ❓         |          ❌          |   ❓    |             ❓             |          ❓           |       ❓        |        ✅         |
+| **Qwen-Image-2.1**       |     ❌     |     ❌      |   ✅ (Ulysses only)    |       ✅        |         ❓         |          ❌          |   ❓    |             ❓             |      ✅ (decode)      |       ❓        |        ✅         |
 | **SenseNova-U1 / U1.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Stable-Diffusion-XL**  |     ❌     |     ❌      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Stable-Diffusion3.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |

examples/offline_inference/image_to_image/qwen_image_21_edit.py

@@ -144,6 +144,11 @@ def parse_args() -> argparse.Namespace:
         choices=[1, 2],
         help="Number of GPUs used for classifier free guidance parallel size.",
     )
+    parser.add_argument(
+        "--vae-use-tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization (also tiles condition-image encoding).",
+    )
     parser.add_argument(
         "--enforce-eager",
         action="store_true",
@@ -174,6 +179,7 @@ def main():
         "model": args.model,
         "ulysses_degree": args.ulysses_degree,
         "cfg_parallel_size": args.cfg_parallel_size,
+        "vae_use_tiling": args.vae_use_tiling,
     }
     if args.enforce_eager is not None:
         omni_kwargs["enforce_eager"] = args.enforce_eager

recipes/Qwen/Qwen-Image-2.1.md

@@ -146,3 +146,8 @@ pipeline) requires matching the cache setting on both sides.
   `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.
+- 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.

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py

@@ -15,6 +15,7 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
 )
 from vllm_omni.diffusion.models.qwen_image_21.autoencoder_kl_qwenimage21 import (
     AutoencoderKLQwenImage21,
+    _cudnn_deterministic,
     _unpatchify,
 )
 
@@ -89,11 +90,12 @@ class DistributedAutoencoderKLQwenImage21(AutoencoderKLQwenImage21, DistributedV
         """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)
+        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
 

vllm_omni/diffusion/models/qwen_image_21/__init__.py

@@ -1,18 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Qwen Image 2.1 diffusion model components."""
+"""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,
-)
+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",
@@ -21,3 +20,31 @@ __all__ = [
     "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)

vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py

@@ -15,6 +15,9 @@
 # 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
@@ -32,6 +35,32 @@ 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,
@@ -1150,7 +1179,17 @@ class AutoencoderKLQwenImage21(ModelMixin, AutoencoderMixin, ConfigMixin, FromOr
             is_residual=is_residual,
         )
 
-        self.spatial_compression_ratio = scale_factor_spatial
+        # 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.
@@ -1161,13 +1200,19 @@ class AutoencoderKLQwenImage21(ModelMixin, AutoencoderMixin, ConfigMixin, FromOr
         # intermediate tiles together, the memory requirement can be lowered.
         self.use_tiling = False
 
-        # The minimal tile height and width for spatial tiling to be used
-        self.tile_sample_min_height = 256
-        self.tile_sample_min_width = 256
+        # The minimal 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
-        self.tile_sample_stride_height = 192
-        self.tile_sample_stride_width = 192
+        # 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 = {
@@ -1219,6 +1264,7 @@ class AutoencoderKLQwenImage21(ModelMixin, AutoencoderMixin, ConfigMixin, FromOr
         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
 
@@ -1273,13 +1319,14 @@ class AutoencoderKLQwenImage21(ModelMixin, AutoencoderMixin, ConfigMixin, FromOr
             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(z, return_dict=return_dict)
+            return self._tiled_decode_adaptive(z, return_dict=return_dict)
 
         self.clear_cache()
         x = self.post_quant_conv(z)
@@ -1345,6 +1392,7 @@ class AutoencoderKLQwenImage21(ModelMixin, AutoencoderMixin, ConfigMixin, FromOr
             )
         return b
 
+    @_with_cudnn_deterministic
     def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
         r"""Encode a batch of images using a tiled encoder.
 
@@ -1417,6 +1465,49 @@ class AutoencoderKLQwenImage21(ModelMixin, AutoencoderMixin, ConfigMixin, FromOr
         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:
+            (
+                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.

vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py

@@ -86,11 +86,14 @@ def _read_vae_scale_factor(od_config: OmniDiffusionConfig) -> int:
     vae_config_path = os.path.join(model_path, "vae/config.json")
     with open(vae_config_path) as f:
         vae_config = json.load(f)
-    if vae_config.get("scale_factor_spatial"):
-        return int(vae_config["scale_factor_spatial"])
-    # 2.1 compresses 16x spatially (four downsampling stages).
+    # 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