diff --git a/tests/diffusion/distributed/test_wan_chunked_decode.py b/tests/diffusion/distributed/test_wan_chunked_decode.py new file mode 100644 index 000000000..87f3ddff7 --- /dev/null +++ b/tests/diffusion/distributed/test_wan_chunked_decode.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import weakref +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from diffusers.models.autoencoders import AutoencoderKLWan + +from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( + DistributedAutoencoderKLWan, + OmniAutoencoderKLWan, +) +from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.fixture +def vae(): + # Real causal convolutions and caches, with small random weights for CPU CI. + return OmniAutoencoderKLWan( + base_dim=4, + z_dim=2, + dim_mult=[1, 2], + num_res_blocks=1, + temperal_downsample=[True], + scale_factor_spatial=2, + scale_factor_temporal=2, + ).eval() + + +@pytest.mark.parametrize("patch_size", [None, 2]) +@pytest.mark.parametrize("return_dict", [False, True]) +@torch.no_grad() +def test_chunked_decode_matches_diffusers(vae, patch_size, return_dict): + if patch_size is not None: + vae = OmniAutoencoderKLWan.from_config(vae.config, patch_size=patch_size, out_channels=12).eval() + z = torch.randn(1, 2, 3, 2, 2) + expected = AutoencoderKLWan._decode(vae, z).sample + actual = vae.decode(z, return_dict=return_dict) + torch.testing.assert_close(actual.sample if return_dict else actual[0], expected, rtol=0, atol=0) + chunks: list[torch.Tensor] = [] + assert supports_chunked_vae_decode(vae) + assert not supports_chunked_vae_decode(object()) + assert vae.decode_with_chunks(z, on_chunk=chunks.append) is None + assert [chunk.shape[2] for chunk in chunks] == [1, 2, 2] + torch.testing.assert_close(torch.cat(chunks, dim=2), expected, rtol=0, atol=0) + assert all(entry is None for entry in vae._feat_map) + + +@torch.no_grad() +def test_slicing_and_forward_hooks(vae): + z = torch.randn(2, 2, 2, 2, 2) + vae.enable_slicing() + batches = [] + handle = vae.post_quant_conv.register_forward_pre_hook(lambda module, args: batches.append(args[0].shape[0])) + vae._hf_hook = SimpleNamespace(pre_forward=Mock()) + try: + assert vae.decode(z).sample.shape[0] == 2 + assert batches == [1, 1] + vae._hf_hook.pre_forward.assert_called_once_with(vae) + batches.clear() + with pytest.raises(ValueError, match="use_slicing"): + vae.decode_with_chunks(z, on_chunk=Mock()) + assert batches == [] + vae._hf_hook.pre_forward.reset_mock() + vae.decode_with_chunks(z[:1], on_chunk=Mock()) + assert batches == [1] + vae._hf_hook.pre_forward.assert_called_once_with(vae) + finally: + handle.remove() + + +@torch.no_grad() +def test_callback_failure_drains_and_clears_cache(vae): + z = torch.randn(1, 2, 4, 2, 2) + decoded = Mock(return_value=None) + handle = vae.decoder.register_forward_hook(lambda *args: decoded()) + error = RuntimeError("consumer failed") + callback = Mock(side_effect=error) + try: + with pytest.raises(RuntimeError, match="consumer failed") as exc: + vae.decode_with_chunks(z, on_chunk=callback) + assert exc.value is error + assert callback.call_count == 1 + assert decoded.call_count == z.shape[2] + assert all(entry is None for entry in vae._feat_map) + finally: + handle.remove() + # The same instance can decode again after the failure. + chunks: list[torch.Tensor] = [] + vae.decode_with_chunks(z, on_chunk=chunks.append) + torch.testing.assert_close(torch.cat(chunks, 2), vae.decode(z).sample, rtol=0, atol=0) + + +def test_chunked_decode_rejects_spatial_tiling(vae): + vae.enable_tiling(tile_sample_min_height=2, tile_sample_min_width=2) + with pytest.raises(ValueError, match="spatial tiling"): + vae.decode_with_chunks(torch.zeros(1, 2, 2, 2, 2), on_chunk=Mock()) + + +@pytest.mark.parametrize("rank", [0, 1]) +@pytest.mark.parametrize("fail_callback", [False, True]) +def test_spatial_callback_ownership_and_bounded_retention(monkeypatch, rank, fail_callback): + # Isolate the temporal loop from halo kernels; weakrefs detect retention of + # decoded GPU-sized tensors even after the consumer has failed. + monkeypatch.setattr(wan_spatial_shard, "install_wan_spatial_shard_decode", lambda *args, **kwargs: None) + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (rank, 2)) + references: list[weakref.ReferenceType[torch.Tensor]] = [] + + def decode(x, **kwargs): + assert sum(ref() is not None for ref in references) <= 1 + chunk = x.clone() + references.append(weakref.ref(chunk)) + return chunk + + vae = SimpleNamespace( + config=SimpleNamespace(patch_size=None), + clear_cache=Mock(), + post_quant_conv=lambda z: z, + decoder=decode, + _feat_map=[], + ) + callback = Mock(side_effect=RuntimeError("consumer failed") if fail_callback else None) + z = torch.zeros(1, 3, 6, 2, 2) + if rank == 0 and fail_callback: + with pytest.raises(RuntimeError, match="consumer failed"): + wan_spatial_shard.spatial_shard_decode(vae, z, group=None, on_chunk=callback) + else: + assert wan_spatial_shard.spatial_shard_decode(vae, z, group=None, on_chunk=callback) is None + assert len(references) == z.shape[2] + assert callback.call_count == (0 if rank else 1 if fail_callback else z.shape[2]) + assert vae.clear_cache.call_count == 2 + + +@pytest.mark.parametrize("rank", [0, 1]) +@torch.no_grad() +def test_distributed_small_latent_callback_ownership(vae, monkeypatch, rank): + vae.is_distributed_enabled = lambda: True + vae.distributed_executor = SimpleNamespace(group=None) + monkeypatch.setattr(torch.distributed, "get_rank", lambda **kwargs: rank) + callback = Mock() + vae.decode_with_chunks(torch.zeros(1, 2, 2, 2, 2), on_chunk=callback) + assert callback.call_count == (2 if rank == 0 else 0) + + +@pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width", "tile"]) +def test_distributed_chunked_routing(monkeypatch, mode): + vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan) + torch.nn.Module.__init__(vae) + vae.use_tiling = True + vae.use_slicing = False + vae.tile_sample_min_height = vae.tile_sample_min_width = 2 + vae.spatial_compression_ratio = 2 + vae.distributed_executor = SimpleNamespace(group=None, parallel_mode=mode, parallel_size=2) + monkeypatch.setattr(vae, "is_distributed_enabled", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda **kwargs: 2) + decode = Mock() + monkeypatch.setattr(wan_spatial_shard, "spatial_shard_decode", decode) + callback = Mock() + z = torch.zeros(1, 2, 2, 2, 2) + if mode == "tile": + with pytest.raises(ValueError, match="require spatial-shard"): + vae.decode_with_chunks(z, on_chunk=callback) + decode.assert_not_called() + else: + assert vae.decode_with_chunks(z, on_chunk=callback) is None + decode.assert_called_once_with( + vae, z, group=None, return_dict=True, split_dim=mode.removeprefix("spatial_shard_"), on_chunk=callback + ) diff --git a/tests/diffusion/utils/test_chunked_mp4.py b/tests/diffusion/utils/test_chunked_mp4.py new file mode 100644 index 000000000..e7ef465f3 --- /dev/null +++ b/tests/diffusion/utils/test_chunked_mp4.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import hashlib + +import av +import numpy as np +import pytest + +from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder, mux_av_video_audio_bytes + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _frames() -> np.ndarray: + rng = np.random.default_rng(6872) + return rng.integers(0, 256, size=(8, 16, 24, 3), dtype=np.uint8) + + +def test_chunked_mp4_is_byte_identical_to_whole_mux() -> None: + frames = _frames() + baseline = mux_av_video_audio_bytes( + (av.VideoFrame.from_ndarray(frame, format="rgb24") for frame in frames), + width=24, + height=16, + fps=24, + video_codec_options={"threads": "1"}, + ) + encoder = ChunkedMP4Encoder( + width=24, + height=16, + fps=24, + max_pending=2, + video_codec_options={"threads": "1"}, + ) + for start in (0, 1, 5): + stop = {0: 1, 1: 5, 5: len(frames)}[start] + encoder.push(np.ascontiguousarray(frames[start:stop])) + chunked = encoder.finish() + assert chunked == baseline + assert hashlib.sha256(chunked).digest() == hashlib.sha256(baseline).digest() + + audio = np.zeros((2, 320), dtype=np.float32) + baseline_audio = mux_av_video_audio_bytes( + (av.VideoFrame.from_ndarray(frame, format="rgb24") for frame in frames), + width=24, + height=16, + fps=24, + audio_waveform=audio, + audio_sample_rate=32000, + video_codec_options={"threads": "1"}, + ) + encoder_audio = ChunkedMP4Encoder( + width=24, + height=16, + fps=24, + audio_waveform=audio, + audio_sample_rate=32000, + video_codec_options={"threads": "1"}, + ) + encoder_audio.push(frames) + assert encoder_audio.finish() == baseline_audio + + +def test_chunked_mp4_abort_joins_worker() -> None: + encoder = ChunkedMP4Encoder(width=24, height=16, fps=24) + encoder.push(_frames()[:1]) + encoder.abort() + assert not encoder._thread.is_alive() + with pytest.raises(RuntimeError, match="already closed"): + encoder.push(_frames()[:1]) + + +def test_chunked_mp4_validates_shape_and_dtype() -> None: + encoder = ChunkedMP4Encoder(width=24, height=16, fps=24) + with pytest.raises(ValueError, match="shape"): + encoder.push(np.zeros((16, 24, 3), dtype=np.uint8)) + with pytest.raises(ValueError, match="dtype"): + encoder.push(np.zeros((1, 16, 24, 3), dtype=np.float32)) + encoder.abort() + + +def test_chunked_mp4_close_returns_encoded_bytes() -> None: + import io + + encoder = ChunkedMP4Encoder(width=24, height=16, fps=24) + encoder.push(_frames()) + result = encoder.close() + assert result == encoder.finish() == encoder.close() + with av.open(io.BytesIO(result)) as container: + assert len(list(container.decode(video=0))) == len(_frames()) + + +def test_chunked_mp4_failure_during_push(monkeypatch) -> None: + from threading import Event + + from vllm_omni.diffusion.utils import media_utils + + fail = Event() + draining = Event() + error = RuntimeError("encoding failed") + + def mux(*args, **kwargs): + assert fail.wait(5) + raise error + + monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", mux) + drain = ChunkedMP4Encoder._drain_until_done + + def notified_drain(self): + draining.set() + drain(self) + + monkeypatch.setattr(ChunkedMP4Encoder, "_drain_until_done", notified_drain) + encoder = ChunkedMP4Encoder(width=24, height=16, fps=24) + put = encoder._queue.put + + def fail_before_put(item): + fail.set() + assert draining.wait(5) + put(item) + + monkeypatch.setattr(encoder._queue, "put", fail_before_put) + try: + with pytest.raises(RuntimeError, match="encoding failed") as exc: + encoder.push(_frames()) + assert exc.value is error + with pytest.raises(RuntimeError, match="encoding failed"): + encoder.finish() + finally: + fail.set() + encoder.abort() + assert not encoder._thread.is_alive() + + +def test_chunked_mp4_flush_failure_does_not_hang(monkeypatch) -> None: + from threading import Thread + + from vllm_omni.diffusion.utils import media_utils + + error = RuntimeError("flush failed") + + def mux(frames, **kwargs): + list(frames) + raise error + + monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", mux) + encoder = ChunkedMP4Encoder(width=24, height=16, fps=24) + encoder.push(_frames()) + errors = [] + + def finish(): + try: + encoder.finish() + except RuntimeError as exc: + errors.append(exc) + + closer = Thread(target=finish, daemon=True) + closer.start() + closer.join(5) + try: + assert not closer.is_alive(), "finish hung after the worker consumed its sentinel" + assert errors == [error] + finally: + if closer.is_alive(): + encoder._send_done() + closer.join(5) + + +def test_chunked_mp4_abort_discards_pending_chunks(monkeypatch) -> None: + from threading import Event, Thread + + from vllm_omni.diffusion.utils import media_utils + + entered = Event() + release = Event() + done_sent = Event() + encoded = [] + + def mux(frames, **kwargs): + encoded.append(next(frames)) + entered.set() + assert release.wait(5) + encoded.extend(frames) + return b"unused" + + monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", mux) + encoder = ChunkedMP4Encoder(width=24, height=16, fps=24, max_pending=1) + send_done = encoder._send_done + + def notified_done(): + send_done() + done_sent.set() + + monkeypatch.setattr(encoder, "_send_done", notified_done) + encoder.push(_frames()[:1]) + assert entered.wait(5) + encoder.push(_frames()[1:]) + closer = Thread(target=encoder.abort, daemon=True) + closer.start() + try: + assert done_sent.wait(5), "abort waited for encoding before discarding the pending queue" + finally: + release.set() + closer.join(5) + assert not closer.is_alive() + assert len(encoded) == 1 + with pytest.raises(RuntimeError, match="aborted"): + encoder.finish() diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index 44c6efe66..eddcb1ad6 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from contextlib import nullcontext from typing import Any @@ -9,6 +9,7 @@ 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 @@ -18,6 +19,7 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor GridSpec, TileTask, ) +from vllm_omni.diffusion.models.interface import DecodedChunkConsumer from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -45,9 +47,133 @@ class OmniAutoencoderKLWan(AutoencoderKLWan): 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) + + @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.""" + 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 @@ -78,7 +204,7 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): blend_height = self.tile_sample_min_height - tile_sample_stride_height blend_width = self.tile_sample_min_width - tile_sample_stride_width - tiletask_list = [] + 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 = [] @@ -151,7 +277,7 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): blend_height = tile_latent_min_height - tile_latent_stride_height blend_width = tile_latent_min_width - tile_latent_stride_width - tiletask_list = [] + 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): @@ -307,8 +433,10 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): if not self.is_distributed_enabled(): return super().tiled_decode(z, return_dict=return_dict) - if self._spatial_shard_decode_enabled(z): + 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, @@ -329,6 +457,27 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): 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. diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py index 07edad694..ac5ea472b 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py +++ b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # # The halo-exchange spatial-parallel decode here is adapted from SGLang's # spatial-parallel VAE decode @@ -24,7 +24,7 @@ from contextlib import nullcontext from contextvars import ContextVar from dataclasses import dataclass from types import MethodType -from typing import Any +from typing import Any, cast import torch import torch.distributed as dist @@ -34,6 +34,8 @@ from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify from diffusers.models.autoencoders.vae import DecoderOutput from vllm.logger import init_logger +from vllm_omni.diffusion.models.interface import DecodedChunkConsumer + logger = init_logger(__name__) @@ -644,7 +646,7 @@ def _replace_child( ) return if isinstance(child, nn.ZeroPad2d): - padding = tuple(int(p) for p in child.padding) + padding = cast(tuple[int, int, int, int], tuple(int(p) for p in child.padding)) module_padding = padding if parent.__class__.__name__ == "Sequential": # Let the following WanDistConv2d account for global after-edge @@ -794,7 +796,8 @@ def spatial_shard_decode( group: dist.ProcessGroup, return_dict: bool = True, split_dim: str = "height", -) -> DecoderOutput | tuple[torch.Tensor]: + on_chunk: DecodedChunkConsumer | None = None, +) -> DecoderOutput | tuple[torch.Tensor] | None: install_wan_spatial_shard_decode(vae, group, split_dim=split_dim) if z.shape[2] == 0: @@ -806,6 +809,7 @@ def spatial_shard_decode( produce_output = world_size <= 1 or rank == 0 vae.clear_cache() + callback_error: BaseException | None = None try: context = vae._execution_context() if hasattr(vae, "_execution_context") else nullcontext() with context: @@ -820,18 +824,33 @@ def spatial_shard_decode( first_chunk=(i == 0), ) if produce_output: - decoded_chunks.append(chunk) + if on_chunk is not None and callback_error is None: + try: + if vae.config.patch_size is not None: + chunk = unpatchify(chunk, patch_size=vae.config.patch_size) + on_chunk(torch.clamp(chunk, min=-1.0, max=1.0)) + except BaseException as exc: + # Keep all ranks in the temporal collective loop; + # surface the callback failure only after decode. + callback_error = exc + elif on_chunk is None: + decoded_chunks.append(chunk) + del chunk - if produce_output: + if produce_output and on_chunk is None: out = torch.cat(decoded_chunks, dim=2) if vae.config.patch_size is not None: out = unpatchify(out, patch_size=vae.config.patch_size) out = torch.clamp(out, min=-1.0, max=1.0) - else: - out = z.new_zeros(0) finally: vae.clear_cache() + if callback_error is not None: + raise callback_error + if on_chunk is not None: + return None + if not produce_output: + out = z.new_zeros(0) if not return_dict: return (out,) return DecoderOutput(sample=out) diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index 75cc5a0d2..dac9a0b1d 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass from typing import ( TYPE_CHECKING, @@ -9,6 +10,7 @@ from typing import ( ClassVar, Literal, Protocol, + TypeGuard, runtime_checkable, ) @@ -46,6 +48,34 @@ class SupportAudioOutput(Protocol): support_audio_output: ClassVar[bool] = True +DecodedChunkConsumer = Callable[["torch.Tensor"], None] + + +@runtime_checkable +class SupportsChunkedVAEDecode(Protocol): + """Optional capability for VAEs that stream decoded temporal chunks. + + Implementations must drain decoding before surfacing callback errors. In a + distributed VAE, every rank invokes the method so collectives stay in + lockstep, while ``on_chunk`` is called only on the rank that owns output. + """ + + def decode_with_chunks( + self, + z: torch.Tensor, + *, + on_chunk: DecodedChunkConsumer, + ) -> None: + """Decode ``z`` and synchronously deliver committed chunks in order.""" + ... + + +def supports_chunked_vae_decode(vae: object) -> TypeGuard[SupportsChunkedVAEDecode]: + """Return whether ``vae`` exposes the optional chunked decode capability.""" + + return isinstance(vae, SupportsChunkedVAEDecode) + + @runtime_checkable class SupportsStepExecution(Protocol): """State-driven step-level execution protocol for diffusion pipelines. diff --git a/vllm_omni/diffusion/utils/media_utils.py b/vllm_omni/diffusion/utils/media_utils.py index df67f9d99..615df090f 100644 --- a/vllm_omni/diffusion/utils/media_utils.py +++ b/vllm_omni/diffusion/utils/media_utils.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Video/audio muxing utilities using PyAV (no ffmpeg binary dependency).""" from __future__ import annotations import io +import queue +import threading from collections.abc import Iterable from fractions import Fraction from typing import Any, cast @@ -12,6 +14,153 @@ from typing import Any, cast import av import numpy as np +_CHUNKED_MP4_DONE = object() + + +def _validate_video_chunk(chunk: np.ndarray, *, width: int, height: int) -> None: + """Validate a ``(T, H, W, 3)`` uint8 RGB chunk against a session's frame size.""" + if chunk.ndim != 4 or chunk.shape[-1] != 3: + raise ValueError("video chunk must have shape (T, H, W, 3)") + if chunk.dtype != np.uint8: + raise ValueError("video chunk must have dtype uint8") + if chunk.shape[1] != height or chunk.shape[2] != width: + raise ValueError("video chunks in a session must use a consistent frame size") + + +class ChunkedMP4Encoder: + """Encode temporal video chunks while the producer is still decoding. + + A bounded queue and one muxing worker provide ordered backpressure while + keeping host memory bounded by ``max_pending`` chunks. Chunks use the same + ``(T, H, W, 3)`` uint8 contract as :func:`mux_video_audio_bytes`. + """ + + def __init__( + self, + *, + width: int, + height: int, + fps: float, + audio_waveform: np.ndarray | None = None, + audio_sample_rate: int | None = None, + max_pending: int = 2, + video_codec: str = "h264", + audio_codec: str = "aac", + crf: str = "18", + video_codec_options: dict[str, str] | None = None, + ) -> None: + if max_pending <= 0: + raise ValueError("max_pending must be positive") + if width <= 0 or height <= 0: + raise ValueError("width and height must be positive") + self.width = int(width) + self.height = int(height) + self._queue: queue.Queue[object] = queue.Queue(maxsize=max_pending) + self._result: bytes | None = None + self._error: BaseException | None = None + self._closed = False + self._aborted = False + self._input_done = False + self._state_lock = threading.Lock() + + def run() -> None: + try: + self._result = mux_av_video_audio_bytes( + self._frames(), + width=self.width, + height=self.height, + audio_waveform=audio_waveform, + fps=fps, + audio_sample_rate=audio_sample_rate, + video_codec=video_codec, + audio_codec=audio_codec, + crf=crf, + video_codec_options=video_codec_options, + ) + except BaseException as exc: + self._error = exc + self._drain_until_done() + + self._thread = threading.Thread(target=run, name="chunked-mp4", daemon=True) + self._thread.start() + + def _drain_until_done(self) -> None: + # A flush/mux failure can occur after _frames consumed the sentinel. + if not self._input_done: + while self._queue.get() is not _CHUNKED_MP4_DONE: + pass + + def _frames(self): + while True: + chunk = self._queue.get() + if chunk is _CHUNKED_MP4_DONE: + self._input_done = True + return + assert isinstance(chunk, np.ndarray) + for frame_data in chunk: + yield av.VideoFrame.from_ndarray(frame_data, format="rgb24") + + def _validate_chunk(self, chunk: np.ndarray) -> None: + _validate_video_chunk(chunk, width=self.width, height=self.height) + + def _raise_if_failed(self) -> None: + if self._error is not None: + raise self._error + if self._closed: + raise RuntimeError("ChunkedMP4Encoder is already closed") + + def push(self, chunk: np.ndarray) -> None: + """Queue a chunk; asynchronous failures surface here or at finish().""" + self._validate_chunk(chunk) + self._raise_if_failed() + self._queue.put(chunk) + self._raise_if_failed() + + def _send_done(self) -> None: + self._queue.put(_CHUNKED_MP4_DONE) + + def finish(self) -> bytes: + """Flush the muxer and return complete progressive MP4 bytes.""" + with self._state_lock: + if not self._closed: + self._closed = True + self._send_done() + self._thread.join() + if self._error is not None: + raise self._error + if self._aborted: + raise RuntimeError("ChunkedMP4Encoder was aborted") + assert self._result is not None + return self._result + + def abort(self) -> None: + """Discard pending chunks and join the worker after cancellation. + + Any chunk already being encoded must finish before the worker exits. + """ + with self._state_lock: + if not self._closed: + self._closed = True + self._aborted = True + while True: + try: + self._queue.get_nowait() + except queue.Empty: + break + self._send_done() + self._thread.join() + + close = finish + + def __enter__(self) -> ChunkedMP4Encoder: + return self + + def __exit__(self, exc_type, exc, tb) -> None: + if exc_type is None: + self.finish() + else: + self.abort() + class FragmentedMP4Muxer: """Incrementally mux video frames into one fragmented MP4 byte stream.""" @@ -57,12 +206,7 @@ class FragmentedMP4Muxer: """Mux a batch of ``uint8`` RGB frames and return newly written MP4 bytes.""" if self._closed: raise RuntimeError("Cannot mux frames after FragmentedMP4Muxer.close().") - if video_frames.ndim != 4 or video_frames.shape[-1] != 3: - raise ValueError("video_frames must have shape (T, H, W, 3).") - if video_frames.dtype != np.uint8: - raise ValueError("video_frames must be uint8.") - if video_frames.shape[1] != self._stream.height or video_frames.shape[2] != self._stream.width: - raise ValueError("All fragmented MP4 chunks in a session must use the same frame size.") + _validate_video_chunk(video_frames, width=self._stream.width, height=self._stream.height) for frame_data in video_frames: frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24")