# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from contextlib import nullcontext from typing import Any, ClassVar import torch import torch.distributed as dist from diffusers.models.autoencoders import AutoencoderKLWan from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify from diffusers.models.autoencoders.vae import DecoderOutput from diffusers.utils.accelerate_utils import apply_forward_hook from vllm.logger import init_logger from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( DistributedOperator, DistributedVaeMixin, GridSpec, TileTask, ) from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import decode_frames, is_installed from vllm_omni.diffusion.models.interface import DecodedChunkConsumer from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) class OmniAutoencoderKLWan(AutoencoderKLWan): def _execution_context(self): try: first_param = next(self.parameters()) except StopIteration: return nullcontext() dtype = first_param.dtype if dtype not in (torch.float16, torch.bfloat16): return nullcontext() return current_omni_platform.create_autocast_context( device_type=first_param.device.type, dtype=dtype, enabled=True, ) def encode(self, x: torch.Tensor, return_dict: bool = True): with self._execution_context(): return super().encode(x, return_dict=return_dict) def decode(self, z: torch.Tensor, return_dict: bool = True): """Decode a Wan latent using the Diffusers-compatible full-tensor API.""" with self._execution_context(): return super().decode(z, return_dict=return_dict) def _decode(self, z: torch.Tensor, return_dict: bool = True): """Decode one latent batch, preserving the parent Diffusers contract.""" _, _, _, 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 getattr(self, "use_tiling", False) and (width > tile_latent_min_width or height > tile_latent_min_height): return self.tiled_decode(z, return_dict=return_dict) return self._decode_temporal(z, return_dict=return_dict) # Chunks are published straight from the decoder, clamped to the # checkpoint's output range. chunk_value_range: ClassVar[tuple[float, float]] = (-1.0, 1.0) @apply_forward_hook def decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) -> None: """Decode ``z`` while synchronously delivering temporal chunks. The ordinary :meth:`decode` contract remains the Diffusers contract. Chunk callbacks deliberately reject a sliced batch: Diffusers slices each sample before decoding, while a streaming consumer needs a stable batch ownership contract. Callers can disable ``use_slicing`` or invoke this capability once per sample. Every rank participating in distributed VAE execution must invoke this method. Spatial-sharded execution invokes ``on_chunk`` only on the output-owning rank. Chunks have shape ``(B, C, T, H, W)`` and values clamped to ``[-1, 1]``. After a callback failure, remaining chunks are decoded and discarded before the exception is re-raised. """ if not callable(on_chunk): raise TypeError("on_chunk must be callable") if getattr(self, "use_slicing", False) and z.shape[0] > 1: raise ValueError( "Wan chunked decode does not support use_slicing for batch size > 1; " "disable VAE slicing or decode one sample at a time" ) with self._execution_context(): self._decode_with_chunks(z, on_chunk=on_chunk) def _decode_temporal( self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer | None = None, return_dict: bool = True, ): """Decode the non-tiled temporal path, optionally streaming chunks.""" if on_chunk is None and is_installed(self) and not torch.is_grad_enabled(): try: out = decode_frames(self, z) finally: self.clear_cache() if not return_dict: return (out,) return DecoderOutput(sample=out) self.clear_cache() callback_error: BaseException | None = None try: x = self.post_quant_conv(z) decoded_chunks: list[torch.Tensor] = [] for i in range(z.shape[2]): self._conv_idx = [0] if i == 0: chunk = self.decoder( x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=True, ) else: chunk = self.decoder( x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, ) if on_chunk is None: decoded_chunks.append(chunk) continue if callback_error is None: try: if self.config.patch_size is not None: chunk = unpatchify(chunk, patch_size=self.config.patch_size) on_chunk(torch.clamp(chunk, min=-1.0, max=1.0)) except BaseException as exc: # A previously installed spatial-shard decoder still # runs collectives, even for latents below the tiling # threshold. Drain before surfacing consumer failures. callback_error = exc del chunk if callback_error is not None: raise callback_error if on_chunk is not None: return None out = torch.cat(decoded_chunks, dim=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) if not return_dict: return (out,) return DecoderOutput(sample=out) finally: self.clear_cache() @staticmethod def _discard_decoded_chunk(_chunk: torch.Tensor) -> None: """Consume a non-owner rank's chunk without retaining its tensor.""" def _decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) -> None: """Run the temporal streaming primitive without API or hook handling.""" distributed_enabled = False if hasattr(self, "is_distributed_enabled"): distributed_enabled = self.is_distributed_enabled() if getattr(self, "use_tiling", False): 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 exceeds_tile = z.shape[-1] > tile_latent_min_width or z.shape[-2] > tile_latent_min_height if exceeds_tile: if distributed_enabled: self._tiled_decode_with_chunks(z, on_chunk=on_chunk) return raise ValueError("Wan chunk callbacks are unsupported for spatial tiling") callback = on_chunk if distributed_enabled and dist.get_rank(group=self.distributed_executor.group) != 0: callback = self._discard_decoded_chunk self._decode_temporal(z, on_chunk=callback) class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, 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]: _, _, 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 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): self._conv_idx = [0] 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 self._execution_context(): 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 encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]: _, _, num_frames, height, width = x.shape encode_spatial_compression_ratio = self.spatial_compression_ratio # Scale tile parameters for patchified coordinate system tile_sample_min_height = self.tile_sample_min_height tile_sample_min_width = self.tile_sample_min_width 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: assert encode_spatial_compression_ratio % self.config.patch_size == 0 encode_spatial_compression_ratio = self.spatial_compression_ratio // self.config.patch_size # When input is patchified, scale tile parameters accordingly tile_sample_min_height = tile_sample_min_height // self.config.patch_size tile_sample_min_width = tile_sample_min_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 latent_height = height // encode_spatial_compression_ratio latent_width = width // encode_spatial_compression_ratio tile_latent_min_height = tile_sample_min_height // encode_spatial_compression_ratio tile_latent_min_width = tile_sample_min_width // encode_spatial_compression_ratio tile_latent_stride_height = tile_sample_stride_height // encode_spatial_compression_ratio tile_latent_stride_width = 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 tiletask_list: list[TileTask] = [] temporal_compression = self.config.scale_factor_temporal for i in range(0, height, tile_sample_stride_height): for j in range(0, width, tile_sample_stride_width): time_list = [] frame_range = 1 + (num_frames - 1) // temporal_compression for k in range(frame_range): if k == 0: tile = x[:, :, :1, i : i + tile_sample_min_height, j : j + tile_sample_min_width] else: tile = x[ :, :, 1 + temporal_compression * (k - 1) : 1 + temporal_compression * k, i : i + tile_sample_min_height, j : j + tile_sample_min_width, ] time_list.append(tile) tiletask_list.append( TileTask( len(tiletask_list), (i // tile_sample_stride_height, j // tile_sample_stride_width), time_list, workload=time_list[0].shape[3] * time_list[0].shape[4], ) ) 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={ "latent_height": latent_height, "latent_width": latent_width, "blend_height": blend_height, "blend_width": blend_width, "tile_latent_stride_height": tile_latent_stride_height, "tile_latent_stride_width": tile_latent_stride_width, }, output_dtype=self.dtype, ) return tiletask_list, grid_spec def encode_tile_exec(self, task: TileTask) -> torch.Tensor: """Encode a single sample tile into latent space.""" self.clear_cache() time = [] for k, tile in enumerate(task.tensor): self._enc_conv_idx = [0] encoded = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) encoded = self.quant_conv(encoded) time.append(encoded) result = torch.cat(time, dim=2) self.clear_cache() return result def encode_tile_merge( self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec ) -> torch.Tensor: """Merge encoded tiles into a full latent tensor.""" grid_h, grid_w = grid_spec.grid_shape 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_latent_stride_height"], : grid_spec.tile_spec["tile_latent_stride_width"], ] ) result_rows.append(torch.cat(result_row, dim=-1)) enc = torch.cat(result_rows, dim=3)[ :, :, :, : grid_spec.tile_spec["latent_height"], : grid_spec.tile_spec["latent_width"] ] return enc 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 result_rows = [] self.clear_cache() 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 _spatial_shard_decode_split_dim(self) -> str | None: mode = self.distributed_executor.parallel_mode if mode == "spatial_shard_width": return "width" if mode == "spatial_shard_height": return "height" return None def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool: split_dim = self._spatial_shard_decode_split_dim() if split_dim is None: return False if z.ndim != 5: logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.") return False if not self.is_distributed_enabled(): return False group = self.distributed_executor.group world_size = dist.get_world_size(group=group) requested_size = int(self.distributed_executor.parallel_size) if requested_size != world_size: logger.warning( "Wan VAE spatial sharded decode currently requires vae_patch_parallel_size " "to match the DIT group size; falling back to tiled decode. " "requested=%s dit_group=%s split_dim=%s", requested_size, world_size, split_dim, ) return False return True 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) spatial_shard_enabled = self._spatial_shard_decode_enabled(z) if spatial_shard_enabled: split_dim = self._spatial_shard_decode_split_dim() assert split_dim is not None logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim) return wan_spatial_shard.spatial_shard_decode( self, z, group=self.distributed_executor.group, return_dict=return_dict, split_dim=split_dim, ) 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=False, ) if not return_dict: return (result,) return DecoderOutput(sample=result) def _tiled_decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) -> None: """Decode tiled latents and deliver temporal chunks to ``on_chunk``.""" if not self.is_distributed_enabled(): raise ValueError("Wan chunk callbacks are unsupported for spatial tiling") spatial_shard_enabled = self._spatial_shard_decode_enabled(z) if not spatial_shard_enabled: raise ValueError("Wan chunk callbacks require spatial-shard decode when tiling is enabled") split_dim = self._spatial_shard_decode_split_dim() assert split_dim is not None logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim) wan_spatial_shard.spatial_shard_decode( self, z, group=self.distributed_executor.group, return_dict=True, split_dim=split_dim, on_chunk=on_chunk, ) def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: """ Encode using distributed VAE executor. Note: x is already patchified by parent's _encode() before calling this method. """ if not self.is_distributed_enabled(): return super().tiled_encode(x) logger.debug("Encode running with distributed executor") self.clear_cache() result = self.distributed_executor.execute( x, DistributedOperator( split=self.encode_tile_split, exec=self.encode_tile_exec, merge=self.encode_tile_merge, ), broadcast_result=True, ) self.clear_cache() return result