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 index ba2142261..fcd090131 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py @@ -3,6 +3,7 @@ from __future__ import annotations +import hashlib import multiprocessing as mp import time from pathlib import Path @@ -14,6 +15,11 @@ import torch import torch.distributed as dist import torch.nn as nn +import vllm_omni.diffusion.models.minimax_h3.temporal_chunks as temporal_chunks_module +import vllm_omni.diffusion.models.minimax_h3.vae as vae_module +from vllm_omni.diffusion.models.minimax_h3.temporal_chunks import ( + MINIMAX_H3_RELEASED_BUNDLE_SHA256, +) from vllm_omni.diffusion.models.minimax_h3.vae import ( MiniMaxH3ChunkCallbackPeerError, MiniMaxH3ChunkedDecodeUnsupportedError, @@ -29,81 +35,145 @@ class _IdentityProcessor: return tensor -class _FakeVendorTemporalModel(nn.Module): - temporal_chunk_callback_api_version = 1 +class _FakeCompatTemporalModel(nn.Module): + use_3d_conv = True + clip_length = 17 + tokens_chunk_size = 5 + token_overlap = 2 + token_drop = 3 + vae_ratio_t = 4 + frame_pre_padding = 3 + frame_overlap = 5 + isolated_first_frame = False + isolated_last_frame = False vae_ratio = 1 def __init__( self, + latent_t: int, *, collective: bool = False, - metadata_mode: str = "valid", + tile_count: int = 2, ) -> None: super().__init__() - values = torch.linspace(0, 1, 124).view(1, 1, 124, 1, 1) - self.expected = values.expand(1, 3, -1, 1, 1).contiguous() - self.processor = _IdentityProcessor() self.parallel_tiling = collective self.collective = collective - self.metadata_mode = metadata_mode + self.tile_count = tile_count + pseudo_tokens = latent_t + self.token_drop + remainder = pseudo_tokens % self.tokens_chunk_size + self.pad_tokens = 0 if remainder == 0 else self.tokens_chunk_size - remainder + self.num_chunks = (pseudo_tokens + self.pad_tokens) // self.tokens_chunk_size - 1 + ideal_frames = self.num_chunks * 17 + 5 + output_frames = { + 33: 111, + 34: 115, + 35: 119, + 36: 120, + }.get(latent_t, ideal_frames) + values = torch.linspace(0, 1, ideal_frames).view(1, 1, ideal_frames, 1, 1) + self.ideal = values.expand(1, 3, -1, 1, 1).contiguous() + self.expected = self.ideal[:, :, :output_frames].clone() + self.processor = _IdentityProcessor() self.decode_base_calls = 0 - self.vendor_callback_calls = 0 + self.adaptive_decode_calls = 0 + self.blend_calls = 0 + self.eval() - @staticmethod - def split_tiles(size: int, tiled: bool): + def split_tiles(self, size: int, tiled: bool): del size, tiled - return [0, 1], None, None + return list(range(self.tile_count)), None, None - def decode_base( - self, - latent: torch.Tensor, - *, - temporal_chunk_callback=None, - ) -> torch.Tensor: + def decode_base(self, latent: torch.Tensor) -> torch.Tensor: del latent self.decode_base_calls += 1 - output = self.expected.clone() - if temporal_chunk_callback is None: - return output + return self.expected.clone() - frame_start = 0 - for chunk_index, chunk_size in enumerate([17] * 7 + [5]): - if self.collective: - collective_probe = torch.tensor([chunk_index], dtype=torch.int32) - dist.all_reduce(collective_probe) - self.vendor_callback_calls += 1 - metadata_index = chunk_index - metadata_total = 8 - metadata_start = frame_start - is_final = chunk_index == 7 - if self.metadata_mode == "gap" and chunk_index == 1: - metadata_start += 1 - elif self.metadata_mode == "bad_total" and chunk_index == 1: - metadata_total += 1 - elif self.metadata_mode == "no_final": - is_final = False - temporal_chunk_callback( - self.expected[:, :, frame_start : frame_start + chunk_size].clone(), - chunk_index=metadata_index, - total_chunks=metadata_total, - frame_start=metadata_start, - is_final=is_final, - ) - frame_start += chunk_size - return output + def _adaptive_decode(self, latent: torch.Tensor) -> torch.Tensor: + del latent + call_index = self.adaptive_decode_calls % self.num_chunks + self.adaptive_decode_calls += 1 + if self.collective and self.parallel_tiling: + collective_probe = torch.tensor([call_index], dtype=torch.int32) + dist.all_reduce(collective_probe) + decoded = torch.full((1, 3, 28, 1, 1), -1.0) + body_start = call_index * 17 + decoded[:, :, 3:20].copy_(self.ideal[:, :, body_start : body_start + 17]) + if call_index == self.num_chunks - 1: + decoded[:, :, 23:28].copy_(self.ideal[:, :, -5:]) + return decoded + + def _decode_temporal_output_frame_plan( + self, + latent: torch.Tensor, + head: torch.Tensor | None, + tail: torch.Tensor | None, + num_chunks: int, + pad_tokens: int, + ) -> tuple[int, int, int]: + del latent + assert head is None + assert tail is None + assert num_chunks == self.num_chunks + assert pad_tokens == self.pad_tokens + total_frames = num_chunks * 17 + 5 + pad_frames = total_frames - int(self.expected.shape[2]) + return total_frames, pad_frames, int(self.expected.shape[2]) + + def _decode_temporal_pad_frames( + self, + latent: torch.Tensor, + pad_tokens: int, + ) -> int: + del latent + assert pad_tokens == self.pad_tokens + return int(self.ideal.shape[2] - self.expected.shape[2]) + + def blend( + self, + previous: torch.Tensor, + current: torch.Tensor, + extent: int, + *, + dim: int, + ) -> torch.Tensor: + del previous + assert extent == self.frame_overlap + assert dim == -3 + self.blend_calls += 1 + return current -def _vae( +def _install_released_source_fingerprint() -> None: + temporal_chunks_module._source_bundle_sha256 = lambda model: ( + "fake/video_vae", + MINIMAX_H3_RELEASED_BUNDLE_SHA256, + ) + + +def _allow_released_source(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + temporal_chunks_module, + "_source_bundle_sha256", + lambda model: ( + "fake/video_vae", + MINIMAX_H3_RELEASED_BUNDLE_SHA256, + ), + ) + + +def _compat_vae( + latent_t: int, *, parallel_size: int = 1, collective: bool = False, - metadata_mode: str = "valid", -) -> tuple[MiniMaxH3VideoVAE, _FakeVendorTemporalModel]: + tile_count: int = 2, +) -> tuple[MiniMaxH3VideoVAE, _FakeCompatTemporalModel]: vae = object.__new__(MiniMaxH3VideoVAE) nn.Module.__init__(vae) - model = _FakeVendorTemporalModel( + model = _FakeCompatTemporalModel( + latent_t, collective=collective, - metadata_mode=metadata_mode, + tile_count=tile_count, ) vae.model = model vae.config_dict = { @@ -115,22 +185,57 @@ def _vae( return vae, model -def _latent() -> torch.Tensor: - return torch.zeros(1, 3, 37, 1, 1) +def _latent(latent_t: int = 37) -> torch.Tensor: + return torch.zeros(1, 3, latent_t, 1, 1) -def test_default_decode_path_does_not_request_vendor_chunks() -> None: - vae, model = _vae() +def test_default_decode_path_does_not_request_chunks() -> None: + vae, model = _compat_vae(37) output = vae.decode_latent(_latent()) torch.testing.assert_close(output, model.expected, rtol=0, atol=0) assert model.decode_base_calls == 1 - assert model.vendor_callback_calls == 0 + assert model.adaptive_decode_calls == 0 -def test_callback_chunks_are_owned_and_reconstruct_the_full_decode() -> None: - vae, model = _vae() +def test_default_decode_bypasses_compatibility_source_checks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + vae, model = _compat_vae(37) + monkeypatch.setattr( + temporal_chunks_module, + "_source_bundle_sha256", + lambda model: pytest.fail("default decode fingerprinted remote code"), + ) + + output = vae.decode_latent(_latent()) + + torch.testing.assert_close(output, model.expected, rtol=0, atol=0) + assert model.decode_base_calls == 1 + assert model.adaptive_decode_calls == 0 + + +@pytest.mark.parametrize( + ("latent_t", "expected_chunk_sizes"), + [ + (37, [17] * 7 + [5]), + (62, [17] * 12 + [5]), + (72, [17] * 14 + [5]), + (107, [17] * 21 + [5]), + (36, [17] * 7 + [1]), + (35, [17] * 7), + (34, [17] * 6 + [13]), + (33, [17] * 6 + [9]), + ], +) +def test_released_compatibility_chunks_reconstruct_full_decode( + monkeypatch: pytest.MonkeyPatch, + latent_t: int, + expected_chunk_sizes: list[int], +) -> None: + _allow_released_source(monkeypatch) + vae, model = _compat_vae(latent_t) chunks: list[torch.Tensor] = [] snapshots: list[torch.Tensor] = [] metadata: list[tuple[int, int, int, bool]] = [] @@ -150,18 +255,35 @@ def test_callback_chunks_are_owned_and_reconstruct_the_full_decode() -> None: snapshots.append(frames.clone()) metadata.append((chunk_index, total_chunks, frame_start, is_final)) - output = vae.decode_latent_with_chunks(_latent(), consume) + output = vae.decode_latent_with_chunks(_latent(latent_t), consume) torch.testing.assert_close(output, model.expected, rtol=0, atol=0) torch.testing.assert_close(torch.cat(chunks, dim=2), output, rtol=0, atol=0) + assert [int(chunk.shape[2]) for chunk in chunks] == expected_chunk_sizes + expected_metadata = [] + frame_start = 0 + for chunk_index, chunk_size in enumerate(expected_chunk_sizes): + expected_metadata.append( + ( + chunk_index, + len(expected_chunk_sizes), + frame_start, + chunk_index == len(expected_chunk_sizes) - 1, + ) + ) + frame_start += chunk_size + assert metadata == expected_metadata for chunk, snapshot in zip(chunks, snapshots, strict=True): torch.testing.assert_close(chunk, snapshot, rtol=0, atol=0) assert len({chunk.untyped_storage().data_ptr() for chunk in chunks}) == len(chunks) - assert metadata == [(index, 8, index * 17, index == 7) for index in range(8)] + assert model.adaptive_decode_calls == model.num_chunks -def test_callback_mutation_does_not_change_returned_full_output() -> None: - vae, model = _vae() +def test_compat_callback_mutation_does_not_change_returned_full_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _allow_released_source(monkeypatch) + vae, model = _compat_vae(37) def mutate(frames: torch.Tensor, **metadata: Any) -> None: del metadata @@ -172,29 +294,77 @@ def test_callback_mutation_does_not_change_returned_full_output() -> None: torch.testing.assert_close(output, model.expected, rtol=0, atol=0) -def test_callback_failure_is_deferred_until_vendor_decode_finishes() -> None: - vae, model = _vae() - callback_calls = 0 +def test_compat_source_fingerprint_is_cached_after_success( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = 0 - def fail(*args, **kwargs) -> None: - nonlocal callback_calls - del args, kwargs - callback_calls += 1 - raise LookupError("sink failed") + def fingerprint(model): + nonlocal calls + del model + calls += 1 + return "fake/video_vae", MINIMAX_H3_RELEASED_BUNDLE_SHA256 - with pytest.raises(LookupError, match="sink failed"): - vae.decode_latent_with_chunks(_latent(), fail) + monkeypatch.setattr( + temporal_chunks_module, + "_source_bundle_sha256", + fingerprint, + ) + vae, model = _compat_vae(37) - assert callback_calls == 1 - assert model.vendor_callback_calls == 8 - torch.testing.assert_close(vae.decode_latent(_latent()), model.expected, rtol=0, atol=0) + for _ in range(2): + output = vae.decode_latent_with_chunks( + _latent(), + lambda *args, **kwargs: None, + ) + torch.testing.assert_close(output, model.expected, rtol=0, atol=0) + + assert calls == 1 -def test_terminal_callback_failure_preserves_the_original_exception() -> None: - vae, model = _vae() - seen: list[tuple[int, int, int, bool]] = [] +def test_compat_callback_publishes_after_temporal_blending( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _allow_released_source(monkeypatch) + vae, model = _compat_vae(37) + chunks: list[torch.Tensor] = [] - def fail_final( + def mark_overlap( + previous: torch.Tensor, + current: torch.Tensor, + extent: int, + *, + dim: int, + ) -> torch.Tensor: + del previous + assert extent == 5 + assert dim == -3 + blended = current.clone() + blended[:, :, :extent].fill_(0.25) + return blended + + monkeypatch.setattr(model, "blend", mark_overlap) + output = vae.decode_latent_with_chunks( + _latent(), + lambda frames, **metadata: chunks.append(frames.clone()), + ) + + expected_overlap = torch.full_like(chunks[1][:, :, :5], 0.25) + torch.testing.assert_close(chunks[1][:, :, :5], expected_overlap, rtol=0, atol=0) + torch.testing.assert_close(output[:, :, 17:22], expected_overlap, rtol=0, atol=0) + torch.testing.assert_close(torch.cat(chunks, dim=2), output, rtol=0, atol=0) + + +@pytest.mark.parametrize("fail_final", [False, True]) +def test_compat_callback_failure_finishes_decode_and_recovers( + monkeypatch: pytest.MonkeyPatch, + fail_final: bool, +) -> None: + _allow_released_source(monkeypatch) + vae, model = _compat_vae(37) + seen: list[tuple[int, bool]] = [] + + def fail( frames: torch.Tensor, *, chunk_index: int, @@ -202,48 +372,165 @@ def test_terminal_callback_failure_preserves_the_original_exception() -> None: frame_start: int, is_final: bool, ) -> None: - del frames - seen.append((chunk_index, total_chunks, frame_start, is_final)) - if is_final: - raise LookupError("final sink failed") + del frames, total_chunks, frame_start + seen.append((chunk_index, is_final)) + if is_final == fail_final: + raise LookupError("compat sink failed") - with pytest.raises(LookupError, match="final sink failed"): - vae.decode_latent_with_chunks(_latent(), fail_final) + with pytest.raises(LookupError, match="compat sink failed"): + vae.decode_latent_with_chunks(_latent(), fail) - assert seen == [(index, 8, index * 17, index == 7) for index in range(8)] - assert model.vendor_callback_calls == 8 + assert model.adaptive_decode_calls == model.num_chunks + if fail_final: + assert seen == [(index, index == 7) for index in range(8)] + else: + assert seen == [(0, False)] + + recovered_chunks: list[torch.Tensor] = [] + recovered = vae.decode_latent_with_chunks( + _latent(), + lambda frames, **metadata: recovered_chunks.append(frames), + ) + torch.testing.assert_close(recovered, model.expected, rtol=0, atol=0) + torch.testing.assert_close( + torch.cat(recovered_chunks, dim=2), + model.expected, + rtol=0, + atol=0, + ) + assert model.adaptive_decode_calls == model.num_chunks * 2 + + +@pytest.mark.parametrize("drift", ["changed", "missing"]) +def test_compat_source_drift_fails_closed_before_decode( + monkeypatch: pytest.MonkeyPatch, + drift: str, +) -> None: + _allow_released_source(monkeypatch) + if drift == "changed": + monkeypatch.setattr( + temporal_chunks_module, + "_source_bundle_sha256", + lambda model: ("fake/video_vae", "0" * 64), + ) + else: + monkeypatch.setattr( + temporal_chunks_module, + "_source_bundle_sha256", + lambda model: (None, None), + ) + vae, model = _compat_vae(37) + + with pytest.raises( + MiniMaxH3ChunkedDecodeUnsupportedError, + match="42ed227e", + ): + vae.decode_latent_with_chunks(_latent(), lambda *args, **kwargs: None) + + assert model.adaptive_decode_calls == 0 + torch.testing.assert_close(vae.decode_latent(_latent()), model.expected, rtol=0, atol=0) + + +def test_uninspectable_compat_source_maps_to_unsupported( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + temporal_chunks_module.inspect, + "getsourcefile", + lambda object_type: (_ for _ in ()).throw(TypeError("no source")), + ) + vae, model = _compat_vae(37) + + with pytest.raises(MiniMaxH3ChunkedDecodeUnsupportedError, match="bundle=None"): + vae.decode_latent_with_chunks(_latent(), lambda *args, **kwargs: None) + + assert model.adaptive_decode_calls == 0 + + +def test_source_bundle_manifest_is_deterministic_and_excludes_init( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + package = tmp_path / "video_vae" + package.mkdir() + (package / "z.py").write_text("z = 1\n", encoding="utf-8") + (package / "klvae.py").write_text("k = 2\n", encoding="utf-8") + (package / "a.py").write_text("a = 3\n", encoding="utf-8") + (package / "__init__.py").write_text("ignored = 1\n", encoding="utf-8") + monkeypatch.setattr( + temporal_chunks_module.inspect, + "getsourcefile", + lambda object_type: str(package / "klvae.py"), + ) + + expected = hashlib.sha256() + for filename in ("a.py", "klvae.py", "z.py"): + expected.update(filename.encode()) + expected.update(b"\0") + expected.update((package / filename).read_bytes()) + expected.update(b"\0") + _, first = temporal_chunks_module._source_bundle_sha256(nn.Linear(1, 1)) + assert first == expected.hexdigest() + + (package / "__init__.py").write_text("ignored = 2\n", encoding="utf-8") + _, after_init = temporal_chunks_module._source_bundle_sha256(nn.Linear(1, 1)) + assert after_init == first + (package / "a.py").write_text("a = 4\n", encoding="utf-8") + _, after_change = temporal_chunks_module._source_bundle_sha256(nn.Linear(1, 1)) + assert after_change != first @pytest.mark.parametrize( - ("metadata_mode", "message"), + ("attribute", "value"), [ - ("gap", "discontinuous chunk metadata"), - ("bad_total", "discontinuous chunk metadata"), - ("no_final", "discontinuous chunk metadata"), + ("tokens_chunk_size", 6), + ("token_overlap", 1), + ("isolated_first_frame", True), + ("isolated_last_frame", True), ], ) -def test_adapter_rejects_a_broken_vendor_callback_contract( - metadata_mode: str, - message: str, +def test_compat_config_drift_fails_closed( + monkeypatch: pytest.MonkeyPatch, + attribute: str, + value: object, ) -> None: - vae, model = _vae(metadata_mode=metadata_mode) + _allow_released_source(monkeypatch) + vae, model = _compat_vae(37) + setattr(model, attribute, value) - with pytest.raises(RuntimeError, match=message): - vae.decode_latent_with_chunks(_latent(), lambda *args, **kwargs: None) - - assert model.vendor_callback_calls == 8 - - -def test_callback_request_fails_closed_without_the_vendor_api() -> None: - vae, model = _vae() - model.temporal_chunk_callback_api_version = -1 - - with pytest.raises(MiniMaxH3ChunkedDecodeUnsupportedError, match="official MiniMax"): + with pytest.raises( + MiniMaxH3ChunkedDecodeUnsupportedError, + match="released temporal configuration", + ): vae.decode_latent_with_chunks(_latent(), lambda *args, **kwargs: None) + assert model.adaptive_decode_calls == 0 torch.testing.assert_close(vae.decode_latent(_latent()), model.expected, rtol=0, atol=0) +@pytest.mark.parametrize("invalid", ["rank", "empty_time", "training", "non_3d"]) +def test_compat_structural_invalids_fail_before_decode( + monkeypatch: pytest.MonkeyPatch, + invalid: str, +) -> None: + _allow_released_source(monkeypatch) + vae, model = _compat_vae(37) + latent = _latent() + if invalid == "rank": + latent = latent[:, :, 0] + elif invalid == "empty_time": + latent = latent[:, :, :0] + elif invalid == "training": + model.train() + else: + model.use_3d_conv = False + + with pytest.raises(MiniMaxH3ChunkedDecodeUnsupportedError): + vae.decode_latent_with_chunks(latent, lambda *args, **kwargs: None) + + assert model.adaptive_decode_calls == 0 + + def _distributed_worker( rank: int, init_file: str, @@ -256,11 +543,105 @@ def _distributed_worker( world_size=2, ) try: - vae, model = _vae(parallel_size=2, collective=True) - vae._native_parallel_state = MethodType( - lambda self: {"sp_process_group": dist.group.WORLD}, - vae, + vae_module.get_data_parallel_world_size = lambda: 1 + vae_module.model_parallel_is_initialized = lambda: True + + def attach_group(vae: MiniMaxH3VideoVAE) -> None: + vae._native_parallel_state = MethodType( + lambda self: {"sp_process_group": dist.group.WORLD}, + vae, + ) + + _install_released_source_fingerprint() + source_vae, source_model = _compat_vae( + 37, + parallel_size=2, + collective=True, ) + attach_group(source_vae) + if rank == 1: + temporal_chunks_module._source_bundle_sha256 = lambda model: ( + "fake/video_vae", + "0" * 64, + ) + try: + source_vae.decode_latent_with_chunks( + _latent(), + (lambda *args, **kwargs: None) if rank == 0 else None, + ) + except BaseException as exc: # noqa: BLE001 + source_mismatch_type = type(exc).__name__ + else: + source_mismatch_type = "none" + source_decode_calls = source_model.adaptive_decode_calls + + _install_released_source_fingerprint() + planner_vae, planner_model = _compat_vae( + 37, + parallel_size=2, + collective=True, + ) + attach_group(planner_vae) + if rank == 1: + + def fail_plan(self, *args, **kwargs): + del self, args, kwargs + raise RuntimeError("rank-local planner failed") + + setattr( + planner_model, + "_decode_temporal_output_frame_plan", + MethodType(fail_plan, planner_model), + ) + try: + planner_vae.decode_latent_with_chunks( + _latent(), + (lambda *args, **kwargs: None) if rank == 0 else None, + ) + except BaseException as exc: # noqa: BLE001 + planner_failure_type = type(exc).__name__ + else: + planner_failure_type = "none" + planner_decode_calls = planner_model.adaptive_decode_calls + + dtype_vae, dtype_model = _compat_vae( + 37, + parallel_size=2, + collective=True, + ) + attach_group(dtype_vae) + temporal_chunks_module.resolve_minimax_h3_temporal_cat_dtype = ( + (lambda model: torch.float16) if rank == 0 else (lambda model: None) + ) + try: + dtype_vae.decode_latent_with_chunks( + _latent(), + (lambda *args, **kwargs: None) if rank == 0 else None, + ) + except BaseException as exc: # noqa: BLE001 + dtype_mismatch_type = type(exc).__name__ + else: + dtype_mismatch_type = "none" + dtype_decode_calls = dtype_model.adaptive_decode_calls + temporal_chunks_module.resolve_minimax_h3_temporal_cat_dtype = lambda model: None + + vae_module.get_data_parallel_world_size = lambda: 2 + dp_vae, dp_model = _compat_vae(37, parallel_size=2, collective=True) + attach_group(dp_vae) + try: + dp_vae.decode_latent_with_chunks( + _latent(), + (lambda *args, **kwargs: None) if rank == 0 else None, + ) + except BaseException as exc: # noqa: BLE001 + dp_vae_type = type(exc).__name__ + else: + dp_vae_type = "none" + dp_decode_calls = dp_model.adaptive_decode_calls + vae_module.get_data_parallel_world_size = lambda: 1 + + vae, model = _compat_vae(37, parallel_size=2, collective=True) + attach_group(vae) try: vae.decode_latent_with_chunks( @@ -298,7 +679,26 @@ def _distributed_worker( succeed if rank == 0 else None, ) - local_vae, local_model = _vae(parallel_size=1) + fallback_vae, fallback_model = _compat_vae( + 37, + parallel_size=2, + collective=True, + tile_count=1, + ) + attach_group(fallback_vae) + fallback_calls = 0 + + def consume_fallback(*args, **kwargs) -> None: + nonlocal fallback_calls + del args, kwargs + fallback_calls += 1 + + fallback_output = fallback_vae.decode_latent_with_chunks( + _latent(), + consume_fallback if rank == 0 else None, + ) + + local_vae, local_model = _compat_vae(37, parallel_size=1) local_calls = 0 def consume_local(*args, **kwargs) -> None: @@ -316,11 +716,26 @@ def _distributed_worker( result_queue.put( { "rank": rank, + "source_mismatch_type": source_mismatch_type, + "source_decode_calls": source_decode_calls, + "planner_failure_type": planner_failure_type, + "planner_decode_calls": planner_decode_calls, + "dtype_mismatch_type": dtype_mismatch_type, + "dtype_decode_calls": dtype_decode_calls, + "dp_vae_type": dp_vae_type, + "dp_decode_calls": dp_decode_calls, "invalid_owner_type": invalid_owner_type, "failure_type": failure_type, - "vendor_calls": model.vendor_callback_calls, + "adaptive_calls": model.adaptive_decode_calls, "success_calls": success_calls, "output_matches": torch.equal(output, model.expected), + "fallback_calls": fallback_calls, + "fallback_decode_calls": fallback_model.adaptive_decode_calls, + "fallback_output_matches": torch.equal( + fallback_output, + fallback_model.expected, + ), + "fallback_tiling_restored": fallback_model.parallel_tiling, "local_calls": local_calls, "local_output_matches": torch.equal(local_output, local_model.expected), "recovery": int(recovery.item()), @@ -367,13 +782,37 @@ def test_two_rank_callback_failure_is_symmetric_and_group_recovers( "LookupError", MiniMaxH3ChunkCallbackPeerError.__name__, ] + assert [result["source_mismatch_type"] for result in results] == [ + MiniMaxH3ChunkedDecodeUnsupportedError.__name__, + MiniMaxH3ChunkedDecodeUnsupportedError.__name__, + ] + assert [result["source_decode_calls"] for result in results] == [0, 0] + assert [result["planner_failure_type"] for result in results] == [ + MiniMaxH3ChunkedDecodeUnsupportedError.__name__, + "RuntimeError", + ] + assert [result["planner_decode_calls"] for result in results] == [0, 0] + assert [result["dtype_mismatch_type"] for result in results] == [ + MiniMaxH3ChunkedDecodeUnsupportedError.__name__, + MiniMaxH3ChunkedDecodeUnsupportedError.__name__, + ] + assert [result["dtype_decode_calls"] for result in results] == [0, 0] + assert [result["dp_vae_type"] for result in results] == [ + MiniMaxH3ChunkedDecodeUnsupportedError.__name__, + MiniMaxH3ChunkedDecodeUnsupportedError.__name__, + ] + assert [result["dp_decode_calls"] for result in results] == [0, 0] assert [result["invalid_owner_type"] for result in results] == [ "ValueError", "ValueError", ] - assert [result["vendor_calls"] for result in results] == [16, 16] + assert [result["adaptive_calls"] for result in results] == [14, 14] assert [result["success_calls"] for result in results] == [8, 0] assert all(result["output_matches"] for result in results) + assert [result["fallback_calls"] for result in results] == [8, 0] + assert [result["fallback_decode_calls"] for result in results] == [7, 7] + assert all(result["fallback_output_matches"] for result in results) + assert all(result["fallback_tiling_restored"] for result in results) # With VAE parallelism disabled, the caller still designates one request owner. assert [result["local_calls"] for result in results] == [8, 0] assert all(result["local_output_matches"] for result in results) 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..27f6d3457 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/temporal_chunks.py @@ -0,0 +1,408 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Source-gated temporal chunk compatibility for the released MiniMax H3 VAE.""" + +from __future__ import annotations + +import hashlib +import importlib +import inspect +from dataclasses import dataclass +from pathlib import Path +from typing import Protocol + +import torch +import torch.nn as nn + +MINIMAX_H3_RELEASED_REVISION = "42ed227ee7df40d41602854ae760620d6eb651fe" +# Compatibility pins detect upstream code drift; they are not a trust or +# security boundary for remote code execution. +MINIMAX_H3_RELEASED_BUNDLE_SHA256 = "30e64afbaa940696bf8bfe2c86003a4b1666c84b22fa3aa12403a3e6a03f705d" + +_RELEASED_TEMPORAL_CONFIG = { + "clip_length": 17, + "tokens_chunk_size": 5, + "token_overlap": 2, + "token_drop": 3, + "vae_ratio_t": 4, + "frame_pre_padding": 3, + "frame_overlap": 5, + "isolated_first_frame": False, + "isolated_last_frame": False, +} + + +class MiniMaxH3TemporalChunkCompatibilityError(RuntimeError): + """Raised when remote H3 code cannot use the local compatibility path.""" + + +class _RawTemporalChunkCallback(Protocol): + def __call__( + self, + frames: torch.Tensor, + *, + chunk_index: int, + total_chunks: int, + frame_start: int, + is_final: bool, + ) -> None: ... + + +@dataclass(slots=True) +class MiniMaxH3TemporalDecodePlan: + latent: torch.Tensor + num_chunks: int + total_frames: int + pad_frames: int + output_frames: int + chunk_frames: tuple[int, ...] + temporal_dtype: torch.dtype | None + + +class _RawChunkAssembly: + def __init__( + self, + *, + callback: _RawTemporalChunkCallback, + plan: MiniMaxH3TemporalDecodePlan, + ) -> None: + self.callback = callback + self.plan = plan + self.frames: torch.Tensor | None = None + self.write_pos = 0 + self.logical_frames = 0 + self.dropped_frames = 0 + self.chunk_index = 0 + self.pending_final: tuple[torch.Tensor, int, int] | None = None + self.contract_error: RuntimeError | None = None + + def write(self, part: torch.Tensor) -> None: + part_frames = int(part.shape[2]) + if part_frames <= 0: + return + self.logical_frames += part_frames + remaining = self.plan.output_frames - self.write_pos + copy_frames = min(part_frames, max(0, remaining)) + if copy_frames <= 0: + self.dropped_frames += part_frames + return + if self.chunk_index >= len(self.plan.chunk_frames): + self.contract_error = RuntimeError("MiniMax-H3 compatibility decode emitted an extra chunk") + else: + expected_frames = self.plan.chunk_frames[self.chunk_index] + if copy_frames != expected_frames and self.contract_error is None: + self.contract_error = RuntimeError( + "MiniMax-H3 compatibility chunk size mismatch: " + f"chunk={self.chunk_index}, frames={copy_frames}/{expected_frames}" + ) + + emitted = part[:, :, :copy_frames] + if self.frames is None: + output_shape = list(emitted.shape) + output_shape[2] = self.plan.output_frames + self.frames = torch.empty( + output_shape, + dtype=emitted.dtype, + device=emitted.device, + ) + frame_start = self.write_pos + frame_stop = frame_start + copy_frames + self.frames[:, :, frame_start:frame_stop].copy_(emitted) + if frame_stop == self.plan.output_frames: + self.pending_final = (emitted, self.chunk_index, frame_start) + elif self.contract_error is None: + self.callback( + emitted, + chunk_index=self.chunk_index, + total_chunks=len(self.plan.chunk_frames), + frame_start=frame_start, + is_final=False, + ) + self.chunk_index += 1 + self.write_pos = frame_stop + self.dropped_frames += part_frames - copy_frames + + def finish(self) -> torch.Tensor: + if self.contract_error is not None: + raise self.contract_error + if self.frames is None or self.pending_final is None: + raise RuntimeError("MiniMax-H3 compatibility decode produced no terminal chunk") + if ( + self.logical_frames != self.plan.total_frames + or self.dropped_frames != self.plan.pad_frames + or self.write_pos != self.plan.output_frames + or self.chunk_index != len(self.plan.chunk_frames) + ): + raise RuntimeError( + "MiniMax-H3 compatibility frame plan mismatch: " + f"logical={self.logical_frames}/{self.plan.total_frames}, " + f"dropped={self.dropped_frames}/{self.plan.pad_frames}, " + f"written={self.write_pos}/{self.plan.output_frames}, " + f"chunks={self.chunk_index}/{len(self.plan.chunk_frames)}" + ) + + final_frames, final_index, final_start = self.pending_final + self.callback( + final_frames, + chunk_index=final_index, + total_chunks=len(self.plan.chunk_frames), + frame_start=final_start, + is_final=True, + ) + return self.frames + + +def _source_bundle_sha256(model: nn.Module) -> tuple[str | None, str | None]: + try: + source = inspect.getsourcefile(type(model)) + except TypeError: + return None, None + if source is None: + return None, None + package_dir = Path(source).parent + files = sorted(path for path in package_dir.glob("*.py") if path.name != "__init__.py") + digest = hashlib.sha256() + try: + for path in files: + digest.update(path.name.encode()) + digest.update(b"\0") + digest.update(path.read_bytes()) + digest.update(b"\0") + except OSError: + return str(package_dir), None + return str(package_dir), digest.hexdigest() + + +def _validate_released_sources(model: nn.Module) -> None: + source_dir, actual = _source_bundle_sha256(model) + if actual != MINIMAX_H3_RELEASED_BUNDLE_SHA256: + raise MiniMaxH3TemporalChunkCompatibilityError( + "The vLLM-owned MiniMax-H3 chunk adapter supports only remote code " + f"from MiniMaxAI/MiniMax-H3@{MINIMAX_H3_RELEASED_REVISION}: " + f"bundle={actual!r} ({source_dir}). " + "The normal complete-output decode remains available." + ) + + +def _validate_released_config(model: nn.Module) -> None: + actual = {name: getattr(model, name, None) for name in _RELEASED_TEMPORAL_CONFIG} + if actual != _RELEASED_TEMPORAL_CONFIG: + raise MiniMaxH3TemporalChunkCompatibilityError( + "The vLLM-owned MiniMax-H3 chunk adapter requires the released temporal " + f"configuration {_RELEASED_TEMPORAL_CONFIG}, found {actual}. " + "The normal complete-output decode remains available." + ) + + +def _validate_compat_contract(model: nn.Module, latent: torch.Tensor) -> None: + if latent.ndim != 5 or int(latent.shape[2]) <= 0: + raise MiniMaxH3TemporalChunkCompatibilityError( + f"MiniMax-H3 video latent must have positive rank-5 time, got {tuple(latent.shape)}" + ) + if not bool(getattr(model, "use_3d_conv", False)): + raise MiniMaxH3TemporalChunkCompatibilityError("MiniMax-H3 temporal chunk decode requires the 3D decoder") + if bool(getattr(model, "training", False)): + raise MiniMaxH3TemporalChunkCompatibilityError("MiniMax-H3 temporal chunk decode is inference-only") + required_methods = ( + "_adaptive_decode", + "_decode_temporal_output_frame_plan", + "_decode_temporal_pad_frames", + "blend", + ) + missing = [name for name in required_methods if not callable(getattr(model, name, None))] + if missing: + raise MiniMaxH3TemporalChunkCompatibilityError( + "Loaded MiniMax-H3 VAE lacks required released-code methods: " + ", ".join(missing) + ) + + +def _truncate_chunk_frames( + logical_chunk_frames: list[int], + output_frames: int, +) -> tuple[int, ...]: + remaining = output_frames + result: list[int] = [] + for logical_frames in logical_chunk_frames: + emitted_frames = min(logical_frames, max(0, remaining)) + if emitted_frames > 0: + result.append(emitted_frames) + remaining -= emitted_frames + if remaining != 0: + raise RuntimeError(f"MiniMax-H3 compatibility chunk plan cannot cover output: remaining={remaining}") + return tuple(result) + + +def _build_decode_plan( + model: nn.Module, + latent: torch.Tensor, + temporal_dtype: torch.dtype | None, +) -> MiniMaxH3TemporalDecodePlan: + pseudo_total_tokens = int(latent.shape[2]) + int(model.token_drop) # type: ignore[attr-defined] + tokens_chunk_size = int(model.tokens_chunk_size) # type: ignore[attr-defined] + remainder = pseudo_total_tokens % tokens_chunk_size + pad_tokens = 0 if remainder == 0 else tokens_chunk_size - remainder + pseudo_total_tokens += pad_tokens + num_chunks = pseudo_total_tokens // tokens_chunk_size - int(model.token_drop > 0) # type: ignore[attr-defined] + if num_chunks <= 0: + raise MiniMaxH3TemporalChunkCompatibilityError( + f"MiniMax-H3 temporal chunk plan is empty for latent T={latent.shape[2]}" + ) + + temporal_latent = latent + if pad_tokens: + pad = temporal_latent[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1) + temporal_latent = torch.cat((temporal_latent, pad), dim=2) + total_frames, pad_frames, output_frames = model._decode_temporal_output_frame_plan( # type: ignore[attr-defined] + temporal_latent, + None, + None, + num_chunks, + pad_tokens, + ) + if output_frames <= 0: + raise MiniMaxH3TemporalChunkCompatibilityError( + f"MiniMax-H3 temporal chunk plan has {output_frames} output frames" + ) + + chunk_decoded_frames = tokens_chunk_size * int(model.vae_ratio_t) # type: ignore[attr-defined] + split_count = int(model.token_drop > 0) + 1 # type: ignore[attr-defined] + logical_chunk_frames: list[int] = [] + final_overlap_frames = 0 + for temporal_index in range(num_chunks): + token_start = temporal_index * tokens_chunk_size + token_stop = token_start + tokens_chunk_size + int(model.token_overlap) # type: ignore[attr-defined] + clip_tokens = max( + 0, + min(token_stop, int(temporal_latent.shape[2])) - min(token_start, int(temporal_latent.shape[2])), + ) + clip_frames = clip_tokens * int(model.vae_ratio_t) # type: ignore[attr-defined] + for split_index in range(split_count): + frame_start = split_index * chunk_decoded_frames + frame_stop = min(frame_start + chunk_decoded_frames, clip_frames) + part_frames = max( + 0, + frame_stop - frame_start - int(model.frame_pre_padding), # type: ignore[attr-defined] + ) + if split_index == 0: + if part_frames > 0: + logical_chunk_frames.append(part_frames) + else: + final_overlap_frames = part_frames + if final_overlap_frames > 0: + logical_chunk_frames.append(final_overlap_frames) + if sum(logical_chunk_frames) != int(total_frames): + raise RuntimeError( + "MiniMax-H3 compatibility planner disagrees with released remote code: " + f"frames={sum(logical_chunk_frames)}/{total_frames}" + ) + chunk_frames = _truncate_chunk_frames(logical_chunk_frames, int(output_frames)) + return MiniMaxH3TemporalDecodePlan( + latent=temporal_latent, + num_chunks=num_chunks, + total_frames=int(total_frames), + pad_frames=int(pad_frames), + output_frames=int(output_frames), + chunk_frames=chunk_frames, + temporal_dtype=temporal_dtype, + ) + + +def resolve_minimax_h3_temporal_cat_dtype(model: nn.Module) -> torch.dtype | None: + module = importlib.import_module(type(model).__module__) + resolver = getattr(module, "_resolve_temporal_cat_dtype", None) + if resolver is None: + return None + if not callable(resolver): + raise MiniMaxH3TemporalChunkCompatibilityError( + "Loaded MiniMax-H3 VAE has a non-callable temporal dtype resolver" + ) + try: + dtype = resolver() + except ValueError as exc: + raise MiniMaxH3TemporalChunkCompatibilityError(str(exc)) from None + if dtype is not None and not isinstance(dtype, torch.dtype): + raise MiniMaxH3TemporalChunkCompatibilityError( + f"Loaded MiniMax-H3 VAE returned invalid temporal dtype {dtype!r}" + ) + return dtype + + +def decode_minimax_h3_temporal_chunks_compat( + model: nn.Module, + plan: MiniMaxH3TemporalDecodePlan, + callback: _RawTemporalChunkCallback, +) -> torch.Tensor: + """Mirror the released Apache-2.0 H3 temporal assembler with callbacks. + + Keep this implementation and the source/config fingerprints above in sync + with ``MiniMaxAI/MiniMax-H3@42ed227e``. Unknown remote code fails closed + before this function is called. + """ + + chunk_decoded_frames = int(model.tokens_chunk_size) * int(model.vae_ratio_t) # type: ignore[attr-defined] + split_count = int(model.token_drop > 0) + 1 # type: ignore[attr-defined] + assembly = _RawChunkAssembly(callback=callback, plan=plan) + overlap: torch.Tensor | None = None + + for temporal_index in range(plan.num_chunks): + token_start = temporal_index * int(model.tokens_chunk_size) # type: ignore[attr-defined] + token_stop = token_start + int(model.tokens_chunk_size) + int(model.token_overlap) # type: ignore[attr-defined] + chunk_latent = plan.latent[:, :, token_start:token_stop] + decoded = model._adaptive_decode(chunk_latent) # type: ignore[attr-defined] + if plan.temporal_dtype is not None and decoded.dtype != plan.temporal_dtype: + decoded = decoded.to(plan.temporal_dtype) + if decoded.device != plan.latent.device: + decoded = decoded.to(plan.latent.device) + + for split_index in range(split_count): + frame_start = split_index * chunk_decoded_frames + frame_stop = min(frame_start + chunk_decoded_frames, int(decoded.shape[2])) + part = decoded[:, :, frame_start:frame_stop] + part = part[:, :, int(model.frame_pre_padding) :] # type: ignore[attr-defined] + if split_index == 0: + if overlap is not None: + part = model.blend( # type: ignore[attr-defined] + overlap, + part, + int(model.frame_overlap), # type: ignore[attr-defined] + dim=-3, + ) + overlap = None + assembly.write(part) + else: + overlap = part.contiguous() + del part + + if temporal_index == plan.num_chunks - 1 and overlap is not None: + assembly.write(overlap) + overlap = None + del decoded, chunk_latent + + return assembly.finish() + + +def prepare_minimax_h3_temporal_chunks_compat( + model: nn.Module, + latent: torch.Tensor, + *, + validate_sources: bool = True, +) -> MiniMaxH3TemporalDecodePlan: + """Build a validated plan without entering decoder collectives.""" + + if validate_sources: + _validate_released_sources(model) + _validate_released_config(model) + _validate_compat_contract(model, latent) + temporal_dtype = resolve_minimax_h3_temporal_cat_dtype(model) + return _build_decode_plan(model, latent, temporal_dtype) + + +__all__ = [ + "MINIMAX_H3_RELEASED_REVISION", + "MINIMAX_H3_RELEASED_BUNDLE_SHA256", + "MiniMaxH3TemporalDecodePlan", + "MiniMaxH3TemporalChunkCompatibilityError", + "decode_minimax_h3_temporal_chunks_compat", + "prepare_minimax_h3_temporal_chunks_compat", + "resolve_minimax_h3_temporal_cat_dtype", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index 7f658d8ae..5a2e08318 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -21,7 +21,11 @@ from vllm.logger import init_logger from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( DistributedVaeMixin, ) -from vllm_omni.diffusion.distributed.parallel_state import get_world_group +from vllm_omni.diffusion.distributed.parallel_state import ( + get_data_parallel_world_size, + get_world_group, + model_parallel_is_initialized, +) from vllm_omni.diffusion.offloader.module_residency import ( BoundedAllocatorCache, PinnedModuleStager, @@ -29,6 +33,12 @@ from vllm_omni.diffusion.offloader.module_residency import ( from .ops import install_h3_vae_optimizations from .packed_tokens import minimax_h3_patchify_video_latent +from .temporal_chunks import ( + MiniMaxH3TemporalChunkCompatibilityError, + MiniMaxH3TemporalDecodePlan, + decode_minimax_h3_temporal_chunks_compat, + prepare_minimax_h3_temporal_chunks_compat, +) MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42 MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000 @@ -60,7 +70,7 @@ class MiniMaxH3VideoChunkCallback(Protocol): class MiniMaxH3ChunkedDecodeUnsupportedError(RuntimeError): - """Raised when the loaded MiniMax H3 VAE lacks the vendor chunk API.""" + """Raised when loaded MiniMax H3 code cannot honor the chunk contract.""" class MiniMaxH3ChunkCallbackPeerError(RuntimeError): @@ -338,6 +348,10 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): def _chunk_process_group(self) -> dist.ProcessGroup | None: if not self.is_distributed_enabled(): return None + if model_parallel_is_initialized() and get_data_parallel_world_size() > 1: + raise MiniMaxH3ChunkedDecodeUnsupportedError( + "MiniMax-H3 temporal chunks do not yet support combining data parallelism with VAE patch parallelism" + ) 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") @@ -395,14 +409,48 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): raise callback_error raise MiniMaxH3ChunkCallbackPeerError("MiniMax-H3 video chunk callback failed on this VAE group's output rank") - def _validate_vendor_chunk_api(self) -> None: - version = getattr(self.model, "temporal_chunk_callback_api_version", None) - if version != 1: - raise MiniMaxH3ChunkedDecodeUnsupportedError( - "Loaded MiniMax-H3 VAE does not provide temporal_chunk_callback API v1 " - f"(found {version!r}); install a model snapshot containing the official " - "MiniMax temporal chunk callback" + @staticmethod + def _synchronize_chunk_preflight_error( + preflight_error: Exception | None, + temporal_dtype: torch.dtype | None, + *, + group: dist.ProcessGroup | None, + device: torch.device, + ) -> None: + if group is None: + if preflight_error is not None: + raise preflight_error + return + failed = torch.tensor( + [int(preflight_error is not None)], + dtype=torch.int32, + device=device, + ) + dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=group) + if int(failed.item()) == 0: + dtype_codes = { + None: 0, + torch.float16: 1, + torch.bfloat16: 2, + torch.float32: 3, + } + code = torch.tensor( + [dtype_codes.get(temporal_dtype, -1)], + dtype=torch.int32, + device=device, ) + lowest = code.clone() + highest = code.clone() + dist.all_reduce(lowest, op=dist.ReduceOp.MIN, group=group) + dist.all_reduce(highest, op=dist.ReduceOp.MAX, group=group) + if int(lowest.item()) == int(highest.item()) and int(lowest.item()) >= 0: + return + raise MiniMaxH3ChunkedDecodeUnsupportedError( + "MiniMax-H3 VAE ranks selected incompatible temporal concat dtypes" + ) + if preflight_error is not None: + raise preflight_error + raise MiniMaxH3ChunkedDecodeUnsupportedError("A peer VAE rank rejected the MiniMax-H3 temporal chunk preflight") @torch.inference_mode() def encode_image(self, image: Image.Image) -> torch.Tensor: @@ -506,11 +554,11 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): latent: torch.Tensor, chunk_callback: MiniMaxH3VideoChunkCallback | None, ) -> torch.Tensor: - """Decode a video while publishing vendor-committed temporal chunks. + """Decode a video while publishing committed temporal chunks. Every rank collaborating in the configured H3 VAE tile group must call this method, with ``chunk_callback`` supplied only on group rank 0. A - callback failure is deferred until the vendor decoder finishes its + callback failure is deferred until the compatibility decoder finishes its remaining collectives, then propagated to every group rank. With ``parallel_size=1`` the runtime designates each request's output owner by supplying a callback there and ``None`` on replica peers. @@ -522,8 +570,41 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): requested-size crop as the complete H3 pipeline output. """ - self._validate_vendor_chunk_api() group = self._chunk_process_group() + plan: MiniMaxH3TemporalDecodePlan | None = None + preflight_error: Exception | None = None + if latent.ndim != 5: + preflight_error = MiniMaxH3ChunkedDecodeUnsupportedError( + f"MiniMax-H3 video latent must be rank 5, got {tuple(latent.shape)}" + ) + else: + try: + denormalized = self._denormalize_latent(latent) + plan = prepare_minimax_h3_temporal_chunks_compat( + self.model, + denormalized, + validate_sources=not getattr( + self, + "_minimax_h3_chunk_sources_validated", + False, + ), + ) + self._minimax_h3_chunk_sources_validated = True + del denormalized + except MiniMaxH3TemporalChunkCompatibilityError as exc: + preflight_error = MiniMaxH3ChunkedDecodeUnsupportedError(str(exc)) + except Exception as exc: # noqa: BLE001 + # Every VAE rank must reach the preflight reduction even when + # one node's private planner or environment fails locally. + exc.__traceback__ = None + preflight_error = exc + self._synchronize_chunk_preflight_error( + preflight_error, + plan.temporal_dtype if plan is not None else None, + group=group, + device=latent.device, + ) + assert plan is not None is_output_rank = self._is_chunk_output_rank( chunk_callback, group=group, @@ -560,7 +641,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): or is_final != (chunk_index + 1 == total_chunks) ): raise RuntimeError( - "MiniMax-H3 vendor callback emitted discontinuous chunk metadata: " + "MiniMax-H3 temporal backend emitted discontinuous chunk metadata: " f"index={chunk_index}/{next_chunk_index}, " f"total={total_chunks}/{expected_total_chunks}, " f"frame_start={frame_start}/{next_frame_start}, " @@ -570,7 +651,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): frames = frames.clone(memory_format=torch.contiguous_format) frame_count = int(frames.shape[2]) if frame_count <= 0: - raise RuntimeError("MiniMax-H3 vendor callback emitted an empty chunk") + raise RuntimeError("MiniMax-H3 temporal backend emitted an empty chunk") next_chunk_index += 1 next_frame_start += frame_count saw_final = is_final @@ -589,9 +670,10 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): callback_error = exc with self._decode_tiling_context(latent): - decoded = self.model.decode_base( - self._denormalize_latent(latent), - temporal_chunk_callback=publish, + decoded = decode_minimax_h3_temporal_chunks_compat( + self.model, + plan, + publish, ) frames = self._normalize_decoded_frames(decoded) if ( @@ -600,7 +682,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): and (not saw_final or next_chunk_index != expected_total_chunks or next_frame_start != int(frames.shape[2])) ): callback_error = RuntimeError( - "MiniMax-H3 vendor callback did not reconstruct the full decode: " + "MiniMax-H3 temporal backend did not reconstruct the full decode: " f"final={saw_final}, chunks={next_chunk_index}/{expected_total_chunks}, " f"frames={next_frame_start}/{frames.shape[2]}" )