# 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)