From b2b6c6c7a9d1a338bb4d5b4a05a6c2ced64a2629 Mon Sep 17 00:00:00 2001 From: ooooooye <53851298+brandneway@users.noreply.github.com> Date: Sun, 20 Sep 2026 10:56:54 +0800 Subject: [PATCH] [Perf][MiniMax-H3] Cut video VAE encode/decode memory peaks for low-memory serving (#7241) Signed-off-by: brandneway Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com> --- tests/config/test_environment_variables.py | 4 +- .../minimax_h3/test_minimax_h3_offload.py | 6 +- .../test_minimax_h3_vae_decode_revert.py | 139 ++++++++ .../test_minimax_h3_vae_encode_prep.py | 243 +++++++++++++ .../test_minimax_h3_vae_split_residency.py | 218 ++++++++++++ .../test_minimax_h3_vae_temporal_patches.py | 335 ++++++++++++++++++ .../config/environment_variable_inventory.py | 2 + .../models/minimax_h3/pipeline_minimax_h3.py | 109 +++++- vllm_omni/diffusion/models/minimax_h3/vae.py | 286 ++++++++++++++- .../models/minimax_h3/vae_temporal.py | 271 ++++++++++++++ 10 files changed, 1584 insertions(+), 29 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_vae_decode_revert.py create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_vae_encode_prep.py create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_vae_split_residency.py create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_vae_temporal_patches.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/vae_temporal.py diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index eb073f0c7..ec3db0501 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PUBLIC_OMNI: 28, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, - EnvironmentVariableCategory.MODEL_SPECIFIC: 65, + EnvironmentVariableCategory.MODEL_SPECIFIC: 67, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20, EnvironmentVariableCategory.INTERNAL: 2, } @@ -182,7 +182,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6, ModelEnvironmentVariableDisposition.EXTERNAL: 0, ModelEnvironmentVariableDisposition.INTERNALIZE: 16, - ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5, + ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 7, } diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py index 24206952d..824ad2c63 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py @@ -86,6 +86,7 @@ def test_decode_releases_only_audio_before_return(monkeypatch): audio_vae=Mock(decode_latent=Mock(return_value=audio)), _component_on_device=lambda component: nullcontext(), _offload_model_cpu_stage_output=Mock(side_effect=lambda value: value), + _release_stage_cache=Mock(), ) monkeypatch.setattr( module.current_omni_platform, @@ -152,6 +153,7 @@ def test_forward_releases_video_after_quantization_and_before_next_seed(monkeypa diffuse=diffuse, decode=decode, _offload_model_cpu_stage_output=_recording_release(events, releases), + _release_stage_cache=lambda: None, ) request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace()) @@ -201,6 +203,7 @@ def test_forward_drops_the_decoded_video_before_the_next_seed_diffuses(monkeypat diffuse=diffuse, decode=decode, _offload_model_cpu_stage_output=lambda value: value.clone(), + _release_stage_cache=lambda: None, ) request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace()) @@ -231,6 +234,7 @@ def test_post_decode_releases_video_after_quantization(monkeypatch): _unpack_denoised_rows=lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1)), decode=decode, _offload_model_cpu_stage_output=_recording_release(events, releases), + _release_stage_cache=lambda: None, ) state = SimpleNamespace( latents=torch.zeros(1), @@ -529,7 +533,7 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp torch.nn.Module.__init__(pipeline) pipeline.audio_vae = Mock() pipeline._model_cpu_offload_modules = [pipeline.audio_vae] - events = [] + events: list[tuple] = [] @contextmanager def record_component(value): diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_decode_revert.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_decode_revert.py new file mode 100644 index 000000000..805770c19 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_decode_revert.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU tests for in-place decode revert and stage-boundary cache release. + +``_revert_decoded_inplace`` must be numerically identical to the checkpoint's +``revert_tensor`` (denormalize → clamp → rearrange) with zero whole-video +copies, and ``_release_stage_cache`` must force the allocator-cache release +(or fall back to a plain empty_cache when no component cache exists). +""" + +import pytest +import torch + +import vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 as pipeline_module +from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline +from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +DENORM_MEAN = (-2.11628412, -2.03571429, -1.80444444) # -imagenet_mean/std +DENORM_STD = (4.36681223, 4.46428571, 4.44444444) # 1/imagenet_std + + +class _FakeDenormalize: + def __init__(self, mean, std): + self.mean = mean + self.std = std + + +class _FakeProcessor: + def __init__(self): + self.transform_rev = _FakeDenormalize(DENORM_MEAN, DENORM_STD) + self.use_3d_conv = True + self.revert_calls = 0 + + def revert_tensor(self, tensor): + self.revert_calls += 1 + return _legacy_revert(tensor) + + +class _FakeCheckpointModel: + def __init__(self): + self.processor = _FakeProcessor() + + +def _vae(model=None): + vae = object.__new__(MiniMaxH3VideoVAE) + vae.model = model if model is not None else _FakeCheckpointModel() + return vae + + +def _legacy_revert(tensor: torch.Tensor) -> torch.Tensor: + """The checkpoint's revert_tensor, op for op, on the CPU (B == 1).""" + if tensor.ndim == 4: + tensor = tensor.unsqueeze(2) + assert tensor.ndim == 5 + # Normalize runs on the (b t) c h w rearrangement; at B == 1 that is a + # view of (B, C, T, H, W), so the per-channel broadcast is equivalent. + mean = torch.tensor(DENORM_MEAN).view(1, 3, 1, 1, 1) + std = torch.tensor(DENORM_STD).view(1, 3, 1, 1, 1) + tensor = (tensor - mean) / std + return tensor.clamp(0, 1) + + +def test_revert_inplace_matches_legacy_bitwise(): + decoded = torch.randn(1, 3, 7, 6, 8).clamp_(-3, 3) + expected = _legacy_revert(decoded.clone()) + out = _vae()._revert_decoded_inplace(decoded.clone()) + assert out.shape == expected.shape == (1, 3, 7, 6, 8) + assert torch.equal(out, expected) + + +def test_revert_inplace_mutates_without_extra_copies(): + decoded = torch.randn(1, 3, 5, 4, 4) + out = _vae()._revert_decoded_inplace(decoded) + assert out is decoded # in-place: no whole-video copy materializes + assert out.is_contiguous() + + +def test_revert_inplace_unsqueezes_4d_input(): + # 4-D decoded tensors follow the (B, C, H, W) image contract: unsqueeze(2) + # lands T=1, matching the checkpoint's revert_tensor branch. + decoded = torch.randn(1, 3, 6, 8) + expected = _legacy_revert(decoded.clone()) + out = _vae()._revert_decoded_inplace(decoded.clone()) + assert torch.equal(out, expected) + + +def test_revert_inplace_falls_back_without_contract(): + vae = _vae() + vae.model.processor.transform_rev = type("Empty", (), {})() + decoded = torch.randn(1, 3, 5, 4, 4) + out = vae._revert_decoded_inplace(decoded) + assert vae.model.processor.revert_calls == 1 + assert torch.equal(out, _legacy_revert(decoded)) + + +def test_revert_inplace_falls_back_on_bad_arity(): + vae = _vae() + vae.model.processor.transform_rev = _FakeDenormalize((0.5, 0.5), (1.0, 1.0, 1.0)) + vae._revert_decoded_inplace(torch.randn(1, 3, 5, 4, 4)) + assert vae.model.processor.revert_calls == 1 + + +class _RecordingCache: + def __init__(self): + self.calls = [] + + def release_if_needed(self, *, force=False): + self.calls.append(force) + return True + + +def _pipeline(): + pipeline = object.__new__(MiniMaxH3Pipeline) + pipeline._dlo_component_cache = None + return pipeline + + +def test_stage_cache_release_forces_component_cache(): + pipeline = _pipeline() + cache = _RecordingCache() + pipeline._dlo_component_cache = cache + pipeline._release_stage_cache() + assert cache.calls == [True] + + +def test_stage_cache_release_falls_back_to_empty_cache(monkeypatch): + released = [] + monkeypatch.setattr( + pipeline_module.current_omni_platform, + "empty_cache", + lambda: released.append(True), + raising=True, + ) + pipeline = _pipeline() + pipeline._release_stage_cache() + assert released == [True] diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_encode_prep.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_encode_prep.py new file mode 100644 index 000000000..3e597f09a --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_encode_prep.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU tests for the streaming uint8 preparation in the MiniMax-H3 video VAE. + +``_stream_prepare_video_tensor`` must be numerically identical to the +checkpoint's whole-video numpy path (``convert_numpy_to_tensor`` → +``transform_tensor`` → ``transpose``) while padding on the host, so the +device-side copies never materialize. +""" + +import numpy as np +import pytest +import torch + +from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +IMAGENET_MEAN = (0.485, 0.456, 0.406) +IMAGENET_STD = (0.229, 0.224, 0.225) +CLIP_LENGTH = 4 + + +class _FakeNormalize: + def __init__(self, mean, std): + self.mean = mean + self.std = std + + +class _FakeProcessor: + # Read by the stream prep via getattr(..., False); annotation only so + # tests can attach it per case. + isolated_last_frame: bool + + def __init__(self): + self.transform = _FakeNormalize(IMAGENET_MEAN, IMAGENET_STD) + + +class _FakeCheckpointModel: + def __init__(self, **overrides): + self.clip_length = CLIP_LENGTH + self.isolated_first_frame = False + self.frame_pre_padding = 0 + self.processor = _FakeProcessor() + for key, value in overrides.items(): + setattr(self, key, value) + + +def _vae(model=None): + vae = object.__new__(MiniMaxH3VideoVAE) + vae.model = model if model is not None else _FakeCheckpointModel() + return vae + + +def _legacy_reference(frames: np.ndarray) -> torch.Tensor: + """The checkpoint's numpy path, op for op, on the CPU. + + convert_numpy_to_tensor: astype(fp32) → permute → ÷255 + transform_tensor: Normalize = (x - mean) / std + encode_videos: transpose(0, 1) + """ + tensor = torch.from_numpy(frames.astype(np.float32)).permute(0, 3, 1, 2) + tensor = tensor / 255.0 + mean = torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1) + std = torch.tensor(IMAGENET_STD).view(1, 3, 1, 1) + tensor = (tensor - mean) / std + return tensor.transpose(0, 1).contiguous() + + +def _video(num_frames, height=6, width=8, seed=0): + rng = np.random.default_rng(seed) + return rng.integers(0, 256, size=(num_frames, height, width, 3), dtype=np.uint8) + + +def _pad_frames(frames: np.ndarray, pad: int) -> np.ndarray: + if not pad: + return frames + return np.concatenate([frames, np.repeat(frames[-1:], pad, axis=0)]) + + +def test_stream_prep_matches_legacy_path_bitwise(): + frames = _video(10) + out = _vae()._stream_prepare_video_tensor(frames, torch.device("cpu")) + assert out is not None + assert out.shape == (3, 12, 6, 8) # pad = (0 - 10) % 4 == 2 + assert out.dtype == torch.float32 + assert torch.equal(out, _legacy_reference(_pad_frames(frames, 2))) + + +def test_stream_prep_isolated_first_frame_without_last_frame_skips_pad(): + # Asymmetric checkpoint (isolated first frame, no isolated last frame): + # trim-stable lengths are k*clip + 0, which never satisfy the + # encode_temporal offset of 1, so the remote re-pads through its own + # whole-video cat regardless. Padding here would only upload frames the + # remote trim discards, so the legacy frame count is kept. + model = _FakeCheckpointModel(isolated_first_frame=True) + frames = _video(10) + out = _vae(model)._stream_prepare_video_tensor(frames, torch.device("cpu")) + assert out is not None + assert out.shape[1] == 10 + assert torch.equal(out, _legacy_reference(frames)) + + +def test_stream_prep_symmetric_isolated_frames_pad_to_trim_stable(): + # isolated first + last frames: offset == tail % clip == 1, so padding + # to the next trim-stable length makes the remote trim a no-op and keeps + # encode_temporal from re-padding: T=10, clip=4, tail=1 → 3*4+1 == 13. + model = _FakeCheckpointModel(isolated_first_frame=True) + model.processor.isolated_last_frame = True + frames = _video(10) + out = _vae(model)._stream_prepare_video_tensor(frames, torch.device("cpu")) + assert out is not None + assert out.shape[1] == 13 + assert out.shape[1] % 4 == 1 + assert torch.equal(out, _legacy_reference(_pad_frames(frames, 3))) + + +def test_stream_prep_uses_processor_align_when_available(): + calls = [] + + class _AlignedProcessor(_FakeProcessor): + def align_video_length(self, video_length, mode="pad", granularity="chunk"): + calls.append((video_length, mode, granularity)) + return 5 + + model = _FakeCheckpointModel() + model.processor = _AlignedProcessor() + frames = _video(10) + out = _vae(model)._stream_prepare_video_tensor(frames, torch.device("cpu")) + assert out is not None + assert calls == [(10, "pad", "chunk")] + assert out.shape[1] == 15 + assert torch.equal(out, _legacy_reference(_pad_frames(frames, 5))) + + +def test_stream_prep_pads_with_last_frame(): + frames = _video(2) # pad = (0 - 2) % 4 == 2 + out = _vae()._stream_prepare_video_tensor(frames, torch.device("cpu")) + padded = np.concatenate([frames, np.repeat(frames[-1:], 2, axis=0)]) + assert torch.equal(out, _legacy_reference(padded)) + + +def test_stream_prep_aligned_input_needs_no_pad(): + frames = _video(8) # 8 % 4 == 0 + out = _vae()._stream_prepare_video_tensor(frames, torch.device("cpu")) + assert out.shape[1] == 8 + assert torch.equal(out, _legacy_reference(frames)) + + +@pytest.mark.parametrize( + "frames", + [ + _video(5).astype(np.float32), # not uint8 + _video(5)[:, :, :, :1], # single channel + np.zeros((5, 6, 8, 4), dtype=np.uint8), # four channels + np.zeros((5, 6, 8), dtype=np.uint8), # missing channel axis + np.zeros((0, 6, 8, 3), dtype=np.uint8), # empty + ], +) +def test_stream_prep_rejects_unsupported_inputs(frames): + assert _vae()._stream_prepare_video_tensor(frames, torch.device("cpu")) is None + + +@pytest.mark.parametrize( + "model", + [ + _FakeCheckpointModel(clip_length=None), + _FakeCheckpointModel(clip_length=0), + type("NoProcessor", (), {"clip_length": CLIP_LENGTH})(), + type( + "NoTransform", + (), + {"clip_length": CLIP_LENGTH, "processor": type("P", (), {})()}, + )(), + ], +) +def test_stream_prep_rejects_missing_checkpoint_contract(model): + frames = _video(5) + assert _vae(model)._stream_prepare_video_tensor(frames, torch.device("cpu")) is None + + +def test_stream_prep_rejects_bad_norm_arity(): + model = _FakeCheckpointModel() + model.processor.transform = _FakeNormalize((0.5, 0.5), IMAGENET_STD) + assert _vae(model)._stream_prepare_video_tensor(_video(5), torch.device("cpu")) is None + + +class _RecordingModel(_FakeCheckpointModel): + """Feeds encode_video a well-shaped latent and records its input.""" + + def __init__(self): + super().__init__() + self.received = None + self.latent = torch.randn(1, 4, 3, 4, 4) # (B, C, T', H', W'), even H/W + + def encode_videos(self, videos, use_fp16_latent=False, **kwargs): + self.received = videos + assert use_fp16_latent is True + return [self.latent] + + +def _vae_for_encode_video(): + vae = object.__new__(MiniMaxH3VideoVAE) + vae.model = _RecordingModel() + vae.device_module = torch + vae.config_dict = { + "latent_channels": 4, + "latents_mean": [0.0] * 4, + "latents_std": [1.0] * 4, + } + # A plain tensor (not nn.Parameter): nn.Module.__setattr__ refuses + # Parameters on instances built via object.__new__. + vae._dummy_parameter = torch.zeros(1, dtype=torch.float32) + vae.parameters = lambda: iter([vae._dummy_parameter]) + return vae + + +def test_encode_video_streams_uint8_numpy(monkeypatch): + monkeypatch.delenv("VLLM_OMNI_VAE_ENCODE_LEGACY_PREP", raising=False) + vae = _vae_for_encode_video() + rows, shape = vae.encode_video(_video(10)) + assert isinstance(vae.model.received, list) + assert isinstance(vae.model.received[0], torch.Tensor) + assert vae.model.received[0].shape == (3, 12, 6, 8) + assert rows.shape[0] == shape[0] * shape[1] * shape[2] // 4 # patchify rows + + +def test_encode_video_legacy_env_escapes_to_checkpoint_path(monkeypatch): + monkeypatch.setenv("VLLM_OMNI_VAE_ENCODE_LEGACY_PREP", "1") + vae = _vae_for_encode_video() + frames = _video(10) + vae.encode_video(frames) + assert vae.model.received is frames # untouched numpy, checkpoint converts + + +def test_encode_video_falls_back_when_preparation_unsupported(monkeypatch): + monkeypatch.delenv("VLLM_OMNI_VAE_ENCODE_LEGACY_PREP", raising=False) + vae = _vae_for_encode_video() + vae.model.clip_length = None # contract not discoverable + frames = _video(10) + vae.encode_video(frames) + assert vae.model.received is frames diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_split_residency.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_split_residency.py new file mode 100644 index 000000000..166f43538 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_split_residency.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU tests for the split encoder/decoder residency of the MiniMax-H3 VAE. + +The encode half (CNN encoder + quant_conv) and the decode half +(post_quant_conv + ViT decoder) stage independently: loading one half must +rebind only its own parameters while the other half stays on its pinned +CPU master. +""" + +from unittest.mock import Mock + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.models.minimax_h3.vae import ( + MiniMaxH3VideoVAE, + _VideoVAEPartProxy, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _FakeRemoteModel(nn.Module): + def __init__(self): + super().__init__() + self.encoder = nn.Linear(8, 8) + self.quant_conv = nn.Linear(4, 4) + self.post_quant_conv = nn.Linear(4, 4) + self.decoder = nn.Linear(16, 16) + + +class _FakeRemote(nn.Module): + def __init__(self): + super().__init__() + self.model = _FakeRemoteModel() + + +def _vae() -> MiniMaxH3VideoVAE: + vae = object.__new__(MiniMaxH3VideoVAE) + nn.Module.__init__(vae) # bare instance: enable Module attribute assignment + vae._device_target = torch.device("cpu") + vae.remote = _FakeRemote() + vae.model = vae.remote.model + vae._stager = None + vae._encoder_stager = None + vae._decoder_stager = None + vae._build_residency_stagers(torch.device("cpu")) + vae.encoder_component = _VideoVAEPartProxy(vae, "encoder") + vae.decoder_component = _VideoVAEPartProxy(vae, "decoder") + return vae + + +def _ptr(param: torch.nn.Parameter) -> int: + return param.untyped_storage().data_ptr() + + +def test_split_stagers_snapshot_disjoint_halves(): + vae = _vae() + enc_master = _ptr(vae.model.encoder.weight) + dec_master = _ptr(vae.model.decoder.weight) + pqc_master = _ptr(vae.model.post_quant_conv.weight) + assert enc_master != dec_master + + vae._load_part_to_device("encoder") + # The encoder half is rebound to freshly allocated staging storage... + assert _ptr(vae.model.encoder.weight) != enc_master + assert _ptr(vae.model.quant_conv.weight) != enc_master + # ...while every decode-half tensor stays on its own CPU master. + assert _ptr(vae.model.decoder.weight) == dec_master + assert _ptr(vae.model.post_quant_conv.weight) == pqc_master + + vae._offload_part_to_cpu("encoder") + assert _ptr(vae.model.encoder.weight) == enc_master + + +def test_part_load_is_idempotent_and_independent(): + vae = _vae() + vae._load_part_to_device("decoder") + staged_ptr = _ptr(vae.model.decoder.weight) + vae._load_part_to_device("decoder") + assert _ptr(vae.model.decoder.weight) == staged_ptr + vae._offload_part_to_cpu("decoder") + + +def test_whole_component_load_and_offload_cover_both_halves(): + vae = _vae() + enc_master = _ptr(vae.model.encoder.weight) + dec_master = _ptr(vae.model.decoder.weight) + vae.load_to_device() + assert _ptr(vae.model.encoder.weight) != enc_master + assert _ptr(vae.model.decoder.weight) != dec_master + vae.offload_to_cpu() + assert _ptr(vae.model.encoder.weight) == enc_master + assert _ptr(vae.model.decoder.weight) == dec_master + + +def test_part_values_survive_a_load_offload_cycle(): + vae = _vae() + expected = vae.model.encoder.weight.detach().clone() + vae._load_part_to_device("encoder") + assert torch.equal(vae.model.encoder.weight.detach(), expected) + vae._offload_part_to_cpu("encoder") + assert torch.equal(vae.model.encoder.weight.detach(), expected) + + +def test_proxy_routes_to_the_matching_half(): + vae = _vae() + enc_master = _ptr(vae.model.encoder.weight) + dec_master = _ptr(vae.model.decoder.weight) + vae.encoder_component.load_to_device() + assert _ptr(vae.model.encoder.weight) != enc_master + assert _ptr(vae.model.decoder.weight) == dec_master + vae.encoder_component.offload_to_cpu() + assert _ptr(vae.model.encoder.weight) == enc_master + + +def test_proxy_does_not_register_the_adapter_as_submodule(): + vae = _vae() + proxy = vae.encoder_component + assert proxy._vae is vae + # The back-reference must not recurse through module traversal. + assert all(child is not vae for child in proxy.children()) + assert len(list(proxy.parameters())) == 0 + + +def test_component_on_device_unwraps_proxy_for_sequential_offload(monkeypatch): + """Model-level CPU offload hooks the real video_vae, never the part + proxies: _component_on_device must hand the hooked module to the + sequential context, or _get_sequential_offload_hook raises on the + hook-less proxy on the first encode/decode of every request.""" + from contextlib import contextmanager + + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + + vae = _vae() + seen: list[nn.Module] = [] + + @contextmanager + def fake_sequential_offload_component(component): + seen.append(component) + yield + + monkeypatch.setattr(pipeline_module, "sequential_offload_component", fake_sequential_offload_component) + + pipeline = object.__new__(MiniMaxH3Pipeline) + nn.Module.__init__(pipeline) + pipeline._model_cpu_offload_modules = [vae] + + with pipeline._component_on_device(vae.encoder_component): + pass + with pipeline._component_on_device(vae.decoder_component): + pass + # A plain component passes through unchanged. + with pipeline._component_on_device(vae): + pass + # Mocks auto-create any attribute, so the unwrap must be a type check: + # a Mock component must reach the sequential context as itself. + mock_vae = Mock() + pipeline._model_cpu_offload_modules = [mock_vae] + with pipeline._component_on_device(mock_vae): + pass + + assert seen == [vae, vae, vae, mock_vae] + + +def test_sequential_offload_target_is_the_hooked_module(): + from vllm_omni.diffusion.offloader.sequential_backend import _get_sequential_offload_hook + + vae = _vae() + proxy = vae.encoder_component + # The pre-fix failure mode: the proxy carries no hook registry. + with pytest.raises(RuntimeError, match="no sequential offload hook"): + _get_sequential_offload_hook(proxy) + assert proxy.sequential_offload_target is vae + + +def test_cache_retention_forwarded_to_both_stagers(): + vae = _vae() + cache = object() + vae.set_omni_component_cache(cache) + assert vae._encoder_stager.cache_retention is cache + assert vae._decoder_stager.cache_retention is cache + + +def test_unknown_part_raises(): + vae = _vae() + with pytest.raises(ValueError, match="unknown video VAE part"): + vae._part_stager("middle") + + +def test_missing_checkpoint_structure_falls_back_to_whole_stager(): + vae = object.__new__(MiniMaxH3VideoVAE) + nn.Module.__init__(vae) + vae._device_target = torch.device("cpu") + + class BareRemote(nn.Module): + def __init__(self): + super().__init__() + self.model = nn.Linear(4, 4) # no encoder/decoder submodules + + vae.remote = BareRemote() + vae.model = vae.remote.model + vae._stager = None + vae._encoder_stager = None + vae._decoder_stager = None + vae._build_residency_stagers(torch.device("cpu")) + assert vae._stager is not None + assert vae._encoder_stager is None + assert vae._decoder_stager is None + # Part loads fall back to the whole-module stager. + vae._load_part_to_device("decoder") + assert vae._stager.loaded + vae._offload_part_to_cpu("decoder") + assert not vae._stager.loaded diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_temporal_patches.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_temporal_patches.py new file mode 100644 index 000000000..8ba8199a1 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_temporal_patches.py @@ -0,0 +1,335 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU tests for the temporal VAE forks and the uint8 output transfer. + +The encode fork must feed bit-identical frames to the checkpoint's chunk +loop (padding moves from a whole-video cat into the tail chunk), the decode +fork must produce the same video as the legacy fp32-buffer chain +(revert -> clamp -> *255 -> round -> uint8), and the output transfer must +match the previous ``.to(uint8, contiguous)`` conversion. +""" + +from typing import Any, cast + +import pytest +import torch + +from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _prepare_minimax_h3_video_output, +) +from vllm_omni.diffusion.models.minimax_h3.vae_temporal import ( + _decode_temporal_streaming_uint8, + _encode_temporal_tail_pad, + install_temporal_stream_patches, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +CLIP = 4 + + +# --------------------------------------------------------------------------- +# encode_temporal fork +# --------------------------------------------------------------------------- + + +class _EncodeModel: + """Minimal stand-in implementing the checkpoint's encode_temporal.""" + + # Attached per test case as capability probes for the install checks; + # annotations only, the fork must keep working when they are absent. + encode_temporal: Any + _decode_temporal_streaming: Any + processor: Any + + def __init__(self, *, isolated_first=False, isolated_last=False, key_frame=False, frame_drop=0): + self.clip_length = CLIP + self.isolated_first_frame = isolated_first + self.frame_pre_padding = 0 + self.isolated_key_frame = key_frame + self.token_drop = 0 + self.isolated_last_frame = isolated_last + self.frame_drop = frame_drop + self.encoded = [] + + def _adaptive_encode(self, clip): + self.encoded.append(clip.clone()) + return clip # identity: z_list entries are the clips themselves + + def reference_encode_temporal(self, x): + """The checkpoint's original: whole-video pad, then chunk.""" + offset_frame = 1 if self.isolated_first_frame and self.frame_pre_padding == 0 else 0 + if x.shape[2] % self.clip_length != offset_frame: + pad_size = (offset_frame - x.shape[2]) % self.clip_length + pad_frames = x[:, :, -1:].repeat(1, 1, pad_size, 1, 1) + x = torch.cat([x, pad_frames], dim=2) + num_chunks = (x.shape[2] - offset_frame) // self.clip_length + z_list = [] + for i in range(num_chunks): + start_idx = i * self.clip_length + offset_frame + end_idx = (i + 1) * self.clip_length + offset_frame + clip_x = x[:, :, start_idx:end_idx, :, :] + if self.isolated_key_frame: + key_frame = clip_x[:, :, :1, :, :] + z_key = self._adaptive_encode(key_frame) + if clip_x.shape[2] > 1: + z_video = self._adaptive_encode(clip_x[:, :, 1:, :, :]) + z = torch.cat([z_key, z_video], dim=2) + else: + z = z_key + else: + z = self._adaptive_encode(clip_x) + z_list.append(z) + z = torch.cat(z_list, dim=2) + if self.token_drop > 0: + z = z[:, :, : -self.token_drop] + if self.isolated_first_frame: + z_first = self._adaptive_encode(x[:, :, :1, :, :]) + z = torch.cat([z_first, z], dim=2) + if self.isolated_last_frame: + frame_num = x.shape[2] + last_idx = frame_num - self.frame_drop + offset_frame + z_last = self._adaptive_encode(x[:, :, last_idx : last_idx + 1, :, :]) + z = torch.cat([z, z_last], dim=2) + return z + + +def _video_x(t, c=3, h=4, w=4): + return torch.arange(t * c * h * w, dtype=torch.float32).reshape(1, c, t, h, w) / 255.0 + + +@pytest.mark.parametrize("num_frames", [8, 10, 11, 13]) +@pytest.mark.parametrize( + "flags", + [ + {}, + {"isolated_first": True}, + {"isolated_last": True}, + {"isolated_first": True, "isolated_last": True}, + {"key_frame": True}, + ], +) +def test_encode_fork_matches_whole_video_pad_bitwise(num_frames, flags): + x = _video_x(num_frames) + ref_model = _EncodeModel(**flags) + expected = ref_model.reference_encode_temporal(x.clone()) + # Empty-slice encodes (out-of-range isolated-last indexes are no-ops in + # the checkpoint) carry no content; compare only real calls. + ref_calls = [c.clone() for c in ref_model.encoded if c.shape[2] > 0] + + fork_model = _EncodeModel(**flags) + out = _encode_temporal_tail_pad(fork_model, x.clone()) + + assert out.shape == expected.shape + assert torch.equal(out, expected) + assert len(fork_model.encoded) == len(ref_calls) + for got, want in zip(fork_model.encoded, ref_calls): + assert torch.equal(got, want) + + +def test_encode_fork_avoids_whole_video_pad_for_misaligned_input(): + # 375 frames with an asymmetric config (isolated last frame only): + # the checkpoint always pads (375 % 17 != 0); the fork must encode the + # tail chunk with an in-chunk pad instead of concatenating the video. + model = _EncodeModel(isolated_last=True, frame_drop=16) + model.clip_length = 17 + x = _video_x(375, c=1, h=2, w=2) + _encode_temporal_tail_pad(model, x) + # 375 % 17 == 1 → pad 16 inside the tail chunk; the isolated last frame + # resolves to the final real frame (pad frames replicate it). + assert len(model.encoded) == (375 // 17) + 1 + 1 + assert torch.equal(model.encoded[-1].flatten(), x[:, :, -1:, :, :].flatten()) + + +# --------------------------------------------------------------------------- +# _decode_temporal_streaming uint8 fork +# --------------------------------------------------------------------------- + +DENORM_MEAN = (-2.11628412, -2.03571429, -1.80444444) +DENORM_STD = (4.36681223, 4.46428571, 4.44444444) + + +class _FakeDenorm: + def __init__(self): + self.mean = DENORM_MEAN + self.std = DENORM_STD + + +class _FakeProcessor: + def __init__(self): + self.transform_rev = _FakeDenorm() + + +class _DecodeModel: + """Minimal stand-in for the checkpoint's streaming decode contract.""" + + def __init__(self, *, token_drop=0, frame_overlap=0): + self.tokens_chunk_size = 2 + self.vae_ratio_t = 1 + self.token_overlap = 1 if token_drop else 0 + self.token_drop = token_drop + self.frame_pre_padding = 0 + self.frame_overlap = frame_overlap + self.processor = _FakeProcessor() + + def _adaptive_decode(self, clip_z): + # Deterministic pseudo-decode: map latent tokens to frames 1:1. + return torch.sin(clip_z * 3.0) + + def blend(self, a, b, overlap, dim=-3): + return (a + b) / 2.0 + + def _decode_temporal_pad_frames(self, z, pad_tokens): + return int(pad_tokens) + + def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens): + total = num_chunks * self.tokens_chunk_size * self.vae_ratio_t + if z_head is not None: + total += 1 + if z_tail is not None: + total += 1 + if self.token_drop > 0: + total += self.frame_overlap + return int(total), int(pad_tokens), int(total - pad_tokens) + + def legacy_streaming(self, z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype): + """The checkpoint's original: accumulate in the cat dtype, quantize later.""" + total_frames, _, output_frames = self._decode_temporal_output_frame_plan( + z, z_head, z_tail, num_chunks, pad_tokens + ) + chunk_dec = self.tokens_chunk_size * self.vae_ratio_t + split_count = int(self.token_drop > 0) + 1 + dec = None + dec_overlap = None + write_pos = 0 + + def write_part(part): + nonlocal dec, write_pos + if dec is None: + out_shape = list(part.shape) + out_shape[2] = output_frames + dec = torch.empty(out_shape, dtype=part.dtype) + n = min(part.shape[2], dec.shape[2] - write_pos) + if n > 0: + dec[:, :, write_pos : write_pos + n].copy_(part[:, :, :n]) + write_pos += n + + for i in range(num_chunks): + t_start = i * self.tokens_chunk_size + t_end = t_start + self.tokens_chunk_size + self.token_overlap + clip_z = z[:, :, t_start:t_end] + clip_dec = self._adaptive_decode(clip_z) + if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype: + clip_dec = clip_dec.to(temporal_cat_dtype) + for j in range(split_count): + chunk = clip_dec[:, :, j * chunk_dec : min(j * chunk_dec + chunk_dec, clip_dec.shape[2])] + if j == 0: + if dec_overlap is not None: + chunk = self.blend(dec_overlap, chunk, self.frame_overlap, dim=-3) + dec_overlap = None + write_part(chunk) + else: + dec_overlap = chunk.contiguous() + if i == num_chunks - 1 and dec_overlap is not None: + write_part(dec_overlap) + dec_overlap = None + # Legacy post-decode chain: revert in place in the decoded dtype, then + # the adapter's ``frames.float()`` before the pipeline's quantizer. + assert dec is not None + mean = torch.tensor(DENORM_MEAN, dtype=dec.dtype).view(1, 3, 1, 1, 1) + std = torch.tensor(DENORM_STD, dtype=dec.dtype).view(1, 3, 1, 1, 1) + dec.sub_(mean).div_(std).clamp_(0.0, 1.0) + dec = dec.float() + dec.mul_(255.0).round_() + return dec.to(torch.uint8) + + +@pytest.mark.parametrize("token_drop,frame_overlap", [(0, 0), (1, 1)]) +@pytest.mark.parametrize("temporal_cat_dtype", [None, torch.float16, torch.bfloat16]) +def test_decode_fork_matches_legacy_chain_bitwise(token_drop, frame_overlap, temporal_cat_dtype): + # The quantizer must see FP32 no matter the cat dtype: the legacy adapter + # returned frames.float() before the pipeline multiplied and rounded, so + # an fp16/bf16 cat dtype only narrows denormalization, never the rounding. + torch.manual_seed(0) + tokens = 2 * 2 + (1 if token_drop else 0) # num_chunks * tokens_chunk_size + z = torch.randn(1, 3, tokens, 4, 4) + legacy = _DecodeModel(token_drop=token_drop, frame_overlap=frame_overlap) + expected = legacy.legacy_streaming(z.clone(), None, None, 2, 0, temporal_cat_dtype) + + fork = _DecodeModel(token_drop=token_drop, frame_overlap=frame_overlap) + out = _decode_temporal_streaming_uint8(fork, z.clone(), None, None, 2, 0, temporal_cat_dtype) + assert out.dtype == torch.uint8 + assert out.shape == expected.shape + assert torch.equal(out, expected) + + +def test_decode_fork_rejects_empty_plan(): + model = _DecodeModel() + model.tokens_chunk_size = 0 + with pytest.raises(ValueError, match="non-positive output_frames"): + _decode_temporal_streaming_uint8(model, torch.randn(1, 3, 0, 4, 4), None, None, 0, 0, None) + + +# --------------------------------------------------------------------------- +# install + escape hatch +# --------------------------------------------------------------------------- + + +def test_install_binds_both_forks(): + model = _EncodeModel() + # cast(Any, ...) keeps mypy from narrowing the probes to the lambda type, + # which has no __func__ for the bound-method assertions below. + model.encode_temporal = cast(Any, lambda x: x) # placeholder for the capability check + model._decode_temporal_streaming = cast(Any, lambda *a: None) + model.processor = _FakeProcessor() + install_temporal_stream_patches(model) + assert model.encode_temporal.__func__ is _encode_temporal_tail_pad + assert model._decode_temporal_streaming.__func__ is _decode_temporal_streaming_uint8 + + +def test_install_skips_decode_fork_without_denorm_contract(): + model = _EncodeModel() + model.encode_temporal = cast(Any, lambda x: x) + model._decode_temporal_streaming = cast(Any, lambda *a: None) + install_temporal_stream_patches(model) # no processor -> keep original + assert model.encode_temporal.__func__ is _encode_temporal_tail_pad + assert not hasattr(model._decode_temporal_streaming, "__func__") + + +def test_install_respects_legacy_env(monkeypatch): + monkeypatch.setenv("VLLM_OMNI_VAE_LEGACY_TEMPORAL", "1") + model = _EncodeModel() + install_temporal_stream_patches(model) + assert not hasattr(getattr(model, "encode_temporal", None), "__func__") + + +# --------------------------------------------------------------------------- +# _prepare_minimax_h3_video_output +# --------------------------------------------------------------------------- + + +def test_prepare_output_matches_previous_conversion(): + video = torch.randn(1, 3, 4, 6, 8).clamp_(-1, 2) + legacy = ( + video.detach() + .float() + .clamp(0, 1) + .mul(255) + .round() + .permute(0, 2, 3, 4, 1) + .to(dtype=torch.uint8, memory_format=torch.contiguous_format) + ) + out = _prepare_minimax_h3_video_output(video.clone()) + assert out.dtype == torch.uint8 + assert out.shape == (1, 4, 6, 8, 3) + assert out.is_contiguous() + assert torch.equal(out, legacy) + + +def test_prepare_output_uint8_passthrough(): + video = torch.randint(0, 256, (1, 3, 4, 6, 8), dtype=torch.uint8) + out = _prepare_minimax_h3_video_output(video) + assert out.dtype == torch.uint8 + assert out.is_contiguous() + assert torch.equal(out, video.permute(0, 2, 3, 4, 1).contiguous()) diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 27b1a2ac2..f57d7ffef 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -214,6 +214,8 @@ _MODEL_INTERNALIZE = ( _MODEL_DEPRECATE_REMOVE = ( "HIGGS_AUDIO_V2_TOKENIZER_PATH", "VLLM_OMNI_LINGBOT_ACTION_ROOT", + "VLLM_OMNI_VAE_ENCODE_LEGACY_PREP", + "VLLM_OMNI_VAE_LEGACY_TEMPORAL", "VLLM_OMNI_VOXCPM_CODE_PATH", "XCODEC1_PATH", "model_stage", diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 1a8395e5c..61c0ceecd 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -8,7 +8,7 @@ import json import math import os from collections.abc import Iterable, Mapping, Sequence -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from dataclasses import fields, replace from itertools import groupby from pathlib import Path @@ -138,9 +138,11 @@ from .time_request import ( MINIMAX_H3_SHAPE_PLANNER, minimax_h3_time_shift_sigmas, ) -from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE +from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE, _VideoVAEPartProxy if TYPE_CHECKING: + from PIL import Image + from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.utils import StepRequestState @@ -376,12 +378,20 @@ def _minimax_h3_post_process(output, output_type: str = "np"): def _prepare_minimax_h3_video_output(video: torch.Tensor) -> torch.Tensor: """Quantize decoded frames in place before worker-to-engine transfer.""" - video = video.detach().float() + video = video.detach() + if video.dtype == torch.uint8: + # Streaming decode already quantized and clamped; only the transfer + # layout remains. + return video.permute(0, 2, 3, 4, 1).contiguous() + video = video.float() video.clamp_(0, 1).mul_(255).round_() - return video.permute(0, 2, 3, 4, 1).to( - dtype=torch.uint8, - memory_format=torch.contiguous_format, - ) + permuted = video.permute(0, 2, 3, 4, 1) + out = torch.empty(permuted.shape, dtype=torch.uint8, device=video.device) + # copy_ fuses the layout change and the cast into one kernel; + # ``.to(dtype=uint8, memory_format=contiguous_format)`` materializes a + # contiguous FP32 intermediate first (~4.2GB for a 15s clip). + out.copy_(permuted) + return out def _register_dlo_component_cache(cache: BoundedAllocatorCache, *components: Any) -> None: @@ -1342,6 +1352,15 @@ class MiniMaxH3Pipeline( @contextmanager def _component_on_device(self, component: nn.Module): if getattr(self, "_model_cpu_offload_modules", None): + # Sequential offload hooks whole modules (enable_omni_model_cpu_offload + # registers them on the discovered components, e.g. the real + # video_vae). Split-residency proxies carry no hook, so unwrap to the + # hooked module before entering the context — whole-module movement + # has no half-residency benefit anyway. The check is a type check, not + # getattr: Mock components auto-create any attribute, which would + # unwrap them to a child mock and break scope tracking. + if isinstance(component, _VideoVAEPartProxy): + component = component.sequential_offload_target with sequential_offload_component(component): yield return @@ -1381,6 +1400,50 @@ class MiniMaxH3Pipeline( """Whether this rank's output is returned by the diffusion executor.""" return not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0 + def _encode_visual_conditions( + self, + images: list[Image.Image], + prepared_videos: list[dict[str, Any]] | None, + *, + video_count: int, + ) -> tuple[torch.Tensor | None, list[tuple[int, int, int]]]: + rows: list[torch.Tensor] = [] + shapes: list[tuple[int, int, int]] = [] + _, rank, _ = _dit_rank_world() + # Keep image and video references in one residency window when both + # appear in a request; otherwise the video branch would reload the VAE. + # Encoding touches only the CNN encoder half, so the 9GB ViT decoder + # stays off the device for the whole window. + needs_video_vae = video_count > 0 or (rank == 0 and bool(images)) + video_vae_context = ( + self._component_on_device(self.video_vae.encoder_component) if needs_video_vae else nullcontext() + ) + with video_vae_context: + if images: + image_rows = None + if rank == 0: + image_rows = torch.cat([self.video_vae.encode_image(image) for image in images]) + rows.append( + _broadcast_tensor( + image_rows, + dtype=torch.float32, + device=self.device, + ) + ) + shapes.extend((1, image.height // 16, image.width // 16) for image in images) + if video_count: + video_rows, video_shapes = self._encode_video_conditions_resident( + prepared_videos, + count=video_count, + ) + rows.append(video_rows) + shapes.extend(video_shapes) + # The latents are extracted; the encode's input/staging pages are idle + # and must not stay mapped through the denoise and decode peaks. + if needs_video_vae: + self._release_stage_cache() + return (torch.cat(rows) if rows else None), shapes + def _offload_model_cpu_stage_output(self, tensor: torch.Tensor) -> torch.Tensor: """Release a decoded output's storage before a later seed reloads the DiT. @@ -1768,6 +1831,26 @@ class MiniMaxH3Pipeline( raise ValueError("MiniMax H3 chunked MP4 encoding currently expects one output per decoder") return videos[0] + def _release_stage_cache(self) -> None: + """Return idle allocator pages to the device at stage boundaries. + + The bounded component cache releases only past its idle-cache bound + (>25% of device capacity), which on large devices lets a finished + stage's freed activations -- the DiT's denoise buffers, an encode + input, a decoded frame tensor -- stay physically mapped across the + next stage's peak. Forcing the release at a boundary is exact (only + free pages are returned; live tensors are untouched) and costs one + remap per later allocation. + """ + cache = getattr(self, "_dlo_component_cache", None) + if cache is not None: + try: + cache.release_if_needed(force=True) + except BaseException: + logger.exception("Failed to release retained allocator cache at stage boundary") + return + current_omni_platform.empty_cache() + def decode( self, video_latent: torch.Tensor, @@ -1776,7 +1859,11 @@ class MiniMaxH3Pipeline( height: int, width: int, ) -> tuple[torch.Tensor, torch.Tensor]: - with self._component_on_device(self.video_vae): + # Denoise just ended: its freed activation pages must not stay mapped + # through the VAE decode peak. Decoding needs only the ViT decoder + # half of the VAE, so the CNN encoder stays off the device. + self._release_stage_cache() + with self._component_on_device(self.video_vae.decoder_component): with current_omni_platform.create_autocast_context( device_type=self.device.type, dtype=torch.float16, @@ -2199,6 +2286,11 @@ class MiniMaxH3Pipeline( # loop rebinds it. post_decode() rebinds for the same reason. video = self._offload_model_cpu_stage_output(_prepare_minimax_h3_video_output(video)) videos.append(video) + # The FP32 decoded frames are quantized into the appended uint8 + # tensor; drop the reference and return the idle pages instead of + # holding them through the next output's denoise/decode. + del video + self._release_stage_cache() audios.append(audio) if videos and isinstance(videos[0], bytes): video = videos[0] if len(videos) == 1 else videos @@ -2518,6 +2610,7 @@ class MiniMaxH3Pipeline( width=shape["width"], ) video = self._offload_model_cpu_stage_output(_prepare_minimax_h3_video_output(video)) + self._release_stage_cache() return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index d02649229..ed28029bb 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -6,11 +6,13 @@ from __future__ import annotations import importlib import json +import os from collections.abc import Iterator from contextlib import AbstractContextManager, contextmanager, nullcontext from pathlib import Path from typing import Any, ClassVar +import numpy as np import torch import torch.distributed as dist import torch.nn as nn @@ -33,11 +35,15 @@ from vllm_omni.diffusion.offloader.module_residency import ( from .chunked_decode import decode_h3_chunks from .ops import install_h3_vae_optimizations from .packed_tokens import minimax_h3_patchify_video_latent +from .vae_temporal import install_temporal_stream_patches MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42 MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000 MINIMAX_H3_AUDIO_CHANNELS = 2 +# Escape hatch back to the checkpoint's whole-video numpy preparation path. +_VAE_ENCODE_LEGACY_PREP_ENV = "VLLM_OMNI_VAE_ENCODE_LEGACY_PREP" + logger = init_logger(__name__) @@ -63,6 +69,11 @@ def _minimax_h3_keyframe_encode_context( yield +def _legacy_encode_prep_enabled() -> bool: + value = os.environ.get(_VAE_ENCODE_LEGACY_PREP_ENV, "0") + return value.strip().lower() not in ("", "0", "false", "off") + + def _load_component_config(component_path: str) -> dict[str, Any]: config_path = Path(component_path) / "config.json" config = json.loads(config_path.read_text(encoding="utf-8")) @@ -264,6 +275,49 @@ class _AudioVAEDeterminismContext(AbstractContextManager): return False +class _VideoVAEPartProxy(nn.Module): + """Residency proxy for one half of the split video VAE. + + The pipeline's ``_component_on_device`` drives whatever object it is + given through ``load_to_device``/``offload_to_cpu``. The checkpoint's + ViT decoder holds ~9GB of FP32 weights while the CNN encoder holds + ~0.7GB, and each is used by exactly one direction, so routing those + calls per half keeps the decoder off the device while reference videos + encode and the encoder off while latents decode. ``object.__setattr__`` + keeps the back-reference out of the module registry (registering the + adapter as a submodule would recurse through ``parameters()``). + """ + + def __init__(self, vae: MiniMaxH3VideoVAE, part: str) -> None: + super().__init__() + object.__setattr__(self, "_vae", vae) + object.__setattr__(self, "_part", part) + + def load_to_device(self) -> None: + self._vae._load_part_to_device(self._part) + + def offload_to_cpu(self) -> None: + self._vae._offload_part_to_cpu(self._part) + + def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None: + self._vae.set_omni_component_cache(cache) + + @property + def sequential_offload_target(self) -> MiniMaxH3VideoVAE: + """The module model-level CPU offload actually hooks. + + ``enable_omni_model_cpu_offload`` registers the sequential hook on the + real ``video_vae``, and the hook moves the whole module through + ``parameters()`` — entering the sequential context through this proxy + would find no ``_hook_registry`` and raise, and whole-module movement + has no half-residency benefit anyway. ``_component_on_device`` unwraps + through this property before entering its sequential-offload branch; + the manual per-half staging above stays proxy-driven, where split + residency is the whole point. + """ + return self._vae + + class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): """Adapter around the checkpoint's native parallel-tiled video VAE.""" @@ -303,14 +357,15 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): decoder, device=device, ) - self._stager = None - if initial_device.type == "cpu" and device.type not in ("cpu", "meta"): - self._stager = PinnedModuleStager( - self.remote, - device, - pin_memory=True, - ) + install_temporal_stream_patches(self.remote.model) self.model = self.remote.model + self._stager = None + self._encoder_stager = None + self._decoder_stager = None + if initial_device.type == "cpu" and device.type not in ("cpu", "meta"): + self._build_residency_stagers(device) + self.encoder_component = _VideoVAEPartProxy(self, "encoder") + self.decoder_component = _VideoVAEPartProxy(self, "decoder") self.use_tiling = True self.use_slicing = False self.parallel_size = 1 @@ -321,27 +376,97 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): if self._tile_gather_reuse_enabled(): self._install_persistent_tile_gather() - def load_to_device(self) -> None: + def _build_residency_stagers(self, device: torch.device) -> None: + """Stage the encode and decode halves of the remote separately. + + The encode path touches only ``encoder`` + ``quant_conv`` and the + decode path only ``post_quant_conv`` + ``decoder`` (verified against + the checkpoint's ``AutoencoderKLLegacy``: the one cross reference, + ``tiled_decode``'s ``getattr(self.encoder, "mask_enabled")``, is + short-circuited by ``self.training`` at inference). No storage is + shared across the two groups, so independent staging is exact. When + the checkpoint's structure is not discoverable, fall back to + whole-module staging (the previous single-stager behavior). + """ + part_names = ("encoder", "quant_conv", "post_quant_conv", "decoder") + if all(isinstance(getattr(self.model, name, None), nn.Module) for name in part_names): + self._encoder_stager = PinnedModuleStager( + [self.model.encoder, self.model.quant_conv], + device, + pin_memory=True, + ) + self._decoder_stager = PinnedModuleStager( + [self.model.post_quant_conv, self.model.decoder], + device, + pin_memory=True, + ) + return + self._stager = PinnedModuleStager( + self.remote, + device, + pin_memory=True, + ) + + def _part_stager(self, part: str) -> PinnedModuleStager | None: + if part == "encoder": + return self._encoder_stager + if part == "decoder": + return self._decoder_stager + raise ValueError(f"unknown video VAE part {part!r}") + + def _load_part_to_device(self, part: str) -> None: + stager = self._part_stager(part) + if stager is not None: + stager.load() + elif self._stager is not None: + self._stager.load() + else: + # No staged residency (the component lives on the device already + # or is fully CPU-resident): fall back to whole-module placement. + self.remote.to(self._device_target) + + def _offload_part_to_cpu(self, part: str) -> None: + stager = self._part_stager(part) + if stager is not None: + stager.offload() + return if self._stager is not None: + self._stager.offload() + return + self.remote.to("cpu") + self._release_component_cache() + + def _release_component_cache(self) -> None: + cache = getattr(self, "_omni_component_cache", None) + if cache is None: + torch.accelerator.empty_cache() + else: + cache.release_if_needed() + + def load_to_device(self) -> None: + if self._encoder_stager is not None: + self._encoder_stager.load() + self._decoder_stager.load() + elif self._stager is not None: self._stager.load() else: self.remote.to(self._device_target) def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None: self._omni_component_cache = cache - if self._stager is not None: - self._stager.set_cache_retention(cache) + for stager in (self._encoder_stager, self._decoder_stager, self._stager): + if stager is not None: + stager.set_cache_retention(cache) def offload_to_cpu(self) -> None: - if self._stager is not None: + if self._encoder_stager is not None: + self._encoder_stager.offload() + self._decoder_stager.offload() + elif self._stager is not None: self._stager.offload() else: self.remote.to("cpu") - cache = getattr(self, "_omni_component_cache", None) - if cache is None: - torch.accelerator.empty_cache() - else: - cache.release_if_needed() + self._release_component_cache() def set_parallel_size( self, @@ -650,6 +775,91 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): patch_size=(1, 2, 2), ).float() + def _stream_prepare_video_tensor( + self, + frames: np.ndarray, + device: torch.device, + ) -> torch.Tensor | None: + """Upload a uint8 ``(T, H, W, 3)`` video as one normalized FP32 tensor. + + The checkpoint's numpy path uploads the whole video as FP32, then + ``transform_tensor`` materializes a second normalized copy, then + ``encode_temporal`` pads through a full-video ``torch.cat``: three + resident pixel-scale copies on the device. Here the padding is + replicated on the host, a single preallocated ``(3, T', H, W)`` tensor + receives clip-by-clip uploads, and the ``÷255 → (x-mean)/std`` chain + runs in place on each clip's staging buffer in the same op order as + ``convert_numpy_to_tensor`` → ``transform_tensor``, so peak device + memory is the output tensor plus one clip instead of three copies. + + Returns ``None`` whenever the checkpoint contract this mirrors (uint8 + frames, ``clip_length`` alignment, processor ``transform`` constants) + is not discoverable; callers fall back to the legacy path. + """ + if frames.dtype != np.uint8 or frames.ndim != 4 or frames.shape[-1] != 3: + return None + if int(frames.shape[0]) == 0: + return None + model = self.model + clip_length = getattr(model, "clip_length", None) + transform = getattr(getattr(model, "processor", None), "transform", None) + mean = getattr(transform, "mean", None) + std = getattr(transform, "std", None) + if not isinstance(clip_length, int) or clip_length <= 0: + return None + if mean is None or std is None or len(mean) != 3 or len(std) != 3: + return None + # Mirror the checkpoint's temporal alignment so the device-side + # ``get_suitable_video_length`` trim is a no-op and the + # ``encode_temporal`` padding ``torch.cat`` never triggers. A + # trim-stable length is ``k * clip_length + tail`` (tail = frame + # overlap plus the isolated last frame); it also satisfies + # encode_temporal only when ``tail % clip_length == offset_frame``. + # With an asymmetric checkpoint (isolated last frame but no isolated + # first frame) no trim-stable length can satisfy encode_temporal, so + # the remote re-pads through a whole-video cat regardless and any + # host-side pad would be trimmed away before that -- keep the legacy + # frame count there instead of uploading dead frames. + isolated_first_frame = bool(getattr(model, "isolated_first_frame", False)) + frame_pre_padding = int(getattr(model, "frame_pre_padding", 0) or 0) + offset = 1 if isolated_first_frame and frame_pre_padding == 0 else 0 + processor = getattr(model, "processor", None) + tail = int(getattr(processor, "frame_overlap", 0) or 0) + if bool(getattr(processor, "isolated_last_frame", False)): + tail += 1 + num_frames = int(frames.shape[0]) + pad = 0 + if tail % clip_length == offset: + align = getattr(processor, "align_video_length", None) + if callable(align): + pad = max(0, int(align(num_frames, mode="pad", granularity="chunk"))) + else: + chunks = -(-(num_frames - tail) // clip_length) + pad = max(max(chunks, 1) * clip_length + tail - num_frames, 0) + if pad: + frames = np.concatenate([frames, np.repeat(frames[-1:], pad, axis=0)]) + num_frames += pad + mean_t = torch.as_tensor(mean, dtype=torch.float32, device=device).view(3, 1, 1, 1) + std_t = torch.as_tensor(std, dtype=torch.float32, device=device).view(3, 1, 1, 1) + height, width = int(frames.shape[1]), int(frames.shape[2]) + out = torch.empty( + (3, num_frames, height, width), + dtype=torch.float32, + device=device, + ) + for start in range(0, num_frames, clip_length): + end = min(start + clip_length, num_frames) + # Frame slices of a C-contiguous (T, H, W, 3) array stay + # contiguous, so each clip uploads at uint8 width (4x less host + # traffic than the legacy FP32 upload) before the in-place + # normalization chain. permute lands on the checkpoint's + # (3, T, H, W) layout directly. + chunk = torch.from_numpy(frames[start:end]).to(device=device) + chunk = chunk.permute(3, 0, 1, 2).to(torch.float32).div_(255.0) + chunk = chunk.sub_(mean_t).div_(std_t) + out[:, start:end].copy_(chunk) + return out + @torch.inference_mode() def encode_video( self, @@ -679,6 +889,14 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): for device in devices: with self.device_module.device(device): self.device_module.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED) + prepared = None + if isinstance(frames, np.ndarray) and not _legacy_encode_prep_enabled(): + prepared = self._stream_prepare_video_tensor(frames, parameter.device) + if prepared is not None: + # Tensor inputs already carry the checkpoint's expected + # (3, T, H, W) normalized-FP32 contract, so encode_videos + # skips its own convert/transform/pad whole-video copies. + frames = [prepared] latent = self.model.encode_videos( frames, use_fp16_latent=True, @@ -738,7 +956,12 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): raise RuntimeError("MiniMax H3 encode-only video VAE cannot decode latents") 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)) + if decoded.dtype == torch.uint8: + # The streaming uint8 write-back already ran the revert and the + # output quantizer inside write_part; nothing remains but the + # shape contract. + return self._normalize_decoded_frames(decoded) + return self._normalize_decoded_frames(self._revert_decoded_inplace(decoded)) def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor: channels = int(self.config_dict["latent_channels"]) @@ -755,7 +978,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): frames = frames.unsqueeze(0).transpose(1, 2) if frames.ndim != 5: raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}") - return frames.float() + return frames if frames.dtype == torch.uint8 else frames.float() # Chunks are reverted through the checkpoint's processor, which # denormalizes and clamps into the unit interval. @@ -791,6 +1014,33 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): with self._decode_tiling_context(z): decode_h3_chunks(self, z, on_chunk, group=group) + def _revert_decoded_inplace(self, decoded: torch.Tensor) -> torch.Tensor: + """In-place counterpart of the processor's ``revert_tensor``. + + The checkpoint version materializes a denormalized copy, a clamped + copy, and a contiguous copy of the whole decoded video -- three + pixel-scale tensors resident at the decode peak. Decoding owns the + tensor here, so the same op order (torchvision ``Normalize`` is + ``(x - mean) / std``, then ``clamp(0, 1)``) runs in place on the + original ``(B, C, T, H, W)`` layout, which is bit-identical + elementwise to normalizing the ``(b t) c h w`` rearrangement and + returns with zero whole-video copies. + + Falls back to ``processor.revert_tensor`` when the checkpoint's + denormalization constants are not discoverable. + """ + processor = getattr(self.model, "processor", None) + transform_rev = getattr(processor, "transform_rev", None) + mean = getattr(transform_rev, "mean", None) + std = getattr(transform_rev, "std", None) + if mean is None or std is None or len(mean) != 3 or len(std) != 3: + return processor.revert_tensor(decoded) + if bool(getattr(processor, "use_3d_conv", True)) and decoded.ndim == 4: + decoded = decoded.unsqueeze(2) + mean_t = torch.as_tensor(mean, dtype=decoded.dtype, device=decoded.device).view(1, 3, 1, 1, 1) + std_t = torch.as_tensor(std, dtype=decoded.dtype, device=decoded.device).view(1, 3, 1, 1, 1) + return decoded.sub_(mean_t).div_(std_t).clamp_(0.0, 1.0) + class MiniMaxH3AudioVAE(nn.Module): def __init__( diff --git a/vllm_omni/diffusion/models/minimax_h3/vae_temporal.py b/vllm_omni/diffusion/models/minimax_h3/vae_temporal.py new file mode 100644 index 000000000..6606ee5ef --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/vae_temporal.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Memory-optimized forks of the checkpoint's temporal VAE loops. + +Two method replacements installed on the checkpoint's ``AutoencoderKLLegacy`` +at adapter construction: + +* ``encode_temporal`` pads misaligned frame counts by concatenating repeated + last frames onto the *whole* video before the chunk loop (~4GB copy for a + 15s clip). The fork pads inside the tail chunk instead: the chunks see + bit-identical frames while the copy shrinks to a single clip. + +* ``_decode_temporal_streaming`` accumulates the whole decoded video in a + floating-point buffer that is denormalized, clamped, and quantized to + uint8 only after the final chunk lands. The fork fuses the adapter's + in-place revert and the output quantizer into ``write_part`` and + accumulates directly into a uint8 buffer -- the same per-element op order + (``(x - mean) / std -> clamp(0, 1) -> *255 -> round -> uint8``), so the + delivered video is bit-identical while the resident decode buffer shrinks + 4x. + +Set ``VLLM_OMNI_VAE_LEGACY_TEMPORAL=1`` to keep the checkpoint's methods. +""" + +from __future__ import annotations + +import os +from types import MethodType + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_LEGACY_TEMPORAL_ENV = "VLLM_OMNI_VAE_LEGACY_TEMPORAL" + + +def _legacy_temporal_enabled() -> bool: + value = os.environ.get(_LEGACY_TEMPORAL_ENV, "0") + return value.strip().lower() not in ("", "0", "false", "off") + + +def _encode_temporal_tail_pad(self, x): + """Drop-in fork of the checkpoint's ``encode_temporal``. + + Frame handling is identical to the original; only the padding site + moves from a whole-video ``torch.cat`` into the tail chunk. + """ + offset_frame = 1 if self.isolated_first_frame and self.frame_pre_padding == 0 else 0 + + orig_frames = int(x.shape[2]) + pad_size = (offset_frame - orig_frames) % self.clip_length + padded_frames = orig_frames + pad_size + num_chunks = (padded_frames - offset_frame) // self.clip_length + + z_list = [] + for i in range(num_chunks): + start_idx = i * self.clip_length + offset_frame + end_idx = (i + 1) * self.clip_length + offset_frame + clip_x = x[:, :, start_idx : min(end_idx, orig_frames), :, :] + if end_idx > orig_frames: + # The original pads the whole video up front by repeating the + # last frame; repeating it inside the tail chunk feeds the + # identical frames to the identical encodes. + clip_x = torch.cat( + [clip_x, clip_x[:, :, -1:].repeat(1, 1, end_idx - orig_frames, 1, 1)], + dim=2, + ) + + if self.isolated_key_frame: + key_frame = clip_x[:, :, :1, :, :] + z_key = self._adaptive_encode(key_frame) + + if clip_x.shape[2] > 1: + video_frames = clip_x[:, :, 1:, :, :] + z_video = self._adaptive_encode(video_frames) + z = torch.cat([z_key, z_video], dim=2) + else: + z = z_key + else: + z = self._adaptive_encode(clip_x) + + z_list.append(z) + + z = torch.cat(z_list, dim=2) + if self.token_drop > 0: + z = z[:, :, : -self.token_drop] + + if self.isolated_first_frame: + input_first_frame = x[:, :, :1, :, :] + z_first_frame = self._adaptive_encode(input_first_frame) + + if self.frame_pre_padding == 0: + z = torch.cat([z_first_frame, z], dim=2) + else: + z = torch.cat([z_first_frame, z[:, :, 1:, :, :]], dim=2) + + if self.isolated_last_frame: + # The original indexes into the padded video: pad frames replicate + # the last real frame, out-of-range indexes produce an empty slice + # (a no-op), and negative indexes wrap. Reproduce that mapping + # against the unpadded tensor instead of materializing the pad. + last_frame_idx = padded_frames - self.frame_drop + offset_frame + if last_frame_idx < 0: + last_frame_idx += padded_frames + if 0 <= last_frame_idx < padded_frames: + src_idx = min(last_frame_idx, orig_frames - 1) + input_last_frame = x[:, :, src_idx : src_idx + 1, :, :] + z_last_frame = self._adaptive_encode(input_last_frame) + z = torch.cat([z, z_last_frame], dim=2) + + return z + + +def _decode_temporal_streaming_uint8(self, z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype): + """Drop-in fork of the checkpoint's ``_decode_temporal_streaming``. + + The chunk loop, overlap blending, frame plan, and assertions are + identical to the original. ``write_part`` additionally runs the + denormalize -> clamp -> quantize chain in place (the same op order the + pipeline applies after decode) and accumulates into a uint8 buffer. + """ + total_frames, pad_frames, output_frames = self._decode_temporal_output_frame_plan( + z, z_head, z_tail, num_chunks, pad_tokens + ) + if output_frames <= 0: + raise ValueError( + f"decode_temporal streaming planned non-positive output_frames={output_frames} " + f"total_frames={total_frames} pad_frames={pad_frames}" + ) + + chunk_dec = self.tokens_chunk_size * self.vae_ratio_t + split_count = int(self.token_drop > 0) + 1 + dec = None + dec_overlap = None + norm = None + write_pos = 0 + logical_frames = 0 + dropped_frames = 0 + decoded_count = 0 + + def write_part(part): + nonlocal dec, dec_overlap, norm, write_pos, logical_frames, dropped_frames + part_frames = int(part.shape[2]) + if part_frames <= 0: + return + logical_frames += part_frames + if dec is None: + out_shape = list(part.shape) + out_shape[2] = output_frames + dec = torch.empty(out_shape, dtype=torch.uint8, device=part.device) + transform_rev = self.processor.transform_rev + mean = torch.as_tensor(transform_rev.mean, dtype=part.dtype, device=part.device).view(1, 3, 1, 1, 1) + std = torch.as_tensor(transform_rev.std, dtype=part.dtype, device=part.device).view(1, 3, 1, 1, 1) + norm = (mean, std) + # revert_tensor ((x - mean) / std, clamp) followed by the output + # quantizer (*255, round, uint8), in place on the part: the blended + # raw values are consumed before this point, and every write into + # ``dec`` goes through here, so the final buffer matches the legacy + # post-decode chain element for element. Denormalization and clamping + # run in the decoded dtype exactly like the legacy adapter; the + # quantizer then runs in FP32 because the legacy adapter returned + # ``frames.float()`` before the pipeline multiplied and rounded — + # rounding fp16/bf16 directly flips borderline pixels by one. + mean_t, std_t = norm + part.sub_(mean_t).div_(std_t).clamp_(0.0, 1.0) + part = part.float() + part.mul_(255.0).round_() + + 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 + dropped_frames += part_frames - copy_frames + + for i in range(num_chunks): + t_start_idx = i * self.tokens_chunk_size + t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap + clip_z = z[:, :, t_start_idx:t_end_idx, :, :] + + 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 = self._adaptive_decode(clip_z) + decoded_count += 1 + if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype: + clip_dec = clip_dec.to(temporal_cat_dtype) + if clip_dec.device != z.device: + clip_dec = clip_dec.to(z.device) + + dec_tail = None + if i == 0 and z_head is not None: + write_part(clip_dec[:, :, self.vae_ratio_t - 1 : self.vae_ratio_t, :, :]) + clip_dec = clip_dec[:, :, self.vae_ratio_t :, :, :] + + if i == num_chunks - 1 and z_tail is not None: + dec_tail = clip_dec[:, :, -1:, :, :] + clip_dec = clip_dec[:, :, : -self.vae_ratio_t, :, :] + + for j in range(split_count): + f_start_idx = j * chunk_dec + f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2]) + clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :] + clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding :, :, :] + + if j == 0: + if dec_overlap is not None: + clip_dec_chunk = self.blend(dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3) + dec_overlap = None + write_part(clip_dec_chunk) + else: + # Break the view's reference to the full decoded clip so earlier + # temporal chunks can be released before the final output exists. + dec_overlap = clip_dec_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) + + del clip_dec, clip_z + + if dec is None: + raise RuntimeError("decode_temporal streaming produced no output tensor") + if logical_frames != total_frames or dropped_frames != pad_frames or write_pos != output_frames: + raise RuntimeError( + "decode_temporal streaming frame plan mismatch: " + f"logical_frames={logical_frames} total_frames={total_frames} " + f"dropped_frames={dropped_frames} pad_frames={pad_frames} " + f"write_pos={write_pos} output_frames={output_frames}" + ) + + return dec + + +def _processor_denorm_available(model) -> bool: + transform_rev = getattr(getattr(model, "processor", None), "transform_rev", None) + mean = getattr(transform_rev, "mean", None) + std = getattr(transform_rev, "std", None) + return mean is not None and std is not None and len(mean) == 3 and len(std) == 3 + + +def install_temporal_stream_patches(model) -> None: + """Install the memory-optimized temporal forks on the checkpoint model. + + Each fork is installed only when the checkpoint contract it mirrors is + discoverable; otherwise the checkpoint's own method keeps running (the + adapter's post-decode revert already covers that path). + """ + if _legacy_temporal_enabled(): + logger.info("MiniMax-H3 VAE temporal stream patches disabled by %s", _LEGACY_TEMPORAL_ENV) + return + if callable(getattr(model, "encode_temporal", None)) and hasattr(model, "clip_length"): + model.encode_temporal = MethodType(_encode_temporal_tail_pad, model) + logger.info("MiniMax-H3 VAE encode_temporal tail-padding patch installed") + else: + logger.warning("MiniMax-H3 VAE encode_temporal contract not found; keeping checkpoint method") + if callable(getattr(model, "_decode_temporal_streaming", None)) and _processor_denorm_available(model): + model._decode_temporal_streaming = MethodType(_decode_temporal_streaming_uint8, model) + logger.info("MiniMax-H3 VAE streaming uint8 decode patch installed") + else: + logger.warning("MiniMax-H3 VAE streaming decode contract not found; keeping checkpoint method") + + +__all__ = ["install_temporal_stream_patches"]