[1/N] Stream Wan VAE chunks to the media consumer

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

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