diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py new file mode 100644 index 000000000..fa7410b46 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py @@ -0,0 +1,357 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import gc +import tempfile +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist + +from vllm_omni.diffusion.models.minimax_h3.chunked_decode import decode_h3_chunks +from vllm_omni.diffusion.models.minimax_h3.temporal_chunks import decode_temporal_chunks + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.fixture(scope="module") +def cpu_process_group(): + if dist.is_initialized(): + yield dist.group.WORLD + return + + with tempfile.NamedTemporaryFile(prefix="h3_chunk_dist_") as rendezvous: + init_method = f"file://{rendezvous.name}" + dist.init_process_group("gloo", rank=0, world_size=1, init_method=init_method) + try: + yield dist.group.WORLD + finally: + dist.destroy_process_group() + gc.collect() + + +class _FakeTemporalModel: + use_3d_conv = True + token_drop = 3 + tokens_chunk_size = 5 + token_overlap = 2 + vae_ratio_t = 4 + frame_pre_padding = 3 + frame_overlap = 5 + isolated_first_frame = False + isolated_last_frame = False + + def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens): + del z, z_head, z_tail, num_chunks, pad_tokens + return 35, 0, 35 + + def _adaptive_decode(self, clip): + value = float(clip[:, :, 0].mean()) + return torch.full((1, 1, 24, 2, 2), value) + + @staticmethod + def blend(overlap, part, frame_overlap, dim): + del overlap, frame_overlap, dim + return part + + +class _FakeHost: + def __init__(self): + self.model = _FakeTemporalModel() + + @staticmethod + def _denormalize_latent(latent): + return latent + + @staticmethod + def _normalize_decoded_frames(frames): + return frames.float() + + +def test_temporal_chunks_emit_ordered_frames_and_collect_when_unconsumed(): + model = _FakeTemporalModel() + latent = torch.arange(8, dtype=torch.float32).view(1, 1, 8, 1, 1) + chunks = [] + marker = decode_temporal_chunks(model, latent, chunks.append) + + assert marker.shape == (0,) + assert [chunk.shape[2] for chunk in chunks] == [17, 17, 1] + assert torch.equal(torch.cat(chunks, dim=2), decode_temporal_chunks(model, latent, None)) + + +def test_h3_callback_failure_is_deferred_until_temporal_decode_finishes(): + host = _FakeHost() + latent = torch.zeros(1, 1, 8, 1, 1) + seen = [] + + def fail_once(frames): + seen.append(frames.shape[2]) + raise RuntimeError("sink failed") + + with pytest.raises(RuntimeError, match="sink failed"): + decode_h3_chunks(host, latent, fail_once, group=None) + assert seen == [17] + + +def test_h3_without_callback_is_a_plain_full_decode_on_a_vae_group(cpu_process_group): + host = _FakeHost() + latent = torch.arange(8, dtype=torch.float32).view(1, 1, 8, 1, 1) + + full = decode_h3_chunks(host, latent, None, group=cpu_process_group) + + chunks = [] + decode_h3_chunks(host, latent, chunks.append, group=cpu_process_group) + assert torch.equal(torch.cat(chunks, dim=2), full) + + +class _SimulatedGroup: + """Stand in for a VAE group so one process can check every rank's verdict. + + ``all_reduce`` is replaced with the sum the real collective would produce + for a group of ``world_size`` ranks, ``supplied`` of which handed in a + callback. That sum is what makes a rank-local verdict observable from a + single process. + """ + + def __init__(self, world_size: int, supplied: int): + self.world_size = world_size + self.supplied = supplied + + +def _patch_simulated_ranks(monkeypatch, simulated: _SimulatedGroup, rank: int) -> None: + from vllm_omni.diffusion.models.minimax_h3 import chunked_decode as mod + + def fake_all_reduce(tensor, group=None): + del group # the simulated census below stands in for the real reduction + tensor.copy_(torch.tensor([simulated.supplied], dtype=tensor.dtype)) + + monkeypatch.setattr(mod.dist, "get_rank", lambda _group: rank) + monkeypatch.setattr(mod.dist, "get_world_size", lambda _group: simulated.world_size) + monkeypatch.setattr(mod.dist, "all_reduce", fake_all_reduce) + monkeypatch.setattr(mod.dist, "broadcast", lambda tensor, src=0, group=None: None) + monkeypatch.setattr(mod.dist, "get_global_rank", lambda _group, _rank: 0) + + +@pytest.mark.parametrize("rank", [0, 1, 2]) +def test_h3_every_rank_rejects_a_partially_supplied_callback(monkeypatch, rank): + """Ranks that disagree must all refuse, not let one decode alone. + + A rank-local verdict let a peer fall through into the decoder collectives + while its peers raised, hanging the stage instead of failing it. + """ + host = _FakeHost() + latent = torch.zeros(1, 1, 8, 1, 1) + # Two ranks of three handed in a callback. + group = _SimulatedGroup(world_size=3, supplied=2) + _patch_simulated_ranks(monkeypatch, group, rank) + + with pytest.raises(ValueError, match="every rank"): + decode_h3_chunks(host, latent, lambda _frames: None, group=object()) + + +@pytest.mark.parametrize("rank", [0, 1, 2]) +def test_h3_only_rank_zero_publishes_when_every_rank_supplies(monkeypatch, rank): + """Ownership is positional: all ranks stream, rank 0 publishes.""" + host = _FakeHost() + latent = torch.zeros(1, 1, 8, 1, 1) + group = _SimulatedGroup(world_size=3, supplied=3) + _patch_simulated_ranks(monkeypatch, group, rank) + + seen: list[int] = [] + decode_h3_chunks(host, latent, lambda f: seen.append(int(f.shape[2])), group=object()) + + # Peers run the loop to stay in the collectives but publish nothing. + assert (seen != []) is (rank == 0) + + +def _stub_video_vae(*, parallel_size: int, tile_count: int): + """A MiniMaxH3VideoVAE with only the pieces the decode entry points read.""" + from contextlib import contextmanager + + from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE + + vae = object.__new__(MiniMaxH3VideoVAE) + vae.parallel_size = parallel_size + vae.model = SimpleNamespace(_adaptive_decode=lambda clip: clip) + vae.entered_rank_local = False + + @contextmanager + def rank_local(): + vae.entered_rank_local = True + yield + + vae._decoder_tile_count = lambda latent: tile_count + vae._rank_local_tiling = rank_local + vae.is_distributed_enabled = lambda: False + return vae + + +def test_chunked_decode_falls_back_when_tiles_are_fewer_than_ranks(monkeypatch): + """Without the fallback, tileless ranks raise while their peers block.""" + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_mod + + vae = _stub_video_vae(parallel_size=4, tile_count=1) + monkeypatch.setattr(vae_mod, "decode_h3_chunks", lambda *args, **kwargs: torch.zeros(1)) + + vae.decode_with_chunks(torch.zeros(1, 1, 8, 1, 1), on_chunk=lambda _frames: None) + + assert vae.entered_rank_local, "chunked decode must share decode_latent's too-few-tiles fallback" + + +def test_chunked_decode_keeps_shared_tiling_when_every_rank_has_a_tile(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_mod + + vae = _stub_video_vae(parallel_size=4, tile_count=8) + monkeypatch.setattr(vae_mod, "decode_h3_chunks", lambda *args, **kwargs: torch.zeros(1)) + + vae.decode_with_chunks(torch.zeros(1, 1, 8, 1, 1), on_chunk=lambda _frames: None) + + assert not vae.entered_rank_local + + +def test_h3_vae_declares_the_chunked_decode_capability(): + """A consumer must be able to detect the capability and the pixel range.""" + from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode + from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE + + vae = object.__new__(MiniMaxH3VideoVAE) + + assert supports_chunked_vae_decode(vae) + # H3 reverts through the checkpoint's processor, which lands in [0, 1]; + # a Wan VAE publishes [-1, 1], so the range cannot be assumed. + assert vae.chunk_value_range == (0.0, 1.0) + + +def _native_decode_temporal_streaming(model, z, z_head, z_tail, num_chunks, pad_tokens): + """The released ``_decode_temporal_streaming`` loop, as the parity oracle. + + Transcribed from the checkpoint's ``AutoencoderKLLegacy`` so the fork can be + compared against the control flow it mirrors, including the isolated + head/tail extraction. + """ + total_frames, pad_frames, output_frames = model._decode_temporal_output_frame_plan( + z, z_head, z_tail, num_chunks, pad_tokens + ) + chunk_dec = model.tokens_chunk_size * model.vae_ratio_t + split_count = int(model.token_drop > 0) + 1 + dec = None + dec_overlap = None + write_pos = 0 + + def write_part(part): + nonlocal dec, write_pos + part_frames = int(part.shape[2]) + if part_frames <= 0: + return + if dec is None: + out_shape = list(part.shape) + out_shape[2] = output_frames + dec = torch.empty(out_shape, dtype=part.dtype, device=part.device) + remaining = int(dec.shape[2]) - write_pos + copy_frames = min(part_frames, max(0, remaining)) + if copy_frames > 0: + dec[:, :, write_pos : write_pos + copy_frames].copy_(part[:, :, :copy_frames]) + write_pos += copy_frames + + for i in range(num_chunks): + t_start = i * model.tokens_chunk_size + clip_z = z[:, :, t_start : t_start + model.tokens_chunk_size + model.token_overlap] + if i == 0 and z_head is not None: + clip_z = torch.cat([z_head, clip_z], dim=2) + if i == num_chunks - 1 and z_tail is not None: + clip_z = torch.cat([clip_z, z_tail], dim=2) + + clip_dec = model._adaptive_decode(clip_z) + + dec_tail = None + if i == 0 and z_head is not None: + write_part(clip_dec[:, :, model.vae_ratio_t - 1 : model.vae_ratio_t]) + clip_dec = clip_dec[:, :, model.vae_ratio_t :] + if i == num_chunks - 1 and z_tail is not None: + dec_tail = clip_dec[:, :, -1:] + clip_dec = clip_dec[:, :, : -model.vae_ratio_t] + + for j in range(split_count): + f_start = j * chunk_dec + f_end = min(f_start + chunk_dec, clip_dec.shape[2]) + chunk = clip_dec[:, :, f_start:f_end][:, :, model.frame_pre_padding :] + if j == 0: + if dec_overlap is not None: + chunk = model.blend(dec_overlap, chunk, model.frame_overlap, dim=-3) + dec_overlap = None + write_part(chunk) + else: + dec_overlap = chunk.contiguous() + + if i == num_chunks - 1: + if dec_overlap is not None: + write_part(dec_overlap) + dec_overlap = None + if dec_tail is not None: + write_part(dec_tail) + + return dec + + +class _IsolatedFrameModel(_FakeTemporalModel): + """Fake decoder with the isolated boundaries the released config can set.""" + + frame_pre_padding = 0 + + def __init__(self, *, first: bool, last: bool): + self.isolated_first_frame = first + self.isolated_last_frame = last + + def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens): + del pad_tokens + frames = num_chunks * self.tokens_chunk_size * self.vae_ratio_t + frames += int(z_head is not None) + int(z_tail is not None) + del z + return frames, 0, frames + + def _adaptive_decode(self, clip): + # Every latent token decodes to a distinct, ordered block, so a shifted + # or reordered emission is visible in the values themselves. + base = float(clip[:, :, 0].mean()) + frames = int(clip.shape[2]) * self.vae_ratio_t + return torch.arange(frames, dtype=torch.float32).view(1, 1, frames, 1, 1) + base + + +@pytest.mark.parametrize( + ("first", "last"), + [(True, False), (False, True), (True, True)], +) +def test_temporal_chunks_match_the_native_loop_at_isolated_boundaries(first, last): + """The fork must extract the isolated head/tail the way the native loop does.""" + model = _IsolatedFrameModel(first=first, last=last) + latent = torch.arange(12, dtype=torch.float32).view(1, 1, 12, 1, 1) + + chunks: list[torch.Tensor] = [] + decode_temporal_chunks(model, latent, chunks.append) + streamed = torch.cat(chunks, dim=2) + + # Feed the oracle exactly what the released ``decode_temporal`` would: + # split the isolated tokens off, then pad the remainder to a whole chunk. + isolated = int(first) + int(last) + pseudo_tokens = int(latent.shape[2]) - isolated + model.token_drop + remainder = pseudo_tokens % model.tokens_chunk_size + pad_tokens = (model.tokens_chunk_size - remainder) if remainder else 0 + num_chunks = (pseudo_tokens + pad_tokens) // model.tokens_chunk_size - int(model.token_drop > 0) + + body = latent + z_head = None + if first: + z_head = body[:, :, :1] + body = body[:, :, 1:] + z_tail = None + if last: + z_tail = body[:, :, -1:] + body = body[:, :, :-1] + if pad_tokens: + body = torch.cat([body, body[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2) + + expected = _native_decode_temporal_streaming(model, body, z_head, z_tail, num_chunks, pad_tokens) + + assert streamed.shape == expected.shape + assert torch.equal(streamed, expected) diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index eddcb1ad6..9ab28d09c 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from contextlib import nullcontext -from typing import Any +from typing import Any, ClassVar import torch import torch.distributed as dist @@ -62,6 +62,10 @@ class OmniAutoencoderKLWan(AutoencoderKLWan): 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. diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index dac9a0b1d..92216d19e 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -58,8 +58,16 @@ class SupportsChunkedVAEDecode(Protocol): 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. + + ``chunk_value_range`` is the closed interval the delivered floats occupy, + which differs per checkpoint: a Wan VAE emits ``(-1.0, 1.0)`` while a + MiniMax-H3 VAE reverts through its processor to ``(0.0, 1.0)``. A consumer + cannot quantize chunks to pixels without it, so it is part of the + capability rather than knowledge each consumer hard-codes per model. """ + chunk_value_range: ClassVar[tuple[float, float]] + def decode_with_chunks( self, z: torch.Tensor, diff --git a/vllm_omni/diffusion/models/minimax_h3/chunked_decode.py b/vllm_omni/diffusion/models/minimax_h3/chunked_decode.py new file mode 100644 index 000000000..08ce0e64b --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/chunked_decode.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Frames-only callback coordination for MiniMax-H3 VAE decoding.""" + +from __future__ import annotations + +from collections.abc import Callable + +import torch +import torch.distributed as dist + +from .temporal_chunks import decode_temporal_chunks + +MiniMaxH3VideoChunkCallback = Callable[[torch.Tensor], None] + + +def decode_h3_chunks( + host, + latent: torch.Tensor, + callback: MiniMaxH3VideoChunkCallback | None, + *, + group: dist.ProcessGroup | None, +) -> torch.Tensor: + """Run the local temporal loop, synchronizing callback errors across ranks.""" + streaming = callback is not None + owner = streaming + if group is not None: + rank = dist.get_rank(group) + world_size = dist.get_world_size(group) + # Ownership is positional: every rank supplies the callback and rank 0 + # publishes. Reduce only whether the ranks agree that this decode + # streams, so the verdict is identical everywhere -- a rank-local + # verdict would let a peer fall through and enter the decoder + # collectives alone, hanging the stage instead of failing it. + census = torch.tensor([int(streaming)], dtype=torch.int64, device=latent.device) + dist.all_reduce(census, group=group) + supplied = int(census.item()) + if supplied not in (0, world_size): + raise ValueError( + "MiniMax-H3 chunk callback must be supplied on every rank of the VAE " + f"group or on none of them; got {supplied} of {world_size}" + ) + # No rank supplied a callback: every rank runs a plain full decode. + streaming = supplied == world_size + owner = streaming and rank == 0 + + error: BaseException | None = None + + def publish(raw: torch.Tensor) -> None: + nonlocal error + if not owner or error is not None: + return + try: + processor = getattr(host.model, "processor", None) + decoded = raw if processor is None else processor.revert_tensor(raw) + frames = host._normalize_decoded_frames(decoded).contiguous() + assert callback is not None + callback(frames) + except BaseException as exc: # noqa: BLE001 + error = exc + + # Peer ranks run the temporal loop without materializing a second video. + sink = publish if owner else ((lambda _frames: None) if streaming else None) + result = decode_temporal_chunks( + host.model, + host._denormalize_latent(latent), + sink, + ) + + if group is not None: + failed = torch.tensor([int(rank == 0 and error is not None)], dtype=torch.int32, device=latent.device) + dist.broadcast(failed, src=dist.get_global_rank(group, 0), group=group) + if int(failed.item()): + if rank == 0: + assert error is not None + raise error + raise RuntimeError("MiniMax-H3 video chunk callback failed on rank zero") + elif error is not None: + raise error + + if not result.numel(): + return result + processor = getattr(host.model, "processor", None) + decoded = result if processor is None else processor.revert_tensor(result) + return host._normalize_decoded_frames(decoded) diff --git a/vllm_omni/diffusion/models/minimax_h3/temporal_chunks.py b/vllm_omni/diffusion/models/minimax_h3/temporal_chunks.py new file mode 100644 index 000000000..a5154525d --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/temporal_chunks.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Small compatibility loop for the released MiniMax-H3 temporal decoder.""" + +from __future__ import annotations + +from collections.abc import Callable + +import torch +import torch.nn as nn + + +def decode_temporal_chunks( + model: nn.Module, + latent: torch.Tensor, + callback: Callable[[torch.Tensor], None] | None, +) -> torch.Tensor: + """Decode H3 temporal clips and synchronously publish committed frames. + + The released VAE exposes the temporal clip primitives but materializes the + final tensor. This is the same loop with the write point replaced by a + frames-only callback. ``callback=None`` retains the complete-output path. + """ + if latent.ndim != 5 or not bool(getattr(model, "use_3d_conv", False)): + raise ValueError("MiniMax-H3 temporal chunk decode requires a rank-5 3D latent") + + token_drop = int(model.token_drop) + chunk_size = int(model.tokens_chunk_size) + overlap_tokens = int(model.token_overlap) + ratio_t = int(model.vae_ratio_t) + pre_padding = int(model.frame_pre_padding) + + isolated_first = bool(model.isolated_first_frame and pre_padding == 0) + isolated_last = bool(model.isolated_last_frame) + z_head = latent[:, :, :1] if isolated_first else None + z_tail = latent[:, :, -1:] if isolated_last else None + start = int(isolated_first) + stop = int(latent.shape[2]) - int(isolated_last) + latent = latent[:, :, start:stop] + + pseudo_tokens = int(latent.shape[2]) + token_drop + pad_tokens = (-pseudo_tokens) % chunk_size + if pad_tokens: + latent = torch.cat((latent, latent[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)), dim=2) + num_chunks = (pseudo_tokens + pad_tokens) // chunk_size - int(token_drop > 0) + if num_chunks <= 0: + raise ValueError("MiniMax-H3 temporal chunk plan is empty") + + total_frames, pad_frames, output_frames = model._decode_temporal_output_frame_plan( + latent, z_head, z_tail, num_chunks, pad_tokens + ) + output_frames = int(output_frames) + if output_frames <= 0: + raise ValueError("MiniMax-H3 temporal chunk plan has no output frames") + + collected: list[torch.Tensor] = [] + overlap: torch.Tensor | None = None + written = 0 + + def emit(part: torch.Tensor) -> None: + nonlocal written + frames = int(part.shape[2]) + if frames <= 0 or written >= output_frames: + return + frames = min(frames, output_frames - written) + part = part[:, :, :frames] + if callback is None: + collected.append(part) + else: + callback(part) + written += frames + + chunk_frames = chunk_size * ratio_t + split_count = int(token_drop > 0) + 1 + for index in range(num_chunks): + begin = index * chunk_size + end = begin + chunk_size + overlap_tokens + clip = latent[:, :, begin:end] + if index == 0 and z_head is not None: + clip = torch.cat((z_head, clip), dim=2) + if index == num_chunks - 1 and z_tail is not None: + clip = torch.cat((clip, z_tail), dim=2) + + decoded = model._adaptive_decode(clip) + + # The isolated boundary latents were prepended/appended to the clip, so + # the decoder returns their frames inside this clip's block. Publish + # them the way the released loop does and drop their temporal block + # before splitting, or every later frame shifts by one block. + decoded_tail = None + if index == 0 and z_head is not None: + emit(decoded[:, :, ratio_t - 1 : ratio_t]) + decoded = decoded[:, :, ratio_t:] + if index == num_chunks - 1 and z_tail is not None: + decoded_tail = decoded[:, :, -1:] + decoded = decoded[:, :, :-ratio_t] + + for split in range(split_count): + begin = split * chunk_frames + end = min(begin + chunk_frames, int(decoded.shape[2])) + part = decoded[:, :, begin:end] + part = part[:, :, pre_padding:] + if split == 0: + if overlap is not None: + part = model.blend(overlap, part, int(model.frame_overlap), dim=-3) + overlap = None + emit(part) + else: + overlap = part.contiguous() + + if index == num_chunks - 1: + if overlap is not None: + emit(overlap) + overlap = None + if decoded_tail is not None: + emit(decoded_tail) + + if written != output_frames: + raise RuntimeError(f"MiniMax-H3 temporal decode emitted {written}/{output_frames} frames") + if callback is not None: + return latent.new_empty((0,)) + return torch.cat(collected, dim=2) if collected else latent.new_empty((0,)) diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index e98659409..ca8359904 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -9,7 +9,7 @@ import json from collections.abc import Iterator from contextlib import AbstractContextManager, contextmanager, nullcontext from pathlib import Path -from typing import Any +from typing import Any, ClassVar import torch import torch.distributed as dist @@ -22,11 +22,13 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor DistributedVaeMixin, ) from vllm_omni.diffusion.distributed.parallel_state import get_world_group +from vllm_omni.diffusion.models.interface import DecodedChunkConsumer from vllm_omni.diffusion.offloader.module_residency import ( BoundedAllocatorCache, PinnedModuleStager, ) +from .chunked_decode import decode_h3_chunks from .ops import install_h3_vae_optimizations from .packed_tokens import minimax_h3_patchify_video_latent @@ -377,25 +379,16 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): ).float() return rows, shape - @torch.inference_mode() - def decode_latent(self, latent: torch.Tensor) -> torch.Tensor: - channels = int(self.config_dict["latent_channels"]) - mean = torch.tensor( - self.config_dict["latents_mean"], - device=latent.device, - dtype=latent.dtype, - ).view(1, channels, 1, 1, 1) - std = torch.tensor( - self.config_dict["latents_std"], - device=latent.device, - dtype=latent.dtype, - ).view(1, channels, 1, 1, 1) - # The checkpoint hands rank r the tiles ``range(r, num_tiles, sp_size)`` - # and then rejects an empty share inside the gather. A rank with no - # tiles raises and leaves the collective while the others block in it - # forever, so too few tiles hangs the whole stage rather than failing - # it. Tile count depends only on the latent shape, so every rank takes - # this branch together. + def _decode_tiling_context(self, latent: torch.Tensor) -> AbstractContextManager: + """Pick the tiling mode a decode of ``latent`` can safely use. + + The checkpoint hands rank r the tiles ``range(r, num_tiles, sp_size)`` + and then rejects an empty share inside the gather. A rank with no + tiles raises and leaves the collective while the others block in it + forever, so too few tiles hangs the whole stage rather than failing + it. Tile count depends only on the latent shape, so every rank takes + this branch together. + """ num_tiles = self._decoder_tile_count(latent) if self.parallel_size > 1 and num_tiles < self.parallel_size: logger.warning_once( @@ -405,19 +398,64 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): num_tiles, self.parallel_size, ) - tiling_context: AbstractContextManager = self._rank_local_tiling() - else: - tiling_context = nullcontext() + return self._rank_local_tiling() + return nullcontext() - with tiling_context: - decoded = self.model.decode_base(latent * std + mean) - frames = self.model.processor.revert_tensor(decoded) + @torch.inference_mode() + def decode_latent(self, latent: torch.Tensor) -> torch.Tensor: + with self._decode_tiling_context(latent): + decoded = self.model.decode_base(self._denormalize_latent(latent)) + return self._normalize_decoded_frames(self.model.processor.revert_tensor(decoded)) + + def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor: + channels = int(self.config_dict["latent_channels"]) + mean = torch.tensor(self.config_dict["latents_mean"], device=latent.device, dtype=latent.dtype) + std = torch.tensor(self.config_dict["latents_std"], device=latent.device, dtype=latent.dtype) + shape = (1, channels, 1, 1, 1) + return latent * std.view(shape) + mean.view(shape) + + @staticmethod + def _normalize_decoded_frames(decoded: torch.Tensor) -> torch.Tensor: + """Canonicalize remote H3 decoder output to [B,C,T,H,W].""" + frames = decoded if frames.ndim == 4: frames = frames.unsqueeze(0).transpose(1, 2) if frames.ndim != 5: raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}") return frames.float() + # Chunks are reverted through the checkpoint's processor, which + # denormalizes and clamps into the unit interval. + chunk_value_range: ClassVar[tuple[float, float]] = (0.0, 1.0) + + @torch.inference_mode() + def decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) -> None: + """Decode temporal clips and synchronously publish frames-only chunks. + + Implements :class:`SupportsChunkedVAEDecode`. Every rank participating + in distributed VAE execution must invoke this method with a callback so + the temporal collectives stay in lockstep; ``on_chunk`` is called only + on the rank that owns output. Chunks arrive as ``[B, C, T, H, W]`` + float frames, normalized through the checkpoint's processor to match + the complete decode path. After a callback failure, the 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 not callable(getattr(self.model, "_adaptive_decode", None)): + raise RuntimeError("Loaded MiniMax-H3 VAE does not expose temporal decode primitives") + group = None + if self.is_distributed_enabled(): + group = self._native_parallel_state().get("sp_process_group") + if group is None or dist.get_world_size(group) != self.parallel_size: + raise RuntimeError("MiniMax-H3 VAE chunk decode has an invalid spatial-parallel group") + # Native H3 tiling performs its own collectives for every temporal clip, + # so this path needs the same too-few-tiles fallback as the complete + # decode: without it a shape that leaves some ranks tileless hangs the + # gather instead of decoding rank-locally. + with self._decode_tiling_context(z): + decode_h3_chunks(self, z, on_chunk, group=group) + class MiniMaxH3AudioVAE(nn.Module): def __init__(