From 3affc7766db7f5dea3e2719c02d3ea87e3ee7c1f Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Fri, 18 Sep 2026 10:02:50 +0800 Subject: [PATCH] Enable VAE tiling for Qwen-Image-2.1 (#5) * init Signed-off-by: Isotr0py * fix Signed-off-by: Isotr0py * test(diffusion): add Qwen-Image-2.1 unit tests Signed-off-by: Isotr0py * 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 * fix cudnn determination Signed-off-by: Isotr0py --------- Signed-off-by: Isotr0py --- docs/user_guide/diffusion_features.md | 2 +- .../image_to_image/qwen_image_21_edit.py | 6 + recipes/Qwen/Qwen-Image-2.1.md | 5 + .../autoencoder_kl_qwenimage21.py | 12 +- .../models/qwen_image_21/__init__.py | 51 +++++++-- .../autoencoder_kl_qwenimage21.py | 107 ++++++++++++++++-- .../qwen_image_21/pipeline_qwen_image_21.py | 9 +- 7 files changed, 163 insertions(+), 29 deletions(-) diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 7eaecc0f0..503f369d8 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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) | ❌ | ❌ | diff --git a/examples/offline_inference/image_to_image/qwen_image_21_edit.py b/examples/offline_inference/image_to_image/qwen_image_21_edit.py index cdc03d39d..e2f81f407 100644 --- a/examples/offline_inference/image_to_image/qwen_image_21_edit.py +++ b/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 diff --git a/recipes/Qwen/Qwen-Image-2.1.md b/recipes/Qwen/Qwen-Image-2.1.md index 4c7aa441b..33ab34743 100644 --- a/recipes/Qwen/Qwen-Image-2.1.md +++ b/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. diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py index f0037186c..771c41ec1 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py +++ b/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 diff --git a/vllm_omni/diffusion/models/qwen_image_21/__init__.py b/vllm_omni/diffusion/models/qwen_image_21/__init__.py index ec77cf534..6aa269033 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/__init__.py +++ b/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) 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 index 0feba0fc9..cbaf0aec8 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py +++ b/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. 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 index 8119da80b..c8dd46e29 100644 --- 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 @@ -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