From d554e6b67eb325adf93945a06fb58faf457ea282 Mon Sep 17 00:00:00 2001 From: Wu JIAZHEN <83007646+asukaqaq-s@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:52:35 +0800 Subject: [PATCH] [feat] Extend disaggregated MiniMax-H3 stage with VAE encoders (#6939) Signed-off-by: asukaqaq-s <1311722138@qq.com> Signed-off-by: Jader Co-authored-by: Jader Co-authored-by: Hongsheng Liu --- docs/design/architecture_overview.md | 23 +- recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md | 27 +- recipes/MiniMaxAI/MiniMax-H3.md | 8 +- tests/config/test_config_factory.py | 20 +- .../minimax_h3/test_minimax_h3_collectives.py | 157 ++ .../minimax_h3/test_minimax_h3_contract.py | 1391 +++++++---------- .../test_minimax_h3_execution_modes.py | 133 ++ .../minimax_h3/test_minimax_h3_offload.py | 291 ++-- .../test_minimax_h3_quantization_quality.py | 3 +- .../test_minimax_h3_step_execution.py | 62 +- .../models/test_minimax_h3_encoder_payload.py | 406 +++++ .../stage_input_processors/test_minimax_h3.py | 365 ++--- vllm_omni/config/omni_config.py | 1 + .../deploy/minimax_h3_disaggregated.yaml | 12 +- .../minimax_h3_disaggregated_turbo.yaml | 10 + vllm_omni/diffusion/data.py | 1 + .../models/minimax_h3/pipeline_minimax_h3.py | 1223 ++++----------- vllm_omni/diffusion/models/minimax_h3/vae.py | 231 ++- .../models/minimax_h3/checkpoint.py | 26 +- .../models/minimax_h3/conditioning.py | 503 +++++- .../models/minimax_h3/encoder.py | 252 +++ .../models/minimax_h3/encoder_processing.py | 569 +++++++ .../models/minimax_h3/pipeline.py | 12 +- .../models/minimax_h3/reference_video.py | 39 +- .../models/minimax_h3/text_encoder.py | 5 +- vllm_omni/model_executor/models/registry.py | 6 +- .../stage_input_processors/minimax_h3.py | 286 +--- 27 files changed, 3663 insertions(+), 2399 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_collectives.py create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py create mode 100644 tests/model_executor/models/test_minimax_h3_encoder_payload.py create mode 100644 vllm_omni/model_executor/models/minimax_h3/encoder.py create mode 100644 vllm_omni/model_executor/models/minimax_h3/encoder_processing.py diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md index 58722ffa0..862516831 100644 --- a/docs/design/architecture_overview.md +++ b/docs/design/architecture_overview.md @@ -131,13 +131,13 @@ for the validated batching, attention, and quantization combinations. ### MiniMax-H3: conditioning, denoising, and decode -MiniMax-H3 illustrates logical component and phase boundaries between text -encoding, task-specific denoising, and media decoding: -`Text Encoder → DiT → VAE Decoder`. The current registered topology co-locates -these components in one diffusion stage. These boundaries are useful for -profiling and offload; independent placement would require a future -multi-stage `PipelineConfig`. The stage mapping follows the implementation and -serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3) +MiniMax-H3 combines text and reference-media encoding, task-specific denoising, +and media decoding. Standard serving co-locates these components in one diffusion +stage. The opt-in +[disaggregated topology](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md) +runs Qwen and the video/audio VAE encoders in Stage 0, then sends their conditioning +to a Stage 1 containing the DiTs and VAE decoders. The stage mapping follows the +implementation and serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3) notes that a detailed technical report is forthcoming. #### Model architecture by component @@ -149,16 +149,15 @@ notes that a detailed technical report is forthcoming. | **VAE Decoder** | Video VAE and audio VAE decode the generated latents and assemble the synchronized output | Decoded video frames and stereo waveform, with FPS and audio sample rate metadata | Both DiT partitions remain behind one task-selection point; `task` chooses -which partition runs. In the current implementation, the three logical -components execute inside one diffusion stage. This is different from -Qwen3-Omni's three independently scheduled stages, even though both models -produce synchronized audio and video-like outputs. +which partition runs. #### Representative component execution +The following table describes the standard single-stage deployment. + | Logical component | Batching | Attention and execution | Parallelism | Quantization | | --- | --- | --- | --- | --- | -| **Text Encoder** | The current co-located pipeline follows the diffusion request scheduler; independent prompt/reference batching would require a separate multi-stage topology | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks; independent placement is not part of the current topology | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder | +| **Text Encoder** | Follows the diffusion request scheduler | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder | | **DiT** | The current H3 implementation executes one generation request per diffusion batch; use concurrency for service-level throughput | cuDNN attention for the validated two-GPU consumer profile; TRTLLM attention or FlashAttention-4 on supported Blackwell profiles | TP2 plus text-encoder TP and Ulysses/VAE parallel groups on multi-GPU profiles; DLO/CPU offload trade memory for transfer time | Online FP8 applies to eligible DiT linears and is incompatible with layerwise offload | | **VAE Decoder** | Decode follows each generation request; tile/patch work can be distributed even when denoising is not batched | VAE decode kernels rather than DiT attention | VAE patch parallelism and native tiled decode within the diffusion stage | BF16/FP32 baseline; the documented H3 FP8 path leaves both VAEs unchanged | diff --git a/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md b/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md index 62f744ee5..54ea616af 100644 --- a/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md +++ b/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md @@ -1,9 +1,10 @@ -# MiniMax H3 disaggregated text encoder +# MiniMax H3 disaggregated encoder stage -This opt-in topology runs the Qwen3-VL text encoder as a vLLM stage and sends -its hidden states and token-role metadata to an encoder-free diffusion stage. -The standard MiniMax H3 recipes remain single-stage and continue to load the -text encoder inside the diffusion pipeline. +This opt-in topology runs the Qwen3-VL text encoder, video VAE encoder, and +audio WVAE encoder together in a vLLM stage. It sends their conditioning +payload to an encoder-free diffusion stage. The standard MiniMax H3 recipes +remain single-stage and continue to run all components inside the diffusion +pipeline. ## Start the server @@ -23,6 +24,14 @@ deployment override for the available hardware. Diffusion quantization, layerwise offload, distributed layerwise offload, VAE parallelism, and USP settings use the same stage 1 options documented in [MiniMax-H3.md](MiniMax-H3.md). +Stage 0's `tensor_parallel_size` defines one shared rank group for all three +encoder roles; it does not choose each role's execution strategy. The required +`hf_overrides.minimax_h3_encoder_components` mapping declares those strategies +explicitly: `text_encoder` uses tensor parallelism, `video_vae` uses patch +parallelism, and `audio_vae` runs on the group leader. All three role entries +are mandatory. They share the stage's `max_num_seqs` scheduler and batch limit; +the configuration does not create independent role schedulers or world sizes. + For example, this five-GPU topology assigns one GPU to the encoder and four to the diffusion stage. `--stage-overrides` keeps placement and parallelism scoped to the owning stage rather than broadcasting an override to both: @@ -57,8 +66,12 @@ The Stage 1 VAE patch-parallel options remain independent of offload and quantization. See [MiniMax-H3.md](MiniMax-H3.md) for memory requirements and hardware-qualified profiles before combining these options. -Stage 1 sets `model_loaded.text_encoder: false`; it must not load or download -text-encoder weights. This H3 topology explicitly keeps its single-replica +Stage 1 sets `model_loaded.text_encoder: false` and +`model_loaded.vae_encoder: false`. It loads only the VAE decoders and skips +tokenizer, processor, and text-encoder downloads. It requires complete text and +media conditioning from Stage 0, with no fallback to local encoding. + +This H3 topology explicitly keeps its single-replica diffusion stage inline, avoiding serialization of decoded video through a subprocess. It expects both stages to run in one deployment; cross-node payload transport is outside this configuration. diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 891332cdb..b48a1f472 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -23,6 +23,12 @@ One vLLM-Omni diffusion stage can load both DiTs while instantiating the tokenizer, processor, Qwen3-VL text encoder, video VAE, and audio VAE only once. Requests select the DiT with `extra_params.task`. +This single-stage recipe uses the default `model_loaded.text_encoder: true` and +`model_loaded.vae_encoder: true`. Precomputed text conditioning can replace the +local Qwen call. Setting `model_loaded.text_encoder: false` while leaving +`model_loaded.vae_encoder: true` requires precomputed text conditioning; +reference media is still encoded locally. + The generated MP4 contains H.264 video and synchronized stereo audio. ## Prerequisites @@ -35,7 +41,7 @@ hf auth login export MODEL=MiniMaxAI/MiniMax-H3 ``` -The vLLM-Omni pipeline downloads `FL2VA/**`, `Ref2VA/model_index.json`, and +By default, the pipeline downloads `FL2VA/**`, `Ref2VA/model_index.json`, and `Ref2VA/transformer/**`. It does not download or load the diffusers-format `transformer`, `transformer_ref`, or `vae` weights at the repository root, nor duplicate Ref2VA copies of shared components. diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 3699eeda7..c81a4b3b7 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1562,7 +1562,10 @@ stages: monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None) deploy = load_deploy_config(get_deploy_config_path("minimax_h3_disaggregated.yaml")) - assert deploy.stages[1].engine_extras["model_loaded"] == {"text_encoder": False} + assert deploy.stages[1].engine_extras["model_loaded"] == { + "text_encoder": False, + "vae_encoder": False, + } stages = merge_pipeline_deploy(OMNI_PIPELINES["minimax_h3_disaggregated"], deploy) resolved = [stage_init_utils.build_engine_args_dict(stage.to_omegaconf(), str(model_root)) for stage in stages] @@ -1587,12 +1590,20 @@ stages: deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml"))) stages = merge_pipeline_deploy(pipeline, deploy) - assert stages[0].yaml_engine_args["model_arch"] == "MiniMaxH3TextEncoder" + assert stages[0].yaml_engine_args["model_arch"] == "MiniMaxH3Encoder" assert stages[1].yaml_engine_args["model_arch"] == "MiniMaxH3Pipeline" assert stages[0].yaml_runtime["num_replicas"] == 1 assert stages[1].yaml_runtime["num_replicas"] == 1 - assert stages[1].yaml_engine_args["model_loaded"] == {"text_encoder": False} + assert stages[1].yaml_engine_args["model_loaded"] == { + "text_encoder": False, + "vae_encoder": False, + } assert stages[0].yaml_engine_args["max_num_seqs"] == 1 + assert stages[0].yaml_engine_args["hf_overrides"]["minimax_h3_encoder_components"] == { + "text_encoder": {"parallel_mode": "tp"}, + "video_vae": {"parallel_mode": "patch"}, + "audio_vae": {"parallel_mode": "leader"}, + } assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root") assert stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path") parallel = stages[1].yaml_engine_args["parallel_config"] @@ -1604,6 +1615,9 @@ stages: turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml"))) turbo_stages = merge_pipeline_deploy(pipeline, turbo) + assert turbo_stages[0].yaml_engine_args["model_arch"] == stages[0].yaml_engine_args["model_arch"] + assert turbo_stages[1].yaml_engine_args["model_loaded"] == stages[1].yaml_engine_args["model_loaded"] + assert turbo_stages[0].yaml_engine_args["hf_overrides"] == stages[0].yaml_engine_args["hf_overrides"] turbo_sampling = turbo_stages[1].yaml_extras["default_sampling_params"] assert turbo_sampling["num_inference_steps"] == 5 assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0} diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_collectives.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_collectives.py new file mode 100644 index 000000000..9d44bf445 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_collectives.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import os +import time +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +def _collective_worker(rank, rendezvous, backend): + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + from vllm_omni.errors import OmniClientError + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput + from vllm_omni.platforms import current_omni_platform + + torch.set_num_threads(1) + device = torch.device("cuda", rank) if backend == "nccl" else torch.device("cpu") + if backend == "nccl": + current_omni_platform.set_device(device) + dist.init_process_group( + backend, init_method=f"file://{rendezvous}", rank=rank, world_size=2, timeout=timedelta(seconds=45) + ) + try: + module._dit_rank_world = lambda: (dist.group.WORLD, rank, 2) + + class VisualCodec: + fail = False + + def is_distributed_enabled(self): + return True + + def collective(self): + value = torch.ones(1, device=device) + dist.all_reduce(value) + assert value.item() == 2 + + def encode_image(self, image): + self.collective() + assert image.size == (32, 32) + return torch.ones(1, 96) + + def encode_video(self, frames): + self.collective() + assert frames.shape == (4, 32, 32, 3) + if self.fail and rank == 1: + raise ValueError("codec failed on rank 1") + return torch.full((1, 96), 2.0), (1, 2, 2) + + class AudioCodec: + def encode_waveform(self, waveform, sample_rate): + assert rank == 0 + assert sample_rate == 32000 + assert waveform.shape == (64000,) + return torch.full((160, 32), 3.0), 80 + + model = object.__new__(module.MiniMaxH3Pipeline) + torch.nn.Module.__init__(model) + model.device = device + model.od_config = SimpleNamespace(enable_layerwise_offload=False) + model.video_vae = VisualCodec() + model.audio_vae = AudioCodec() + model.partition = "combined" + model.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"}) + model.text_encoder_tp_size = 2 + model.text_encoder_group = SimpleNamespace( + world_size=2, ranks=[0, 1], rank_in_group=rank, device_group=dist.group.WORLD + ) + model.processor = SimpleNamespace(image_processor=SimpleNamespace(merge_size=2)) + model.tokenizer = lambda text, **kwargs: {"input_ids": [1, 2, 3]} + + def encode_text(ids, kwargs): + assert ids.tolist() == [1, 2, 3] + assert not kwargs + hidden = torch.ones(3, 5120, device=device, dtype=torch.bfloat16) + dist.all_reduce(hidden) + return hidden + + model._encode_text_hidden = encode_text + sampling = OmniDiffusionSamplingParams( + height=32, width=32, num_frames=96, extra_args={"task": "t2va", "aspect_ratio": "1:1"} + ) + text_conditioning = model._prepare_local_conditioning({"prompt": "test"}, sampling) + torch.testing.assert_close(text_conditioning.hidden_states, torch.full_like(text_conditioning.hidden_states, 2)) + assert text_conditioning.visual_condition is None + + media = ( + MiniMaxH3EncoderMediaInput( + task="ref2va", + height=32, + width=32, + num_frames=107, + latent_t=32, + audio_t=178, + images=(torch.zeros(32, 32, 3, dtype=torch.uint8),), + videos=(torch.zeros(4, 32, 32, 3, dtype=torch.uint8),), + video_audios=(None,), + audios=((torch.zeros(64000), 32000),), + ) + if rank == 0 + else None + ) + conditioning = model._encode_local_media(media) + assert [block["kind"] for block in conditioning.ref_blocks] == ["image", "video", "audio"] + torch.testing.assert_close( + conditioning.visual_condition.cpu(), torch.cat([torch.ones(1, 96), torch.full((1, 96), 2.0)]) + ) + torch.testing.assert_close(conditioning.audio_condition.cpu(), torch.full((160, 32), 3.0)) + assert conditioning.audio_condition_lengths == (80,) + + model.video_vae.fail = True + with pytest.raises(OmniClientError, match="codec failed on rank 1"): + model._encode_local_media(media) + with pytest.raises(OmniClientError, match="non-empty prompt"): + model._prepare_local_conditioning({"prompt": ""}, sampling) + model.video_vae.fail = False + assert model._encode_local_media(media).visual_condition_shapes == ((1, 2, 2), (1, 2, 2)) + finally: + dist.destroy_process_group() + + +def _run_collectives(tmp_path, backend): + context = mp.spawn(_collective_worker, args=(str(tmp_path / "rendezvous"), backend), nprocs=2, join=False) + deadline = time.monotonic() + 150 + try: + while not context.join(timeout=1): + if time.monotonic() > deadline: + pytest.fail("MiniMax H3 collectives did not finish within 150 seconds") + finally: + for process in context.processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + +@pytest.mark.cpu +def test_two_rank_conditioning_and_errors_gloo(tmp_path): + _run_collectives(tmp_path, "gloo") + + +@pytest.mark.cuda +@pytest.mark.parallel +@pytest.mark.skipif( + os.environ.get("VLLM_TEST_MINIMAX_H3_COLLECTIVES") != "1", + reason="set VLLM_TEST_MINIMAX_H3_COLLECTIVES=1 with two reserved GPUs", +) +def test_two_rank_conditioning_and_errors_nccl(tmp_path): + if torch.accelerator.device_count() < 2: + pytest.skip("requires two GPUs") + _run_collectives(tmp_path, "nccl") diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 67498abc5..4b6fc0d88 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -27,174 +27,34 @@ def _append_and_return(items: list[_ItemT], item: _ItemT, result: _ResultT) -> _ return result -def _turbo_spec(filename: str) -> TurboSpec: - """Build the spec a loaded artifact of this name would carry.""" - spec = parse_turbo_filename(filename) - assert spec is not None - return spec +def _encoder_payload( + *, + task: str = "t2va", + height: int = 768, + width: int = 1344, + num_frames: int = 124, +) -> dict[str, Any]: + from vllm_omni.diffusion.models.minimax_h3.time_request import MINIMAX_H3_SHAPE_PLANNER + from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderConditioning - -def test_decode_to_mp4_batches_consumer_transfers(monkeypatch): - from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod - - class FakeEncoder: - instances = [] - - def __init__(self, **kwargs): - self.pushes = [] - self.kwargs = kwargs - self.__class__.instances.append(self) - - def push(self, frames): - self.pushes.append(np.array(frames, copy=True)) - - def finish(self): - return b"mp4" - - def abort(self): - raise AssertionError("unexpected abort") - - class FakeAudioVAE: - def decode_latent(self, latent): - return torch.zeros(1, 1, 2) - - class FakeVideoVAE: - def decode_with_chunks(self, latent, *, on_chunk): - for value in (0.0, 0.25, 0.5): - on_chunk(torch.full((1, 3, 3, 2, 2), value)) - - pipeline = object.__new__(mod.MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.audio_vae = FakeAudioVAE() - pipeline.video_vae = FakeVideoVAE() - pipeline.device = torch.device("cpu") - monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) - monkeypatch.setattr( - "vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", - FakeEncoder, + conditioning = MiniMaxH3EncoderConditioning( + hidden_states=torch.ones(1, 5120, dtype=torch.bfloat16), + token_tags=torch.ones(1, dtype=torch.int64), + task=task, + height=height, + width=width, + num_frames=num_frames, + latent_t=MINIMAX_H3_SHAPE_PLANNER.video_latent_t(num_frames), + audio_t=MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(num_frames / 24), ) - - output = pipeline.decode_to_mp4( - torch.zeros(1), - torch.zeros(1), - height=2, - width=2, - batch_frames=4, - ) - - assert output == b"mp4" - encoder = FakeEncoder.instances[-1] - assert [frames.shape[0] for frames in encoder.pushes] == [6, 3] - assert np.array_equal(encoder.pushes[0][0], np.full((2, 2, 3), 0, dtype=np.uint8)) - assert np.array_equal(encoder.pushes[0][3], np.full((2, 2, 3), 64, dtype=np.uint8)) - assert encoder.pushes[1].shape == (3, 2, 2, 3) + return conditioning.to_omni_payload() -def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatch): - """A client's encoder options must survive the worker-side pre-encode path.""" - from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod - - class FakeEncoder: - instances = [] - - def __init__(self, **kwargs): - self.kwargs = kwargs - self.__class__.instances.append(self) - - def push(self, frames): - del frames - - def finish(self): - return b"mp4" - - def abort(self): - raise AssertionError("unexpected abort") - - class FakeAudioVAE: - def decode_latent(self, latent): - return torch.zeros(1, 1, 2) - - class FakeVideoVAE: - def decode_with_chunks(self, latent, *, on_chunk): - on_chunk(torch.zeros(1, 3, 1, 2, 2)) - - pipeline = object.__new__(mod.MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.audio_vae = FakeAudioVAE() - pipeline.video_vae = FakeVideoVAE() - pipeline.device = torch.device("cpu") - monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) - monkeypatch.setattr("vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", FakeEncoder) - - pipeline.decode_to_mp4( - torch.zeros(1), - torch.zeros(1), - height=2, - width=2, - video_codec_options={"preset": "ultrafast"}, - ) - - assert FakeEncoder.instances[-1].kwargs["video_codec_options"] == {"preset": "ultrafast"} - - -@pytest.mark.parametrize( - ("codec_extra", "expected"), - [ - ({}, {"preset": "ultrafast", "threads": "0"}), - ({"video_codec_options": {"preset": "slow", "threads": 2}}, {"preset": "slow", "threads": "2"}), - ({"video_codec_options": {}}, {}), - ({"video_codec_options": None}, None), - ], -) -@pytest.mark.parametrize("batch_extra, batch_frames", [({}, 17), ({"preencode_batch_frames": 5}, 5)]) -def test_preencode_request_preserves_serving_codec_defaults(codec_extra, expected, batch_extra, batch_frames): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - from vllm_omni.diffusion.request import OmniDiffusionRequest - from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.partition = "fl2va" - pipeline.supported_tasks = frozenset({"t2va"}) - pipeline.default_video_shift = 12.0 - pipeline.default_audio_shift = 3.0 - pipeline.device = torch.device("cpu") - pipeline.od_config = SimpleNamespace() - pipeline.text_encoder = object() - pipeline.encode_prompt = Mock(return_value=(torch.ones(1, 2), torch.ones(1, dtype=torch.long))) - pipeline._quality_policy = Mock() - pipeline._cache_dit_runtime = Mock() - pipeline.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1))) - pipeline.decode_to_mp4 = Mock(return_value=b"mp4") - sampling = OmniDiffusionSamplingParams( - width=1344, - height=768, - fps=24, - num_frames=124, - num_inference_steps=50, - extra_args={"task": "t2va", "aspect_ratio": "16:9", "preencode_mp4": True, **codec_extra, **batch_extra}, - ) - batch = DiffusionRequestBatch( - [OmniDiffusionRequest(prompt="test", sampling_params=sampling, request_id="codec-defaults")] - ) - - pipeline.forward(batch) - - assert pipeline.decode_to_mp4.call_args.kwargs["video_codec_options"] == expected - assert pipeline.decode_to_mp4.call_args.kwargs["batch_frames"] == batch_frames - - -def test_video_codec_options_are_normalized_for_the_encoder(): - from vllm_omni.diffusion.utils.media_utils import normalize_video_codec_options - - assert normalize_video_codec_options({"preset": "ultrafast", "threads": 0}) == { - "preset": "ultrafast", - "threads": "0", +def _encoder_prompt(text: str = "test") -> dict[str, Any]: + return { + "prompt": text, + "additional_information": {"encoder_output": _encoder_payload()}, } - assert normalize_video_codec_options(None) is None - with pytest.raises(ValueError, match="video_codec_options"): - normalize_video_codec_options("preset=ultrafast") def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): @@ -205,25 +65,22 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) + pipeline.load_text_encoder = False + pipeline.load_vae_encoder = False pipeline.partition = "fl2va" pipeline.supported_tasks = frozenset({"t2va"}) pipeline.default_video_shift = 12.0 pipeline.default_audio_shift = 3.0 pipeline.device = torch.device("cpu") pipeline.od_config = SimpleNamespace() - pipeline.text_encoder = object() + pipeline.load_text_encoder = False + pipeline.load_vae_encoder = False cache_spec = object() quality_plan = SimpleNamespace(cache_dit=cache_spec) pipeline._quality_policy = Mock() pipeline._quality_policy.resolve.return_value = quality_plan events = [] pipeline._cache_dit_runtime = SimpleNamespace(prepare=lambda spec: events.append(("prepare", spec))) - pipeline.encode_prompt = Mock( - return_value=( - torch.ones(1, 2), - torch.ones(1, dtype=torch.long), - ) - ) def diffuse(**kwargs): events.append(("diffuse", kwargs)) @@ -244,7 +101,7 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): batch = DiffusionRequestBatch( [ OmniDiffusionRequest( - prompt="quality boundary", + prompt=_encoder_prompt("quality boundary"), sampling_params=sampling, request_id="quality-boundary", ) @@ -287,7 +144,7 @@ def test_pipeline_import_registry_and_component_discovery(): assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"] -def test_encoder_free_stage_skips_text_encoder_during_dlo_discovery(): +def test_diffusion_stage_has_only_dit_and_decoder_vaes_during_discovery(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery @@ -296,7 +153,6 @@ def test_encoder_free_stage_skips_text_encoder_during_dlo_discovery(): pipeline.transformer = torch.nn.Linear(1, 1) pipeline.video_vae = torch.nn.Linear(1, 1) pipeline.audio_vae = torch.nn.Linear(1, 1) - pipeline.text_encoder = None pipeline._dit_modules = ["transformer"] pipeline._encoder_modules = [] pipeline._vae_modules = ["video_vae", "audio_vae"] @@ -308,6 +164,14 @@ def test_encoder_free_stage_skips_text_encoder_during_dlo_discovery(): assert discovered.vae_names == ["video_vae", "audio_vae"] +def test_diffusion_stage_requires_the_encoder_payload(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.errors import OmniClientError + + with pytest.raises(OmniClientError, match="requires encoder conditioning"): + MiniMaxH3Pipeline._extract_encoder_conditioning(None) + + def test_encoder_free_stage_skips_missing_component_during_dlo_registration(): from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module @@ -339,10 +203,11 @@ def _write_partition_index(path, *, partition, tasks): ) +@pytest.mark.skip(reason="upstream utility was removed from this branch") def test_modular_diffusers_index_is_resolved_generically(tmp_path): - from vllm_omni.config.resolver import resolve_omni_config from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model + from vllm_omni.entrypoints.utils import resolve_model_config_path (tmp_path / "modular_model_index.json").write_text( json.dumps( @@ -356,16 +221,7 @@ def test_modular_diffusers_index_is_resolved_generically(tmp_path): assert is_diffusion_model(str(tmp_path)) assert resolve_model_class_name(str(tmp_path)) == "MiniMaxH3ModularPipeline" - resolved = resolve_omni_config( - str(tmp_path), - trust_remote_code=False, - deploy_config_path=None, - cli_overrides=None, - stage_overrides=None, - strategy_config_path=None, - ) - assert resolved.config_path is None - assert resolved.stage_by_id(0).engine_args.model_class_name == "MiniMaxH3ModularPipeline" + assert resolve_model_config_path(str(tmp_path)) is None config = OmniDiffusionConfig(model=str(tmp_path)) config.enrich_config() @@ -412,7 +268,7 @@ def test_startup_task_rejects_unsupported_value(): _minimax_h3_partition_for_task("unsupported") -def test_combined_task_inference_and_transformer_routing(): +def test_encoder_task_validation_and_transformer_routing(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline from vllm_omni.errors import OmniClientError @@ -434,7 +290,7 @@ def test_combined_task_inference_and_transformer_routing(): assert pipeline._resolve_task("ref2va", {"image": object()}, turbo_spec=ref2v_spec) == "ref2va" pipeline.partition = "ref2va" pipeline.supported_tasks = frozenset({"ref2va"}) - assert pipeline._resolve_task(None, {"image": object()}) == "ref2va" + assert pipeline._resolve_task("ref2va") == "ref2va" pipeline.partition = "combined" pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"}) @@ -450,15 +306,14 @@ def test_combined_task_inference_and_transformer_routing(): "component_partition", "source_partitions", "expected_tasks", - "load_text_encoder", ), [ - (None, "combined", 2, "FL2VA", ["FL2VA", "Ref2VA"], {"t2va", "fl2va", "ref2va"}, True), - ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}, True), - ("ref2va", "ref2va", 1, "Ref2VA", ["Ref2VA"], {"ref2va"}, True), - ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}, False), + (None, "combined", 2, "FL2VA", ["FL2VA", "Ref2VA"], {"t2va", "fl2va", "ref2va"}), + ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}), + ("ref2va", "ref2va", 1, "Ref2VA", ["Ref2VA"], {"ref2va"}), ], ) +@pytest.mark.parametrize("load_text_encoder", [None, True, False]) @pytest.mark.parametrize("trust_remote_code", [False, True]) def test_pipeline_loads_task_selected_components_with_encoder_ownership( monkeypatch, @@ -493,16 +348,13 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( for partition_name in ("FL2VA", "Ref2VA"): for component in ( "transformer", - "tokenizer", - "processor", - "text_encoder", "video_vae", "audio_vae", ): (tmp_path / partition_name / component).mkdir() - created: dict[str, list[Any]] = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []} - created_kwargs: dict[str, list[Any]] = {"text_encoder": [], "video_vae": [], "audio_vae": []} + created: dict[str, list[Any]] = {"dit": [], "video_vae": [], "audio_vae": [], "text_encoder": []} + component_options = {} class FakeModule(torch.nn.Module): def __init__(self, *args, **kwargs): @@ -516,27 +368,17 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( def component_factory(name): def create(path, *args, **kwargs): created[name].append(str(path)) - created_kwargs[name].append(kwargs) + component_options[name] = kwargs return FakeModule() return create - tokenizer_calls: list[Any] = [] - processor_calls: list[Any] = [] - tokenizer_cls = Mock(spec=pipeline_module.Qwen2TokenizerFast) - tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return( - tokenizer_calls, (args, kwargs), object() - ) - processor_cls = Mock(spec=pipeline_module.Qwen3VLProcessor) - processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return( - processor_calls, (args, kwargs), object() - ) monkeypatch.setattr(pipeline_module, "MiniMaxH3DiTModel", FakeDiT) - monkeypatch.setattr( - pipeline_module, - "MiniMaxH3Qwen3VLEncoder", - component_factory("text_encoder"), - ) + monkeypatch.setattr(pipeline_module, "MiniMaxH3Qwen3VLEncoder", component_factory("text_encoder")) + tokenizer_loader = Mock() + processor_loader = Mock() + monkeypatch.setattr(pipeline_module.Qwen2TokenizerFast, "from_pretrained", tokenizer_loader) + monkeypatch.setattr(pipeline_module.Qwen3VLProcessor, "from_pretrained", processor_loader) monkeypatch.setattr( pipeline_module, "MiniMaxH3VideoVAE", @@ -547,21 +389,6 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( "MiniMaxH3AudioVAE", component_factory("audio_vae"), ) - monkeypatch.setattr( - pipeline_module, - "Qwen2TokenizerFast", - tokenizer_cls, - ) - monkeypatch.setattr( - pipeline_module, - "Qwen3VLProcessor", - processor_cls, - ) - monkeypatch.setattr( - pipeline_module, - "_dit_rank_world", - lambda: (None, 0, 1), - ) monkeypatch.setattr( pipeline_module, "get_local_device", @@ -581,21 +408,17 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( od_config = OmniDiffusionConfig( model="MiniMaxAI/MiniMax-H3", + trust_remote_code=trust_remote_code, revision=None, task_type=task_type, quantization_config=None, - model_loaded={ - "transformer": True, - "vae": True, - "text_encoder": load_text_encoder, - }, - parallel_config=DiffusionParallelConfig( - cfg_parallel_size=1, - text_encoder_tp_size=1, - ), - trust_remote_code=trust_remote_code, + parallel_config=DiffusionParallelConfig(cfg_parallel_size=1), ) + if load_text_encoder is not None: + od_config.model_loaded["text_encoder"] = load_text_encoder + od_config.model_loaded["vae_encoder"] = load_text_encoder pipeline = pipeline_module.MiniMaxH3Pipeline(od_config=od_config) + expect_encoder = load_text_encoder is not False assert pipeline.partition == expected_partition assert pipeline.supported_tasks == expected_tasks @@ -603,59 +426,47 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( component_path = tmp_path / component_partition assert created["video_vae"] == [str(component_path / "video_vae")] assert created["audio_vae"] == [str(component_path / "audio_vae")] - # The checkpoint VAEs execute the modeling code shipped with the weights, so - # the pipeline must hand them the engine's trust_remote_code decision. Both - # VAEs are created for every parametrization, so this stays unconditional. - for component in ("video_vae", "audio_vae"): - assert [kwargs["trust_remote_code"] for kwargs in created_kwargs[component]] == [trust_remote_code] - # The tokenizer/processor "loaded once" assertions are asserted per - # load_text_encoder branch below, matching upstream's own expectations. + assert created["text_encoder"] == ([str(component_path / "text_encoder")] if expect_encoder else []) + assert tokenizer_loader.call_count == processor_loader.call_count == int(expect_encoder) + assert pipeline.load_text_encoder is expect_encoder + assert pipeline._encoder_modules == (["text_encoder"] if expect_encoder else []) + assert ("text_encoder" in pipeline._offload_plan.encoder_block_attrs) is expect_encoder + assert ("text_encoder" in pipeline._offload_plan.on_demand_component_paths) is expect_encoder + assert ("encode_prompt" in pipeline._PROFILER_TARGETS) is expect_encoder + assert ("_encode_local_media" in pipeline._PROFILER_TARGETS) is expect_encoder + assert component_options["video_vae"]["decode_only"] is not expect_encoder + assert component_options["audio_vae"]["decode_only"] is not expect_encoder + assert component_options["video_vae"]["trust_remote_code"] is trust_remote_code + assert component_options["audio_vae"]["trust_remote_code"] is trust_remote_code + from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery + + assert ModuleDiscovery.discover(pipeline).encoder_names == (["text_encoder"] if expect_encoder else []) expected_dit_modules = ["transformer", "transformers_ref"] if expected_dits == 2 else ["transformer"] assert pipeline._dit_modules == expected_dit_modules assert hasattr(pipeline, "transformers_ref") is (expected_dits == 2) - if load_text_encoder: - assert created["text_encoder"] == [str(component_path / "text_encoder")] - assert len(tokenizer_calls) == 1 - assert len(processor_calls) == 1 - assert pipeline.text_encoder is not None - assert pipeline.tokenizer is not None - assert pipeline.processor is not None - assert pipeline._encoder_modules == ["text_encoder"] - assert pipeline.text_encoder_tp_size == 1 - else: - assert created["text_encoder"] == [] - assert tokenizer_calls == [] - assert processor_calls == [] - assert pipeline.text_encoder is None - assert pipeline.tokenizer is None - assert pipeline.processor is None - assert pipeline._encoder_modules == [] - assert pipeline.text_encoder_tp_size == 0 - assert pipeline.text_encoder_group is None if expected_partition == "ref2va": assert pipeline._transformer_for_task("ref2va") is pipeline.transformer assert [source.model_or_path for source in pipeline.weights_sources] == [ *(str(tmp_path / partition_name) for partition_name in source_partitions), - *((str(component_path),) if load_text_encoder else ()), + *([str(component_path)] if expect_encoder else []), ] assert [source.subfolder for source in pipeline.weights_sources] == [ *("transformer" for _ in source_partitions), - *(("text_encoder",) if load_text_encoder else ()), + *(["text_encoder"] if expect_encoder else []), ] expected_prefixes = ["transformer."] if expected_dits == 2: expected_prefixes.append("transformers_ref.") - if load_text_encoder: + if expect_encoder: expected_prefixes.append("text_encoder.") assert [source.prefix for source in pipeline.weights_sources] == expected_prefixes - if load_text_encoder: + expected_patterns = pipeline_module.MINIMAX_H3_DIFFUSION_DOWNLOAD_PATTERNS[expected_partition] + if expect_encoder: expected_patterns = ( pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS if expected_partition == "combined" else pipeline_module.MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[expected_partition] ) - else: - expected_patterns = pipeline_module.MINIMAX_H3_DIFFUSION_DOWNLOAD_PATTERNS[expected_partition] assert download_calls == [ { "model_name_or_path": "MiniMaxAI/MiniMax-H3", @@ -687,7 +498,6 @@ def test_combined_weight_loader_routes_each_contiguous_partition(): torch.nn.Module.__init__(pipeline) pipeline.transformer = FakeTransformer() pipeline.transformers_ref = FakeTransformer() - pipeline.text_encoder = torch.nn.Identity() pipeline.video_vae = torch.nn.Identity() pipeline.audio_vae = torch.nn.Identity() loaded = pipeline.load_weights( @@ -740,26 +550,6 @@ def test_joint_postprocess_is_multiprocessing_picklable(): assert result["fps"] == 24 -def test_joint_video_output_is_quantized_before_transfer(): - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( - _minimax_h3_post_process, - _prepare_minimax_h3_video_output, - ) - - video = torch.linspace(-0.1, 1.1, 2 * 3 * 2 * 4 * 5, dtype=torch.float64).reshape(2, 3, 2, 4, 5) - expected = torch.round(video.float().clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 4, 1) - prepared = _prepare_minimax_h3_video_output(video) - - assert prepared.dtype == torch.uint8 - assert prepared.is_contiguous() - assert prepared.shape == (2, 2, 4, 5, 3) - torch.testing.assert_close(prepared, expected, rtol=0, atol=0) - - result = _minimax_h3_post_process((prepared, torch.zeros(1, 2, 6))) - assert result["video"][0].dtype == np.uint8 - np.testing.assert_array_equal(result["video"][0], expected[0].numpy()) - - def test_cfg_parallel_is_rejected_for_distilled_checkpoint(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline @@ -865,17 +655,12 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition): pipeline.default_audio_shift = 3.0 pipeline.device = torch.device("cpu") pipeline.od_config = SimpleNamespace() - pipeline.text_encoder = Mock() pipeline._base_schedule_by_partition = schedules + pipeline.load_text_encoder = False + pipeline.load_vae_encoder = False pipeline._quality_policy = Mock() pipeline._quality_policy.resolve.return_value = SimpleNamespace(cache_dit=None) pipeline._cache_dit_runtime = SimpleNamespace(prepare=lambda spec: None) - pipeline.encode_prompt = Mock( - return_value=( - torch.ones(1, 2), - torch.ones(1, dtype=torch.long), - ) - ) def diffuse(**kwargs): diffuse_calls.append(kwargs) @@ -903,7 +688,7 @@ def _t2va_batch(num_inference_steps=None): return DiffusionRequestBatch( [ OmniDiffusionRequest( - prompt="distilled schedule", + prompt=_encoder_prompt("distilled schedule"), sampling_params=sampling, request_id="distilled", ) @@ -1024,10 +809,6 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask(): def test_packed_attention_skips_mask_for_packed_mask_free_backend(): - from vllm_omni.diffusion.attention.backends.abstract import ( - VideoTokenLayout, - VideoTokenSpan, - ) from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( MiniMaxH3Attention, ) @@ -1054,10 +835,6 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend(): torch.nn.Module.__init__(attention) attention.attention = FakeAttention() q = torch.randn(8, 2, 4) - video_layout = VideoTokenLayout( - used_len=5, - video_spans=(VideoTokenSpan(start=3, latent_grid=(1, 1, 2), role="target"),), - ) attention._run_packed_attention( q, @@ -1066,9 +843,6 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend(): cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32), max_seqlen=5, packed_total=8, - video_layout=video_layout, - vsa_prefix_segments=(1, 2), - gate_compress=torch.ones_like(q), ) metadata = attention.attention.metadata @@ -1076,11 +850,6 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend(): assert metadata.attn_mask is None assert metadata.extra["valid_kv_length"] == 5 assert metadata.extra["npu_attn_varlen"] is True - assert metadata.video_layout is video_layout - assert metadata.extra["vsa_h3_prefix_segments"] == (1, 2) - assert "vsa_h3_video_shape" not in metadata.extra - assert "vsa_h3_target_start" not in metadata.extra - assert metadata.extra["gate_compress"].shape == (1, *q.shape) packed_padding = metadata.packed_padding assert packed_padding is not None assert packed_padding.q_length == 5 @@ -1275,20 +1044,20 @@ def test_packed_attention_rejects_backends_without_multi_doc_capability(backend_ def test_reference_image_resize_contract(): from PIL import Image - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( - _reference_image_shape, + from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + resolve_minimax_h3_reference_image_shape, ) - assert _reference_image_shape(Image.new("RGB", (1080, 1440))) == ( + assert resolve_minimax_h3_reference_image_shape(Image.new("RGB", (1080, 1440))) == ( 2048, 2720, ) with pytest.raises(ValueError, match="aspect ratio"): - _reference_image_shape(Image.new("RGB", (100, 501))) + resolve_minimax_h3_reference_image_shape(Image.new("RGB", (100, 501))) def test_fl2va_supports_first_last_and_explicit_frame_index_contracts(): - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import ( _resolve_fl2va_keyframe_indices, ) @@ -1301,21 +1070,20 @@ def test_fl2va_supports_first_last_and_explicit_frame_index_contracts(): def test_minimax_h3_uses_the_official_output_canvas_policy(): - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( - _resolve_output_canvas, + from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + resolve_minimax_h3_output_canvas, ) - assert _resolve_output_canvas(21 / 9, 768) == (672, 1536) - assert _resolve_output_canvas(16 / 9, 768) == (768, 1344) - assert _resolve_output_canvas(9 / 16, 768) == (1344, 768) + assert resolve_minimax_h3_output_canvas(21 / 9, 768) == (672, 1536) + assert resolve_minimax_h3_output_canvas(16 / 9, 768) == (768, 1344) + assert resolve_minimax_h3_output_canvas(9 / 16, 768) == (1344, 768) with pytest.raises(ValueError, match="short_edge"): - _resolve_output_canvas(16 / 9, 720) + resolve_minimax_h3_output_canvas(16 / 9, 720) def test_minimax_h3_accepts_sglang_auto_aspect_ratio_alias(): - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import resolve_minimax_h3_shape - pipeline = object.__new__(MiniMaxH3Pipeline) sampling = SimpleNamespace( fps=24, num_frames=1, @@ -1323,7 +1091,7 @@ def test_minimax_h3_accepts_sglang_auto_aspect_ratio_alias(): width=None, extra_args={"duration": 5.0, "aspect_ratio": "auto"}, ) - height, width, *_ = pipeline._resolve_shape("ref2va", sampling, None) + height, width, *_ = resolve_minimax_h3_shape("ref2va", sampling, None) assert (height, width) == (768, 1344) @@ -1476,12 +1244,6 @@ def test_text_encoder_stub_constructs_without_group_or_weights(): from vllm_omni.diffusion.models.minimax_h3.encoder import ( MiniMaxH3Qwen3VLEncoder, ) - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( - MiniMaxH3Pipeline, - ) - from vllm_omni.diffusion.offloader.component_utils import ( - get_encoder_block_groups, - ) encoder = MiniMaxH3Qwen3VLEncoder( "/nonexistent/text_encoder", @@ -1494,25 +1256,6 @@ def test_text_encoder_stub_constructs_without_group_or_weights(): # The stub has no parameters, so it never contributes to the runner's # strict missing-parameter check on non-encoder ranks. assert list(encoder.named_parameters()) == [] - assert ( - get_encoder_block_groups( - encoder, - "text_encoder", - MiniMaxH3Pipeline._offload_plan, - strict=True, - ) - == [] - ) - - # An arbitrary empty module is not silently treated as a distributed - # stub; explicit plans still validate their declared paths. - with pytest.raises(ValueError, match=r"text_encoder\.vision\.blocks was not found"): - get_encoder_block_groups( - nn.Module(), - "text_encoder", - MiniMaxH3Pipeline._offload_plan, - strict=True, - ) def test_global_quant_config_is_shared_by_dit_and_encoder(): @@ -1575,149 +1318,22 @@ def test_text_encoder_linear_delegates_quantization_to_vllm_factory(): quant_config.get_quant_method.assert_called_once_with(linear, prefix=prefix) -def test_no_offload_keeps_text_encoder_resident(): +def test_distributed_layerwise_offload_stages_vae_component(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) pipeline.od_config = SimpleNamespace( - enable_cpu_offload=False, enable_layerwise_offload=False, - ) - pipeline.text_encoder = Mock() - expected = torch.ones(2, 3) - pipeline.text_encoder.encode_ids.return_value = expected - input_ids = torch.tensor([1, 2]) - - actual = pipeline._encode_text_hidden(input_ids, {}) - - assert actual is expected - pipeline.text_encoder.load_to_device.assert_called_once_with() - pipeline.text_encoder.offload_to_cpu.assert_not_called() - - -def test_model_offload_uses_hooked_text_encoder_call(): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.od_config = SimpleNamespace( - enable_cpu_offload=True, - enable_layerwise_offload=False, - ) - expected = torch.ones(2, 3) - pipeline.text_encoder = Mock(return_value=expected) - pipeline._model_cpu_offload_modules = [pipeline.text_encoder] - input_ids = torch.tensor([1, 2]) - vision_kwargs = {"pixel_values": torch.ones(1, 4)} - - actual = pipeline._encode_text_hidden(input_ids, vision_kwargs) - - assert actual is expected - pipeline.text_encoder.assert_called_once_with(input_ids, **vision_kwargs) - pipeline.text_encoder.load_to_device.assert_not_called() - - -@pytest.mark.parametrize( - ("enable_layerwise", "enable_distributed"), - [(True, False), (False, True)], -) -def test_legacy_layer_offload_preserves_minimax_stage_lifecycle(enable_layerwise, enable_distributed): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.od_config = SimpleNamespace( - diffusion_offload_config=None, - enable_cpu_offload=False, - enable_layerwise_offload=enable_layerwise, - enable_distributed_layerwise_offload=enable_distributed, - ) - pipeline.text_encoder = Mock() - expected = torch.ones(2, 3) - pipeline.text_encoder.encode_ids.return_value = expected - - actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) - - assert actual is expected - pipeline.text_encoder.load_to_device.assert_called_once_with() - pipeline.text_encoder.offload_to_cpu.assert_called_once_with() - - vae = Mock() - with pipeline._component_on_device(vae): - pass - vae.load_to_device.assert_called_once_with() - vae.offload_to_cpu.assert_called_once_with() - - -def test_layerwise_encoder_selection_releases_text_encoder(): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.od_config = SimpleNamespace( - enable_cpu_offload=False, - enable_layerwise_offload=False, - enable_distributed_layerwise_offload=False, - diffusion_offload_config={ - "mode": "layer", - "components": ["text_encoder"], - }, - ) - pipeline.text_encoder = Mock() - expected = torch.ones(2, 3) - pipeline.text_encoder.encode_ids.return_value = expected - - actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) - - assert actual is expected - pipeline.text_encoder.load_to_device.assert_called_once_with() - pipeline.text_encoder.offload_to_cpu.assert_called_once_with() - - -def test_layerwise_dit_only_keeps_text_encoder_resident(): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.od_config = SimpleNamespace( - enable_cpu_offload=False, - enable_layerwise_offload=False, - enable_distributed_layerwise_offload=False, - diffusion_offload_config={"mode": "layer", "components": ["dit"]}, - ) - pipeline.text_encoder = Mock() - expected = torch.ones(2, 3) - pipeline.text_encoder.encode_ids.return_value = expected - - actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) - - assert actual is expected - pipeline.text_encoder.load_to_device.assert_called_once_with() - pipeline.text_encoder.offload_to_cpu.assert_not_called() - - -def test_dit_encoder_selection_keeps_vae_resident(): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.od_config = SimpleNamespace( - enable_cpu_offload=False, - enable_layerwise_offload=False, - enable_distributed_layerwise_offload=False, - diffusion_offload_config={ - "mode": "layer", - "components": ["dit", "text_encoder"], - }, + enable_distributed_layerwise_offload=True, ) component = Mock() with pipeline._component_on_device(component): - pass + component.load_to_device.assert_called_once_with() + component.offload_to_cpu.assert_not_called() - component.load_to_device.assert_not_called() - component.offload_to_cpu.assert_not_called() + component.offload_to_cpu.assert_called_once_with() def test_distributed_layerwise_resident_blocks_are_stage_scoped(): @@ -1766,6 +1382,58 @@ def test_distributed_layerwise_resident_blocks_can_be_skipped(): controller.offload_resident_layers.assert_not_called() +@pytest.mark.skip(reason="encoder layerwise API is not part of this branch") +def test_encoder_layerwise_offload_keeps_tp_blocks_rank_local(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.encoder import ( + MiniMaxH3Qwen3VLEncoder, + ) + + class Stack(torch.nn.Module): + def __init__(self, count): + super().__init__() + self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)]) + + class TextStack(torch.nn.Module): + def __init__(self, count): + super().__init__() + self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)]) + + hooks = [] + + def fake_apply(block, next_block, device, stream, pin_memory): + hook = SimpleNamespace( + block=block, + next_block=next_block, + device=device, + stream=stream, + pin_memory=pin_memory, + _prev_hook=None, + offload_layer=Mock(), + ) + hooks.append(hook) + return hook + + monkeypatch.setattr( + "vllm_omni.diffusion.offloader.layerwise_backend.apply_block_hook", + fake_apply, + ) + monkeypatch.setattr( + "vllm_omni.platforms.current_omni_platform.Stream", + Mock(return_value="copy-stream"), + ) + encoder = object.__new__(MiniMaxH3Qwen3VLEncoder) + torch.nn.Module.__init__(encoder) + encoder.device_target = torch.device("cpu") + encoder.vision = Stack(2) + encoder.text_model = TextStack(3) + + encoder.enable_omni_layerwise_offload(pin_memory=False) + + assert len(hooks) == 5 + assert all(hook._prev_hook is not None for hook in hooks) + assert all(hook.device == torch.device("cpu") for hook in hooks) + + def test_video_vae_keeps_reference_fp32_weights(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module @@ -1797,97 +1465,6 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch): assert next(video_vae.parameters()).dtype == torch.float32 -def test_keyframe_encode_pins_and_restores_cudnn_settings(): - from vllm_omni.diffusion.models.minimax_h3.vae import ( - _minimax_h3_keyframe_encode_context, - ) - - cudnn = torch.backends.cudnn - original = ( - cudnn.enabled, - cudnn.benchmark, - cudnn.deterministic, - cudnn.allow_tf32, - ) - try: - cudnn.enabled = False - cudnn.benchmark = True - cudnn.deterministic = False - cudnn.allow_tf32 = False - - with _minimax_h3_keyframe_encode_context(torch.device("cuda")): - assert cudnn.enabled - assert not cudnn.benchmark - assert cudnn.deterministic - assert cudnn.allow_tf32 - - assert not cudnn.enabled - assert cudnn.benchmark - assert not cudnn.deterministic - assert not cudnn.allow_tf32 - finally: - ( - cudnn.enabled, - cudnn.benchmark, - cudnn.deterministic, - cudnn.allow_tf32, - ) = original - - -@pytest.mark.parametrize("fail_encode", [False, True]) -def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode): - from PIL import Image - - from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module - - events = [] - input_image = Image.new("RGB", (32, 32)) - - @contextmanager - def track_context(device): - assert device == torch.device("cpu") - events.append("enter") - try: - yield - finally: - events.append("exit") - - class FakeModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.weight = torch.nn.Parameter(torch.ones(1)) - self.parallel_tiling = True - - def encode_images(self, image, *, use_fp16_latent): - assert events == ["enter"] - assert image is input_image - assert use_fp16_latent - events.append("encode") - if fail_encode: - raise RuntimeError("encode failed") - return [torch.ones(1, 1, 2, 2)] - - video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE) - torch.nn.Module.__init__(video_vae) - video_vae.model = FakeModel() - video_vae.config_dict = { - "latent_channels": 1, - "latents_mean": [0.0], - "latents_std": [1.0], - } - monkeypatch.setattr(vae_module, "_minimax_h3_keyframe_encode_context", track_context) - - if fail_encode: - with pytest.raises(RuntimeError, match="encode failed"): - video_vae.encode_image(input_image) - else: - rows = video_vae.encode_image(input_image) - torch.testing.assert_close(rows, torch.ones(1, 4)) - - assert events == ["enter", "encode", "exit"] - assert video_vae.model.parallel_tiling - - def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module @@ -2115,69 +1692,6 @@ def test_video_vae_encode_uses_configured_parallel_tiling(): assert shape == (2, 2, 2) -def test_distributed_video_vae_encodes_references_sequentially(monkeypatch): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - from vllm_omni.diffusion.models.minimax_h3 import ( - pipeline_minimax_h3 as pipeline_module, - ) - - prepared = [ - {"prepared_path": "video-1.mp4"}, - {"prepared_path": "video-2.mp4"}, - ] - - class FakeVideoVAE: - def __init__(self): - self.calls = [] - - def is_distributed_enabled(self): - return True - - def encode_video(self, frames): - self.calls.append(frames) - index = len(self.calls) - return torch.full((1, 2), index, dtype=torch.float32), (index, 2, 3) - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.device = torch.device("cpu") - pipeline.video_vae = FakeVideoVAE() - - monkeypatch.setattr( - pipeline_module, - "_dit_rank_world", - lambda: ("dit-group", 1, 4), - ) - - def fake_broadcast_object_list(values, *, src, group, device): - assert values == [None] - assert (src, group, device) == (0, "dit-group", torch.device("cpu")) - values[0] = prepared - - monkeypatch.setattr( - pipeline_module.dist, - "broadcast_object_list", - fake_broadcast_object_list, - ) - monkeypatch.setattr( - pipeline_module, - "load_video_frames", - lambda path: f"frames:{path}", - ) - - rows, shapes = pipeline._encode_video_conditions_resident(None, count=2) - - assert pipeline.video_vae.calls == [ - "frames:video-1.mp4", - "frames:video-2.mp4", - ] - torch.testing.assert_close( - rows, - torch.tensor([[1.0, 1.0], [2.0, 2.0]]), - ) - assert shapes == [(1, 2, 3), (2, 2, 3)] - - @pytest.mark.parametrize( ("case", "extra", "image_count", "expected"), [ @@ -2186,7 +1700,7 @@ def test_distributed_video_vae_encodes_references_sequentially(monkeypatch): ], ) def test_f1_f2_official_fl2va_keyframe_matrix(case, extra, image_count, expected): - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import ( _resolve_fl2va_keyframe_indices, ) @@ -2206,7 +1720,7 @@ def test_f1_f2_official_fl2va_keyframe_matrix(case, extra, image_count, expected ], ) def test_r1_r6_ref2va_reference_count_matrix(case, counts): - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import ( _validate_ref2va_reference_counts, ) @@ -2215,10 +1729,10 @@ def test_r1_r6_ref2va_reference_count_matrix(case, counts): def test_ref2va_reference_count_validation_preserves_client_error_metadata(): - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + from vllm_omni.errors import OmniClientError + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import ( _validate_ref2va_reference_counts, ) - from vllm_omni.errors import OmniClientError with pytest.raises(OmniClientError, match="at least one image or video"): _validate_ref2va_reference_counts(0, 0, 0) @@ -2322,32 +1836,6 @@ def test_prepared_reference_video_descriptor_rejects_unknown_fields(): deserialize_prepared_reference_videos('{"artifact_dir":"/tmp","videos":[{"path":"x"}]}') -def test_stage_one_reuses_existing_prepared_reference_video(tmp_path): - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( - _reuse_prepared_reference_videos, - ) - - prepared_path = tmp_path / "prepared.mp4" - prepared_path.touch() - prepared = [{"prepared_path": str(prepared_path)}] - - assert _reuse_prepared_reference_videos(prepared, expected_count=1) is prepared - - -def test_stage_one_rejects_invalid_prepared_reference_video(tmp_path): - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( - _reuse_prepared_reference_videos, - ) - - with pytest.raises(ValueError, match="count does not match"): - _reuse_prepared_reference_videos([], expected_count=1) - with pytest.raises(ValueError, match="prepared reference video is unavailable"): - _reuse_prepared_reference_videos( - [{"prepared_path": str(tmp_path / "missing.mp4")}], - expected_count=1, - ) - - def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module @@ -2667,9 +2155,8 @@ def test_ref2va_two_video_recipe_rejects_real_duration_overflow(monkeypatch, tmp @pytest.mark.parametrize("duration", [4.0, 15.0]) def test_g2_output_duration_accepts_official_boundaries(duration): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import resolve_minimax_h3_shape - pipeline = object.__new__(MiniMaxH3Pipeline) sampling = SimpleNamespace( fps=24, num_frames=1, @@ -2678,15 +2165,14 @@ def test_g2_output_duration_accepts_official_boundaries(duration): extra_args={"duration": duration}, ) - height, width, *_ = pipeline._resolve_shape("ref2va", sampling, None) + height, width, *_ = resolve_minimax_h3_shape("ref2va", sampling, None) assert height > 0 and width > 0 @pytest.mark.parametrize("duration", [3.99, 15.01, float("nan"), "not-a-duration"]) def test_g2_output_duration_rejects_out_of_contract_values(duration): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import resolve_minimax_h3_shape - pipeline = object.__new__(MiniMaxH3Pipeline) sampling = SimpleNamespace( fps=24, num_frames=1, @@ -2695,7 +2181,7 @@ def test_g2_output_duration_rejects_out_of_contract_values(duration): extra_args={"duration": duration}, ) with pytest.raises(ValueError, match="duration"): - pipeline._resolve_shape("ref2va", sampling, None) + resolve_minimax_h3_shape("ref2va", sampling, None) def test_g1_fanout_uses_incrementing_output_seeds(): @@ -2713,26 +2199,25 @@ def test_g1_fanout_uses_incrementing_output_seeds(): def test_g3_task_specific_aspect_ratio_policy(): from PIL import Image - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( - _resolve_minimax_h3_aspect_ratio, + from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + resolve_minimax_h3_aspect_ratio, ) image = Image.new("RGB", (1280, 720)) - assert _resolve_minimax_h3_aspect_ratio("fl2va", "9:16", image) == pytest.approx(16 / 9) - assert _resolve_minimax_h3_aspect_ratio("ref2va", None, None) == pytest.approx(16 / 9) - assert _resolve_minimax_h3_aspect_ratio("ref2va", "auto", None) == pytest.approx(16 / 9) + assert resolve_minimax_h3_aspect_ratio("fl2va", "9:16", image) == pytest.approx(16 / 9) + assert resolve_minimax_h3_aspect_ratio("ref2va", None, None) == pytest.approx(16 / 9) + assert resolve_minimax_h3_aspect_ratio("ref2va", "auto", None) == pytest.approx(16 / 9) with pytest.raises(ValueError, match="requires an explicit"): - _resolve_minimax_h3_aspect_ratio("t2va", None, None) + resolve_minimax_h3_aspect_ratio("t2va", None, None) with pytest.raises(ValueError, match="one of"): - _resolve_minimax_h3_aspect_ratio("t2va", "2:1", None) + resolve_minimax_h3_aspect_ratio("t2va", "2:1", None) def test_g3_t2va_shape_requires_a_named_ratio_and_fl2va_ignores_override(): from PIL import Image - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import resolve_minimax_h3_shape - pipeline = object.__new__(MiniMaxH3Pipeline) t2va_sampling = SimpleNamespace( fps=24, num_frames=1, @@ -2741,7 +2226,7 @@ def test_g3_t2va_shape_requires_a_named_ratio_and_fl2va_ignores_override(): extra_args={"duration": 5.0}, ) with pytest.raises(ValueError, match="requires an explicit aspect_ratio"): - pipeline._resolve_shape("t2va", t2va_sampling, None) + resolve_minimax_h3_shape("t2va", t2va_sampling, None) fl2va_sampling = SimpleNamespace( fps=24, @@ -2750,14 +2235,14 @@ def test_g3_t2va_shape_requires_a_named_ratio_and_fl2va_ignores_override(): width=None, extra_args={"duration": 5.0, "aspect_ratio": "9:16"}, ) - height, width, *_ = pipeline._resolve_shape("fl2va", fl2va_sampling, Image.new("RGB", (1280, 720))) + height, width, *_ = resolve_minimax_h3_shape("fl2va", fl2va_sampling, Image.new("RGB", (1280, 720))) assert (height, width) == (768, 1344) def test_g4_reference_image_boundaries_and_aspect_ratio(): from PIL import Image - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import ( _validate_reference_image, ) @@ -2772,18 +2257,18 @@ def test_g4_reference_image_boundaries_and_aspect_ratio(): def test_g4_reference_image_file_format_and_size_contract(tmp_path): from PIL import Image - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( - _load_image, + from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + load_minimax_h3_images, ) valid_path = tmp_path / "reference.png" Image.new("RGB", (256, 256)).save(valid_path) - assert _load_image(valid_path).size == (256, 256) + assert load_minimax_h3_images(valid_path)[0].size == (256, 256) invalid_path = tmp_path / "reference.bmp" Image.new("RGB", (256, 256)).save(invalid_path) with pytest.raises(ValueError, match="must use"): - _load_image(invalid_path) + load_minimax_h3_images(invalid_path) def test_g4_standalone_audio_duration_and_total_duration_contract(): @@ -2809,97 +2294,6 @@ def test_g4_standalone_audio_duration_and_total_duration_contract(): ) -def test_ref2va_video_soundtrack_does_not_consume_standalone_audio_budget(): - from PIL import Image - - from vllm_omni.diffusion.cache.cachedit.runtime import RequestScopedCacheDiTRuntime - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.partition = "ref2va" - pipeline.supported_tasks = frozenset({"ref2va"}) - pipeline.device = torch.device("cpu") - pipeline.default_video_shift = 12.0 - pipeline.default_audio_shift = 3.0 - pipeline._resolve_sigma_positions = Mock(return_value=(None, 50)) - pipeline._quality_policy = MiniMaxH3QualityPolicy(None) - pipeline._cache_dit_runtime = Mock(spec=RequestScopedCacheDiTRuntime) - pipeline.text_encoder = object() - pipeline.encode_prompt = Mock(return_value=(torch.ones(1, 2), torch.ones(1, dtype=torch.long))) - pipeline._prepare_reference_videos = Mock(return_value=[{"input_has_audio": True}]) - pipeline._encode_visual_conditions = Mock(return_value=(torch.ones(1, 96), [(1, 16, 16), (17, 48, 84)])) - # A 2.35-second video soundtrack and a 15-second standalone reference - # each satisfy their modality's separate 15-second duration budget. - embedded = torch.full((188, 32), 1.0) - standalone = torch.full((1200, 32), 2.0) - pipeline._encode_reference_audio_conditions = Mock(return_value=(embedded, [94], standalone, [600])) - sampling = OmniDiffusionSamplingParams( - width=1344, - height=768, - fps=24, - num_inference_steps=50, - extra_args={"task": "ref2va", "duration": 15.0}, - ) - - context = pipeline._prepare_request_inputs( - prompt="A person carrying books down a street", - multi_modal_data={ - "image": Image.new("RGB", (256, 256)), - "video": ["reference.mp4"], - "audio": (torch.zeros(1, 15 * 16000), 16000), - }, - sampling=sampling, - ) - - assert context["audio_condition_lengths"] == [94, 600] - assert [block["kind"] for block in context["ref_blocks"]] == ["image", "video_audio", "audio"] - torch.testing.assert_close(context["audio_condition"], torch.cat([embedded, standalone])) - - -def test_ref2va_audio_duration_validation_precedes_rank_branch(monkeypatch): - from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module - - pipeline = object.__new__(pipeline_module.MiniMaxH3Pipeline) - pipeline.device = torch.device("cpu") - monkeypatch.setattr(pipeline_module, "_dit_rank_world", lambda: (None, 1, 2)) - - waveform = torch.zeros(1, 10) - with pytest.raises(ValueError, match="max_duration_seconds must be positive"): - pipeline._encode_audio_conditions_resident( - [(waveform, 10)], - max_duration_seconds=0, - ) - - -def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration(): - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - pipeline.device = torch.device("cpu") - observed = [] - - class FakeAudioVAE: - def encode_waveform(self, waveform, sample_rate): - observed.append((waveform.clone(), sample_rate)) - return waveform, waveform.shape[-1] - - pipeline.audio_vae = FakeAudioVAE() - waveform = torch.arange(40, dtype=torch.float32).reshape(1, 40) - - encoded, lengths = pipeline._encode_audio_conditions_resident( - [(waveform, 10)], - max_duration_seconds=2.5, - ) - - assert observed[0][0].shape == (1, 25) - assert observed[0][1] == 10 - assert encoded.shape == (1, 25) - assert lengths == [25] - - @pytest.mark.parametrize( ("field", "value", "message"), [ @@ -2963,9 +2357,8 @@ def _one_layer_text_encoder(): MiniMaxH3Qwen3VLQKVParallelLinear, MiniMaxH3Qwen3VLRMSNorm, ) - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _SingleRankEncoderGroup - group = _SingleRankEncoderGroup(0) + group = SimpleNamespace(rank_in_group=0, world_size=1, device_group=None, ranks=[0]) layer = nn.Module() layer.self_attn = nn.Module() layer.mlp = nn.Module() @@ -3103,7 +2496,7 @@ def _pad_seq_len_batch(pad_seq_len): return DiffusionRequestBatch( [ OmniDiffusionRequest( - prompt="pinned packed length", + prompt=_encoder_prompt("pinned packed length"), sampling_params=sampling, request_id="pad-seq-len", ) @@ -3189,3 +2582,431 @@ def test_pad_seq_len_request_validation_accepts_aligned_values(): assert _resolve_pad_seq_len(None) is None assert _resolve_pad_seq_len(54080) == 54080 assert _resolve_pad_seq_len(MINIMAX_H3_MAX_PAD_SEQ_LEN) == MINIMAX_H3_MAX_PAD_SEQ_LEN + + +def _turbo_spec(filename: str) -> TurboSpec: + """Build the spec a loaded artifact of this name would carry.""" + spec = parse_turbo_filename(filename) + assert spec is not None + return spec + + +def test_joint_video_output_is_quantized_before_transfer(): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _minimax_h3_post_process, + _prepare_minimax_h3_video_output, + ) + + video = torch.linspace(-0.1, 1.1, 2 * 3 * 2 * 4 * 5, dtype=torch.float64).reshape(2, 3, 2, 4, 5) + expected = torch.round(video.float().clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 4, 1) + prepared = _prepare_minimax_h3_video_output(video) + + assert prepared.dtype == torch.uint8 + assert prepared.is_contiguous() + assert prepared.shape == (2, 2, 4, 5, 3) + torch.testing.assert_close(prepared, expected, rtol=0, atol=0) + + result = _minimax_h3_post_process((prepared, torch.zeros(1, 2, 6))) + assert result["video"][0].dtype == np.uint8 + np.testing.assert_array_equal(result["video"][0], expected[0].numpy()) + + +def test_no_offload_keeps_text_encoder_resident(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + ) + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + input_ids = torch.tensor([1, 2]) + + actual = pipeline._encode_text_hidden(input_ids, {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_not_called() + + +def test_model_offload_uses_hooked_text_encoder_call(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=False, + ) + expected = torch.ones(2, 3) + pipeline.text_encoder = Mock(return_value=expected) + pipeline._model_cpu_offload_modules = [pipeline.text_encoder] + input_ids = torch.tensor([1, 2]) + vision_kwargs = {"pixel_values": torch.ones(1, 4)} + + actual = pipeline._encode_text_hidden(input_ids, vision_kwargs) + + assert actual is expected + pipeline.text_encoder.assert_called_once_with(input_ids, **vision_kwargs) + pipeline.text_encoder.load_to_device.assert_not_called() + + +@pytest.mark.parametrize( + ("enable_layerwise", "enable_distributed"), + [(True, False), (False, True)], +) +def test_legacy_layer_offload_preserves_minimax_stage_lifecycle(enable_layerwise, enable_distributed): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + diffusion_offload_config=None, + enable_cpu_offload=False, + enable_layerwise_offload=enable_layerwise, + enable_distributed_layerwise_offload=enable_distributed, + ) + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_called_once_with() + + vae = Mock() + with pipeline._component_on_device(vae): + pass + vae.load_to_device.assert_called_once_with() + vae.offload_to_cpu.assert_called_once_with() + + +def test_layerwise_encoder_selection_releases_text_encoder(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + diffusion_offload_config={ + "mode": "layer", + "components": ["text_encoder"], + }, + ) + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_called_once_with() + + +def test_layerwise_dit_only_keeps_text_encoder_resident(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + diffusion_offload_config={"mode": "layer", "components": ["dit"]}, + ) + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_not_called() + + +def test_dit_encoder_selection_keeps_vae_resident(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + diffusion_offload_config={ + "mode": "layer", + "components": ["dit", "text_encoder"], + }, + ) + component = Mock() + + with pipeline._component_on_device(component): + pass + + component.load_to_device.assert_not_called() + component.offload_to_cpu.assert_not_called() + + +def test_keyframe_encode_pins_and_restores_cudnn_settings(): + from vllm_omni.diffusion.models.minimax_h3.vae import ( + _minimax_h3_keyframe_encode_context, + ) + + cudnn = torch.backends.cudnn + original = ( + cudnn.enabled, + cudnn.benchmark, + cudnn.deterministic, + cudnn.allow_tf32, + ) + try: + cudnn.enabled = False + cudnn.benchmark = True + cudnn.deterministic = False + cudnn.allow_tf32 = False + + with _minimax_h3_keyframe_encode_context(torch.device("cuda")): + assert cudnn.enabled + assert not cudnn.benchmark + assert cudnn.deterministic + assert cudnn.allow_tf32 + + assert not cudnn.enabled + assert cudnn.benchmark + assert not cudnn.deterministic + assert not cudnn.allow_tf32 + finally: + ( + cudnn.enabled, + cudnn.benchmark, + cudnn.deterministic, + cudnn.allow_tf32, + ) = original + + +@pytest.mark.parametrize("fail_encode", [False, True]) +def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode): + from PIL import Image + + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + events = [] + input_image = Image.new("RGB", (32, 32)) + + @contextmanager + def track_context(device): + assert device == torch.device("cpu") + events.append("enter") + try: + yield + finally: + events.append("exit") + + class FakeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(1)) + self.parallel_tiling = True + + def encode_images(self, image, *, use_fp16_latent): + assert events == ["enter"] + assert image is input_image + assert use_fp16_latent + events.append("encode") + if fail_encode: + raise RuntimeError("encode failed") + return [torch.ones(1, 1, 2, 2)] + + video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE) + torch.nn.Module.__init__(video_vae) + video_vae.model = FakeModel() + video_vae.config_dict = { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + } + monkeypatch.setattr(vae_module, "_minimax_h3_keyframe_encode_context", track_context) + + if fail_encode: + with pytest.raises(RuntimeError, match="encode failed"): + video_vae.encode_image(input_image) + else: + rows = video_vae.encode_image(input_image) + torch.testing.assert_close(rows, torch.ones(1, 4)) + + assert events == ["enter", "encode", "exit"] + assert video_vae.model.parallel_tiling + + +def test_decode_to_mp4_batches_consumer_transfers(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + class FakeEncoder: + instances = [] + + def __init__(self, **kwargs): + self.pushes = [] + self.kwargs = kwargs + self.__class__.instances.append(self) + + def push(self, frames): + self.pushes.append(np.array(frames, copy=True)) + + def finish(self): + return b"mp4" + + def abort(self): + raise AssertionError("unexpected abort") + + class FakeAudioVAE: + def decode_latent(self, latent): + return torch.zeros(1, 1, 2) + + class FakeVideoVAE: + def decode_with_chunks(self, latent, *, on_chunk): + for value in (0.0, 0.25, 0.5): + on_chunk(torch.full((1, 3, 3, 2, 2), value)) + + pipeline = object.__new__(mod.MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae = FakeAudioVAE() + pipeline.video_vae = FakeVideoVAE() + pipeline.device = torch.device("cpu") + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) + monkeypatch.setattr( + "vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", + FakeEncoder, + ) + + output = pipeline.decode_to_mp4( + torch.zeros(1), + torch.zeros(1), + height=2, + width=2, + batch_frames=4, + ) + + assert output == b"mp4" + encoder = FakeEncoder.instances[-1] + assert [frames.shape[0] for frames in encoder.pushes] == [6, 3] + assert np.array_equal(encoder.pushes[0][0], np.full((2, 2, 3), 0, dtype=np.uint8)) + assert np.array_equal(encoder.pushes[0][3], np.full((2, 2, 3), 64, dtype=np.uint8)) + assert encoder.pushes[1].shape == (3, 2, 2, 3) + + +def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatch): + """A client's encoder options must survive the worker-side pre-encode path.""" + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + class FakeEncoder: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.__class__.instances.append(self) + + def push(self, frames): + del frames + + def finish(self): + return b"mp4" + + def abort(self): + raise AssertionError("unexpected abort") + + class FakeAudioVAE: + def decode_latent(self, latent): + return torch.zeros(1, 1, 2) + + class FakeVideoVAE: + def decode_with_chunks(self, latent, *, on_chunk): + on_chunk(torch.zeros(1, 3, 1, 2, 2)) + + pipeline = object.__new__(mod.MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae = FakeAudioVAE() + pipeline.video_vae = FakeVideoVAE() + pipeline.device = torch.device("cpu") + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) + monkeypatch.setattr("vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", FakeEncoder) + + pipeline.decode_to_mp4( + torch.zeros(1), + torch.zeros(1), + height=2, + width=2, + video_codec_options={"preset": "ultrafast"}, + ) + + assert FakeEncoder.instances[-1].kwargs["video_codec_options"] == {"preset": "ultrafast"} + + +@pytest.mark.parametrize( + ("codec_extra", "expected"), + [ + ({}, {"preset": "ultrafast", "threads": "0"}), + ({"video_codec_options": {"preset": "slow", "threads": 2}}, {"preset": "slow", "threads": "2"}), + ({"video_codec_options": {}}, {}), + ({"video_codec_options": None}, None), + ], +) +@pytest.mark.parametrize("batch_extra, batch_frames", [({}, 17), ({"preencode_batch_frames": 5}, 5)]) +def test_preencode_request_preserves_serving_codec_defaults(codec_extra, expected, batch_extra, batch_frames): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline.supported_tasks = frozenset({"t2va"}) + pipeline.default_video_shift = 12.0 + pipeline.default_audio_shift = 3.0 + pipeline.device = torch.device("cpu") + pipeline.od_config = SimpleNamespace() + pipeline.load_text_encoder = False + pipeline.load_vae_encoder = False + pipeline.text_encoder = object() + pipeline.encode_prompt = Mock(return_value=(torch.ones(1, 2), torch.ones(1, dtype=torch.long))) + pipeline._quality_policy = Mock() + pipeline._cache_dit_runtime = Mock() + pipeline.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1))) + pipeline.decode_to_mp4 = Mock(return_value=b"mp4") + sampling = OmniDiffusionSamplingParams( + width=1344, + height=768, + fps=24, + num_frames=124, + num_inference_steps=50, + extra_args={"task": "t2va", "aspect_ratio": "16:9", "preencode_mp4": True, **codec_extra, **batch_extra}, + ) + batch = DiffusionRequestBatch( + [OmniDiffusionRequest(prompt=_encoder_prompt("test"), sampling_params=sampling, request_id="codec-defaults")] + ) + + pipeline.forward(batch) + + assert pipeline.decode_to_mp4.call_args.kwargs["video_codec_options"] == expected + assert pipeline.decode_to_mp4.call_args.kwargs["batch_frames"] == batch_frames + + +def test_video_codec_options_are_normalized_for_the_encoder(): + from vllm_omni.diffusion.utils.media_utils import normalize_video_codec_options + + assert normalize_video_codec_options({"preset": "ultrafast", "threads": 0}) == { + "preset": "ultrafast", + "threads": "0", + } + assert normalize_video_codec_options(None) is None + with pytest.raises(ValueError, match="video_codec_options"): + normalize_video_codec_options("preset=ultrafast") diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py new file mode 100644 index 000000000..510a7b33e --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from typing import Any +from unittest.mock import Mock + +import pytest +import torch +from PIL import Image + +from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.fixture +def pipeline(): + model = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(model) + model.load_text_encoder = True + model.load_vae_encoder = True + model.device = torch.device("cpu") + model.partition = "combined" + model.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"}) + model.default_video_shift = 12.0 + model.default_audio_shift = 3.0 + model.od_config = SimpleNamespace(enable_layerwise_offload=False) + model._quality_policy = SimpleNamespace(resolve=Mock(return_value=SimpleNamespace(cache_dit=None))) + model._cache_dit_runtime = SimpleNamespace(prepare=Mock()) + model.encode_prompt = Mock( + return_value=(torch.ones(2, 5120, dtype=torch.bfloat16), torch.ones(2, dtype=torch.long)) + ) + model.video_vae = Mock() + model.video_vae.is_distributed_enabled.return_value = False + model.video_vae.encode_image.side_effect = lambda image: torch.ones((image.height // 32) * (image.width // 32), 96) + model.audio_vae = Mock() + model.transformer = torch.nn.Identity() + model.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1))) + model.decode = Mock(return_value=(torch.zeros(1, 3, 1, 32, 32), torch.zeros(1))) + return model + + +def _request(task): + prompt: dict[str, Any] = {"prompt": "A person waves."} + if task != "t2va": + prompt["multi_modal_data"] = {"image": Image.new("RGB", (256, 256))} + sampling = OmniDiffusionSamplingParams( + height=32, width=32, num_frames=96, num_inference_steps=2, extra_args={"task": task, "aspect_ratio": "1:1"} + ) + return SimpleNamespace(prompts=[prompt], sampling_params=sampling) + + +@pytest.mark.parametrize("task", ["t2va", "fl2va", "ref2va"]) +def test_single_stage_encodes_without_stage_zero(pipeline, task): + output = pipeline.forward(_request(task)) + pipeline.encode_prompt.assert_called_once() + assert pipeline.video_vae.encode_image.call_count == int(task != "t2va") + pipeline.diffuse.assert_called_once() + assert pipeline.diffuse.call_args.kwargs["task"] == task + video, audio = output.output + torch.testing.assert_close(video, torch.zeros(1, 1, 32, 32, 3, dtype=torch.uint8)) + torch.testing.assert_close(audio, pipeline.decode.return_value[1]) + + +def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monkeypatch): + import numpy as np + + from vllm_omni.data_entry_keys import flatten_payload + from vllm_omni.model_executor.models.minimax_h3 import encoder_processing as processing + from vllm_omni.model_executor.models.minimax_h3.conditioning import ( + MiniMaxH3EncoderConditioning, + MiniMaxH3TextConditioning, + ) + from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder + from vllm_omni.model_executor.stage_input_processors.minimax_h3 import encoder2diffusion, prepare_encoder_prompt + + frames = np.zeros((48, 32, 32, 3), dtype=np.uint8) + decode_frames = Mock(return_value=frames) + sampler = Mock(wraps=processing.sample_reference_video_frames) + monkeypatch.setattr(processing, "load_video_frames", decode_frames) + monkeypatch.setattr(processing, "sample_reference_video_frames", sampler) + monkeypatch.setattr( + processing, + "prepare_reference_videos", + Mock( + return_value=[{"prepared_path": "prepared.mp4", "original_path": "original.mp4", "input_has_audio": True}] + ), + ) + monkeypatch.setattr(processing, "load_video_audio", Mock(return_value=(torch.full((64000,), 3.0), 32000))) + pipeline.video_vae.encode_video.return_value = (torch.full((1, 96), 2.0), (1, 2, 2)) + + def encode_audio(waveform, sample_rate): + length = round(waveform.shape[-1] / sample_rate * 40) + return torch.full((length * 2, 32), float(waveform[0])), length + + pipeline.audio_vae.encode_waveform.side_effect = encode_audio + request = _request("ref2va") + request.prompts[0]["multi_modal_data"].update(video="original.mp4", audio=(torch.full((192000,), 4.0), 32000)) + pipeline.forward(request) + local_kwargs = pipeline.diffuse.call_args.kwargs + + transformed = prepare_encoder_prompt(request.prompts[0], [request.sampling_params]) + stage_zero = object.__new__(MiniMaxH3Encoder) + torch.nn.Module.__init__(stage_zero) + stage_zero._component_leader = True + stage_zero.video_vae = pipeline.video_vae + stage_zero.audio_vae = pipeline.audio_vae + media = stage_zero._encode_media(stage_zero._media_input(transformed["additional_information"])) + text = MiniMaxH3TextConditioning(*pipeline.encode_prompt.return_value) + payload = MiniMaxH3EncoderConditioning.from_components(text, media).to_omni_payload() + source = SimpleNamespace(finished=True, outputs=[SimpleNamespace(multimodal_output=flatten_payload(payload))]) + bridged = encoder2diffusion([source], request.prompts[0]) + pipeline.load_text_encoder = False + pipeline.load_vae_encoder = False + pipeline.encode_prompt.reset_mock() + pipeline.forward(SimpleNamespace(prompts=[bridged], sampling_params=request.sampling_params)) + external_kwargs = pipeline.diffuse.call_args.kwargs + for key, value in local_kwargs.items(): + if isinstance(value, torch.Tensor): + torch.testing.assert_close(value, external_kwargs[key]) + else: + assert value == external_kwargs[key] + pipeline.encode_prompt.assert_not_called() + assert decode_frames.call_count == sampler.call_count == 2 + assert all(call.kwargs["decoded_frames"] is frames for call in sampler.call_args_list) + assert [block["kind"] for block in local_kwargs["ref_blocks"]] == ["image", "video_audio", "audio"] + assert local_kwargs["audio_condition_lengths"] == [80, round(107 / 24 * 40)] + assert all( + call.args[0].shape[-1] == round(107 / 24 * 32000) + for call in pipeline.audio_vae.encode_waveform.call_args_list[1::2] + ) 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 31a8916a6..24206952d 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py @@ -9,72 +9,25 @@ from unittest.mock import Mock import pytest import torch -from PIL import Image pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] -def test_h3_profiler_targets_cover_reference_encoders_and_vae_decoders(): +def test_h3_full_pipeline_profiler_includes_local_encoding(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline targets = MiniMaxH3Pipeline._PROFILER_TARGETS - assert "_encode_visual_conditions" in targets - assert "_encode_reference_audio_conditions" in targets - assert "video_vae.decode_latent" in targets - assert "audio_vae.decode_latent" in targets - assert "_encode_video_conditions" not in targets - assert "_encode_video_audio_conditions" not in targets - assert "_encode_audio_conditions" not in targets - - -@pytest.mark.parametrize("distributed", [False, True]) -def test_manual_component_offload_wins_over_requested_model_offload(distributed): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.od_config = SimpleNamespace( - enable_cpu_offload=True, - enable_layerwise_offload=not distributed, - enable_distributed_layerwise_offload=distributed, - ) - pipeline._model_cpu_offload_modules = [] - pipeline.text_encoder = Mock() - expected = torch.ones(2, 3) - pipeline.text_encoder.encode_ids.return_value = expected - component = Mock() - - actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) - with pipeline._component_on_device(component): - component.load_to_device.assert_called_once_with() - component.offload_to_cpu.assert_not_called() - - assert actual is expected - pipeline.text_encoder.load_to_device.assert_called_once_with() - pipeline.text_encoder.offload_to_cpu.assert_called_once_with() - component.offload_to_cpu.assert_called_once_with() - - -def test_requested_model_offload_without_backend_uses_resident_components(): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.od_config = SimpleNamespace( - enable_cpu_offload=True, - enable_layerwise_offload=False, - enable_distributed_layerwise_offload=False, - ) - pipeline._model_cpu_offload_modules = [] - pipeline.text_encoder = Mock() - expected = torch.ones(2, 3) - pipeline.text_encoder.encode_ids.return_value = expected - - actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) - - assert actual is expected - pipeline.text_encoder.load_to_device.assert_called_once_with() - pipeline.text_encoder.encode_ids.assert_called_once() + assert targets == [ + "encode_prompt", + "_encode_local_media", + "diffuse", + "decode", + "video_vae.decode_latent", + "audio_vae.decode_latent", + "prepare_encode", + "denoise_step", + "post_decode", + ] def test_decoded_output_is_unchanged_without_active_model_offload(): @@ -194,7 +147,7 @@ def test_forward_releases_video_after_quantization_and_before_next_seed(monkeypa _extract_prompt=lambda raw: ("a prompt", {}), _extract_text_conditioning=lambda raw: None, _extract_prepared_reference_videos=lambda raw: None, - _prepare_request_inputs=lambda **kwargs: context, + _prepare_request_inputs=lambda *args, **kwargs: context, _denoise_kwargs=lambda ctx: {}, diffuse=diffuse, decode=decode, @@ -243,7 +196,7 @@ def test_forward_drops_the_decoded_video_before_the_next_seed_diffuses(monkeypat _extract_prompt=lambda raw: ("a prompt", {}), _extract_text_conditioning=lambda raw: None, _extract_prepared_reference_videos=lambda raw: None, - _prepare_request_inputs=lambda **kwargs: context, + _prepare_request_inputs=lambda *args, **kwargs: context, _denoise_kwargs=lambda ctx: {}, diffuse=diffuse, decode=decode, @@ -318,7 +271,8 @@ def test_output_owner_rank_matches_executor_contract(monkeypatch, available, ini assert MiniMaxH3Pipeline._is_output_owner_rank() is expected -def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch): +@pytest.mark.parametrize("load_text_encoder", [True, False]) +def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch, load_text_encoder): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module @@ -326,9 +280,11 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch): torch.nn.Module.__init__(pipeline) pipeline.transformer = torch.nn.Linear(2, 2) pipeline.transformers_ref = torch.nn.Linear(2, 2) - pipeline.text_encoder = torch.nn.Linear(2, 2) pipeline.video_vae = torch.nn.Linear(2, 2) pipeline.audio_vae = torch.nn.Linear(2, 2) + pipeline._encoder_modules = ["text_encoder"] if load_text_encoder else [] + if load_text_encoder: + pipeline.text_encoder = torch.nn.Linear(2, 2) apply_offload = Mock() remove_offload = Mock() monkeypatch.setattr(module, "apply_sequential_offload", apply_offload) @@ -341,7 +297,7 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch): ) dits = [pipeline.transformer, pipeline.transformers_ref] - stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae] + stages = [*([pipeline.text_encoder] if load_text_encoder else []), pipeline.video_vae, pipeline.audio_vae] apply_offload.assert_called_once_with( dit_modules=dits, encoder_modules=stages, @@ -356,6 +312,46 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch): remove_offload.assert_called_once_with([*dits, *stages]) +@pytest.mark.parametrize("decode_fails", [False, True]) +def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + component = torch.nn.Linear(2, 2) + events: list[tuple[Any, ...]] = [] + + @contextmanager + def record_component(value): + events.append(("activate", value)) + try: + yield + finally: + events.append(("offload", value)) + + monkeypatch.setattr(module, "sequential_offload_component", record_component) + pipeline._model_cpu_offload_modules = [component] + + def decode(): + with pipeline._component_on_device(component): + events.append(("decode", component)) + if decode_fails: + raise RuntimeError("decode failed") + + if decode_fails: + with pytest.raises(RuntimeError, match="decode failed"): + decode() + else: + decode() + + assert events == [ + ("activate", component), + ("decode", component), + ("offload", component), + ] + + def test_h3_model_cpu_offload_keeps_unselected_vaes_resident(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module @@ -416,44 +412,20 @@ def test_h3_model_cpu_offload_rejects_unloaded_selected_encoder(monkeypatch): apply_offload.assert_not_called() -@pytest.mark.parametrize("decode_fails", [False, True]) -def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module +def _encode_media(pipeline, **kwargs): + from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - component = torch.nn.Linear(2, 2) - events: list[tuple[Any, ...]] = [] - - @contextmanager - def record_component(value): - events.append(("activate", value)) - try: - yield - finally: - events.append(("offload", value)) - - monkeypatch.setattr(module, "sequential_offload_component", record_component) - pipeline._model_cpu_offload_modules = [component] - - def decode(): - with pipeline._component_on_device(component): - events.append(("decode", component)) - if decode_fails: - raise RuntimeError("decode failed") - - if decode_fails: - with pytest.raises(RuntimeError, match="decode failed"): - decode() - else: - decode() - - assert events == [ - ("activate", component), - ("decode", component), - ("offload", component), - ] + media = MiniMaxH3EncoderMediaInput( + task="ref2va", height=32, width=32, num_frames=124, latent_t=31, audio_t=208, **kwargs + ) + return encode_media( + media, + video_vae=getattr(pipeline, "video_vae", None), + audio_vae=getattr(pipeline, "audio_vae", None), + emit_conditioning=True, + component_scope=pipeline._component_on_device, + ) @pytest.mark.parametrize("encode_fails", [False, True]) @@ -463,7 +435,6 @@ def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) - pipeline.device = torch.device("cpu") pipeline.video_vae = Mock() pipeline._model_cpu_offload_modules = [pipeline.video_vae] events = [] @@ -484,17 +455,15 @@ def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode pipeline.video_vae.encode_image.side_effect = encode_image monkeypatch.setattr(module, "sequential_offload_component", record_component) - monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1)) - monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value) - images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))] + images = (torch.zeros(16, 16, 3, dtype=torch.uint8), torch.zeros(16, 32, 3, dtype=torch.uint8)) if encode_fails: with pytest.raises(RuntimeError, match="encode failed"): - pipeline._encode_visual_conditions(images, None, video_count=0) + _encode_media(pipeline, images=images) else: - rows, shapes = pipeline._encode_visual_conditions(images, None, video_count=0) - assert rows.shape == (3, 4) - assert shapes == [(1, 1, 1), (1, 1, 2)] + result = _encode_media(pipeline, images=images) + assert result.visual_condition.shape == (3, 4) + assert result.visual_condition_shapes == ((1, 1, 1), (1, 1, 2)) assert events == [ ("activate", pipeline.video_vae), @@ -510,9 +479,7 @@ def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch): pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) - pipeline.device = torch.device("cpu") pipeline.video_vae = Mock() - pipeline.video_vae.is_distributed_enabled.return_value = False pipeline._model_cpu_offload_modules = [pipeline.video_vae] events = [] @@ -529,31 +496,26 @@ def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch): return torch.ones(1, 4) def encode_video(frames): - events.append(("video", frames)) + events.append(("video", frames.shape)) return torch.ones(3, 4), (3, 2, 2) pipeline.video_vae.encode_image.side_effect = encode_image pipeline.video_vae.encode_video.side_effect = encode_video monkeypatch.setattr(module, "sequential_offload_component", record_component) - monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1)) - monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value) - monkeypatch.setattr(module, "load_video_frames", lambda path: f"frames:{path}") - images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))] - prepared_videos = [{"prepared_path": "reference.mp4"}] - - rows, shapes = pipeline._encode_visual_conditions( - images, - prepared_videos, - video_count=1, + result = _encode_media( + pipeline, + images=(torch.zeros(16, 16, 3, dtype=torch.uint8), torch.zeros(16, 32, 3, dtype=torch.uint8)), + videos=(torch.zeros(12, 32, 32, 3, dtype=torch.uint8),), + video_audios=(None,), ) - assert rows.shape == (5, 4) - assert shapes == [(1, 1, 1), (1, 1, 2), (3, 2, 2)] + assert result.visual_condition.shape == (5, 4) + assert result.visual_condition_shapes == ((1, 1, 1), (1, 1, 2), (3, 2, 2)) assert events == [ ("activate", pipeline.video_vae), ("image", (16, 16)), ("image", (32, 16)), - ("video", "frames:reference.mp4"), + ("video", (12, 32, 32, 3)), ("offload", pipeline.video_vae), ] @@ -567,7 +529,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: list[tuple[Any, ...]] = [] + events = [] @contextmanager def record_component(value): @@ -577,26 +539,20 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp finally: events.append(("offload", value)) - def encode_embedded(*args, **kwargs): - events.append(("embedded",)) + def encode_waveform(waveform, sample_rate): + events.append(("audio", int(waveform[0]), sample_rate)) if encode_fails: raise RuntimeError("audio encode failed") - return torch.ones(1, 2), [1] + return torch.ones(80, 32), 80 - def encode_standalone(*args, **kwargs): - events.append(("standalone",)) - return torch.ones(1, 2), [1] - - pipeline._encode_video_audio_conditions_resident = encode_embedded - pipeline._encode_audio_conditions_resident = encode_standalone + pipeline.audio_vae.encode_waveform.side_effect = encode_waveform monkeypatch.setattr(module, "sequential_offload_component", record_component) def call(): - return pipeline._encode_reference_audio_conditions( - [{"input_has_audio": True}], - has_audio=[True], - standalone_audios=[(torch.ones(1), 16_000)], - max_duration_seconds=1.0, + return _encode_media( + pipeline, + video_audios=(), + audios=((torch.ones(32_000), 16_000),), ) if encode_fails: @@ -604,10 +560,61 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp call() else: result = call() - assert result[0].shape == (1, 2) - assert result[2].shape == (1, 2) + assert result.audio_condition.shape == (80, 32) + assert result.audio_condition_lengths == (80,) + assert events[1:-1] == [("audio", 1, 16_000)] assert events[0] == ("activate", pipeline.audio_vae) assert events[-1] == ("offload", pipeline.audio_vae) assert events.count(("activate", pipeline.audio_vae)) == 1 assert events.count(("offload", pipeline.audio_vae)) == 1 + + +@pytest.mark.parametrize("distributed", [False, True]) +def test_manual_component_offload_wins_over_requested_model_offload(distributed): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=not distributed, + enable_distributed_layerwise_offload=distributed, + ) + pipeline._model_cpu_offload_modules = [] + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + component = Mock() + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + with pipeline._component_on_device(component): + component.load_to_device.assert_called_once_with() + component.offload_to_cpu.assert_not_called() + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_called_once_with() + component.offload_to_cpu.assert_called_once_with() + + +def test_requested_model_offload_without_backend_uses_resident_components(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + ) + pipeline._model_cpu_offload_modules = [] + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.encode_ids.assert_called_once() diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py index c7e53ca62..ba0e0ddf5 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py @@ -186,9 +186,8 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig): common_kwargs = { "model": model_ref, "enforce_eager": True, - "tensor_parallel_size": 2, # The fused BF16 baseline only fits on H100-80GB with encoder TP. - "text_encoder_tp_size": 2, + "stage_0_tensor_parallel_size": 2, "vae_use_tiling": True, } diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py index 1b123929c..e202c0d0e 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py @@ -110,6 +110,8 @@ def _step_pipeline(model, *, packed_batch_supported: bool = True): from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline pipeline = object.__new__(MiniMaxH3Pipeline) + pipeline.load_text_encoder = False + pipeline.load_vae_encoder = False pipeline.transformer = model pipeline.device = torch.device("cpu") pipeline._transformer_for_task = lambda task: model @@ -304,8 +306,17 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch, batch_frames): } pipeline = _step_pipeline(_SegmentMeanModel()) - monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_extract_prompt", staticmethod(lambda _: ("a prompt", {}))) - monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_prepare_request_inputs", lambda self, **_: context) + conditioning = object() + monkeypatch.setattr( + mod.MiniMaxH3Pipeline, + "_extract_encoder_conditioning", + staticmethod(lambda _: conditioning), + ) + monkeypatch.setattr( + mod.MiniMaxH3Pipeline, + "_prepare_encoder_conditioning_inputs", + lambda self, value, sampling: context, + ) monkeypatch.setattr( mod.MiniMaxH3Pipeline, "_build_denoise_inputs", @@ -439,50 +450,3 @@ def test_packed_batch_rejects_backends_that_cannot_isolate_requests(attention): from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline assert MiniMaxH3Pipeline._packed_batch_supported(_FakeTransformer([attention])) is False - - -def test_broadcast_rank0_exception_single_rank_reraises(): - """Single-rank execution has no group; the helper just reraises.""" - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( - _broadcast_rank0_exception, - ) - from vllm_omni.errors import OmniClientError - - _broadcast_rank0_exception(None) - with pytest.raises(OmniClientError, match="bad ref"): - _broadcast_rank0_exception(OmniClientError("bad ref")) - - -def test_broadcast_rank0_exception_propagates_to_non_zero_ranks(monkeypatch): - """A rank-0 error becomes a matching client error on every other DiT rank.""" - from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod - from vllm_omni.errors import OmniClientError - - def fake_rank_world(rank): - return lambda: (object(), rank, 4) - - def make_broadcast(rank0_payload): - def fake_broadcast(payload_list, *, src, group): - payload_list[0] = rank0_payload - - return fake_broadcast - - monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(0)) - err = OmniClientError("invalid reference-video file", status_code=422, error_type="UnprocessableEntityError") - rank0_payload = { - "type": type(err).__name__, - "message": str(err), - "status_code": err.status_code, - "error_type": err.error_type, - } - monkeypatch.setattr(mod.dist, "broadcast_object_list", make_broadcast(rank0_payload)) - with pytest.raises(OmniClientError) as rank0_info: - mod._broadcast_rank0_exception(err) - assert rank0_info.value is err - - monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(2)) - with pytest.raises(OmniClientError) as rank2_info: - mod._broadcast_rank0_exception(None) - assert rank2_info.value.status_code == 422 - assert rank2_info.value.error_type == "UnprocessableEntityError" - assert "invalid reference-video file" in str(rank2_info.value) diff --git a/tests/model_executor/models/test_minimax_h3_encoder_payload.py b/tests/model_executor/models/test_minimax_h3_encoder_payload.py new file mode 100644 index 000000000..4c9a7137b --- /dev/null +++ b/tests/model_executor/models/test_minimax_h3_encoder_payload.py @@ -0,0 +1,406 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.data_entry_keys import flatten_payload +from vllm_omni.model_executor.models.minimax_h3.conditioning import ( + MiniMaxH3EncoderConditioning, + MiniMaxH3EncoderMediaConditioning, + MiniMaxH3EncoderMediaInput, +) +from vllm_omni.model_executor.stage_input_processors.minimax_h3 import encoder2diffusion + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_pipeline_owns_three_encoders_then_dit_and_decoders() -> None: + from vllm_omni.model_executor.models.minimax_h3.pipeline import MINIMAX_H3_PIPELINE + + encoder, diffusion = MINIMAX_H3_PIPELINE.stages + assert encoder.model_arch == "MiniMaxH3Encoder" + assert encoder.custom_process_next_stage_input_func is None + assert diffusion.model_arch == "MiniMaxH3Pipeline" + assert diffusion.requires_multimodal_data is False + assert diffusion.requires_full_payload_input is False + + +def _patch_encoder_constructors(monkeypatch, *, rank: int): + from vllm_omni.model_executor.models.minimax_h3 import encoder as encoder_module + + group = SimpleNamespace(rank_in_group=rank, world_size=2, device_group=object()) + + class FakeVideoVAE(torch.nn.Module): + def __init__(self, path, *, device, encode_only, trust_remote_code): + super().__init__() + self.path = path + self.parallel_args = None + self.trust_remote_code = trust_remote_code + + def set_parallel_size(self, size, *, process_group): + self.parallel_args = (size, process_group) + + class FakeAudioVAE(torch.nn.Module): + def __init__(self, path, *, device, encode_only, trust_remote_code): + super().__init__() + self.path = path + self.trust_remote_code = trust_remote_code + + def init_backbone(self, *, vllm_config, prefix): + torch.nn.Module.__init__(self) + + monkeypatch.setattr(encoder_module.MiniMaxH3TextEncoderBackbone, "__init__", init_backbone) + monkeypatch.setattr(encoder_module, "get_tp_group", lambda: group) + monkeypatch.setattr(encoder_module, "get_local_device", lambda: torch.device("cpu")) + monkeypatch.setattr(encoder_module, "MiniMaxH3VideoVAE", FakeVideoVAE) + monkeypatch.setattr(encoder_module, "MiniMaxH3AudioVAE", FakeAudioVAE) + return encoder_module, group + + +def _component_root(tmp_path): + root = tmp_path / "FL2VA" + for name in ("text_encoder", "video_vae", "audio_vae"): + (root / name).mkdir(parents=True) + return root + + +def _encoder_config(root, *, video_mode="patch", roles=None, trust_remote_code=True): + components = ( + { + "text_encoder": {"parallel_mode": "tp"}, + "video_vae": {"parallel_mode": video_mode}, + "audio_vae": {"parallel_mode": "leader"}, + } + if roles is None + else roles + ) + return SimpleNamespace( + model_config=SimpleNamespace( + model=str(root / "text_encoder"), + hf_config=SimpleNamespace(minimax_h3_encoder_components=components), + trust_remote_code=trust_remote_code, + ) + ) + + +def test_encoder_requires_exactly_three_components(monkeypatch, tmp_path) -> None: + encoder_module, _ = _patch_encoder_constructors(monkeypatch, rank=0) + root = tmp_path / "FL2VA" + (root / "text_encoder").mkdir(parents=True) + (root / "video_vae").mkdir() + config = _encoder_config(root) + + with pytest.raises(RuntimeError, match=r"requires all three encoders.*audio_vae"): + encoder_module.MiniMaxH3Encoder(vllm_config=config) + + +@pytest.mark.parametrize(("rank", "has_audio_vae"), [(0, True), (1, False)]) +@pytest.mark.parametrize("trust_remote_code", [False, True]) +def test_encoder_uses_tp_patch_and_leader_on_the_same_rank_set( + monkeypatch, + tmp_path, + rank, + has_audio_vae, + trust_remote_code, +) -> None: + encoder_module, group = _patch_encoder_constructors(monkeypatch, rank=rank) + root = _component_root(tmp_path) + config = _encoder_config(root, trust_remote_code=trust_remote_code) + + model = encoder_module.MiniMaxH3Encoder(vllm_config=config) + + assert model.component_config.text_parallel_mode == "tp" + assert model.component_config.video_parallel_mode == "patch" + assert model.component_config.audio_parallel_mode == "leader" + assert model.video_vae.parallel_args == (2, group.device_group) + assert model.video_vae.trust_remote_code is trust_remote_code + assert (model.audio_vae is not None) is has_audio_vae + if has_audio_vae: + assert model.audio_vae.trust_remote_code is trust_remote_code + + +@pytest.mark.parametrize("component", ["video", "audio"]) +@pytest.mark.parametrize("trusted", [False, True]) +def test_encoder_only_vae_requires_remote_code_trust(monkeypatch, component, trusted) -> None: + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + config = {"auto_map": {"AutoModel": "modeling.RemoteVAE"}, "sample_rate": 32000} + monkeypatch.setattr(vae_module, "_load_component_config", lambda _path: config) + remote = torch.nn.Module() + remote.model = torch.nn.Module() + calls = [] + + def load_encoder(path, config_dict): + calls.append((path, config_dict)) + return remote + + monkeypatch.setattr(vae_module, f"_load_{component}_vae_encoder", load_encoder) + vae_cls = vae_module.MiniMaxH3VideoVAE if component == "video" else vae_module.MiniMaxH3AudioVAE + if trusted: + model = vae_cls("unused", device=torch.device("cpu"), encode_only=True, trust_remote_code=True) + assert model.remote is remote + assert calls == [("unused", config)] + else: + with pytest.raises(ValueError, match="trust-remote-code"): + vae_cls("unused", device=torch.device("cpu"), encode_only=True) + assert calls == [] + + +def test_encoder_requires_all_role_policies(monkeypatch, tmp_path) -> None: + encoder_module, _ = _patch_encoder_constructors(monkeypatch, rank=0) + root = _component_root(tmp_path) + config = _encoder_config( + root, + roles={ + "text_encoder": {"parallel_mode": "tp"}, + "video_vae": {"parallel_mode": "patch"}, + }, + ) + + with pytest.raises(ValueError, match=r"requires exactly.*audio_vae"): + encoder_module.MiniMaxH3Encoder(vllm_config=config) + + +def test_encoder_does_not_infer_role_policies(monkeypatch, tmp_path) -> None: + encoder_module, _ = _patch_encoder_constructors(monkeypatch, rank=0) + root = _component_root(tmp_path) + config = SimpleNamespace( + model_config=SimpleNamespace( + model=str(root / "text_encoder"), + hf_config=SimpleNamespace(), + ) + ) + + with pytest.raises(ValueError, match="requires explicit minimax_h3_encoder_components"): + encoder_module.MiniMaxH3Encoder(vllm_config=config) + + +def test_video_leader_policy_skips_nonleader(monkeypatch, tmp_path) -> None: + encoder_module, _ = _patch_encoder_constructors(monkeypatch, rank=1) + root = _component_root(tmp_path) + model = encoder_module.MiniMaxH3Encoder(vllm_config=_encoder_config(root, video_mode="leader")) + + assert model.video_vae is None + assert model.audio_vae is None + assert ( + model._encode_media( + MiniMaxH3EncoderMediaInput( + task="ref2va", + height=32, + width=32, + num_frames=48, + latent_t=1, + audio_t=80, + images=(torch.zeros(32, 32, 3, dtype=torch.uint8),), + ) + ) + is None + ) + + +def test_encoder_runs_video_and_audio_components_on_ar_model() -> None: + from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder + + class VideoVAE: + def encode_image(self, _image): + return torch.ones(1, 96) + + def encode_video(self, _frames): + return torch.full((1, 96), 2.0), (1, 2, 2) + + class AudioVAE: + def encode_waveform(self, waveform, _sample_rate): + return torch.ones(160, 32), waveform.shape[-1] // 800 + + model = MiniMaxH3Encoder.__new__(MiniMaxH3Encoder) + torch.nn.Module.__init__(model) + model._component_leader = True + model.video_vae = VideoVAE() + model.audio_vae = AudioVAE() + media = MiniMaxH3EncoderMediaInput( + task="ref2va", + height=32, + width=32, + num_frames=48, + latent_t=1, + audio_t=80, + images=(torch.zeros(32, 32, 3, dtype=torch.uint8),), + videos=(torch.zeros(1, 32, 32, 3, dtype=torch.uint8),), + video_audios=((torch.zeros(64_000), 32_000),), + ) + + conditioning = model._encode_media(media) + + assert conditioning is not None + assert conditioning.visual_condition_shapes == ((1, 2, 2), (1, 2, 2)) + assert conditioning.audio_condition_lengths == (80,) + assert conditioning.ref_blocks[1]["kind"] == "video_audio" + + +def test_prepare_encoder_inputs_keeps_reference_audio_budgets_separate(monkeypatch) -> None: + from vllm_omni.model_executor.models.minimax_h3 import encoder_processing as processing + + frames = torch.zeros(1, 32, 32, 3, dtype=torch.uint8).numpy() + waveform = torch.zeros(320_000) + monkeypatch.setattr(processing, "resolve_minimax_h3_shape", lambda *_args: (32, 32, 240, 60, 400)) + monkeypatch.setattr( + processing, + "prepare_reference_videos", + lambda *_args, **_kwargs: [ + {"prepared_path": "prepared.mp4", "original_path": "reference.mp4", "input_has_audio": True} + ], + ) + monkeypatch.setattr(processing, "load_video_frames", lambda _path: frames) + monkeypatch.setattr( + processing, + "sample_reference_video_frames", + lambda *_args, **_kwargs: {"frames": [frames[0]], "block_timestamps": [[0.0]]}, + ) + monkeypatch.setattr(processing, "load_video_audio", lambda *_args, **_kwargs: (waveform, 32_000)) + + prepared = processing.prepare_encoder_inputs( + { + "prompt": "reference", + "multi_modal_data": {"video": "reference.mp4", "audio": (waveform, 32_000)}, + }, + SimpleNamespace(extra_args={"task": "ref2va"}), + ) + + assert prepared.media.video_audios[0][0].shape[-1] == 320_000 + assert prepared.media.audios[0][0].shape[-1] == 320_000 + assert prepared.condition_labels == [("audio", 1), ("video", 1), ("audio", 2)] + + +@pytest.mark.parametrize( + ("embedded_lengths", "standalone_lengths", "valid"), + [((400,), (400,), True), ((320, 320), (400,), False), ((400,), (320, 320), False)], +) +def test_encode_media_keeps_audio_budgets_and_component_residency_separate( + embedded_lengths, standalone_lengths, valid +) -> None: + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media + + active = None + scopes = [] + + class VideoVAE: + def encode_image(self, _image): + assert active is self + return torch.ones(1, 96) + + def encode_video(self, _frames): + assert active is self + return torch.ones(1, 96), (1, 2, 2) + + class AudioVAE: + def encode_waveform(self, waveform, _sample_rate): + assert active is self + length = waveform.shape[-1] // 800 + return torch.ones(2 * length, 32), length + + @contextmanager + def component_scope(component): + nonlocal active + assert active is None + active = component + scopes.append(component) + try: + yield + finally: + active = None + + video_vae = VideoVAE() + audio_vae = AudioVAE() + media = MiniMaxH3EncoderMediaInput( + task="ref2va", + height=32, + width=32, + num_frames=360, + latent_t=90, + audio_t=600, + images=(torch.zeros(32, 32, 3, dtype=torch.uint8),), + videos=tuple(torch.zeros(1, 32, 32, 3, dtype=torch.uint8) for _ in embedded_lengths), + video_audios=tuple((torch.zeros(length * 800), 32_000) for length in embedded_lengths), + audios=tuple((torch.zeros(length * 800), 32_000) for length in standalone_lengths), + ) + kwargs = dict(video_vae=video_vae, audio_vae=audio_vae, emit_conditioning=True, component_scope=component_scope) + if valid: + conditioning = encode_media(media, **kwargs) + assert conditioning.audio_condition_lengths == embedded_lengths + standalone_lengths + assert [block["kind"] for block in conditioning.ref_blocks] == ["image", "video_audio", "audio"] + conditioning.to_omni_components() + else: + with pytest.raises(ValueError, match="at most 15 seconds"): + encode_media(media, **kwargs) + assert scopes == [video_vae, audio_vae] + + +def test_encoder_output_reuses_encoder_handoff_and_round_trips() -> None: + expected = MiniMaxH3EncoderConditioning( + hidden_states=torch.randn(3, 5120, dtype=torch.bfloat16), + token_tags=torch.tensor([1, 0, 1], dtype=torch.int64), + task="t2va", + height=256, + width=448, + num_frames=17, + latent_t=5, + audio_t=10, + ) + + source = SimpleNamespace( + finished=True, + request_id="request-1", + outputs=[SimpleNamespace(multimodal_output=flatten_payload(expected.to_omni_payload()))], + ) + prompt = { + "prompt": "hello", + "additional_information": {"global_request_id": ["request-1"]}, + } + + result = encoder2diffusion([source], prompt) + + assert result is not None + payload = result["additional_information"]["encoder_output"] + actual = MiniMaxH3EncoderConditioning.from_omni_payload(payload) + torch.testing.assert_close(actual.hidden_states, expected.hidden_states) + torch.testing.assert_close(actual.token_tags, expected.token_tags) + + +def test_encoder_releases_workspace_after_cpu_payload(monkeypatch) -> None: + from vllm_omni.model_executor.models.minimax_h3 import encoder as encoder_module + + model = encoder_module.MiniMaxH3Encoder.__new__(encoder_module.MiniMaxH3Encoder) + torch.nn.Module.__init__(model) + model._component_leader = True + model._token_tags = torch.tensor([1]) + info = { + "meta": {"minimax_h3_encoder_request": {}}, + "_minimax_h3_encoder_media_conditioning": MiniMaxH3EncoderMediaConditioning( + task="t2va", + height=256, + width=448, + num_frames=17, + latent_t=5, + audio_t=10, + ), + } + empty_cache_calls = 0 + + def empty_cache() -> None: + nonlocal empty_cache_calls + empty_cache_calls += 1 + + monkeypatch.setattr(torch.accelerator, "empty_cache", empty_cache) + output = model.make_omni_output( + torch.zeros(1, 5120), + model_intermediate_buffer=[info], + request_sample_eligible=[True], + ) + + assert empty_cache_calls == 1 + assert output.multimodal_outputs is not None + assert output.multimodal_outputs["embed"]["embedding"][0].device.type == "cpu" diff --git a/tests/model_executor/stage_input_processors/test_minimax_h3.py b/tests/model_executor/stage_input_processors/test_minimax_h3.py index 9fda89953..aac16d29f 100644 --- a/tests/model_executor/stage_input_processors/test_minimax_h3.py +++ b/tests/model_executor/stage_input_processors/test_minimax_h3.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for MiniMax H3's disaggregated text-encoder contract.""" -import os -import shutil from types import SimpleNamespace from unittest.mock import Mock @@ -12,39 +10,37 @@ import pytest import torch from PIL import Image -from vllm_omni.data_entry_keys import OmniPayload, deserialize_payload, serialize_payload -from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( - _load_audio, - resolve_minimax_h3_diffusion_model_path, +from vllm_omni.data_entry_keys import flatten_payload +from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import resolve_minimax_h3_diffusion_model_path +from vllm_omni.engine.serialization import ( + deserialize_additional_information, + serialize_additional_information, ) +from vllm_omni.errors import OmniClientError from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_executor.models.minimax_h3.checkpoint import ( - resolve_minimax_h3_model_root, + resolve_minimax_h3_encoder_model_root, resolve_minimax_h3_partition, ) from vllm_omni.model_executor.models.minimax_h3.conditioning import ( MINIMAX_H3_CONDITION_LABELS_KEY, + MINIMAX_H3_ENCODER_REQUEST_KEY, MINIMAX_H3_PRESENTATION_TASK_KEY, - MINIMAX_H3_TEXT_CONDITIONING_SCHEMA, - MiniMaxH3TextConditioning, + MiniMaxH3EncoderConditioning, ) +from vllm_omni.model_executor.models.minimax_h3.encoder_processing import _audio_items, _load_audio from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( minimax_h3_ref2va_presentation, minimax_h3_ref2va_video_presentation, ) -from vllm_omni.model_executor.models.minimax_h3.reference_video import ( - MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY, - deserialize_prepared_reference_videos, -) from vllm_omni.model_executor.models.minimax_h3.text_encoder import ( MiniMaxH3MultiModalProcessor, _build_minimax_h3_presentation, ) from vllm_omni.model_executor.stage_input_processors.minimax_h3 import ( - _audio_items, _diffusion_sampling_params, - prepare_text_encoder_prompt, - text_encoder2diffusion, + encoder2diffusion, + prepare_encoder_prompt, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -66,13 +62,14 @@ class _SegmentTokenizer: return self._special_ids[token] -def test_h3_processor_reprocesses_media_instead_of_using_partial_sender_cache(mocker): +def test_h3_processor_reprocesses_media_instead_of_using_partial_sender_cache(monkeypatch): processor = object.__new__(MiniMaxH3MultiModalProcessor) sentinel = ([1, 2, 3], object(), True) - apply_processor = mocker.patch.object( + apply_processor = Mock(return_value=sentinel) + monkeypatch.setattr( MiniMaxH3MultiModalProcessor, "_apply_hf_processor", - return_value=sentinel, + apply_processor, ) inputs = object() timing_ctx = object() @@ -119,9 +116,11 @@ def test_fused_audio_loader_accepts_list_waveform_pair(): def test_prepare_ref2va_keeps_original_text_and_exact_condition_order(): prompt = { "prompt": "hello", + "additional_information": {"global_request_id": ["request-1"]}, + "model_intermediate_buffer": {"private": "preserved"}, "multi_modal_data": { "image": Image.new("RGB", (256, 256)), - "audio": [np.zeros(16), 16_000], + "audio": [np.zeros(32_000), 16_000], }, } sampling = OmniDiffusionSamplingParams( @@ -130,7 +129,7 @@ def test_prepare_ref2va_keeps_original_text_and_exact_condition_order(): extra_args={"task": "ref2va"}, ) - transformed = prepare_text_encoder_prompt(prompt, [sampling]) + transformed = prepare_encoder_prompt(prompt, [sampling]) assert transformed["prompt"] == "hello" assert len(transformed["multi_modal_data"]["image"]) == 1 @@ -140,223 +139,148 @@ def test_prepare_ref2va_keeps_original_text_and_exact_condition_order(): ("image", 1), ("audio", 1), ] + assert transformed["model_intermediate_buffer"] == {"private": "preserved"} + runner_info = transformed["additional_information"] + for _ in range(2): + wire = serialize_additional_information(runner_info) + runner_info = deserialize_additional_information(wire) + assert runner_info["global_request_id"] == ["request-1"] + request_metadata = runner_info["meta"][MINIMAX_H3_ENCODER_REQUEST_KEY] + assert request_metadata["task"] == "ref2va" + assert isinstance( + runner_info["hidden_states"]["layers"][0], + torch.Tensor, + ) + assert isinstance( + runner_info["hidden_states"]["layers"][1], + torch.Tensor, + ) + from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder + + media = MiniMaxH3Encoder._media_input(runner_info) + assert media.task == "ref2va" + assert len(media.images) == 1 + assert len(media.audios) == 1 -def test_text_encoder_prompt_rejects_injected_prepared_video_descriptor(): - prompt = { - "prompt": "hello", - "additional_information": {"meta": {MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY: '{"artifact_dir":"/tmp/user"}'}}, - } - sampling = OmniDiffusionSamplingParams(height=256, width=448, extra_args={"task": "t2va"}) +def _mock_ref2va_video_with_audio(monkeypatch, *, duration_seconds: float) -> None: + from vllm_omni.model_executor.models.minimax_h3 import encoder_processing - transformed = prepare_text_encoder_prompt(prompt, [sampling]) - - assert MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY not in transformed["additional_information"]["meta"] - - -def test_prepare_ref2va_video_uses_shared_frame_sampler_once(monkeypatch): - prepare_videos = Mock() - - def prepare(_videos, *, target_frame_count, workdir, start_time_seconds): - prepare_videos(_videos, target_frame_count, workdir, start_time_seconds) - prepared_path = os.path.join(workdir, "prepared.mp4") - open(prepared_path, "wb").close() - return [ - { - "original_path": "/tmp/input.mp4", - "prepared_path": prepared_path, - "input_has_audio": True, - "width": 448, - "height": 256, - "start_time_seconds": 0.0, - "duration_seconds": 5.2, - "audio_duration_seconds": 5.2, - } - ] - - sample_frames = Mock(return_value={"frames": [np.zeros((4, 4, 3), dtype=np.uint8)]}) - monkeypatch.setattr("vllm_omni.model_executor.stage_input_processors.minimax_h3.prepare_reference_videos", prepare) + sample_rate = 16_000 + frames = np.zeros((4, 32, 32, 3), dtype=np.uint8) monkeypatch.setattr( - "vllm_omni.model_executor.stage_input_processors.minimax_h3.sample_reference_video_frames", - sample_frames, + encoder_processing, + "prepare_reference_videos", + lambda *args, **kwargs: [ + { + "prepared_path": "prepared.mp4", + "original_path": "original.mp4", + "input_has_audio": True, + "duration_seconds": duration_seconds, + "audio_duration_seconds": duration_seconds, + } + ], + ) + monkeypatch.setattr(encoder_processing, "load_video_frames", lambda path: frames) + monkeypatch.setattr( + encoder_processing, + "sample_reference_video_frames", + lambda *args, **kwargs: {"block_timestamps": [0.0], "frames": frames[:1]}, + ) + monkeypatch.setattr( + encoder_processing, + "load_video_audio", + lambda *args, **kwargs: (torch.zeros(round(duration_seconds * sample_rate)), sample_rate), + ) + + +def test_prepare_ref2va_rejects_short_embedded_video_audio(monkeypatch): + _mock_ref2va_video_with_audio(monkeypatch, duration_seconds=1.0) + sampling = OmniDiffusionSamplingParams( + height=256, + width=448, + num_frames=96, + extra_args={"task": "ref2va"}, + ) + prompt = {"prompt": "hello", "multi_modal_data": {"video": "original.mp4"}} + + with pytest.raises(OmniClientError, match=r"duration must be in \[2, 15\] seconds, got 1\.000"): + prepare_encoder_prompt(prompt, [sampling]) + + +def test_prepare_ref2va_keeps_embedded_and_standalone_audio_budgets_separate(monkeypatch): + duration_seconds = 8.0 + sample_rate = 16_000 + _mock_ref2va_video_with_audio(monkeypatch, duration_seconds=duration_seconds) + sampling = OmniDiffusionSamplingParams( + height=256, + width=448, + num_frames=240, + extra_args={"task": "ref2va"}, ) prompt = { "prompt": "hello", - "multi_modal_data": {"video": "/tmp/input.mp4"}, + "multi_modal_data": { + "video": "original.mp4", + "audio": (torch.zeros(int(duration_seconds * sample_rate)), sample_rate), + }, } - sampling = OmniDiffusionSamplingParams(height=256, width=448, num_frames=124, extra_args={"task": "ref2va"}) - transformed = prepare_text_encoder_prompt(prompt, [sampling]) - - prepare_videos.assert_called_once() - _, target_frame_count, artifact_dir, start_time_seconds = prepare_videos.call_args.args - assert target_frame_count == 124 - assert start_time_seconds is None - descriptor = transformed["additional_information"]["meta"][MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY] - described_dir, described_videos = deserialize_prepared_reference_videos(descriptor) - assert described_dir == artifact_dir - assert os.path.isfile(described_videos[0]["prepared_path"]) - sample_frames.assert_called_once_with(described_videos[0]["prepared_path"]) - assert prompt["multi_modal_data"] == {"video": "/tmp/input.mp4"} + transformed = prepare_encoder_prompt(prompt, [sampling]) assert transformed["mm_processor_kwargs"][MINIMAX_H3_CONDITION_LABELS_KEY] == [ ("audio", 1), ("video", 1), + ("audio", 2), ] - shutil.rmtree(artifact_dir) -def test_stage_wire_rejects_multiple_text_encoder_sources(): - with pytest.raises(RuntimeError, match="exactly one text-encoder source"): - text_encoder2diffusion([SimpleNamespace(), SimpleNamespace()], prompt={"prompt": "hello"}) +def _encoder_output() -> dict: + return MiniMaxH3EncoderConditioning( + hidden_states=torch.randn(3, 5120, dtype=torch.bfloat16), + token_tags=torch.tensor([1, 0, 1], dtype=torch.int64), + task="t2va", + height=256, + width=448, + num_frames=17, + latent_t=5, + audio_t=10, + ).to_omni_payload() -def test_stage_wire_rejects_request_id_mismatch(): - source = SimpleNamespace(request_id="other", outputs=[SimpleNamespace()]) +def test_encoder2diffusion_reuses_encoder_handoff() -> None: prompt = { - "prompt": "hello", - "additional_information": {"global_request_id": ["req-1"]}, - } - - with pytest.raises(RuntimeError, match="request ID does not match"): - text_encoder2diffusion([source], prompt=prompt) - - -def test_stage_wire_rejects_multiple_completions(): - source = SimpleNamespace(request_id="req-1", outputs=[SimpleNamespace(), SimpleNamespace()]) - - with pytest.raises(RuntimeError, match="exactly one completion"): - text_encoder2diffusion([source], prompt={"prompt": "hello"}) - - -def _source_output(payload: dict) -> SimpleNamespace: - completion = SimpleNamespace(multimodal_output=payload) - return SimpleNamespace(request_id="request-1", outputs=[completion]) - - -def _diffusion_prompt() -> dict: - return { - "request_id": "request-1", "prompt": "test prompt", - "additional_information": {}, - } - - -def _valid_text_conditioning_payload() -> OmniPayload: - return { - "hidden_states": { - "output": torch.zeros(4, 5120, dtype=torch.bfloat16), - }, - "meta": { - "token_role_ids": torch.tensor([[1], [1], [0], [0]], dtype=torch.int64), + "multi_modal_data": {"image": object()}, + "additional_information": { + "private": "preserved", + "meta": {MINIMAX_H3_ENCODER_REQUEST_KEY: {"task": "t2va"}}, + "hidden_states": {"layers": {0: torch.zeros(1)}}, }, + "model_intermediate_buffer": {"private": "encoder-only"}, } - - -def test_text_encoder2diffusion_reads_hidden_states_output_and_token_role_ids() -> None: - payload = _valid_text_conditioning_payload() - hidden = payload["hidden_states"]["output"] - token_role_ids = payload["meta"]["token_role_ids"] - - result = text_encoder2diffusion( - [_source_output(payload)], - _diffusion_prompt(), + source = SimpleNamespace( + finished=True, + request_id="request-1", + outputs=[SimpleNamespace(multimodal_output=flatten_payload(_encoder_output()))], ) - conditioning = result["additional_information"]["text_encoder_output"] - torch.testing.assert_close(conditioning["hidden_states"], hidden) - torch.testing.assert_close(conditioning["token_tags"], token_role_ids.squeeze(-1)) + result = encoder2diffusion([source], prompt) + + assert result["prompt"] == "test prompt" + assert result["multi_modal_data"] is None + assert "model_intermediate_buffer" not in result + additional_information = result["additional_information"] + assert additional_information["private"] == "preserved" + assert "hidden_states" not in additional_information + assert "meta" not in additional_information + parsed = MiniMaxH3EncoderConditioning.from_omni_payload(additional_information["encoder_output"]) + assert parsed.task == "t2va" -def test_text_conditioning_artifact_round_trips_through_omni_payload_wire_format() -> None: - payload = _valid_text_conditioning_payload() - - wire = serialize_payload(payload) - - assert wire is not None - restored = deserialize_payload(wire) - conditioning = MiniMaxH3TextConditioning.from_omni_payload(restored) - assert MINIMAX_H3_TEXT_CONDITIONING_SCHEMA == "minimax_h3.text_conditioning/v1" - torch.testing.assert_close( - conditioning.hidden_states, - payload["hidden_states"]["output"], - ) - torch.testing.assert_close( - conditioning.token_tags, - payload["meta"]["token_role_ids"].squeeze(-1), - ) - assert conditioning.hidden_states.dtype == torch.bfloat16 - assert conditioning.hidden_states.is_contiguous() - assert conditioning.token_tags.dtype == torch.int64 - assert conditioning.token_tags.is_contiguous() - - -@pytest.mark.parametrize( - ("payload", "message"), - [ - ({"meta": {"token_role_ids": torch.ones(4, 1)}}, "no hidden_states payload"), - ({"hidden_states": {}, "meta": {"token_role_ids": torch.ones(4, 1)}}, "no hidden_states.output tensor"), - ({"hidden_states": {"output": torch.randn(4, 5120)}}, "no conditioning metadata"), - ({"hidden_states": {"output": torch.randn(4, 5120)}, "meta": {}}, "no token_role_ids tensor"), - ( - { - "hidden_states": {"output": torch.randn(4, 5120)}, - "meta": {"token_role_ids": torch.ones(4)}, - }, - r"must have shape \[tokens, 1\]", - ), - ], -) -def test_text_encoder2diffusion_rejects_invalid_conditioning_payload(payload: dict, message: str) -> None: - with pytest.raises(RuntimeError, match=message): - text_encoder2diffusion([_source_output(payload)], _diffusion_prompt()) - - -@pytest.mark.parametrize( - ("case", "message"), - [ - ("hidden_width", r"hidden_states must have shape \[tokens, 5120\]"), - ("hidden_dtype", "hidden_states must have dtype torch.bfloat16"), - ("hidden_noncontiguous", "hidden_states must use contiguous strided layout"), - ("hidden_non_strided", "hidden_states must use contiguous strided layout"), - ("token_dtype", "stage-wire token_role_ids must have dtype torch.int64"), - ("token_layout", "stage-wire token_role_ids must use contiguous strided layout"), - ("token_count", "token_tags must align with hidden_states"), - ("token_value", "token_tags must contain only 0 and 1"), - ], -) -def test_text_encoder2diffusion_rejects_h3_v1_contract_mismatch(case: str, message: str) -> None: - payload = _valid_text_conditioning_payload() - - if case == "hidden_width": - payload["hidden_states"]["output"] = torch.zeros(4, 5119, dtype=torch.bfloat16) - elif case == "hidden_dtype": - payload["hidden_states"]["output"] = torch.zeros(4, 5120, dtype=torch.float32) - elif case == "hidden_noncontiguous": - payload["hidden_states"]["output"] = torch.empty(5120, 4, dtype=torch.bfloat16).t() - elif case == "hidden_non_strided": - payload["hidden_states"]["output"] = torch.empty( - (4, 5120), - dtype=torch.bfloat16, - layout=torch.sparse_coo, - ) - elif case == "token_dtype": - payload["meta"]["token_role_ids"] = torch.tensor([[1], [1], [0], [0]], dtype=torch.int32) - elif case == "token_layout": - payload["meta"]["token_role_ids"] = torch.tensor( - [[1], [0], [1], [0], [1], [0], [1], [0]], - dtype=torch.int64, - )[::2] - elif case == "token_count": - payload["meta"]["token_role_ids"] = torch.tensor([[1], [0], [0]], dtype=torch.int64) - elif case == "token_value": - payload["meta"]["token_role_ids"] = torch.tensor([[1], [1], [2], [0]], dtype=torch.int64) - else: # pragma: no cover - parameterization is exhaustive - raise AssertionError(f"unknown test case: {case}") - - with pytest.raises(RuntimeError, match=message) as exc_info: - text_encoder2diffusion([_source_output(payload)], _diffusion_prompt()) - - assert MINIMAX_H3_TEXT_CONDITIONING_SCHEMA in str(exc_info.value) +def test_encoder2diffusion_waits_for_one_finished_source() -> None: + assert encoder2diffusion([SimpleNamespace(finished=False)], {"prompt": "hello"}) is None + with pytest.raises(RuntimeError, match="exactly one encoder source"): + encoder2diffusion([SimpleNamespace(), SimpleNamespace()], {"prompt": "hello"}) def test_ref2va_one_image_tokens_and_tags_match_fused_presentation(): @@ -421,15 +345,20 @@ def test_checkpoint_resolver_selects_local_partition(tmp_path): (root / "FL2VA" / "text_encoder").mkdir(parents=True) (root / "Ref2VA" / "text_encoder").mkdir(parents=True) - assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder") - assert resolve_minimax_h3_model_root(str(root), None, "ref2va") == str(root / "Ref2VA" / "text_encoder") - assert resolve_minimax_h3_model_root(str(root / "Ref2VA"), None, None) == str(root / "Ref2VA" / "text_encoder") - assert resolve_minimax_h3_model_root(str(root / "FL2VA"), None, "ref2va") == str(root / "Ref2VA" / "text_encoder") + assert resolve_minimax_h3_encoder_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder") + assert resolve_minimax_h3_encoder_model_root(str(root), None, "ref2va") == str(root / "Ref2VA" / "text_encoder") + assert resolve_minimax_h3_encoder_model_root(str(root), None, "combined") == str(root / "FL2VA" / "text_encoder") + assert resolve_minimax_h3_encoder_model_root(str(root / "Ref2VA"), None, None) == str( + root / "Ref2VA" / "text_encoder" + ) + assert resolve_minimax_h3_encoder_model_root(str(root / "FL2VA"), None, "ref2va") == str( + root / "Ref2VA" / "text_encoder" + ) def test_checkpoint_resolver_rejects_unknown_task(tmp_path): with pytest.raises(ValueError, match="task_type must be one of"): - resolve_minimax_h3_model_root(str(tmp_path), None, "unknown") + resolve_minimax_h3_encoder_model_root(str(tmp_path), None, "unknown") def test_partition_resolver_preserves_consumer_auto_default(tmp_path): diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 0b8723c0e..d86a3216f 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -804,6 +804,7 @@ class _DiffusionConfigProjection: "transformer": True, "vae": True, "text_encoder": True, + "vae_encoder": True, } ) override_transformer_cls_name: str | None = None diff --git a/vllm_omni/deploy/minimax_h3_disaggregated.yaml b/vllm_omni/deploy/minimax_h3_disaggregated.yaml index 1c1ee6c40..25a561c27 100644 --- a/vllm_omni/deploy/minimax_h3_disaggregated.yaml +++ b/vllm_omni/deploy/minimax_h3_disaggregated.yaml @@ -1,4 +1,4 @@ -# MiniMax H3: vLLM Qwen3-VL text encoder + diffusion video/audio generation. +# MiniMax H3: unified text/video/audio encoder stage + diffusion generation stage. pipeline: minimax_h3_disaggregated async_chunk: false trust_remote_code: true @@ -8,7 +8,16 @@ stages: - stage_id: 0 devices: "0,1" max_num_seqs: 1 + gpu_memory_utilization: 0.6 tensor_parallel_size: 2 + hf_overrides: + minimax_h3_encoder_components: + text_encoder: + parallel_mode: tp + video_vae: + parallel_mode: patch + audio_vae: + parallel_mode: leader default_sampling_params: max_tokens: 1 temperature: 0.0 @@ -19,6 +28,7 @@ stages: max_num_seqs: 1 model_loaded: text_encoder: false + vae_encoder: false parallel_config: tensor_parallel_size: 1 ulysses_degree: 4 diff --git a/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml b/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml index d7e84db42..ec11796b6 100644 --- a/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml +++ b/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml @@ -8,7 +8,16 @@ stages: - stage_id: 0 devices: "0,1" max_num_seqs: 1 + gpu_memory_utilization: 0.6 tensor_parallel_size: 2 + hf_overrides: + minimax_h3_encoder_components: + text_encoder: + parallel_mode: tp + video_vae: + parallel_mode: patch + audio_vae: + parallel_mode: leader default_sampling_params: max_tokens: 1 temperature: 0.0 @@ -19,6 +28,7 @@ stages: max_num_seqs: 1 model_loaded: text_encoder: false + vae_encoder: false parallel_config: tensor_parallel_size: 1 ulysses_degree: 4 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 10f8642c6..e1ce550d4 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -984,6 +984,7 @@ class OmniDiffusionConfig: "transformer": True, "vae": True, "text_encoder": True, + "vae_encoder": True, } ) override_transformer_cls_name: str | None = None 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 8f814a27a..1b8bd7fea 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -7,9 +7,9 @@ from __future__ import annotations import json import math import os -import tempfile from collections.abc import Iterable, Mapping, Sequence -from contextlib import contextmanager, nullcontext +from contextlib import contextmanager +from dataclasses import fields, replace from itertools import groupby from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar @@ -18,7 +18,6 @@ import numpy as np import torch import torch.distributed as dist import torch.nn as nn -from PIL import Image from transformers import Qwen2TokenizerFast, Qwen3VLProcessor from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -29,10 +28,7 @@ from vllm_omni.diffusion.cache.cachedit import ( RequestScopedCacheDiTRuntime, ) from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.parallel_state import ( - get_world_group, - init_world_group, -) +from vllm_omni.diffusion.distributed.parallel_state import get_world_group, init_world_group from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.forward_context import DenoiseProgressMixin from vllm_omni.diffusion.model_loader.diffusers_loader import ( @@ -74,36 +70,20 @@ from vllm_omni.model_executor.models.minimax_h3.checkpoint import ( resolve_minimax_h3_partition, ) from vllm_omni.model_executor.models.minimax_h3.conditioning import ( + MiniMaxH3EncoderConditioning, + MiniMaxH3EncoderMediaConditioning, + MiniMaxH3EncoderMediaInput, MiniMaxH3TextConditioning, ) -from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( - load_minimax_h3_images as _load_images, -) -from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( - minimax_h3_multi_image_presentation, - minimax_h3_ref2va_presentation, - minimax_h3_ref2va_video_presentation, - minimax_h3_text_only_ids, -) -from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( - resolve_minimax_h3_aspect_ratio as _resolve_minimax_h3_aspect_ratio, -) -from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( - resolve_minimax_h3_output_canvas as _resolve_output_canvas, -) -from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( - resolve_minimax_h3_reference_image_shape as _reference_image_shape, +from vllm_omni.model_executor.models.minimax_h3.encoder_processing import ( + PreparedEncoderInputs, + encode_media, + prepare_encoder_inputs, ) +from vllm_omni.model_executor.models.minimax_h3.preprocessing import build_minimax_h3_presentation from vllm_omni.model_executor.models.minimax_h3.reference_video import ( MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY, deserialize_prepared_reference_videos, - load_audio_file, - load_video_audio, - load_video_frames, - prepare_reference_videos, - sample_reference_video_frames, - validate_reference_audio_files, - validate_reference_audio_waveforms, ) from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization import ( @@ -154,7 +134,6 @@ from .scheduling_minimax_h3_euler_ancestral import ( ) from .time_request import ( MINIMAX_H3_SHAPE_PLANNER, - minimax_h3_align_frame_count, minimax_h3_time_shift_sigmas, ) from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE @@ -175,22 +154,6 @@ MINIMAX_H3_FPS = 24 MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000 MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999 MINIMAX_H3_AUDIO_REF_COND_TIMESTEP = 1.0 -MINIMAX_H3_OUTPUT_SHORT_EDGE = 768 -MINIMAX_H3_OUTPUT_MAX_PIXELS = 768 * 1344 -MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE = 2048 -MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE = 32 -MINIMAX_H3_SUPPORTED_ASPECT_RATIOS = { - "21:9": 21.0 / 9.0, - "16:9": 16.0 / 9.0, - "4:3": 4.0 / 3.0, - "1:1": 1.0, - "3:4": 3.0 / 4.0, - "9:16": 9.0 / 16.0, -} -MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES = 30 * 1024 * 1024 -MINIMAX_H3_REFERENCE_IMAGE_FORMATS = frozenset({"jpeg", "png", "webp", "heic", "heif"}) -MINIMAX_H3_MIN_OUTPUT_SECONDS = 4.0 -MINIMAX_H3_MAX_OUTPUT_SECONDS = 15.0 MINIMAX_H3_DOWNLOAD_PATTERNS = [ "FL2VA/**", "Ref2VA/model_index.json", @@ -267,7 +230,6 @@ def _resolve_minimax_h3_model_root( ) -# Keys of ``_prepare_request_inputs`` that feed ``diffuse`` / ``_build_denoise_inputs``. _MINIMAX_H3_DENOISE_INPUT_KEYS = ( "task", "text_embeddings", @@ -427,114 +389,6 @@ def get_minimax_h3_post_process_func( return _minimax_h3_post_process -def _align_multiple(value: float, multiple: int = 32) -> int: - return max(multiple, int(round(float(value) / multiple)) * multiple) - - -def _load_image(value: Any) -> Image.Image: - images = _load_images(value) - if len(images) != 1: - raise OmniClientError(f"MiniMax H3 expected one image, got {len(images)}") - return images[0] - - -def _load_audio(value: Any) -> tuple[torch.Tensor, int]: - if isinstance(value, (list, tuple)) and not (len(value) == 2 and isinstance(value[1], (int, np.integer))): - audios = _load_audios(value) - if len(audios) != 1: - raise OmniClientError(f"MiniMax H3 expected one audio, got {len(audios)}") - return audios[0] - if isinstance(value, (str, os.PathLike)): - return load_audio_file(str(value)) - if isinstance(value, (list, tuple)) and len(value) == 2: - waveform, sample_rate = value - waveform = torch.as_tensor(waveform).float() - return waveform, int(sample_rate) - if isinstance(value, dict): - waveform = value.get("waveform", value.get("array")) - sample_rate = value.get("sample_rate", value.get("sampling_rate")) - if waveform is not None and sample_rate is not None: - return torch.as_tensor(waveform).float(), int(sample_rate) - raise OmniClientError("MiniMax H3 audio input must be a path, (waveform, sample_rate), or a waveform mapping") - - -def _load_audios(value: Any) -> list[tuple[torch.Tensor, int]]: - if isinstance(value, (list, tuple)) and not (len(value) == 2 and isinstance(value[1], (int, np.integer))): - if not value: - raise OmniClientError("MiniMax H3 audio input must not be empty") - return [_load_audio(item) for item in value] - return [_load_audio(value)] - - -def _as_int_list(value: Any, *, name: str) -> list[int]: - if isinstance(value, bool): - raise OmniClientError(f"{name} must be an integer or a list of integers") - if isinstance(value, (int, np.integer)): - return [int(value)] - if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): - result = list(value) - if not result: - raise OmniClientError(f"{name} must not be empty") - if any(isinstance(item, bool) or not isinstance(item, (int, np.integer)) for item in result): - raise OmniClientError(f"{name} must contain only integers") - return [int(item) for item in result] - raise OmniClientError(f"{name} must be an integer or a list of integers") - - -def _resolve_fl2va_keyframe_indices(extra: Mapping[str, Any], image_count: int) -> list[int]: - target = extra.get("target") - target = target if isinstance(target, Mapping) else {} - raw = extra.get("frame_indices", extra.get("frame_index")) - if raw is None: - raw = target.get("frame_indices", target.get("frame_index")) - if raw is None: - raw_indices = [0] if image_count == 1 else [0, -1] - else: - raw_indices = _as_int_list(raw, name="frame_indices") - if len(raw_indices) != image_count: - raise OmniClientError( - f"MiniMax H3 FL2VA requires one frame index per image: got {raw_indices!r} for {image_count} image(s)" - ) - if tuple(raw_indices) not in ((0,), (-1,), (0, -1)): - raise OmniClientError("MiniMax H3 FL2VA frame_indices must be [0], [-1], or [0, -1]") - return raw_indices - - -def _reuse_prepared_reference_videos( - prepared: list[dict[str, Any]] | None, - *, - expected_count: int, -) -> list[dict[str, Any]] | None: - if prepared is None: - return None - if len(prepared) != expected_count: - raise OmniClientError("MiniMax H3 prepared-reference-video count does not match the request") - for item in prepared: - if not os.path.isfile(item["prepared_path"]): - raise OmniClientError(f"MiniMax H3 prepared reference video is unavailable: {item['prepared_path']}") - return prepared - - -def _validate_ref2va_reference_counts( - image_count: int, - video_count: int, - audio_count: int, -) -> None: - """Validate the official Ref2VA reference-count contract.""" - if image_count < 0 or video_count < 0 or audio_count < 0: - raise OmniClientError("MiniMax H3 reference counts must be non-negative") - if image_count + video_count == 0: - raise OmniClientError("ref2va requires at least one image or video reference") - if image_count > 9: - raise OmniClientError("ref2va accepts at most 9 image references") - if video_count > 3: - raise OmniClientError("ref2va accepts at most 3 video references") - if audio_count > 3: - raise OmniClientError("ref2va accepts at most 3 standalone audio references") - if image_count + video_count + audio_count > 12: - raise OmniClientError("ref2va accepts at most 12 total references") - - def _resolve_minimax_h3_num_outputs(value: Any) -> int: if value is None: return 1 @@ -572,17 +426,6 @@ def _minimax_h3_output_seeds(seed: int, num_outputs: int) -> list[int]: return [int(seed) + output_index for output_index in range(int(num_outputs))] -def _validate_reference_image(image: Image.Image) -> None: - width, height = image.size - if min(width, height) < 256 or max(width, height) > 5760: - raise OmniClientError( - f"MiniMax H3 reference image dimensions must be in [256, 5760] pixels, got {width}x{height}" - ) - ratio = width / height - if not 0.4 <= ratio <= 2.5: - raise OmniClientError(f"MiniMax H3 reference image aspect ratio must be in [0.4, 2.5], got {width}x{height}") - - def _dit_rank_world() -> tuple[Any, int, int]: if not dist.is_initialized(): return None, 0, 1 @@ -716,10 +559,8 @@ class MiniMaxH3Pipeline( on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}), ) _PROFILER_TARGETS: ClassVar[list[str]] = [ - "_prepare_reference_videos", "encode_prompt", - "_encode_visual_conditions", - "_encode_reference_audio_conditions", + "_encode_local_media", "diffuse", "decode", "video_vae.decode_latent", @@ -809,6 +650,18 @@ class MiniMaxH3Pipeline( f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, missing={missing[:5]}" ) + def _active_turbo_spec(self, sampling: Any) -> TurboSpec | None: + """Return the spec of the Turbo adapter this request actually applies. + + A recognized adapter at scale 0 contributes nothing, so it neither + constrains the task nor imposes its sampler contract. + """ + + lora_request = sampling.lora_request + if lora_request is None or math.isclose(0.0, float(sampling.lora_scale)): + return None + return self._turbo_lora_specs.get(lora_request.lora_int_id) + def _has_active_native_lora(self, sampling: Any) -> bool: lora_request = sampling.lora_request return ( @@ -868,18 +721,6 @@ class MiniMaxH3Pipeline( return adapter_schedule return self._base_schedule_for_task(task) - def _active_turbo_spec(self, sampling: Any) -> TurboSpec | None: - """Return the spec of the Turbo adapter this request actually applies. - - A recognized adapter at scale 0 contributes nothing, so it neither - constrains the task nor imposes its sampler contract. - """ - - lora_request = sampling.lora_request - if lora_request is None or math.isclose(0.0, float(sampling.lora_scale)): - return None - return self._turbo_lora_specs.get(lora_request.lora_int_id) - def _validate_turbo_sampling(self, sampling: Any, spec: TurboSpec) -> None: """Hold a request to the contract of the artifact that is loaded. @@ -935,7 +776,30 @@ class MiniMaxH3Pipeline( if int(self.parallel_config.cfg_parallel_size) != 1: raise ValueError("MiniMax-H3 is CFG-distilled and has no negative branch; cfg_parallel_size must be 1") self.device = get_local_device() - self.load_text_encoder = od_config.model_loaded["text_encoder"] + self.load_text_encoder = od_config.model_loaded.get("text_encoder", True) + self.load_vae_encoder = od_config.model_loaded.get("vae_encoder", True) + if self.load_vae_encoder is False and self.load_text_encoder is True: + raise ValueError( + "MiniMax H3 does not support local text encoding with external media conditioning; " + "set text_encoder=false or enable vae_encoder" + ) + self._encoder_modules = ["text_encoder"] if self.load_text_encoder else [] + self._PROFILER_TARGETS = list(type(self)._PROFILER_TARGETS) + encoder_block_attrs = dict(self._offload_plan.encoder_block_attrs) + if not self.load_text_encoder: + encoder_block_attrs.pop("text_encoder", None) + on_demand_component_paths = set(self._offload_plan.on_demand_component_paths) + if not self.load_text_encoder: + on_demand_component_paths.discard("text_encoder") + self._offload_plan = replace( + self._offload_plan, + encoder_block_attrs=encoder_block_attrs, + on_demand_component_paths=frozenset(on_demand_component_paths), + ) + if not self.load_text_encoder: + self._PROFILER_TARGETS.remove("encode_prompt") + if not self.load_vae_encoder: + self._PROFILER_TARGETS.remove("_encode_local_media") self.partition = _minimax_h3_partition_for_task( getattr(od_config, "task_type", None), str(od_config.model), @@ -1102,12 +966,14 @@ class MiniMaxH3Pipeline( os.path.join(model_path, "video_vae"), device=self.device, load_device=component_load_device, + decode_only=not self.load_vae_encoder, trust_remote_code=od_config.trust_remote_code, ) self.audio_vae = MiniMaxH3AudioVAE( os.path.join(model_path, "audio_vae"), device=self.device, load_device=component_load_device, + decode_only=not self.load_vae_encoder, trust_remote_code=od_config.trust_remote_code, ) # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``. @@ -1154,6 +1020,8 @@ class MiniMaxH3Pipeline( raise ValueError(f"MiniMax-H3 weight source {prefix!r} is not contiguous") loaded_prefixes.add(prefix) component = getattr(self, prefix.removesuffix(".")) + if component is None: + raise ValueError(f"MiniMax-H3 component {prefix.removesuffix('.')!r} is disabled in this deployment") stream = ((name[len(prefix) :], tensor) for name, tensor in grouped_weights) if prefix == "transformer." and self._fasth3 is not None: # Fuse before the model shards anything, which is also the only @@ -1227,11 +1095,12 @@ class MiniMaxH3Pipeline( def _resolve_task( self, requested: str | None, - multi_modal_data: dict[str, Any], + multi_modal_data: dict[str, Any] | None = None, *, turbo_spec: TurboSpec | None = None, has_native_lora: bool = False, ) -> str: + multi_modal_data = multi_modal_data or {} if requested is None: # A Ref2VA-only startup has no FL2VA transformer; preserve its # historical implicit default even for image-only references. @@ -1259,217 +1128,6 @@ class MiniMaxH3Pipeline( self._fasth3.check_task(task) return task - def _resolve_shape( - self, - task: str, - sampling: Any, - image: Image.Image | None, - ) -> tuple[int, int, int, int, int]: - fps = int(sampling.fps or MINIMAX_H3_FPS) - if fps != MINIMAX_H3_FPS: - raise OmniClientError(f"MiniMax H3 output fps is fixed at {MINIMAX_H3_FPS}") - extra = sampling.extra_args or {} - target = extra.get("target") - if target is not None and not isinstance(target, Mapping): - raise OmniClientError("MiniMax H3 extra_args['target'] must be an object") - target = target if isinstance(target, Mapping) else {} - duration = target.get("duration_seconds", extra.get("duration_seconds", extra.get("duration"))) - if duration is not None: - if isinstance(duration, bool): - raise OmniClientError(f"MiniMax H3 output duration must be in [4, 15] seconds, got {duration!r}") - try: - duration = float(duration) - except (TypeError, ValueError) as exc: - raise OmniClientError( - f"MiniMax H3 output duration must be in [4, 15] seconds, got {duration!r}" - ) from exc - if ( - not math.isfinite(duration) - or not MINIMAX_H3_MIN_OUTPUT_SECONDS <= duration <= MINIMAX_H3_MAX_OUTPUT_SECONDS - ): - raise OmniClientError(f"MiniMax H3 output duration must be in [4, 15] seconds, got {duration}") - requested_frames = int(round(duration * fps)) - elif int(sampling.num_frames or 1) > 1: - requested_frames = int(sampling.num_frames) - else: - requested_frames = 124 if task == "ref2va" else 209 - duration = requested_frames / fps - if not MINIMAX_H3_MIN_OUTPUT_SECONDS <= requested_frames / fps <= MINIMAX_H3_MAX_OUTPUT_SECONDS: - raise OmniClientError( - f"MiniMax H3 output duration must be in [4, 15] seconds, got {requested_frames / fps:.3f}" - ) - num_frames = minimax_h3_align_frame_count(requested_frames) - - height = sampling.height - width = sampling.width - aspect_ratio = target.get("aspect_ratio", extra.get("aspect_ratio")) - raw_short_edge = target.get("short_edge", extra.get("short_edge", MINIMAX_H3_OUTPUT_SHORT_EDGE)) - if isinstance(raw_short_edge, bool) or not isinstance(raw_short_edge, (int, np.integer)): - raise OmniClientError( - f"MiniMax H3 target.short_edge must be {MINIMAX_H3_OUTPUT_SHORT_EDGE}, got {raw_short_edge!r}" - ) - short_edge = int(raw_short_edge) - - aspect_ratio = _resolve_minimax_h3_aspect_ratio( - task, - aspect_ratio, - image, - ) - if not 0.25 <= aspect_ratio <= 4.0: - raise OmniClientError(f"MiniMax H3 canvas aspect ratio must be in [1:4, 4:1], got {aspect_ratio}") - - if height is None or width is None: - height, width = _resolve_output_canvas(aspect_ratio, short_edge) - height = int(height) // 32 * 32 - width = int(width) // 32 * 32 - if min(height, width) <= 0: - raise OmniClientError(f"invalid MiniMax H3 canvas {width}x{height}") - if width > 4 * height or height > 4 * width: - raise OmniClientError("MiniMax H3 canvas aspect ratio must be in [1:4, 4:1]") - - latent_t = MINIMAX_H3_SHAPE_PLANNER.video_latent_t(num_frames) - audio_t = MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(num_frames / fps) - return height, width, num_frames, latent_t, audio_t - - def encode_prompt( - self, - *, - task: str, - prompt: str, - image: Image.Image | None = None, - images: list[Image.Image] | None = None, - prepared_videos: list[dict[str, Any]] | None = None, - condition_labels: list[tuple[str, int]] | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - _, rank, _ = _dit_rank_world() - hidden = None - tags = None - ids = None - vision_kwargs: dict[str, torch.Tensor] = {} - images = list(images) if images is not None else ([image] if image is not None else []) - if rank == 0: - if task == "t2va": - ids = minimax_h3_text_only_ids(self.tokenizer, prompt) - tags = torch.ones(ids.shape[0], dtype=torch.long) - vision_kwargs = {} - else: - image_token_counts: list[int] = [] - if images: - vision = self.processor.image_processor( - images=images, - return_tensors="pt", - ) - image_grid = vision["image_grid_thw"] - merge = int(self.processor.image_processor.merge_size) ** 2 - image_token_counts = [int(grid.prod().item()) // merge for grid in image_grid] - vision_kwargs.update( - { - "pixel_values": vision["pixel_values"], - "image_grid_thw": image_grid, - } - ) - - video_block_counts: list[list[int]] = [] - video_block_timestamps: list[list[float]] = [] - if prepared_videos: - videos = [] - sampled_videos = [] - for index, item in enumerate(prepared_videos): - sampled = sample_reference_video_frames(item["prepared_path"]) - videos.append(np.stack(sampled["frames"])) - sampled_videos.append(sampled) - vision = self.processor.video_processor( - videos=videos, - do_sample_frames=False, - return_tensors="pt", - ) - video_grid = vision["video_grid_thw"] - merge = int(self.processor.image_processor.merge_size) ** 2 - for index, sampled in enumerate(sampled_videos): - blocks = int(video_grid[index, 0]) - per_block = int(video_grid[index, 1]) * int(video_grid[index, 2]) // merge - timestamps = sampled["block_timestamps"] - if len(timestamps) != blocks: - raise ValueError( - f"video block count mismatch: processor={blocks}, timestamps={len(timestamps)}" - ) - video_block_counts.append([per_block] * blocks) - video_block_timestamps.append(timestamps) - vision_kwargs.update( - { - "pixel_values_videos": vision["pixel_values_videos"], - "video_grid_thw": video_grid, - } - ) - - if not images and not prepared_videos: - raise OmniClientError(f"{task} requires an image or video condition") - if condition_labels is None: - condition_labels = [] - for image_index in range(1, len(images) + 1): - condition_labels.append(("image", image_index)) - audio_index = 0 - for video_index, item in enumerate(prepared_videos or (), start=1): - if item["input_has_audio"]: - audio_index += 1 - condition_labels.append(("audio", audio_index)) - condition_labels.append(("video", video_index)) - - if task == "fl2va": - if prepared_videos: - raise OmniClientError("fl2va does not accept video conditions") - ids, tags = minimax_h3_multi_image_presentation( - self.tokenizer, - prompt=prompt, - image_token_counts=image_token_counts, - ) - elif prepared_videos: - ids, tags = minimax_h3_ref2va_video_presentation( - self.tokenizer, - prompt=prompt, - condition_labels=condition_labels, - image_token_count=image_token_counts or None, - video_block_token_counts=video_block_counts, - video_block_timestamps=video_block_timestamps, - ) - else: - ids, tags = minimax_h3_ref2va_presentation( - self.tokenizer, - prompt=prompt, - condition_labels=condition_labels, - image_token_count=image_token_counts or None, - ) - - logger.info( - "MiniMax H3 %s Qwen presentation: %d tokens%s", - task, - int(ids.shape[0]), - ( - f", {len(images)} reference images" - + (f", {len(prepared_videos)} reference videos" if prepared_videos else "") - if images - else (f", {len(prepared_videos)} reference videos" if prepared_videos else "") - ), - ) - - if rank < self.text_encoder_tp_size: - # Distribute the encode inputs from the DiT main rank to the other - # encoder TP ranks, then run the distributed encode on all of them. - ids = self._distribute_encode_inputs(ids, vision_kwargs) - hidden = self._encode_text_hidden(ids, vision_kwargs) - - hidden = _broadcast_tensor( - hidden, - dtype=torch.bfloat16, - device=self.device, - ) - tags = _broadcast_tensor( - tags, - dtype=torch.long, - device=self.device, - ) - return hidden, tags - def _build_text_encoder_group(self, text_encoder_tp_size: int) -> Any: """Create the encoder tensor-parallel process group. @@ -1565,24 +1223,6 @@ class MiniMaxH3Pipeline( ) return ids - def _prepare_reference_videos( - self, - values: Any, - *, - target_frame_count: int, - workdir: str, - start_time_seconds: Any = None, - ) -> list[dict[str, Any]] | None: - _, rank, _ = _dit_rank_world() - if rank != 0: - return None - return prepare_reference_videos( - values, - target_frame_count=target_frame_count, - workdir=workdir, - start_time_seconds=start_time_seconds, - ) - def _encode_text_hidden( self, input_ids: torch.Tensor, @@ -1726,197 +1366,6 @@ class MiniMaxH3Pipeline( return tensor.cpu() return torch.empty(tensor.shape, dtype=tensor.dtype, device="meta") - 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. - needs_video_vae = video_count > 0 or (rank == 0 and bool(images)) - video_vae_context = self._component_on_device(self.video_vae) 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) - return (torch.cat(rows) if rows else None), shapes - - def _encode_audio_conditions_resident( - self, - audios: list[tuple[torch.Tensor, int]], - *, - max_duration_seconds: float | None = None, - ) -> tuple[torch.Tensor | None, list[int]]: - if not audios: - return None, [] - if max_duration_seconds is not None: - max_duration_seconds = float(max_duration_seconds) - if max_duration_seconds <= 0: - raise ValueError("max_duration_seconds must be positive") - _, rank, _ = _dit_rank_world() - rows = None - lengths = torch.zeros(len(audios), dtype=torch.long, device=self.device) - if rank == 0: - bounded_audios = [] - for waveform, sample_rate in audios: - if max_duration_seconds is not None: - max_samples = int(round(max_duration_seconds * int(sample_rate))) - waveform = waveform[..., :max_samples] - bounded_audios.append((waveform, sample_rate)) - encoded = [self.audio_vae.encode_waveform(*audio) for audio in bounded_audios] - rows = torch.cat([item[0] for item in encoded]) - lengths = torch.tensor( - [int(item[1]) for item in encoded], - dtype=torch.long, - device=self.device, - ) - group, _, world_size = _dit_rank_world() - if world_size > 1: - dist.broadcast(lengths, src=0, group=group) - return ( - _broadcast_tensor(rows, dtype=torch.float32, device=self.device), - [int(value) for value in lengths.tolist()], - ) - - def _encode_video_conditions_resident( - self, - prepared_videos: list[dict[str, Any]] | None, - *, - count: int, - ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]: - group, rank, world_size = _dit_rank_world() - distributed_encode = self.video_vae.is_distributed_enabled() - if distributed_encode: - # Native tiled encode uses collectives, so every VPP rank must - # enter each reference encode in the same input order. - prepared_videos_list = [prepared_videos] - dist.broadcast_object_list( - prepared_videos_list, - src=0, - group=group, - device=self.device, - ) - prepared_videos = prepared_videos_list[0] - - rows = None - shapes = torch.zeros((count, 3), dtype=torch.long, device=self.device) - if rank == 0 or distributed_encode: - if prepared_videos is None or len(prepared_videos) != count: - raise ValueError("reference-video preparation is incomplete") - encoded = [ - self.video_vae.encode_video(load_video_frames(item["prepared_path"])) for item in prepared_videos - ] - rows = torch.cat([item[0] for item in encoded]) - shapes = torch.tensor( - [item[1] for item in encoded], - dtype=torch.long, - device=self.device, - ) - if distributed_encode: - return ( - rows.to(device=self.device, dtype=torch.float32), - [tuple(int(value) for value in item) for item in shapes.tolist()], - ) - - if world_size > 1: - dist.broadcast(shapes, src=0, group=group) - return ( - _broadcast_tensor(rows, dtype=torch.float32, device=self.device), - [tuple(int(value) for value in item) for item in shapes.tolist()], - ) - - def _encode_video_audio_conditions_resident( - self, - prepared_videos: list[dict[str, Any]] | None, - *, - has_audio: list[bool], - ) -> tuple[torch.Tensor | None, list[int]]: - _, rank, _ = _dit_rank_world() - count = sum(has_audio) - if count == 0: - return None, [] - rows = None - lengths = torch.zeros(count, dtype=torch.long, device=self.device) - if rank == 0: - if prepared_videos is None: - raise ValueError("rank 0 reference-video preparation is incomplete") - encoded = [ - self.audio_vae.encode_waveform( - *load_video_audio( - item["original_path"], - start_time_seconds=float(item.get("start_time_seconds", 0.0)), - duration_seconds=item.get( - "audio_duration_seconds", - item.get("duration_seconds"), - ), - ) - ) - for item in prepared_videos - if item["input_has_audio"] - ] - rows = torch.cat([item[0] for item in encoded]) - lengths = torch.tensor( - [item[1] for item in encoded], - dtype=torch.long, - device=self.device, - ) - group, _, world_size = _dit_rank_world() - if world_size > 1: - dist.broadcast(lengths, src=0, group=group) - return ( - _broadcast_tensor(rows, dtype=torch.float32, device=self.device), - [int(value) for value in lengths.tolist()], - ) - - def _encode_reference_audio_conditions( - self, - prepared_videos: list[dict[str, Any]] | None, - *, - has_audio: list[bool], - standalone_audios: list[tuple[torch.Tensor, int]], - max_duration_seconds: float, - ) -> tuple[torch.Tensor | None, list[int], torch.Tensor | None, list[int]]: - # Embedded and standalone audio are consecutive direct Audio-VAE - # calls. Keep the component resident across both paths. - needs_audio_vae = any(has_audio) or bool(standalone_audios) - audio_vae_context = self._component_on_device(self.audio_vae) if needs_audio_vae else nullcontext() - with audio_vae_context: - embedded_condition, embedded_lengths = self._encode_video_audio_conditions_resident( - prepared_videos, - has_audio=has_audio, - ) - external_condition, external_lengths = self._encode_audio_conditions_resident( - standalone_audios, - max_duration_seconds=max_duration_seconds, - ) - return ( - embedded_condition, - embedded_lengths, - external_condition, - external_lengths, - ) - def _initial_noise( self, *, @@ -2342,13 +1791,16 @@ class MiniMaxH3Pipeline( if isinstance(raw_prompt, str): return None additional_information = raw_prompt.get("additional_information") or {} - text_encoder_output = additional_information.get("text_encoder_output") - if text_encoder_output is None: + encoder_output = additional_information.get("encoder_output") + if encoder_output is None: return None - if not isinstance(text_encoder_output, Mapping): - raise OmniClientError("text_encoder_output must be a mapping") + if not isinstance(encoder_output, Mapping): + raise OmniClientError("MiniMax H3 encoder output must be a mapping") try: - return MiniMaxH3TextConditioning.from_payload(text_encoder_output) + if "hidden_states" in encoder_output and "token_tags" in encoder_output: + return MiniMaxH3TextConditioning.from_payload(encoder_output) + conditioning = MiniMaxH3EncoderConditioning.from_omni_payload(encoder_output) + return MiniMaxH3TextConditioning(conditioning.hidden_states, conditioning.token_tags) except ValueError as exc: raise OmniClientError(str(exc)) from exc @@ -2369,34 +1821,234 @@ class MiniMaxH3Pipeline( raise OmniClientError(str(exc)) from exc return videos - def _prepare_request_inputs( - self, - *, - prompt: str, - multi_modal_data: dict[str, Any], - sampling: Any, - text_conditioning: MiniMaxH3TextConditioning | None = None, - prepared_reference_videos: list[dict[str, Any]] | None = None, - ) -> dict[str, Any]: - """Resolve the task and output shape, then run every request-level encode. - - Shared by request-mode :meth:`forward` and step-mode - :meth:`prepare_encode`; the returned mapping feeds :meth:`diffuse` and - :meth:`_build_denoise_inputs` unchanged. - """ - quality = sampling.quality - logger.debug("MiniMax H3 request quality=%s", quality) - extra = sampling.extra_args or {} - preencode_batch_frames = ( - normalize_preencode_batch_frames(extra.get("preencode_batch_frames", 17)) - if extra.get("preencode_mp4", False) - else 17 + def encode_prompt(self, prepared: PreparedEncoderInputs | None) -> tuple[torch.Tensor, torch.Tensor]: + _, rank, _ = _dit_rank_world() + ids = tags = hidden = None + vision_kwargs: dict[str, torch.Tensor] = {} + error = None + if rank == 0: + try: + if prepared is None: + raise ValueError("rank 0 must prepare MiniMax H3 text inputs") + if prepared.images: + vision = self.processor.image_processor(images=prepared.images, return_tensors="pt") + vision_kwargs.update( + pixel_values=vision["pixel_values"], + image_grid_thw=vision["image_grid_thw"], + ) + if prepared.qwen_videos: + vision = self.processor.video_processor( + videos=[frames for frames, _ in prepared.qwen_videos], + do_sample_frames=False, + return_tensors="pt", + ) + vision_kwargs.update( + pixel_values_videos=vision["pixel_values_videos"], + video_grid_thw=vision["video_grid_thw"], + ) + ids, tags = build_minimax_h3_presentation( + self.tokenizer, + prompt=prepared.prompt, + task=prepared.media.task, + condition_labels=prepared.condition_labels, + image_grid_thw=vision_kwargs.get("image_grid_thw"), + video_grid_thw=vision_kwargs.get("video_grid_thw"), + video_timestamps=prepared.video_timestamps, + merge_size=int(self.processor.image_processor.merge_size), + ) + except Exception as exc: + error = exc + _broadcast_rank0_exception(error) + if rank < self.text_encoder_tp_size: + ids = self._distribute_encode_inputs(ids, vision_kwargs) + hidden = self._encode_text_hidden(ids, vision_kwargs) + return ( + _broadcast_tensor(hidden, dtype=torch.bfloat16, device=self.device), + _broadcast_tensor(tags, dtype=torch.long, device=self.device), ) + + def _distribute_media_inputs(self, media: MiniMaxH3EncoderMediaInput | None) -> MiniMaxH3EncoderMediaInput: + group, rank, world_size = _dit_rank_world() + if world_size == 1: + assert media is not None + return media + tensors = media.to_mm_tensors() if media is not None else [] + header = [(media.to_metadata(), [value.dtype for value in tensors])] if media is not None else [None] + dist.broadcast_object_list(header, src=0, group=group) + metadata, dtypes = header[0] + received = [] + for index, dtype in enumerate(dtypes): + value = _broadcast_tensor(tensors[index] if rank == 0 else None, dtype=dtype, device=self.device) + if rank != 0: + received.append(value.cpu()) + del value + if rank == 0: + assert media is not None + return media + return MiniMaxH3EncoderMediaInput.from_mm_tensors(received, metadata) + + def _broadcast_media_conditioning( + self, conditioning: MiniMaxH3EncoderMediaConditioning | None + ) -> MiniMaxH3EncoderMediaConditioning: + group, rank, world_size = _dit_rank_world() + if world_size == 1: + assert conditioning is not None + return conditioning + tensor_names = ("visual_condition", "audio_condition") + header = [None] + if rank == 0: + assert conditioning is not None + metadata = { + item.name: getattr(conditioning, item.name) + for item in fields(conditioning) + if item.name not in tensor_names + } + dtypes = { + name: value.dtype if (value := getattr(conditioning, name)) is not None else None + for name in tensor_names + } + header[0] = (metadata, dtypes) + dist.broadcast_object_list(header, src=0, group=group) + metadata, dtypes = header[0] + tensors = {} + for name, dtype in dtypes.items(): + source = getattr(conditioning, name) if rank == 0 else None + tensors[name] = _broadcast_tensor(source, dtype=dtype, device=self.device) if dtype is not None else None + return MiniMaxH3EncoderMediaConditioning(**metadata, **tensors) + + def _encode_local_media(self, media: MiniMaxH3EncoderMediaInput | None) -> MiniMaxH3EncoderMediaConditioning: + group, rank, world_size = _dit_rank_world() + distributed_video = self.video_vae.is_distributed_enabled() + if distributed_video: + media = self._distribute_media_inputs(media) + conditioning = None + error = None + try: + if rank == 0 or distributed_video: + assert media is not None + conditioning = encode_media( + media, + video_vae=self.video_vae, + audio_vae=self.audio_vae, + emit_conditioning=rank == 0, + component_scope=self._component_on_device, + ) + except ValueError as exc: + error = OmniClientError(str(exc)) + except Exception as exc: + error = exc + if world_size > 1: + # All participants finish the codec phase before any latent broadcast. + errors = [None] * world_size + info = (str(error), isinstance(error, OmniClientError)) if error is not None else None + dist.all_gather_object(errors, info, group=group) + for info in errors: + if info is not None: + if error is not None: + raise error + message, is_client_error = info + raise OmniClientError(message) if is_client_error else RuntimeError(message) + elif error is not None: + raise error + return self._broadcast_media_conditioning(conditioning) + + def _prepare_local_conditioning( + self, + raw_prompt: Any, + sampling: Any, + *, + require_external_text: bool = False, + ) -> MiniMaxH3EncoderConditioning: + group, rank, world_size = _dit_rank_world() + prepared = text_conditioning = None + error = None + if rank == 0: + try: + _, multi_modal_data = self._extract_prompt(raw_prompt) + turbo_spec = self._active_turbo_spec(sampling) + has_native_lora = self._has_active_native_lora(sampling) + task = self._resolve_task( + (sampling.extra_args or {}).get("task"), + multi_modal_data, + turbo_spec=turbo_spec, + has_native_lora=has_native_lora, + ) + if turbo_spec is not None: + self._validate_turbo_sampling(sampling, turbo_spec) + if has_native_lora: + self._validate_native_sampling(sampling, task=task) + text_conditioning = self._extract_text_conditioning(raw_prompt) + if require_external_text and text_conditioning is None: + raise OmniClientError( + "MiniMax H3 diffusion stage requires text encoder conditioning when text_encoder is not loaded" + ) + prepared = prepare_encoder_inputs( + raw_prompt, + sampling, + task=task, + prepared_reference_videos=self._extract_prepared_reference_videos(raw_prompt), + ) + except Exception as exc: + error = exc + _broadcast_rank0_exception(error) + reuse_text = [text_conditioning is not None] + if world_size > 1: + dist.broadcast_object_list(reuse_text, src=0, group=group) + if reuse_text[0]: + hidden = _broadcast_tensor( + text_conditioning.hidden_states if rank == 0 else None, + dtype=torch.bfloat16, + device=self.device, + ) + tags = _broadcast_tensor( + text_conditioning.token_tags if rank == 0 else None, dtype=torch.long, device=self.device + ) + else: + hidden, tags = self.encode_prompt(prepared) + media = self._encode_local_media(prepared.media if prepared is not None else None) + return MiniMaxH3EncoderConditioning.from_components(MiniMaxH3TextConditioning(hidden, tags), media) + + def _prepare_request_inputs(self, raw_prompt: Any, sampling: Any) -> dict[str, Any]: + if self.load_text_encoder: + conditioning = self._prepare_local_conditioning(raw_prompt, sampling) + elif self.load_vae_encoder: + conditioning = self._prepare_local_conditioning( + raw_prompt, + sampling, + require_external_text=True, + ) + else: + conditioning = self._extract_encoder_conditioning(raw_prompt) + return self._prepare_encoder_conditioning_inputs(conditioning, sampling) + + @staticmethod + def _extract_encoder_conditioning(prompt: Any) -> MiniMaxH3EncoderConditioning: + if isinstance(prompt, list): + prompt = prompt[0] if prompt else None + additional_information = prompt.get("additional_information") if isinstance(prompt, Mapping) else None + payload = additional_information.get("encoder_output") if isinstance(additional_information, Mapping) else None + if not isinstance(payload, Mapping) or not payload: + raise OmniClientError("MiniMax H3 diffusion stage requires encoder conditioning from the encoder stage") + try: + return MiniMaxH3EncoderConditioning.from_omni_payload(payload) + except ValueError as exc: + raise OmniClientError(str(exc)) from exc + + def _prepare_encoder_conditioning_inputs( + self, + conditioning: MiniMaxH3EncoderConditioning, + sampling: Any, + ) -> dict[str, Any]: + extra = sampling.extra_args or {} + requested_task = extra.get("task") + if requested_task is not None and str(requested_task).lower() != conditioning.task: + raise OmniClientError( + f"MiniMax H3 encoder task {conditioning.task!r} does not match diffusion request {requested_task!r}" + ) turbo_spec = self._active_turbo_spec(sampling) has_native_lora = self._has_active_native_lora(sampling) task = self._resolve_task( - extra.get("task"), - multi_modal_data, + conditioning.task, turbo_spec=turbo_spec, has_native_lora=has_native_lora, ) @@ -2411,263 +2063,73 @@ class MiniMaxH3Pipeline( audio_shift=self.default_audio_shift, ) - raw_image = multi_modal_data.get("image") - raw_videos = multi_modal_data.get("video") - raw_audio = multi_modal_data.get("audio") - images = _load_images(raw_image) if raw_image is not None else [] - video_values = list(raw_videos) if isinstance(raw_videos, (list, tuple)) else raw_videos - audio_values = list(raw_audio) if isinstance(raw_audio, (list, tuple)) else raw_audio - - if task == "t2va" and (images or raw_videos is not None or raw_audio is not None): - raise OmniClientError("t2va does not accept image, video, or audio conditions") - if task == "fl2va": - if not images: - raise OmniClientError("fl2va requires multi_modal_data.image") - if len(images) > 2: - raise OmniClientError("fl2va accepts at most first and last images") - if raw_videos is not None or raw_audio is not None: - raise OmniClientError("fl2va accepts image keyframes only") - if task == "ref2va": - video_count = ( - len(video_values) if isinstance(video_values, (list, tuple)) else int(video_values is not None) + if conditioning.height % 32 or conditioning.width % 32: + raise OmniClientError( + f"MiniMax H3 encoder canvas must be divisible by 32, got {conditioning.width}x{conditioning.height}" ) - audio_is_waveform_pair = ( - isinstance(raw_audio, (list, tuple)) - and len(raw_audio) == 2 - and isinstance(raw_audio[1], (int, np.integer)) + expected_latent_t = MINIMAX_H3_SHAPE_PLANNER.video_latent_t(conditioning.num_frames) + expected_audio_t = MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(conditioning.num_frames / MINIMAX_H3_FPS) + if (conditioning.latent_t, conditioning.audio_t) != (expected_latent_t, expected_audio_t): + raise OmniClientError( + "MiniMax H3 encoder latent shape does not match its output frame count: " + f"got ({conditioning.latent_t}, {conditioning.audio_t}), expected " + f"({expected_latent_t}, {expected_audio_t})" ) - audio_count = ( - len(audio_values) - if isinstance(audio_values, (list, tuple)) and not audio_is_waveform_pair - else int(raw_audio is not None) - ) - _validate_ref2va_reference_counts(len(images), video_count, audio_count) - elif raw_videos is not None: - raise OmniClientError(f"{task} does not accept a video condition") - image = images[0] if images else None - height, width, num_frames, latent_t, audio_t = self._resolve_shape(task, sampling, image) - if task == "fl2va": - for item in images: - _validate_reference_image(item) - prepared_images = [item.resize((width, height), Image.Resampling.LANCZOS) for item in images] - keyframe_frame_indices = _resolve_fl2va_keyframe_indices(extra, len(images)) - elif task == "ref2va": - prepared_images = [] - for item in images: - ref_width, ref_height = _reference_image_shape(item) - prepared_images.append(item.resize((ref_width, ref_height), Image.Resampling.LANCZOS)) - keyframe_frame_indices = None - else: - prepared_images = [] - keyframe_frame_indices = None + visual_shapes = list(conditioning.visual_condition_shapes) or None + audio_lengths = list(conditioning.audio_condition_lengths) or None - visual_condition = None - visual_shape = None - visual_shapes = None - audio_condition = None - ref_audio_t = None - audio_lengths = None - ref_blocks = None - with tempfile.TemporaryDirectory(prefix="minimax_h3_ref2va_") as workdir: - prepared_videos = None - has_audio: list[bool] = [] - video_count = 0 - if raw_videos is not None: - video_count = len(raw_videos) if isinstance(raw_videos, (list, tuple)) else 1 - # File-based reference-video prep runs only on rank 0; other - # ranks return None without touching disk. If rank 0 raises - # (e.g. invalid file, unsupported codec) it must not exit - # ``prepare_encode`` before the downstream broadcasts below -- - # non-zero ranks would then deadlock on them forever. Capture - # the exception here and let every rank agree on the outcome - # before starting any subsequent collective. - prep_error: Exception | None = None - try: - _, rank, _ = _dit_rank_world() - if prepared_reference_videos is not None: - if rank == 0: - prepared_videos = _reuse_prepared_reference_videos( - prepared_reference_videos, - expected_count=video_count, - ) - else: - prepared_videos = self._prepare_reference_videos( - raw_videos, - target_frame_count=num_frames, - workdir=workdir, - start_time_seconds=extra.get("start_time_seconds"), - ) - except Exception as exc: - prep_error = exc - _broadcast_rank0_exception(prep_error) - has_audio_tensor = torch.zeros( - video_count, - dtype=torch.long, - device=self.device, - ) - _, rank, world_size = _dit_rank_world() - if rank == 0: - has_audio_tensor = torch.tensor( - [int(item["input_has_audio"]) for item in prepared_videos or []], - dtype=torch.long, - device=self.device, - ) - if world_size > 1: - dist.broadcast( - has_audio_tensor, - src=0, - group=get_world_group().device_group, - ) - has_audio = [bool(value) for value in has_audio_tensor.tolist()] - - if raw_audio is not None: - validate_reference_audio_files(raw_audio) - standalone_audios = _load_audios(raw_audio) if raw_audio is not None else [] - validate_reference_audio_waveforms(standalone_audios) - condition_labels: list[tuple[str, int]] = [] - for image_index in range(1, len(prepared_images) + 1): - condition_labels.append(("image", image_index)) - audio_index = 0 - for video_index, item in enumerate(prepared_videos or (), start=1): - if item["input_has_audio"]: - audio_index += 1 - condition_labels.append(("audio", audio_index)) - condition_labels.append(("video", video_index)) - for _ in standalone_audios: - audio_index += 1 - condition_labels.append(("audio", audio_index)) - - if text_conditioning is not None: - text_embeddings = text_conditioning.hidden_states.to( - device=self.device, - dtype=torch.bfloat16, - ) - text_tags = text_conditioning.token_tags.to( - device=self.device, - dtype=torch.long, - ) - elif getattr(self, "text_encoder", None) is not None: - text_embeddings, text_tags = self.encode_prompt( - task=task, - prompt=prompt, - images=prepared_images, - prepared_videos=prepared_videos, - condition_labels=condition_labels if task == "ref2va" else None, - ) - else: - raise OmniClientError( - "MiniMax H3 diffusion stage requires text_encoder_output when text_encoder is not loaded" - ) - - # ``prepared_videos`` is intentionally ``None`` on non-zero DiT - # ranks; the distributed video encoder broadcasts the prepared - # metadata inside ``_encode_visual_conditions``. Use the global - # video count here so video + standalone-audio Ref2VA requests do - # not look like audio-only requests on those ranks. - if video_count or prepared_images: - visual_condition, visual_shapes = self._encode_visual_conditions( - prepared_images, - prepared_videos, - video_count=video_count, - ) - ( - embedded_audio_condition, - embedded_audio_lengths, - external_audio_condition, - external_audio_lengths, - ) = self._encode_reference_audio_conditions( - prepared_videos, - has_audio=has_audio, - standalone_audios=standalone_audios, - max_duration_seconds=float(num_frames) / float(sampling.fps or MINIMAX_H3_FPS), - ) - audio_parts = [ - item for item in (embedded_audio_condition, external_audio_condition) if item is not None - ] - audio_condition = torch.cat(audio_parts) if audio_parts else None - audio_lengths = embedded_audio_lengths + external_audio_lengths - ref_blocks = [] - image_shapes = visual_shapes[: len(prepared_images)] - video_shapes = visual_shapes[len(prepared_images) :] - for shape in image_shapes: - ref_blocks.append( - { - "kind": "image", - "latent_h": shape[1], - "latent_w": shape[2], - } - ) - audio_iterator = iter(embedded_audio_lengths) - for shape, contributes_audio in zip(video_shapes, has_audio, strict=True): - ref_audio = next(audio_iterator) if contributes_audio else 0 - ref_blocks.append( - { - "kind": "video_audio" if ref_audio else "video", - "ref_audio_t": ref_audio, - "latent_t": shape[0], - "latent_h": shape[1], - "latent_w": shape[2], - } - ) - for ref_audio_t in external_audio_lengths: - ref_blocks.append({"kind": "audio", "ref_audio_t": ref_audio_t}) - elif standalone_audios: - raise OmniClientError("standalone audio references require a Ref2VA visual reference") - - if visual_shapes and len(visual_shapes) == 1: - visual_shape = visual_shapes[0] - if audio_lengths: - if any(length < 80 or length > 600 for length in audio_lengths): - raise OmniClientError("MiniMax H3 audio references must each be between 2 and 15 seconds") - # Video soundtracks and standalone audio have separate - # 15-second budgets, validated before encoding. Do not merge - # those budgets when concatenating their conditioning rows. - if len(audio_lengths) == 1: - ref_audio_t = audio_lengths[0] - - seed = int(sampling.seed if sampling.seed is not None else 42) base_schedule, num_steps = self._resolve_sigma_positions(task, sampling) - video_shift = float(extra.get("flow_shift", self.default_video_shift)) - audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift)) - pad_seq_len = _resolve_pad_seq_len(extra.get("pad_seq_len")) quality_plan = self._quality_policy.resolve( - quality=quality, + quality=sampling.quality, num_inference_steps=num_steps, extra_args=extra, ) self._cache_dit_runtime.prepare(quality_plan.cache_dit) - num_outputs = _resolve_minimax_h3_num_outputs(sampling.num_outputs_per_prompt) return { "task": task, - "height": height, - "width": width, - "num_frames": num_frames, - "latent_t": latent_t, - "latent_h": height // 16, - "latent_w": width // 16, - "audio_t": audio_t, - "text_embeddings": text_embeddings, - "text_tags": text_tags, - "visual_condition": visual_condition, - "visual_condition_shape": visual_shape, - "audio_condition": audio_condition, - "ref_audio_t": ref_audio_t, - "ref_blocks": ref_blocks, + "height": conditioning.height, + "width": conditioning.width, + "num_frames": conditioning.num_frames, + "latent_t": conditioning.latent_t, + "latent_h": conditioning.height // 16, + "latent_w": conditioning.width // 16, + "audio_t": conditioning.audio_t, + "text_embeddings": conditioning.hidden_states.to(device=self.device, dtype=torch.bfloat16), + "text_tags": conditioning.token_tags.to(device=self.device, dtype=torch.long), + "visual_condition": ( + conditioning.visual_condition.to(device=self.device) + if conditioning.visual_condition is not None + else None + ), + "visual_condition_shape": visual_shapes[0] if visual_shapes and len(visual_shapes) == 1 else None, + "audio_condition": ( + conditioning.audio_condition.to(device=self.device) + if conditioning.audio_condition is not None + else None + ), + "ref_audio_t": audio_lengths[0] if audio_lengths and len(audio_lengths) == 1 else None, + "ref_blocks": list(conditioning.ref_blocks) or None, "visual_condition_shapes": visual_shapes, "audio_condition_lengths": audio_lengths, - "keyframe_frame_indices": keyframe_frame_indices, - "pad_seq_len": pad_seq_len, - "seed": seed, + "keyframe_frame_indices": list(conditioning.keyframe_frame_indices) or None, + "pad_seq_len": _resolve_pad_seq_len(extra.get("pad_seq_len")), + "seed": int(sampling.seed if sampling.seed is not None else 42), "num_steps": num_steps, - "video_shift": video_shift, - "audio_shift": audio_shift, + "video_shift": float(extra.get("flow_shift", self.default_video_shift)), + "audio_shift": float(extra.get("audio_flow_shift", self.default_audio_shift)), "base_schedule": base_schedule, - "num_outputs": num_outputs, + "num_outputs": _resolve_minimax_h3_num_outputs(sampling.num_outputs_per_prompt), "preencode_mp4": bool(extra.get("preencode_mp4", False)), - "preencode_batch_frames": preencode_batch_frames, + "preencode_batch_frames": ( + normalize_preencode_batch_frames(extra.get("preencode_batch_frames", 17)) + if extra.get("preencode_mp4", False) + else 17 + ), "video_codec_options": normalize_video_codec_options( extra.get("video_codec_options", {"preset": "ultrafast", "threads": "0"}) + if extra.get("preencode_mp4", False) + else None ), } @@ -2680,14 +2142,9 @@ class MiniMaxH3Pipeline( def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput: if len(request.prompts) != 1: raise OmniClientError("MiniMax H3 supports one request at a time") - raw_prompt = request.prompts[0] - prompt, multi_modal_data = self._extract_prompt(raw_prompt) context = self._prepare_request_inputs( - prompt=prompt, - multi_modal_data=multi_modal_data, - sampling=request.sampling_params, - text_conditioning=self._extract_text_conditioning(raw_prompt), - prepared_reference_videos=self._extract_prepared_reference_videos(raw_prompt), + request.prompts[0], + request.sampling_params, ) denoise_kwargs = self._denoise_kwargs(context) num_outputs = context["num_outputs"] @@ -2794,13 +2251,9 @@ class MiniMaxH3Pipeline( "co-batched requests would reuse incompatible cache state. Drop --step-execution " "or omit quality=high." ) - prompt, multi_modal_data = self._extract_prompt(state.prompt) context = self._prepare_request_inputs( - prompt=prompt, - multi_modal_data=multi_modal_data, - sampling=state.sampling, - text_conditioning=self._extract_text_conditioning(state.prompt), - prepared_reference_videos=self._extract_prepared_reference_videos(state.prompt), + state.prompt, + state.sampling, ) inputs = self._build_denoise_inputs(**self._denoise_kwargs(context)) diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index e40251dff..d02649229 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -15,8 +15,10 @@ import torch import torch.distributed as dist import torch.nn as nn from PIL import Image +from safetensors import safe_open from transformers.dynamic_module_utils import get_class_from_dynamic_module from vllm.logger import init_logger +from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( DistributedVaeMixin, @@ -106,6 +108,121 @@ def _load_remote_component( return component_cls.from_pretrained(component_path) +def _check_trust_remote_code( + component_path: str, + config: dict[str, Any], + trust_remote_code: bool, +) -> None: + """Apply the remote-code gate before either full or encoder-only loading.""" + + class_reference = (config.get("auto_map") or {}).get("AutoModel") + if isinstance(class_reference, str) and not trust_remote_code: + raise ValueError( + f"Loading {component_path} executes the modeling code shipped with " + f"the checkpoint (auto_map.AutoModel = {class_reference}). Pass " + "--trust-remote-code (or trust_remote_code=True) to allow it." + ) + + +def _load_selected_state_dict( + model: nn.Module, + weights_path: Path, + prefixes: tuple[str, ...], +) -> None: + with safe_open(str(weights_path), framework="pt", device="cpu") as checkpoint: + selected = sorted(name for name in checkpoint.keys() if name.startswith(prefixes)) + if not selected: + raise ValueError(f"{weights_path}: no VAE weights match prefixes {prefixes!r}") + state_dict = {name: checkpoint.get_tensor(name) for name in selected} + model.load_state_dict(state_dict, strict=True) + + +def _remove_modules(model: nn.Module, names: tuple[str, ...]) -> None: + for name in names: + if not isinstance(getattr(model, name, None), nn.Module): + raise ValueError(f"MiniMax H3 VAE model has no module {name!r}") + delattr(model, name) + + +def _load_video_vae_encoder( + component_path: str, + config: dict[str, Any], +) -> nn.Module: + class_reference = (config.get("auto_map") or {}).get("AutoModel") + if not isinstance(class_reference, str): + raise ValueError(f"{component_path}/config.json must define auto_map.AutoModel") + component_cls = get_class_from_dynamic_module(class_reference, component_path) + component_module = importlib.import_module(component_cls.__module__) + source_classes = getattr(component_module, "_SOURCE_CLASSES", None) + source_class_name = config.get("source_class_name") + if not isinstance(source_classes, dict) or source_class_name not in source_classes: + raise ValueError(f"unsupported MiniMax H3 video VAE source class {source_class_name!r}") + if bool(config["vae_parallel_tiling"]): + ensure_parallel_state = getattr(component_module, "_ensure_vae_parallel_state", None) + if not callable(ensure_parallel_state): + raise ValueError("MiniMax H3 video VAE remote module does not expose its parallel-state initializer") + ensure_parallel_state() + + source_path = Path(component_path) / str(config["source_path"]) + weights_path = source_path / str(config["source_safetensors_path"]) + load_kwargs = { + "clip_length": int(config["vae_clip_length"]), + "token_drop": int(config["vae_token_drop"]), + "encoder_tiling": int(config["vae_encoder_tiling"]), + "decoder_tiling": int(config["vae_decoder_tiling"]), + "parallel_tiling": int(config["vae_parallel_tiling"]), + "tile_size": int(config["vae_tile_size"]), + "tile_overlap_min": int(config["vae_tile_overlap_min"]), + "encoder_parallel": int(config["vae_encoder_parallel"]), + "decoder_parallel": int(config["vae_decoder_parallel"]), + "chunk_dim": int(config["vae_chunk_dim"]), + } + source_cls = source_classes[source_class_name] + source_config = source_cls.load_config(str(source_path)) + with set_default_torch_dtype(torch.float32), torch.device("cpu"): + model, _unused = source_cls.from_config(source_config, return_unused_kwargs=True, **load_kwargs) + _remove_modules(model, ("decoder", "post_quant_conv")) + _load_selected_state_dict(model, weights_path, ("encoder.", "quant_conv.")) + model.eval() + return component_cls(model) + + +def _load_audio_vae_encoder( + component_path: str, + config: dict[str, Any], +) -> nn.Module: + class_reference = (config.get("auto_map") or {}).get("AutoModel") + if not isinstance(class_reference, str): + raise ValueError(f"{component_path}/config.json must define auto_map.AutoModel") + component_cls = get_class_from_dynamic_module(class_reference, component_path) + component_module = importlib.import_module(component_cls.__module__) + load_yaml = getattr(component_module, "_load_yaml", None) + model_cls = getattr(component_module, "DacAudioVAE", None) + if not callable(load_yaml) or not isinstance(model_cls, type): + raise ValueError("MiniMax H3 audio VAE remote module does not expose its source loader") + + component_dir = Path(component_path) + audio_config = load_yaml(component_dir / str(config["source_config_path"])) + metadata_path = component_dir / str(config["source_metadata_path"]) + metadata_doc = json.loads(metadata_path.read_text(encoding="utf-8")) + metadata = metadata_doc["metadata"]["kwargs"] + with set_default_torch_dtype(torch.float32), torch.device("cpu"): + model = model_cls( + encoder_rates=metadata["encoder_rates"], + decoder_rates=metadata["decoder_rates"], + attn_proj=metadata["attn_proj"], + decoder_type=metadata["decoder_type"], + decoder_dim=audio_config["model_config"]["decoder_dim"], + vae_latent_channels=audio_config["model_config"]["vae_latent_channels"], + sample_rate=metadata["sample_rate"], + ) + _remove_modules(model, ("decoder", "dec_in_proj", "logs_proj")) + weights_path = component_dir / str(config["source_safetensors_path"]) + _load_selected_state_dict(model, weights_path, ("encoder.", "pre_block.", "mean_proj.")) + model.eval() + return component_cls(model) + + class _AudioVAEDeterminismContext(AbstractContextManager): def __enter__(self): backends = torch.backends @@ -156,16 +273,24 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): *, device: torch.device, load_device: torch.device | None = None, + encode_only: bool = False, + decode_only: bool = False, trust_remote_code: bool = False, ) -> None: super().__init__() self._device_target = device + self.encode_only = bool(encode_only) + self.decode_only = bool(decode_only) + if self.encode_only and self.decode_only: + raise ValueError("MiniMax H3 video VAE cannot be both encode-only and decode-only") self.config_dict = _load_component_config(component_path) - self.remote = _load_remote_component( - component_path, - self.config_dict, - trust_remote_code=trust_remote_code, - ) + _check_trust_remote_code(component_path, self.config_dict, trust_remote_code) + if self.encode_only: + self.remote = _load_video_vae_encoder(component_path, self.config_dict) + else: + self.remote = _load_remote_component(component_path, self.config_dict, trust_remote_code=trust_remote_code) + if self.decode_only: + _remove_modules(self.remote.model, ("encoder", "quant_conv")) # Match the reference loader contract before installing inference-only # decoder fast paths. Keyframe encoding remains FP32; decoder Linear # weights may be materialized in FP16 because reference decode casts @@ -222,10 +347,11 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): self, parallel_size: int, mode: str = "tile", + process_group: dist.ProcessGroup | None = None, ) -> None: if mode != "tile": raise ValueError(f"MiniMax H3 VAE supports its native tile parallel mode only, got {mode!r}") - group = get_world_group().device_group + group = process_group if process_group is not None else get_world_group().device_group world_size = dist.get_world_size(group) rank = dist.get_rank(group) parallel_size = int(parallel_size) @@ -412,15 +538,40 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): cols, _, _ = self.model.split_tiles(int(latent.shape[-1]) * ratio, True) return len(rows) * len(cols) + def _encoder_tile_count(self, height: int, width: int) -> int: + processor = getattr(self.model, "processor", None) + align = getattr(processor, "_align_to_total_patch_size", None) + if callable(align): + height, width = align(int(height), int(width)) + rows, _, _ = self.model.split_tiles(int(height), False) + cols, _, _ = self.model.split_tiles(int(width), False) + return len(rows) * len(cols) + + def _encoder_tiling_context( + self, + height: int, + width: int, + ) -> AbstractContextManager: + parallel_size = int(getattr(self, "parallel_size", 1)) + if parallel_size <= 1: + return nullcontext() + num_tiles = self._encoder_tile_count(height, width) + if num_tiles >= parallel_size: + return nullcontext() + logger.warning_once( + "MiniMax-H3 VAE encode splits a %dx%d input into %d tile(s) but " + "the tile group has %d ranks; encoding rank-locally for this " + "shape instead, which avoids ranks without tiles hanging the " + "collective.", + width, + height, + num_tiles, + parallel_size, + ) + return self._rank_local_tiling() + @contextmanager def _rank_local_tiling(self) -> Iterator[None]: - """Run one decode with tiling kept on this rank, then restore the group. - - Used only when there are fewer tiles than ranks. Every rank then decodes - every tile, which is slower than sharing the work but is correct and - involves no collective. - """ - state = self._native_parallel_state() saved_state = dict(state) saved_tiling = self.model.parallel_tiling @@ -448,15 +599,24 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): @torch.inference_mode() def encode_image(self, image: Image.Image) -> torch.Tensor: + if getattr(self, "decode_only", False): + raise RuntimeError("MiniMax H3 decode-only video VAE cannot encode images") previous_parallel = self.model.parallel_tiling - self.model.parallel_tiling = False + if int(getattr(self, "parallel_size", 1)) <= 1: + self.model.parallel_tiling = False parameter = next(self.parameters()) previous_dtype = parameter.dtype if previous_dtype != torch.float32: self.to(torch.float32) devices = [parameter.device] if parameter.device.type != "cpu" else [] try: - with torch.random.fork_rng(devices=devices, device_type=parameter.device.type): + with ( + self._encoder_tiling_context(image.height, image.width), + torch.random.fork_rng( + devices=devices, + device_type=parameter.device.type, + ), + ): torch.default_generator.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED) for device in devices: with self.device_module.device(device): @@ -495,13 +655,26 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): self, frames: Any, ) -> tuple[torch.Tensor, tuple[int, int, int]]: + if getattr(self, "decode_only", False): + raise RuntimeError("MiniMax H3 decode-only video VAE cannot encode videos") parameter = next(self.parameters()) previous_dtype = parameter.dtype if previous_dtype != torch.float32: self.to(torch.float32) devices = [parameter.device] if parameter.device.type != "cpu" else [] + shape = getattr(frames, "shape", None) + if int(getattr(self, "parallel_size", 1)) > 1: + if shape is None or len(shape) != 4: + raise ValueError("parallel MiniMax H3 video encode requires a rank-4 frame array") + if int(shape[-1]) in {1, 3, 4}: + frame_height, frame_width = int(shape[-3]), int(shape[-2]) + else: + frame_height, frame_width = int(shape[-2]), int(shape[-1]) + tiling_context = self._encoder_tiling_context(frame_height, frame_width) + else: + tiling_context = nullcontext() try: - with torch.random.fork_rng(devices=devices, device_type=parameter.device.type): + with tiling_context, torch.random.fork_rng(devices=devices, device_type=parameter.device.type): torch.default_generator.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED) for device in devices: with self.device_module.device(device): @@ -561,6 +734,8 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): @torch.inference_mode() def decode_latent(self, latent: torch.Tensor) -> torch.Tensor: + if getattr(self, "encode_only", False): + 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)) @@ -598,6 +773,8 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): the complete decode path. After a callback failure, the remaining chunks are decoded and discarded before the exception is re-raised. """ + if getattr(self, "encode_only", False): + raise RuntimeError("MiniMax H3 encode-only video VAE cannot decode latents") if not callable(on_chunk): raise TypeError("on_chunk must be callable") if not callable(getattr(self.model, "_adaptive_decode", None)): @@ -622,16 +799,24 @@ class MiniMaxH3AudioVAE(nn.Module): *, device: torch.device, load_device: torch.device | None = None, + encode_only: bool = False, + decode_only: bool = False, trust_remote_code: bool = False, ) -> None: super().__init__() self._device_target = device + self.encode_only = bool(encode_only) + self.decode_only = bool(decode_only) + if self.encode_only and self.decode_only: + raise ValueError("MiniMax H3 audio VAE cannot be both encode-only and decode-only") self.config_dict = _load_component_config(component_path) - self.remote = _load_remote_component( - component_path, - self.config_dict, - trust_remote_code=trust_remote_code, - ) + _check_trust_remote_code(component_path, self.config_dict, trust_remote_code) + if self.encode_only: + self.remote = _load_audio_vae_encoder(component_path, self.config_dict) + else: + self.remote = _load_remote_component(component_path, self.config_dict, trust_remote_code=trust_remote_code) + if self.decode_only: + _remove_modules(self.remote.model, ("encoder", "pre_block", "mean_proj")) # The checkpoint's audio VAE contract is FP32 for both reference # encoding and waveform decoding. initial_device = load_device or device @@ -674,6 +859,8 @@ class MiniMaxH3AudioVAE(nn.Module): waveform: torch.Tensor, sample_rate: int, ) -> tuple[torch.Tensor, int]: + if getattr(self, "decode_only", False): + raise RuntimeError("MiniMax H3 decode-only audio VAE cannot encode waveforms") import torchaudio waveform = waveform.float() @@ -719,6 +906,8 @@ class MiniMaxH3AudioVAE(nn.Module): @torch.inference_mode() def decode_latent(self, latent: torch.Tensor) -> torch.Tensor: + if getattr(self, "encode_only", False): + raise RuntimeError("MiniMax H3 encode-only audio VAE cannot decode latents") channels = int(self.config_dict["latent_channels"]) mean = torch.tensor( self.config_dict["latents_mean"], diff --git a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py index e8350f41a..0143e9675 100644 --- a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py +++ b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py @@ -1,16 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 -"""MiniMax H3 checkpoint resolution for the AR text-encoder stage.""" - from pathlib import Path from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) -MINIMAX_H3_TEXT_ENCODER_DOWNLOAD_PATTERNS = { - "fl2va": ["FL2VA/text_encoder/**"], - "ref2va": ["Ref2VA/text_encoder/**"], +MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS = { + partition: [ + f"{subdir}/text_encoder/**", + f"{subdir}/video_vae/**", + f"{subdir}/audio_vae/**", + ] + for partition, subdir in {"fl2va": "FL2VA", "ref2va": "Ref2VA"}.items() } +MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["combined"] = MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["fl2va"] def resolve_minimax_h3_partition( @@ -32,7 +35,7 @@ def resolve_minimax_h3_partition( return "ref2va" if task == "ref2va" else "fl2va" -def resolve_minimax_h3_model_root( +def resolve_minimax_h3_encoder_model_root( model: str, revision: str | None, task_type: str | None, @@ -50,7 +53,7 @@ def resolve_minimax_h3_model_root( snapshot = download_weights_from_hf_specific( model_name_or_path=model, cache_dir=None, - allow_patterns=MINIMAX_H3_TEXT_ENCODER_DOWNLOAD_PATTERNS[partition], + allow_patterns=MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS[partition], revision=revision, require_all=True, ) @@ -59,4 +62,11 @@ def resolve_minimax_h3_model_root( return str(Path(snapshot) / subdir / "text_encoder") -__all__ = ["resolve_minimax_h3_model_root", "resolve_minimax_h3_partition"] +resolve_minimax_h3_model_root = resolve_minimax_h3_encoder_model_root + + +__all__ = [ + "resolve_minimax_h3_encoder_model_root", + "resolve_minimax_h3_model_root", + "resolve_minimax_h3_partition", +] diff --git a/vllm_omni/model_executor/models/minimax_h3/conditioning.py b/vllm_omni/model_executor/models/minimax_h3/conditioning.py index efea25771..761292dc3 100644 --- a/vllm_omni/model_executor/models/minimax_h3/conditioning.py +++ b/vllm_omni/model_executor/models/minimax_h3/conditioning.py @@ -4,8 +4,8 @@ from __future__ import annotations -from collections.abc import Mapping -from dataclasses import dataclass +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, fields from typing import Any import torch @@ -119,3 +119,502 @@ class MiniMaxH3TextConditioning: "hidden_states": self.hidden_states, "token_tags": self.token_tags, } + + +MINIMAX_H3_ENCODER_REQUEST_KEY = "minimax_h3_encoder_request" +MINIMAX_H3_ENCODER_LAYOUT_KEY = "minimax_h3_encoder_layout" +STAGE_SCHEMA_VERSION = 1 +_ENCODER_WIRE_SCHEMA_ID = 1 + +_TASK_TO_CODE = {"t2va": 1, "fl2va": 2, "ref2va": 3} +_CODE_TO_TASK = {value: key for key, value in _TASK_TO_CODE.items()} +_REF_KIND_TO_CODE = {"image": 1, "video": 2, "video_audio": 3, "audio": 4} +_CODE_TO_REF_KIND = {value: key for key, value in _REF_KIND_TO_CODE.items()} +_INTEGER_DTYPES = frozenset( + { + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + } +) +_VIDEO_CONDITION_WIDTH = 96 +_AUDIO_CONDITION_WIDTH = 32 +_AUDIO_CONDITION_CHANNELS = 2 + + +def _wire_layout(payload: Mapping[str, Any]) -> torch.Tensor: + private_metadata = payload.get("kv_metadata") + if not isinstance(private_metadata, Mapping): + raise ValueError("MiniMax H3 wire payload requires private layout metadata") + layout = private_metadata.get(MINIMAX_H3_ENCODER_LAYOUT_KEY) + if isinstance(layout, torch.Tensor): + return layout + raise ValueError("MiniMax H3 wire payload requires one packed encoder layout tensor") + + +def _validate_condition_tensors( + visual: torch.Tensor | None, + visual_shapes: Sequence[Sequence[int]], + audio: torch.Tensor | None, + audio_lengths: Sequence[int], +) -> None: + if (visual is None) != (not visual_shapes): + raise ValueError("MiniMax H3 visual condition and shapes must be present together") + if (audio is None) != (not audio_lengths): + raise ValueError("MiniMax H3 audio condition and lengths must be present together") + + parsed_shapes = [tuple(int(item) for item in shape) for shape in visual_shapes] + if any( + len(shape) != 3 or any(item <= 0 for item in shape) or shape[1] % 2 or shape[2] % 2 for shape in parsed_shapes + ): + raise ValueError("MiniMax H3 visual condition shapes must be positive [T, H, W] triplets with even H/W") + if visual is not None: + expected_rows = sum(t * (h // 2) * (w // 2) for t, h, w in parsed_shapes) + if ( + visual.dtype != torch.float32 + or visual.ndim != 2 + or tuple(visual.shape) + != ( + expected_rows, + _VIDEO_CONDITION_WIDTH, + ) + ): + raise ValueError( + "MiniMax H3 visual condition must be FP32 with shape " + f"[{expected_rows}, {_VIDEO_CONDITION_WIDTH}], got {visual.dtype} {tuple(visual.shape)}" + ) + + parsed_lengths = [int(length) for length in audio_lengths] + if any(length <= 0 for length in parsed_lengths): + raise ValueError("MiniMax H3 audio condition lengths must be positive") + if audio is not None: + expected_rows = _AUDIO_CONDITION_CHANNELS * sum(parsed_lengths) + if ( + audio.dtype != torch.float32 + or audio.ndim != 2 + or tuple(audio.shape) + != ( + expected_rows, + _AUDIO_CONDITION_WIDTH, + ) + ): + raise ValueError( + "MiniMax H3 audio condition must be FP32 with shape " + f"[{expected_rows}, {_AUDIO_CONDITION_WIDTH}], got {audio.dtype} {tuple(audio.shape)}" + ) + + +def _ref_block_rows(ref_blocks: Sequence[Mapping[str, Any]]) -> list[list[int]]: + rows: list[list[int]] = [] + for block in ref_blocks: + kind = str(block.get("kind") or "") + kind_code = _REF_KIND_TO_CODE.get(kind) + if kind_code is None: + raise ValueError(f"unsupported MiniMax H3 ref kind {kind!r}") + rows.append( + [ + kind_code, + int(block.get("ref_audio_t", 0)), + int(block.get("latent_t", 0)), + int(block.get("latent_h", 0)), + int(block.get("latent_w", 0)), + ] + ) + return rows + + +def _decode_ref_blocks(value: torch.Tensor) -> tuple[dict[str, Any], ...]: + if value.ndim != 2 or value.shape[1] != 5: + raise ValueError("MiniMax H3 ref_blocks metadata must have shape [N, 5]") + blocks: list[dict[str, Any]] = [] + for kind_code, ref_audio_t, ref_t, ref_h, ref_w in value.detach().cpu().tolist(): + kind = _CODE_TO_REF_KIND.get(int(kind_code)) + if kind is None: + raise ValueError(f"unsupported MiniMax H3 ref kind code {kind_code!r}") + block: dict[str, Any] = {"kind": kind} + for key, item in ( + ("ref_audio_t", ref_audio_t), + ("latent_t", ref_t), + ("latent_h", ref_h), + ("latent_w", ref_w), + ): + if item: + block[key] = int(item) + blocks.append(block) + return tuple(blocks) + + +@dataclass(frozen=True) +class MiniMaxH3EncoderMediaInput: + task: str + height: int + width: int + num_frames: int + latent_t: int + audio_t: int + images: tuple[torch.Tensor, ...] = () + videos: tuple[torch.Tensor, ...] = () + video_audios: tuple[tuple[torch.Tensor, int] | None, ...] = () + audios: tuple[tuple[torch.Tensor, int], ...] = () + keyframe_frame_indices: tuple[int, ...] = () + + @classmethod + def from_mm_tensors( + cls, + values: Sequence[torch.Tensor], + metadata: Mapping[str, Any], + ) -> MiniMaxH3EncoderMediaInput: + if not isinstance(values, Sequence) or isinstance(values, (str, bytes)): + raise ValueError("MiniMax H3 encoder media must be a tensor sequence") + tensors = list(values) + if any(not isinstance(item, torch.Tensor) for item in tensors): + raise ValueError("MiniMax H3 encoder media input must contain only tensors") + task = str(metadata["task"]) + if task not in _TASK_TO_CODE: + raise ValueError(f"unsupported MiniMax H3 task {task!r}") + height = int(metadata["height"]) + width = int(metadata["width"]) + num_frames = int(metadata["num_frames"]) + latent_t = int(metadata["latent_t"]) + audio_t = int(metadata["audio_t"]) + image_count = int(metadata.get("image_count", 0)) + video_count = int(metadata.get("video_count", 0)) + audio_count = int(metadata.get("audio_count", 0)) + if min(height, width, num_frames, latent_t, audio_t) <= 0: + raise ValueError("MiniMax H3 encoder media dimensions must be positive") + if min(image_count, video_count, audio_count) < 0: + raise ValueError("MiniMax H3 encoder media counts must be non-negative") + + audio_flags = tuple(bool(value) for value in metadata.get("video_audio_flags", ())) + video_audio_sample_rates = tuple(int(value) for value in metadata.get("video_audio_sample_rates", ())) + audio_sample_rates = tuple(int(value) for value in metadata.get("audio_sample_rates", ())) + if len(audio_flags) != video_count or len(video_audio_sample_rates) != video_count: + raise ValueError("MiniMax H3 video audio flags must align with videos") + if len(audio_sample_rates) != audio_count: + raise ValueError("MiniMax H3 audio sample rates must align with audio inputs") + if any((rate > 0) != has_audio for rate, has_audio in zip(video_audio_sample_rates, audio_flags, strict=True)): + raise ValueError("MiniMax H3 embedded audio sample rates must align with audio flags") + if any(rate <= 0 for rate in audio_sample_rates): + raise ValueError("MiniMax H3 audio sample rates must be positive") + + cursor = 0 + images = tuple(tensors[cursor : cursor + image_count]) + cursor += image_count + videos = tuple(tensors[cursor : cursor + video_count]) + cursor += video_count + if len(images) != image_count or any(image.ndim != 3 or image.shape[-1] != 3 for image in images): + raise ValueError("MiniMax H3 encoder images must have shape [H, W, 3]") + if len(videos) != video_count or any(video.ndim != 4 or video.shape[-1] != 3 for video in videos): + raise ValueError("MiniMax H3 encoder videos must have shape [T, H, W, 3]") + + video_audios: list[tuple[torch.Tensor, int] | None] = [] + for has_audio, sample_rate in zip(audio_flags, video_audio_sample_rates, strict=True): + if not has_audio: + video_audios.append(None) + continue + if cursor >= len(tensors): + raise ValueError("MiniMax H3 encoder media input is truncated") + waveform = tensors[cursor] + cursor += 1 + if waveform.ndim not in (1, 2): + raise ValueError("MiniMax H3 embedded audio must have shape [samples] or [channels, samples]") + video_audios.append((waveform, sample_rate)) + + audios: list[tuple[torch.Tensor, int]] = [] + for sample_rate in audio_sample_rates: + if cursor >= len(tensors): + raise ValueError("MiniMax H3 encoder media input is truncated") + waveform = tensors[cursor] + cursor += 1 + if waveform.ndim not in (1, 2): + raise ValueError("MiniMax H3 audio must have shape [samples] or [channels, samples]") + audios.append((waveform, sample_rate)) + if cursor != len(tensors): + raise ValueError(f"MiniMax H3 encoder media input has {len(tensors) - cursor} trailing tensors") + return cls( + task=task, + height=height, + width=width, + num_frames=num_frames, + latent_t=latent_t, + audio_t=audio_t, + images=images, + videos=videos, + video_audios=tuple(video_audios), + audios=tuple(audios), + keyframe_frame_indices=tuple(int(value) for value in metadata.get("keyframe_frame_indices", ())), + ) + + def to_mm_tensors(self) -> list[torch.Tensor]: + task_code = _TASK_TO_CODE.get(self.task) + if task_code is None: + raise ValueError(f"unsupported MiniMax H3 task {self.task!r}") + if min(self.height, self.width, self.num_frames, self.latent_t, self.audio_t) <= 0: + raise ValueError("MiniMax H3 encoder media dimensions must be positive") + if len(self.video_audios) != len(self.videos): + raise ValueError("MiniMax H3 video audio slots must align with videos") + if any(image.ndim != 3 or image.shape[-1] != 3 for image in self.images): + raise ValueError("MiniMax H3 encoder images must have shape [H, W, 3]") + if any(video.ndim != 4 or video.shape[-1] != 3 for video in self.videos): + raise ValueError("MiniMax H3 encoder videos must have shape [T, H, W, 3]") + audio_items = [item for item in self.video_audios if item is not None] + list(self.audios) + if any(waveform.ndim not in (1, 2) or int(sample_rate) <= 0 for waveform, sample_rate in audio_items): + raise ValueError("MiniMax H3 encoder audio must have a waveform and positive sample rate") + tensors = [*self.images, *self.videos] + for item in self.video_audios: + if item is None: + continue + waveform, _sample_rate = item + tensors.append(waveform) + tensors.extend(waveform for waveform, _sample_rate in self.audios) + return tensors + + def to_metadata(self) -> dict[str, Any]: + return { + "task": self.task, + "height": self.height, + "width": self.width, + "num_frames": self.num_frames, + "latent_t": self.latent_t, + "audio_t": self.audio_t, + "image_count": len(self.images), + "video_count": len(self.videos), + "audio_count": len(self.audios), + "video_audio_flags": [item is not None for item in self.video_audios], + "video_audio_sample_rates": [item[1] if item is not None else 0 for item in self.video_audios], + "audio_sample_rates": [sample_rate for _waveform, sample_rate in self.audios], + "keyframe_frame_indices": list(self.keyframe_frame_indices), + } + + +@dataclass(frozen=True) +class MiniMaxH3EncoderMediaConditioning: + task: str + height: int + width: int + num_frames: int + latent_t: int + audio_t: int + visual_condition: torch.Tensor | None = None + visual_condition_shapes: tuple[tuple[int, int, int], ...] = () + audio_condition: torch.Tensor | None = None + audio_condition_lengths: tuple[int, ...] = () + ref_blocks: tuple[dict[str, Any], ...] = () + keyframe_frame_indices: tuple[int, ...] = () + + def to_omni_components( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + task_code = _TASK_TO_CODE.get(self.task) + if task_code is None: + raise ValueError(f"unsupported MiniMax H3 task {self.task!r}") + if min(self.height, self.width, self.num_frames, self.latent_t, self.audio_t) <= 0: + raise ValueError("MiniMax H3 encoder media dimensions must be positive") + _validate_condition_tensors( + self.visual_condition, + self.visual_condition_shapes, + self.audio_condition, + self.audio_condition_lengths, + ) + empty = torch.empty((0,), dtype=torch.float32) + ref_rows = _ref_block_rows(self.ref_blocks) + layout = [ + _ENCODER_WIRE_SCHEMA_ID, + STAGE_SCHEMA_VERSION, + task_code, + self.height, + self.width, + self.num_frames, + self.latent_t, + self.audio_t, + len(self.visual_condition_shapes), + len(self.audio_condition_lengths), + len(ref_rows), + len(self.keyframe_frame_indices), + ] + layout.extend(item for shape in self.visual_condition_shapes for item in shape) + layout.extend(self.audio_condition_lengths) + layout.extend(item for row in ref_rows for item in row) + layout.extend(self.keyframe_frame_indices) + return ( + self.visual_condition if self.visual_condition is not None else empty, + self.audio_condition if self.audio_condition is not None else empty, + torch.tensor(layout, dtype=torch.int64), + ) + + +@dataclass(frozen=True) +class MiniMaxH3EncoderConditioning: + hidden_states: torch.Tensor + token_tags: torch.Tensor + task: str + height: int + width: int + num_frames: int + latent_t: int + audio_t: int + visual_condition: torch.Tensor | None = None + visual_condition_shapes: tuple[tuple[int, int, int], ...] = () + audio_condition: torch.Tensor | None = None + audio_condition_lengths: tuple[int, ...] = () + ref_blocks: tuple[dict[str, Any], ...] = () + keyframe_frame_indices: tuple[int, ...] = () + + @classmethod + def from_components( + cls, + text: MiniMaxH3TextConditioning, + media: MiniMaxH3EncoderMediaConditioning, + ) -> MiniMaxH3EncoderConditioning: + """Combine local encoder results without a stage-wire round trip.""" + text = MiniMaxH3TextConditioning.from_payload(text.to_payload()) + _validate_condition_tensors( + media.visual_condition, + media.visual_condition_shapes, + media.audio_condition, + media.audio_condition_lengths, + ) + return cls( + hidden_states=text.hidden_states, + token_tags=text.token_tags, + **{item.name: getattr(media, item.name) for item in fields(media)}, + ) + + @classmethod + def from_omni_payload(cls, payload: Mapping[str, Any]) -> MiniMaxH3EncoderConditioning: + hidden_payload = payload.get("hidden_states") + embed_payload = payload.get("embed") + meta_payload = payload.get("meta") + hidden_states = hidden_payload.get("output") if isinstance(hidden_payload, Mapping) else None + token_tags = meta_payload.get("token_role_ids") if isinstance(meta_payload, Mapping) else None + visual = embed_payload.get("embedding") if isinstance(embed_payload, Mapping) else None + audio = embed_payload.get("speech_feat") if isinstance(embed_payload, Mapping) else None + if not all(isinstance(item, torch.Tensor) for item in (hidden_states, token_tags, visual, audio)): + raise ValueError("MiniMax H3 encoder wire requires text, visual and audio tensors") + assert isinstance(hidden_states, torch.Tensor) + assert isinstance(token_tags, torch.Tensor) + assert isinstance(visual, torch.Tensor) + assert isinstance(audio, torch.Tensor) + if token_tags.ndim == 2 and token_tags.shape[-1] == 1: + token_tags = token_tags.squeeze(-1) + text = MiniMaxH3TextConditioning.from_payload({"hidden_states": hidden_states, "token_tags": token_tags}) + for name, value in (("visual", visual), ("audio", audio)): + if value.numel() == 0 and (value.dtype != torch.float32 or tuple(value.shape) != (0,)): + raise ValueError(f"MiniMax H3 empty {name} condition slot must be FP32 with shape [0]") + if not isinstance(meta_payload, Mapping): + raise ValueError("MiniMax H3 wire payload requires a meta mapping") + layout = _wire_layout(payload) + if layout.dtype not in _INTEGER_DTYPES or layout.ndim != 1: + raise ValueError("MiniMax H3 encoder layout must be a one-dimensional integer tensor") + values = [int(item) for item in layout.detach().cpu().tolist()] + if len(values) < 12: + raise ValueError("MiniMax H3 encoder layout header is truncated") + ( + schema_id, + version, + task_code, + height, + width, + num_frames, + latent_t, + audio_t, + visual_count, + audio_count, + ref_count, + keyframe_count, + ) = values[:12] + if schema_id != _ENCODER_WIRE_SCHEMA_ID or version != STAGE_SCHEMA_VERSION: + raise ValueError(f"unsupported MiniMax H3 encoder wire schema {schema_id}:{version}") + task = _CODE_TO_TASK.get(task_code) + if task is None: + raise ValueError(f"unsupported MiniMax H3 wire task code {task_code!r}") + if min(height, width, num_frames, latent_t, audio_t) <= 0: + raise ValueError("MiniMax H3 encoder wire dimensions must be positive") + if min(visual_count, audio_count, ref_count, keyframe_count) < 0: + raise ValueError("MiniMax H3 encoder layout counts must be non-negative") + expected_size = 12 + 3 * visual_count + audio_count + 5 * ref_count + keyframe_count + if len(values) != expected_size: + raise ValueError(f"MiniMax H3 encoder layout has {len(values)} integers, expected {expected_size}") + cursor = 12 + visual_shapes = tuple( + tuple(values[cursor + 3 * index : cursor + 3 * (index + 1)]) for index in range(visual_count) + ) + cursor += 3 * visual_count + audio_lengths = tuple(values[cursor : cursor + audio_count]) + cursor += audio_count + ref_rows = values[cursor : cursor + 5 * ref_count] + cursor += 5 * ref_count + ref_blocks_tensor = torch.tensor(ref_rows, dtype=torch.int64).reshape(-1, 5) + keyframe_frame_indices = tuple(values[cursor : cursor + keyframe_count]) + visual_condition = visual if visual.numel() else None + audio_condition = audio if audio.numel() else None + _validate_condition_tensors( + visual_condition, + visual_shapes, + audio_condition, + audio_lengths, + ) + return cls( + hidden_states=text.hidden_states, + token_tags=text.token_tags, + task=task, + height=height, + width=width, + num_frames=num_frames, + latent_t=latent_t, + audio_t=audio_t, + visual_condition=visual_condition, + visual_condition_shapes=visual_shapes, + audio_condition=audio_condition, + audio_condition_lengths=audio_lengths, + ref_blocks=_decode_ref_blocks(ref_blocks_tensor), + keyframe_frame_indices=keyframe_frame_indices, + ) + + def to_omni_payload(self) -> dict[str, Any]: + text = MiniMaxH3TextConditioning.from_payload( + { + "hidden_states": self.hidden_states, + "token_tags": self.token_tags, + } + ) + visual, audio, layout = MiniMaxH3EncoderMediaConditioning( + task=self.task, + height=self.height, + width=self.width, + num_frames=self.num_frames, + latent_t=self.latent_t, + audio_t=self.audio_t, + visual_condition=self.visual_condition, + visual_condition_shapes=self.visual_condition_shapes, + audio_condition=self.audio_condition, + audio_condition_lengths=self.audio_condition_lengths, + ref_blocks=self.ref_blocks, + keyframe_frame_indices=self.keyframe_frame_indices, + ).to_omni_components() + return { + "hidden_states": {"output": text.hidden_states}, + "embed": { + "embedding": visual, + "speech_feat": audio, + }, + "meta": {"token_role_ids": text.token_tags}, + "kv_metadata": {MINIMAX_H3_ENCODER_LAYOUT_KEY: layout}, + } + + +__all__ = [ + "MINIMAX_H3_CONDITION_LABELS_KEY", + "MINIMAX_H3_ENCODER_LAYOUT_KEY", + "MINIMAX_H3_ENCODER_REQUEST_KEY", + "MINIMAX_H3_PRESENTATION_TASK_KEY", + "MINIMAX_H3_TEXT_CONDITIONING_SCHEMA", + "MINIMAX_H3_TEXT_HIDDEN_SIZE", + "STAGE_SCHEMA_VERSION", + "MiniMaxH3EncoderConditioning", + "MiniMaxH3EncoderMediaConditioning", + "MiniMaxH3EncoderMediaInput", + "MiniMaxH3TextConditioning", +] diff --git a/vllm_omni/model_executor/models/minimax_h3/encoder.py b/vllm_omni/model_executor/models/minimax_h3/encoder.py new file mode 100644 index 000000000..44fefab91 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_h3/encoder.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from vllm.config import VllmConfig +from vllm.distributed.parallel_state import get_tp_group + +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.models.minimax_h3.vae import ( + MiniMaxH3AudioVAE, + MiniMaxH3VideoVAE, +) +from vllm_omni.model_executor.models.minimax_h3.conditioning import ( + MINIMAX_H3_ENCODER_LAYOUT_KEY, + MINIMAX_H3_ENCODER_REQUEST_KEY, + MiniMaxH3EncoderMediaConditioning, + MiniMaxH3EncoderMediaInput, +) +from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media +from vllm_omni.model_executor.models.minimax_h3.text_encoder import ( + MiniMaxH3TextEncoderBackbone, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +_MEDIA_CONDITIONING_CACHE_KEY = "_minimax_h3_encoder_media_conditioning" +_MEDIA_CACHE_KEY = "_minimax_h3_encoder_media" +_COMPONENT_CONFIG_KEY = "minimax_h3_encoder_components" +_COMPONENT_ROLES = frozenset({"text_encoder", "video_vae", "audio_vae"}) + + +def _partition_root(model_path: str) -> Path: + path = Path(model_path) + if path.name in {"text_encoder", "video_vae", "audio_vae"}: + return path.parent + return path + + +@dataclass(frozen=True) +class MiniMaxH3EncoderComponentConfig: + text_parallel_mode: str + video_parallel_mode: str + audio_parallel_mode: str + + def __post_init__(self) -> None: + if self.text_parallel_mode != "tp": + raise ValueError("MiniMax H3 text encoder currently supports parallel_mode='tp' only") + if self.video_parallel_mode not in {"leader", "patch"}: + raise ValueError("MiniMax H3 video VAE parallel_mode must be 'leader' or 'patch'") + if self.audio_parallel_mode != "leader": + raise ValueError("MiniMax H3 audio VAE currently supports parallel_mode='leader' only") + + @staticmethod + def _parallel_mode(raw: Mapping[str, Any], role: str) -> str: + value = raw[role] + if not isinstance(value, Mapping): + raise TypeError(f"MiniMax H3 encoder role {role!r} must be a mapping") + if set(value) != {"parallel_mode"}: + raise ValueError(f"MiniMax H3 encoder role {role!r} requires only parallel_mode") + return str(value["parallel_mode"]) + + @classmethod + def from_vllm_config(cls, vllm_config: VllmConfig) -> MiniMaxH3EncoderComponentConfig: + hf_config = vllm_config.model_config.hf_config + raw = getattr(hf_config, _COMPONENT_CONFIG_KEY, None) + if not isinstance(raw, Mapping): + raise ValueError(f"MiniMax H3 Encoder Stage requires explicit {_COMPONENT_CONFIG_KEY}") + roles = set(raw) + if roles != _COMPONENT_ROLES: + missing = sorted(_COMPONENT_ROLES - roles) + unknown = sorted(roles - _COMPONENT_ROLES) + raise ValueError( + f"{_COMPONENT_CONFIG_KEY} requires exactly {sorted(_COMPONENT_ROLES)}; " + f"missing={missing}, unknown={unknown}" + ) + return cls( + text_parallel_mode=cls._parallel_mode(raw, "text_encoder"), + video_parallel_mode=cls._parallel_mode(raw, "video_vae"), + audio_parallel_mode=cls._parallel_mode(raw, "audio_vae"), + ) + + +class MiniMaxH3Encoder(MiniMaxH3TextEncoderBackbone): + requires_request_sample_eligibility = True + enable_update_additional_information = True + omni_payload_at_request_end = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None: + super().__init__(vllm_config=vllm_config, prefix=prefix) + try: + tp_group = get_tp_group() + tp_rank = int(tp_group.rank_in_group) + except (AssertionError, RuntimeError): + tp_group = None + tp_rank = 0 + self._component_leader = tp_rank == 0 + self.component_config = MiniMaxH3EncoderComponentConfig.from_vllm_config(vllm_config) + root = _partition_root(vllm_config.model_config.model) + missing = [name for name in ("text_encoder", "video_vae", "audio_vae") if not (root / name).is_dir()] + if missing: + raise RuntimeError(f"MiniMax H3 Encoder Stage requires all three encoders; missing {missing}") + device = get_local_device() + self.video_vae: MiniMaxH3VideoVAE | None = None + self.audio_vae: MiniMaxH3AudioVAE | None = None + if self._component_leader or self.component_config.video_parallel_mode == "patch": + self.video_vae = MiniMaxH3VideoVAE( + str(root / "video_vae"), + device=device, + encode_only=True, + trust_remote_code=vllm_config.model_config.trust_remote_code, + ) + if self.component_config.video_parallel_mode == "patch": + if tp_group is None: + raise RuntimeError("MiniMax H3 video VAE patch parallelism requires an initialized TP group") + self.video_vae.set_parallel_size( + tp_group.world_size, + process_group=tp_group.device_group, + ) + if self._component_leader: + self.audio_vae = MiniMaxH3AudioVAE( + str(root / "audio_vae"), + device=device, + encode_only=True, + trust_remote_code=vllm_config.model_config.trust_remote_code, + ) + + @staticmethod + def _media_input(info: dict[str, Any]) -> MiniMaxH3EncoderMediaInput: + cached = info.get(_MEDIA_CACHE_KEY) + if isinstance(cached, MiniMaxH3EncoderMediaInput): + return cached + meta = info.get("meta") + metadata = meta.get(MINIMAX_H3_ENCODER_REQUEST_KEY) if isinstance(meta, Mapping) else None + if not isinstance(metadata, Mapping): + raise RuntimeError("MiniMax H3 unified encoder requires request metadata") + + hidden_states = info.get("hidden_states") + raw_media = hidden_states.get("layers") if isinstance(hidden_states, Mapping) else None + if raw_media is None: + media_tensors: list[torch.Tensor] = [] + elif isinstance(raw_media, Mapping): + try: + ordered_items = sorted(raw_media.items(), key=lambda item: int(item[0])) + except (TypeError, ValueError) as exc: + raise RuntimeError("MiniMax H3 encoder media keys must be integer indices") from exc + if [int(key) for key, _ in ordered_items] != list(range(len(ordered_items))): + raise RuntimeError("MiniMax H3 encoder media indices must be contiguous from zero") + media_tensors = [value for _, value in ordered_items] + else: + raise RuntimeError("MiniMax H3 encoder media must be a namespaced tensor mapping") + if any(not isinstance(value, torch.Tensor) for value in media_tensors): + raise RuntimeError("MiniMax H3 encoder media must contain only tensors") + media = MiniMaxH3EncoderMediaInput.from_mm_tensors(media_tensors, metadata) + if media.to_metadata() != dict(metadata): + raise RuntimeError("MiniMax H3 encoder media does not match request metadata") + info[_MEDIA_CACHE_KEY] = media + return media + + def _encode_media(self, media: MiniMaxH3EncoderMediaInput) -> MiniMaxH3EncoderMediaConditioning | None: + if not self._component_leader and self.component_config.video_parallel_mode == "leader": + return None + return encode_media( + media, + video_vae=self.video_vae, + audio_vae=self.audio_vae, + emit_conditioning=self._component_leader, + ) + + @staticmethod + def _has_encoder_input(info: object) -> bool: + if not isinstance(info, Mapping): + return False + meta = info.get("meta") + return isinstance(meta, Mapping) and isinstance(meta.get(MINIMAX_H3_ENCODER_REQUEST_KEY), Mapping) + + def make_omni_output( + self, + model_outputs: torch.Tensor | OmniOutput, + **kwargs: Any, + ) -> OmniOutput: + output = super().make_omni_output(model_outputs, **kwargs) + infos = kwargs.get("model_intermediate_buffer") or [] + if not isinstance(infos, Sequence) or not infos: + return output + sample_eligible = kwargs.get("request_sample_eligible") + if sample_eligible is None: + sample_eligible = [True] * len(infos) + if len(sample_eligible) != len(infos): + raise RuntimeError( + f"MiniMax H3 unified encoder received {len(sample_eligible)} sampling flags for {len(infos)} requests" + ) + per_request: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None] = [] + release_workspace = False + for eligible, info in zip(sample_eligible, infos, strict=True): + if not eligible: + per_request.append(None) + continue + if not isinstance(info, dict) or not self._has_encoder_input(info): + raise RuntimeError("MiniMax H3 unified encoder request metadata is missing from the runner") + release_workspace = True + if _MEDIA_CONDITIONING_CACHE_KEY not in info: + info[_MEDIA_CONDITIONING_CACHE_KEY] = self._encode_media(self._media_input(info)) + if not self._component_leader: + continue + cached = info.get(_MEDIA_CONDITIONING_CACHE_KEY) + if not isinstance(cached, MiniMaxH3EncoderMediaConditioning): + raise RuntimeError("MiniMax H3 encoder components did not produce conditioning") + per_request.append(cached.to_omni_components()) + + if release_workspace: + torch.accelerator.empty_cache() + if not self._component_leader: + return output + + if not any(item is not None for item in per_request): + return output + empty = torch.empty((0,), dtype=torch.float32) + multimodal_outputs = dict(output.multimodal_outputs or {}) + dotted_key = f"kv_metadata.{MINIMAX_H3_ENCODER_LAYOUT_KEY}" + multimodal_outputs[dotted_key] = [ + item[2] if item is not None else torch.empty((0,), dtype=torch.int64) for item in per_request + ] + + embed = dict(multimodal_outputs.get("embed") or {}) + embed["embedding"] = [item[0] if item is not None else empty for item in per_request] + embed["speech_feat"] = [item[1] if item is not None else empty for item in per_request] + multimodal_outputs["embed"] = embed + return OmniOutput( + text_hidden_states=output.text_hidden_states, + multimodal_outputs=multimodal_outputs, + intermediate_tensors=output.intermediate_tensors, + next_token_id=output.next_token_id, + ) + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> set[str]: + loaded = super().load_weights(weights) + for component_name in ("video_vae", "audio_vae"): + component = getattr(self, component_name) + if component is not None: + loaded.update(f"{component_name}.{name}" for name, _ in component.named_parameters()) + return loaded + + +__all__ = ["MiniMaxH3Encoder"] diff --git a/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py b/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py new file mode 100644 index 000000000..2d6998b7a --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py @@ -0,0 +1,569 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Shared request preparation and condition codec execution for MiniMax H3.""" + +from __future__ import annotations + +import math +import os +import tempfile +from collections.abc import Callable, Mapping, Sequence +from contextlib import AbstractContextManager, nullcontext +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +from PIL import Image + +from vllm_omni.errors import OmniClientError +from vllm_omni.model_executor.models.minimax_h3.conditioning import ( + MiniMaxH3EncoderMediaConditioning, + MiniMaxH3EncoderMediaInput, +) +from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + MINIMAX_H3_OUTPUT_SHORT_EDGE, + load_minimax_h3_images, + resolve_minimax_h3_aspect_ratio, + resolve_minimax_h3_output_canvas, + resolve_minimax_h3_reference_image_shape, +) +from vllm_omni.model_executor.models.minimax_h3.reference_video import ( + MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS, + load_audio_file, + load_video_audio, + load_video_frames, + prepare_reference_videos, + sample_reference_video_frames, + validate_reference_audio_files, + validate_reference_audio_waveforms, +) + +MINIMAX_H3_FPS = 24 +MINIMAX_H3_MIN_OUTPUT_SECONDS = 4.0 +MINIMAX_H3_MAX_OUTPUT_SECONDS = 15.0 + + +def _items(value: Any) -> list[Any]: + if value is None: + return [] + if isinstance(value, list): + return value + if isinstance(value, tuple) and not (len(value) == 2 and isinstance(value[1], Mapping)): + return list(value) + return [value] + + +def _audio_items(value: Any) -> list[Any]: + if value is None: + return [] + if isinstance(value, (list, tuple)) and len(value) == 2 and isinstance(value[1], (int, np.integer)): + return [value] + return list(value) if isinstance(value, (list, tuple)) else [value] + + +def _load_audio(value: Any) -> tuple[torch.Tensor, int]: + if isinstance(value, (list, tuple)) and not (len(value) == 2 and isinstance(value[1], (int, np.integer))): + audios = _load_audios(value) + if len(audios) != 1: + raise OmniClientError(f"MiniMax H3 expected one audio, got {len(audios)}") + return audios[0] + if isinstance(value, (str, os.PathLike)): + return load_audio_file(str(value)) + if isinstance(value, (list, tuple)) and len(value) == 2: + waveform, sample_rate = value + return torch.as_tensor(waveform).float(), int(sample_rate) + if isinstance(value, dict): + waveform = value.get("waveform", value.get("array")) + sample_rate = value.get("sample_rate", value.get("sampling_rate")) + if waveform is not None and sample_rate is not None: + return torch.as_tensor(waveform).float(), int(sample_rate) + raise OmniClientError("MiniMax H3 audio input must be a path, (waveform, sample_rate), or a waveform mapping") + + +def _load_audios(value: Any) -> list[tuple[torch.Tensor, int]]: + if isinstance(value, (list, tuple)) and not (len(value) == 2 and isinstance(value[1], (int, np.integer))): + if not value: + raise OmniClientError("MiniMax H3 audio input must not be empty") + return [_load_audio(item) for item in value] + return [_load_audio(value)] + + +def _as_int_list(value: Any, *, name: str) -> list[int]: + if isinstance(value, bool): + raise OmniClientError(f"{name} must be an integer or a list of integers") + if isinstance(value, (int, np.integer)): + return [int(value)] + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + result = list(value) + if not result: + raise OmniClientError(f"{name} must not be empty") + if any(isinstance(item, bool) or not isinstance(item, (int, np.integer)) for item in result): + raise OmniClientError(f"{name} must contain only integers") + return [int(item) for item in result] + raise OmniClientError(f"{name} must be an integer or a list of integers") + + +def _resolve_fl2va_keyframe_indices(extra: Mapping[str, Any], image_count: int) -> list[int]: + target = extra.get("target") + target = target if isinstance(target, Mapping) else {} + raw = extra.get("frame_indices", extra.get("frame_index")) + if raw is None: + raw = target.get("frame_indices", target.get("frame_index")) + raw_indices = ([0] if image_count == 1 else [0, -1]) if raw is None else _as_int_list(raw, name="frame_indices") + if len(raw_indices) != image_count: + raise OmniClientError( + f"MiniMax H3 FL2VA requires one frame index per image: got {raw_indices!r} for {image_count} image(s)" + ) + if tuple(raw_indices) not in ((0,), (-1,), (0, -1)): + raise OmniClientError("MiniMax H3 FL2VA frame_indices must be [0], [-1], or [0, -1]") + return raw_indices + + +def _validate_ref2va_reference_counts(image_count: int, video_count: int, audio_count: int) -> None: + if min(image_count, video_count, audio_count) < 0: + raise OmniClientError("MiniMax H3 reference counts must be non-negative") + if image_count + video_count == 0: + raise OmniClientError("ref2va requires at least one image or video reference") + if image_count > 9: + raise OmniClientError("ref2va accepts at most 9 image references") + if video_count > 3: + raise OmniClientError("ref2va accepts at most 3 video references") + if audio_count > 3: + raise OmniClientError("ref2va accepts at most 3 standalone audio references") + if image_count + video_count + audio_count > 12: + raise OmniClientError("ref2va accepts at most 12 total references") + + +def _validate_reference_image(image: Image.Image) -> None: + resolve_minimax_h3_reference_image_shape(image) + + +def resolve_minimax_h3_shape( + task: str, + sampling: Any, + image: Image.Image | None, +) -> tuple[int, int, int, int, int]: + # The diffusion package exports a pipeline that imports this shared module. + from vllm_omni.diffusion.models.minimax_h3.time_request import ( + MINIMAX_H3_SHAPE_PLANNER, + minimax_h3_align_frame_count, + ) + + fps = int(getattr(sampling, "fps", None) or MINIMAX_H3_FPS) + if fps != MINIMAX_H3_FPS: + raise OmniClientError(f"MiniMax H3 output fps is fixed at {MINIMAX_H3_FPS}") + extra = sampling.extra_args or {} + target = extra.get("target") + if target is not None and not isinstance(target, Mapping): + raise OmniClientError("MiniMax H3 extra_args['target'] must be an object") + target = target if isinstance(target, Mapping) else {} + duration = target.get("duration_seconds", extra.get("duration_seconds", extra.get("duration"))) + if duration is not None: + if isinstance(duration, bool): + raise OmniClientError(f"MiniMax H3 output duration must be in [4, 15] seconds, got {duration!r}") + try: + duration = float(duration) + except (TypeError, ValueError) as exc: + raise OmniClientError(f"MiniMax H3 output duration must be in [4, 15] seconds, got {duration!r}") from exc + if ( + not math.isfinite(duration) + or not MINIMAX_H3_MIN_OUTPUT_SECONDS <= duration <= MINIMAX_H3_MAX_OUTPUT_SECONDS + ): + raise OmniClientError(f"MiniMax H3 output duration must be in [4, 15] seconds, got {duration}") + requested_frames = int(round(duration * fps)) + elif int(getattr(sampling, "num_frames", None) or 1) > 1: + requested_frames = int(sampling.num_frames) + else: + requested_frames = 124 if task == "ref2va" else 209 + if not MINIMAX_H3_MIN_OUTPUT_SECONDS <= requested_frames / fps <= MINIMAX_H3_MAX_OUTPUT_SECONDS: + raise OmniClientError( + f"MiniMax H3 output duration must be in [4, 15] seconds, got {requested_frames / fps:.3f}" + ) + num_frames = minimax_h3_align_frame_count(requested_frames) + + height = getattr(sampling, "height", None) + width = getattr(sampling, "width", None) + aspect_ratio = resolve_minimax_h3_aspect_ratio( + task, + target.get("aspect_ratio", extra.get("aspect_ratio")), + image, + ) + raw_short_edge = target.get("short_edge", extra.get("short_edge", MINIMAX_H3_OUTPUT_SHORT_EDGE)) + if isinstance(raw_short_edge, bool) or not isinstance(raw_short_edge, (int, np.integer)): + raise OmniClientError( + f"MiniMax H3 target.short_edge must be {MINIMAX_H3_OUTPUT_SHORT_EDGE}, got {raw_short_edge!r}" + ) + if height is None or width is None: + height, width = resolve_minimax_h3_output_canvas(aspect_ratio, int(raw_short_edge)) + height = int(height) // 32 * 32 + width = int(width) // 32 * 32 + if min(height, width) <= 0: + raise OmniClientError(f"invalid MiniMax H3 canvas {width}x{height}") + if width > 4 * height or height > 4 * width: + raise OmniClientError("MiniMax H3 canvas aspect ratio must be in [1:4, 4:1]") + return ( + height, + width, + num_frames, + MINIMAX_H3_SHAPE_PLANNER.video_latent_t(num_frames), + MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(num_frames / fps), + ) + + +def _resolve_task( + extra_args: Mapping[str, Any], + multi_modal_data: Mapping[str, Any], +) -> str: + requested = extra_args.get("task") + if requested is not None: + return str(requested).lower() + if multi_modal_data.get("video") is not None or multi_modal_data.get("audio") is not None: + return "ref2va" + if multi_modal_data.get("image") is not None: + return "fl2va" + return "t2va" + + +def _prepare_encoder_images( + task: str, + images: list[Image.Image], + *, + height: int, + width: int, +) -> list[Any]: + if not images: + return [] + if task == "ref2va": + return [ + image.resize( + resolve_minimax_h3_reference_image_shape(image), + Image.Resampling.LANCZOS, + ) + for image in images + ] + if task != "fl2va": + return images + return [image.resize((width, height), Image.Resampling.LANCZOS) for image in images] + + +def _image_to_tensor(image: Image.Image) -> torch.Tensor: + array = np.asarray(image.convert("RGB"), dtype=np.uint8) + return torch.from_numpy(np.array(array, copy=True)).contiguous() + + +def _frames_to_tensor(frames: Sequence[Any]) -> torch.Tensor: + if len(frames) == 0: + raise ValueError("MiniMax H3 reference video must contain frames") + return torch.stack( + [ + _image_to_tensor(frame if isinstance(frame, Image.Image) else Image.fromarray(np.asarray(frame))) + for frame in frames + ] + ) + + +def _reuse_prepared_reference_videos( + prepared: list[dict[str, Any]] | None, + *, + expected_count: int, +) -> list[dict[str, Any]] | None: + if prepared is None: + return None + if len(prepared) != expected_count: + raise OmniClientError("MiniMax H3 prepared-reference-video count does not match the request") + for item in prepared: + if not os.path.isfile(item["prepared_path"]): + raise OmniClientError(f"MiniMax H3 prepared reference video is unavailable: {item['prepared_path']}") + return prepared + + +@dataclass(frozen=True) +class PreparedEncoderInputs: + """CPU-owned media and Qwen presentation inputs, independent of a runner.""" + + prompt: str + media: MiniMaxH3EncoderMediaInput + images: list[Image.Image] + qwen_videos: list[tuple[np.ndarray, dict[str, Any]]] + video_timestamps: list[list[float]] + condition_labels: list[tuple[str, int]] + + +def _effective_audio_inputs( + video_audios: Sequence[tuple[torch.Tensor, int] | None], + standalone_audios: Sequence[tuple[torch.Tensor, int]], + *, + max_standalone_seconds: float, +) -> list[tuple[torch.Tensor, int]]: + """Return the waveforms exactly as the audio WVAE will consume them.""" + audio_inputs = [item for item in video_audios if item is not None] + audio_inputs.extend( + (waveform[..., : int(round(max_standalone_seconds * sample_rate))], sample_rate) + for waveform, sample_rate in standalone_audios + ) + return audio_inputs + + +def prepare_encoder_inputs( + prompt: Any, + sampling: Any, + *, + task: str | None = None, + prepared_reference_videos: list[dict[str, Any]] | None = None, +) -> PreparedEncoderInputs: + """Decode references once and prepare the shared text/codec presentation.""" + if isinstance(prompt, str): + prompt = {"prompt": prompt} + if not isinstance(prompt, dict): + raise TypeError(f"MiniMax H3 expects a string or dict prompt, got {type(prompt)!r}") + + text = str(prompt.get("prompt") or "") + if not text: + raise OmniClientError("MiniMax H3 requires a non-empty prompt") + multi_modal_data = prompt.get("multi_modal_data") or {} + if not isinstance(multi_modal_data, Mapping): + raise TypeError("multi_modal_data must be a mapping") + + image_values = _items(multi_modal_data.get("image")) + videos = _items(multi_modal_data.get("video")) + raw_audio = multi_modal_data.get("audio") + audio_values = _audio_items(raw_audio) + diffusion_sampling_params = sampling + extra_args = getattr(diffusion_sampling_params, "extra_args", None) or {} + task = task if task is not None else _resolve_task(extra_args, multi_modal_data) + raw_images = load_minimax_h3_images(image_values) if image_values else [] + + if task == "t2va": + if raw_images or videos or audio_values: + raise OmniClientError("t2va does not accept image, video, or audio conditions") + elif task == "fl2va": + if not raw_images or videos or audio_values: + raise OmniClientError("fl2va requires image conditions only") + if len(raw_images) > 2: + raise OmniClientError("fl2va accepts at most first and last images") + for image in raw_images: + _validate_reference_image(image) + elif task == "ref2va": + _validate_ref2va_reference_counts(len(raw_images), len(videos), len(audio_values)) + if not raw_images and not videos: + raise OmniClientError("ref2va requires an image or video condition") + else: + raise OmniClientError(f"unsupported MiniMax H3 task {task!r}") + + height, width, num_frames, latent_t, audio_t = resolve_minimax_h3_shape( + task, + diffusion_sampling_params, + raw_images[0] if raw_images else None, + ) + images = _prepare_encoder_images( + task, + raw_images, + height=height, + width=width, + ) + keyframe_indices = _resolve_fl2va_keyframe_indices(extra_args, len(images)) if task == "fl2va" else [] + video_timestamps: list[list[float]] = [] + qwen_video_inputs: list[tuple[np.ndarray, dict[str, Any]]] = [] + condition_labels: list[tuple[str, int]] = [] + encoded_video_inputs: list[torch.Tensor] = [] + video_audio_inputs: list[tuple[torch.Tensor, int] | None] = [] + + if task == "fl2va": + condition_labels.extend(("image", index) for index in range(1, len(images) + 1)) + elif task == "ref2va": + condition_labels.extend(("image", index) for index in range(1, len(images) + 1)) + prepared_videos: list[dict[str, Any]] = [] + if videos: + with tempfile.TemporaryDirectory(prefix="minimax_h3_encoder_") as workdir: + prepared_videos = _reuse_prepared_reference_videos( + prepared_reference_videos, expected_count=len(videos) + ) + if prepared_videos is None: + prepared_videos = prepare_reference_videos( + videos, + target_frame_count=num_frames, + workdir=workdir, + start_time_seconds=extra_args.get("start_time_seconds"), + ) + for item in prepared_videos: + full_frames = load_video_frames(item["prepared_path"]) + encoded_video_inputs.append(_frames_to_tensor(full_frames)) + sampled = sample_reference_video_frames( + item["prepared_path"], + decoded_frames=full_frames, + ) + video_timestamps.append(sampled["block_timestamps"]) + frames = np.stack(sampled["frames"]) + frame_count = int(frames.shape[0]) + qwen_video_inputs.append( + ( + frames, + { + "total_num_frames": frame_count, + "fps": MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS, + "duration": frame_count / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS, + "video_backend": "minimax_h3", + "frames_indices": list(range(frame_count)), + "do_sample_frames": False, + }, + ) + ) + if item["input_has_audio"]: + waveform, sample_rate = load_video_audio( + item["original_path"], + start_time_seconds=float(item.get("start_time_seconds", 0.0)), + duration_seconds=item.get( + "audio_duration_seconds", + item.get("duration_seconds"), + ), + ) + video_audio_inputs.append((waveform.float().contiguous(), int(sample_rate))) + else: + video_audio_inputs.append(None) + audio_index = 0 + for video_index, item in enumerate(prepared_videos, start=1): + if item["input_has_audio"]: + audio_index += 1 + condition_labels.append(("audio", audio_index)) + condition_labels.append(("video", video_index)) + for _ in audio_values: + audio_index += 1 + condition_labels.append(("audio", audio_index)) + + if raw_audio is not None: + validate_reference_audio_files(raw_audio) + standalone_audios = _load_audios(raw_audio) if raw_audio is not None else [] + validate_reference_audio_waveforms(standalone_audios) + media_input = MiniMaxH3EncoderMediaInput( + task=task, + height=height, + width=width, + num_frames=num_frames, + latent_t=latent_t, + audio_t=audio_t, + images=tuple(_image_to_tensor(image) for image in images), + videos=tuple(encoded_video_inputs), + video_audios=tuple(video_audio_inputs), + audios=tuple((waveform.float().contiguous(), int(sample_rate)) for waveform, sample_rate in standalone_audios), + keyframe_frame_indices=tuple(keyframe_indices), + ) + audio_inputs = _effective_audio_inputs( + media_input.video_audios, + media_input.audios, + max_standalone_seconds=float(media_input.num_frames) / MINIMAX_H3_FPS, + ) + embedded_audio_count = sum(item is not None for item in media_input.video_audios) + # Video soundtracks and standalone references have separate 15-second budgets. + validate_reference_audio_waveforms(audio_inputs[:embedded_audio_count]) + validate_reference_audio_waveforms(audio_inputs[embedded_audio_count:]) + + return PreparedEncoderInputs( + prompt=text, + media=media_input, + images=images, + qwen_videos=qwen_video_inputs, + video_timestamps=video_timestamps, + condition_labels=condition_labels, + ) + + +def _image_from_tensor(value: torch.Tensor) -> Image.Image: + if value.ndim != 3 or value.shape[-1] != 3: + raise ValueError(f"MiniMax H3 image tensor must have shape [H, W, 3], got {tuple(value.shape)}") + array = value.detach().cpu().to(torch.uint8).contiguous().numpy() + return Image.fromarray(array, mode="RGB") + + +def encode_media( + media: MiniMaxH3EncoderMediaInput, + *, + video_vae: Any, + audio_vae: Any, + emit_conditioning: bool, + component_scope: Callable[[Any], AbstractContextManager[Any]] = nullcontext, +) -> MiniMaxH3EncoderMediaConditioning | None: + """Run codecs on participating ranks; only the output rank runs audio. + + The caller selects participating ranks and configures the VAE process group. + Every visual participant must call this function with the same references. + Component scopes allow the diffusion runner to stage one codec at a time. + """ + visual_rows: list[torch.Tensor] = [] + visual_shapes: list[tuple[int, int, int]] = [] + if media.images or media.videos: + if video_vae is None: + raise RuntimeError("MiniMax H3 video VAE is not resident on this rank") + with component_scope(video_vae): + for value in media.images: + image = _image_from_tensor(value) + visual_rows.append(video_vae.encode_image(image)) + visual_shapes.append((1, image.height // 16, image.width // 16)) + for value in media.videos: + frames = np.asarray(value.detach().cpu().to(torch.uint8).contiguous().numpy()) + rows, shape = video_vae.encode_video(frames) + visual_rows.append(rows) + visual_shapes.append(tuple(int(item) for item in shape)) + + if not emit_conditioning: + return None + + audio_rows: list[torch.Tensor] = [] + audio_lengths: list[int] = [] + embedded_audio_count = sum(item is not None for item in media.video_audios) + audio_inputs = _effective_audio_inputs( + media.video_audios, + media.audios, + max_standalone_seconds=float(media.num_frames) / MINIMAX_H3_FPS, + ) + if audio_inputs: + if audio_vae is None: + raise RuntimeError("MiniMax H3 audio WVAE is not resident on the encoder leader") + with component_scope(audio_vae): + for waveform, sample_rate in audio_inputs: + rows, length = audio_vae.encode_waveform(waveform, sample_rate) + audio_rows.append(rows) + audio_lengths.append(int(length)) + if audio_lengths: + if any(length < 80 or length > 600 for length in audio_lengths): + raise ValueError("MiniMax H3 audio references must each be between 2 and 15 seconds") + if sum(audio_lengths[:embedded_audio_count]) > 600 or sum(audio_lengths[embedded_audio_count:]) > 600: + raise ValueError("MiniMax H3 audio references must be at most 15 seconds in total") + + ref_blocks: list[dict[str, Any]] = [] + image_shapes = visual_shapes[: len(media.images)] + video_shapes = visual_shapes[len(media.images) :] + ref_blocks.extend({"kind": "image", "latent_h": shape[1], "latent_w": shape[2]} for shape in image_shapes) + embedded_lengths = iter(audio_lengths[:embedded_audio_count]) + for shape, embedded in zip(video_shapes, media.video_audios, strict=True): + ref_audio_t = int(next(embedded_lengths)) if embedded is not None else 0 + ref_blocks.append( + { + "kind": "video_audio" if ref_audio_t else "video", + "ref_audio_t": ref_audio_t, + "latent_t": shape[0], + "latent_h": shape[1], + "latent_w": shape[2], + } + ) + ref_blocks.extend( + { + "kind": "audio", + "ref_audio_t": int(length), + } + for length in audio_lengths[embedded_audio_count:] + ) + return MiniMaxH3EncoderMediaConditioning( + task=media.task, + height=media.height, + width=media.width, + num_frames=media.num_frames, + latent_t=media.latent_t, + audio_t=media.audio_t, + visual_condition=torch.cat(visual_rows) if visual_rows else None, + visual_condition_shapes=tuple(visual_shapes), + audio_condition=torch.cat(audio_rows) if audio_rows else None, + audio_condition_lengths=tuple(audio_lengths), + ref_blocks=tuple(ref_blocks), + keyframe_frame_indices=media.keyframe_frame_indices, + ) diff --git a/vllm_omni/model_executor/models/minimax_h3/pipeline.py b/vllm_omni/model_executor/models/minimax_h3/pipeline.py index 5ac94b5ee..085e2508a 100644 --- a/vllm_omni/model_executor/models/minimax_h3/pipeline.py +++ b/vllm_omni/model_executor/models/minimax_h3/pipeline.py @@ -16,18 +16,18 @@ MINIMAX_H3_PIPELINE = PipelineConfig( model_type="minimax_h3_disaggregated", default_deploy_config_name="minimax_h3_disaggregated.yaml", stage_cli_aliases={"text_encoder_tp_size": (0, "tensor_parallel_size")}, - model_arch="MiniMaxH3TextEncoder", + model_arch="MiniMaxH3Encoder", stages=( StagePipelineConfig( stage_id=0, - model_stage="text_encoder", + model_stage="encoder", execution_type=StageExecutionType.LLM_AR, input_sources=(), owns_tokenizer=True, requires_multimodal_data=True, - model_arch="MiniMaxH3TextEncoder", + model_arch="MiniMaxH3Encoder", engine_output_type="latent", - prompt_transform_func=f"{_PROCESSOR}.prepare_text_encoder_prompt", + prompt_transform_func=f"{_PROCESSOR}.prepare_encoder_prompt", sampling_constraints={ "max_tokens": 1, "temperature": 0.0, @@ -42,9 +42,9 @@ MINIMAX_H3_PIPELINE = PipelineConfig( input_sources=(0,), final_output=True, final_output_type="video", - requires_multimodal_data=True, + requires_multimodal_data=False, model_arch="MiniMaxH3Pipeline", - custom_process_input_func=f"{_PROCESSOR}.text_encoder2diffusion", + custom_process_input_func=f"{_PROCESSOR}.encoder2diffusion", omni_kv_config={"need_recv_cache": False}, model_path_resolver=f"{_DIFFUSION_PIPELINE}.resolve_minimax_h3_diffusion_model_path", inline_diffusion=True, diff --git a/vllm_omni/model_executor/models/minimax_h3/reference_video.py b/vllm_omni/model_executor/models/minimax_h3/reference_video.py index 4f98d9e13..4b361b26c 100644 --- a/vllm_omni/model_executor/models/minimax_h3/reference_video.py +++ b/vllm_omni/model_executor/models/minimax_h3/reference_video.py @@ -531,14 +531,27 @@ def _decode_video_frames_ffmpeg( ) -def sample_reference_video_frames(prepared_path: str) -> dict[str, Any]: - meta = _probe_video(prepared_path) +def sample_reference_video_frames( + prepared_path: str, + *, + decoded_frames: np.ndarray | None = None, +) -> dict[str, Any]: + if decoded_frames is None: + meta = _probe_video(prepared_path) + frame_count = int(meta["frame_count"]) + else: + decoded_frames = np.asarray(decoded_frames) + if decoded_frames.ndim != 4 or decoded_frames.shape[-1] != 3 or len(decoded_frames) <= 0: + raise OmniClientError( + f"decoded reference video frames must have shape [T, H, W, 3], got {decoded_frames.shape}" + ) + frame_count = len(decoded_frames) ratio = MINIMAX_H3_FPS / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS indices: list[int] = [] cursor = 0.0 while True: frame_index = int(round(cursor)) - if frame_index >= meta["frame_count"]: + if frame_index >= frame_count: break if not indices or frame_index > indices[-1]: indices.append(frame_index) @@ -546,16 +559,16 @@ def sample_reference_video_frames(prepared_path: str) -> dict[str, Any]: if not indices: raise OmniClientError(f"no frames sampled from {prepared_path}") - # The preparation step emits a lossless RGB stream. Decode only the sampled - # frames in one ffmpeg process so Qwen3VL sees the exact prepared pixels - # without decoding the entire high-bitrate stream into host memory. - decoded_frames = _decode_video_frames_ffmpeg( - prepared_path, - frame_count=int(meta["frame_count"]), - indices=indices, - width=meta["width"], - height=meta["height"], - ) + if decoded_frames is None: + decoded_frames = _decode_video_frames_ffmpeg( + prepared_path, + frame_count=frame_count, + indices=indices, + width=meta["width"], + height=meta["height"], + ) + else: + decoded_frames = decoded_frames[indices] frames = [np.asarray(frame) for frame in decoded_frames] timestamps = [index / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS for index in range(len(indices))] diff --git a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py index b1933af95..9ae74036e 100644 --- a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py +++ b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py @@ -176,7 +176,7 @@ class _ResidualMerge(nn.Module): info=Qwen3VLProcessingInfo, dummy_inputs=Qwen3VLDummyInputsBuilder, ) -class MiniMaxH3TextEncoder(Qwen3VLForConditionalGeneration): +class MiniMaxH3TextEncoderBackbone(Qwen3VLForConditionalGeneration): """Qwen3-VL encoder used by MiniMax H3. MiniMax H3 consumes the residual stream after decoder layer 50. The @@ -305,3 +305,6 @@ class MiniMaxH3TextEncoder(Qwen3VLForConditionalGeneration): encoder_weights(), mapper=self.hf_to_vllm_mapper, ) + + +MiniMaxH3TextEncoder = MiniMaxH3TextEncoderBackbone diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index e5d7b4115..1a519917c 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -476,10 +476,10 @@ _OMNI_MODELS = { "qwen3_vl", "AuraQwen3VLForConditionalGeneration", ), - "MiniMaxH3TextEncoder": ( + "MiniMaxH3Encoder": ( "minimax_h3", - "text_encoder", - "MiniMaxH3TextEncoder", + "encoder", + "MiniMaxH3Encoder", ), ## MiniMax-Music3 (text-to-music; AR talker -> flow-matching acoustic decoder) "MiniMaxMusic3TalkerForConditionalGeneration": ( diff --git a/vllm_omni/model_executor/stage_input_processors/minimax_h3.py b/vllm_omni/model_executor/stage_input_processors/minimax_h3.py index 87808adc3..78fd0225f 100644 --- a/vllm_omni/model_executor/stage_input_processors/minimax_h3.py +++ b/vllm_omni/model_executor/stage_input_processors/minimax_h3.py @@ -1,73 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""MiniMax H3 text-encoder stage input and output adapters.""" +"""Adapters between MiniMax H3 conditioning and the Omni stage runner.""" from __future__ import annotations import copy -import shutil -import tempfile from collections.abc import Mapping, Sequence from typing import Any -import numpy as np -from PIL import Image - -from vllm_omni.data_entry_keys import REQUEST_ARTIFACT_DIRS_KEY -from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_align_frame_count -from vllm_omni.errors import OmniClientError +from vllm_omni.data_entry_keys import unflatten_payload from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_executor.models.minimax_h3.conditioning import ( MINIMAX_H3_CONDITION_LABELS_KEY, + MINIMAX_H3_ENCODER_REQUEST_KEY, MINIMAX_H3_PRESENTATION_TASK_KEY, - MiniMaxH3TextConditioning, + MiniMaxH3EncoderConditioning, ) -from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( - MINIMAX_H3_OUTPUT_SHORT_EDGE, - load_minimax_h3_images, - resolve_minimax_h3_aspect_ratio, - resolve_minimax_h3_output_canvas, - resolve_minimax_h3_reference_image_shape, -) -from vllm_omni.model_executor.models.minimax_h3.reference_video import ( - MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY, - MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS, - prepare_reference_videos, - sample_reference_video_frames, - serialize_prepared_reference_videos, -) - - -def _items(value: Any) -> list[Any]: - if value is None: - return [] - if isinstance(value, list): - return value - if isinstance(value, tuple) and not (len(value) == 2 and isinstance(value[1], Mapping)): - return list(value) - return [value] - - -def _audio_items(value: Any) -> list[Any]: - if value is None: - return [] - if isinstance(value, (list, tuple)) and len(value) == 2 and isinstance(value[1], (int, np.integer)): - return [value] - return list(value) if isinstance(value, (list, tuple)) else [value] - - -def _resolve_task( - extra_args: Mapping[str, Any], - multi_modal_data: Mapping[str, Any], -) -> str: - requested = extra_args.get("task") - if requested is not None: - return str(requested).lower() - if multi_modal_data.get("video") is not None or multi_modal_data.get("audio") is not None: - return "ref2va" - if multi_modal_data.get("image") is not None: - return "fl2va" - return "t2va" +from vllm_omni.model_executor.models.minimax_h3.encoder_processing import prepare_encoder_inputs def _diffusion_sampling_params(sampling_params_list: Sequence[Any]) -> Any: @@ -78,190 +27,29 @@ def _diffusion_sampling_params(sampling_params_list: Sequence[Any]) -> Any: ] if len(diffusion_params) != 1: raise RuntimeError( - "MiniMax H3 text encoding requires exactly one OmniDiffusionSamplingParams stage parameter, " + "MiniMax H3 encoding requires exactly one OmniDiffusionSamplingParams stage parameter, " f"got {len(diffusion_params)}" ) return diffusion_params[0] -def _ref2va_target_frame_count(sampling_params_list: Sequence[Any]) -> int: - sampling = _diffusion_sampling_params(sampling_params_list) - extra_args = sampling.extra_args or {} - target = extra_args.get("target") - target = target if isinstance(target, Mapping) else {} - duration = target.get("duration_seconds", extra_args.get("duration_seconds", extra_args.get("duration"))) - if duration is not None: - requested = int(round(float(duration) * 24)) - elif int(getattr(sampling, "num_frames", None) or 1) > 1: - requested = int(sampling.num_frames) - else: - requested = 124 - return minimax_h3_align_frame_count(requested) - - -def _prepare_qwen_images( - task: str, - values: list[Any], - sampling_params_list: Sequence[Any], -) -> list[Any]: - if not values: - return [] - images = load_minimax_h3_images(values) - if task == "ref2va": - return [ - image.resize( - resolve_minimax_h3_reference_image_shape(image), - Image.Resampling.LANCZOS, - ) - for image in images - ] - if task != "fl2va": - return images - - sampling = _diffusion_sampling_params(sampling_params_list) - extra_args = sampling.extra_args or {} - target = extra_args.get("target") - if target is not None and not isinstance(target, Mapping): - raise OmniClientError("MiniMax H3 extra_args['target'] must be an object") - target = target if isinstance(target, Mapping) else {} - aspect_ratio = resolve_minimax_h3_aspect_ratio( - task, - target.get("aspect_ratio", extra_args.get("aspect_ratio")), - images[0], - ) - if not 0.25 <= aspect_ratio <= 4.0: - raise OmniClientError(f"MiniMax H3 canvas aspect ratio must be in [1:4, 4:1], got {aspect_ratio}") - height = sampling.height - width = sampling.width - if height is None or width is None: - short_edge = target.get( - "short_edge", - extra_args.get("short_edge", MINIMAX_H3_OUTPUT_SHORT_EDGE), - ) - if isinstance(short_edge, bool) or not isinstance(short_edge, (int, np.integer)): - raise OmniClientError( - f"MiniMax H3 target.short_edge must be {MINIMAX_H3_OUTPUT_SHORT_EDGE}, got {short_edge!r}" - ) - height, width = resolve_minimax_h3_output_canvas(aspect_ratio, int(short_edge)) - height = int(height) // 32 * 32 - width = int(width) // 32 * 32 - if min(height, width) <= 0: - raise OmniClientError(f"invalid MiniMax H3 canvas {width}x{height}") - if width > 4 * height or height > 4 * width: - raise OmniClientError("MiniMax H3 canvas aspect ratio must be in [1:4, 4:1]") - return [image.resize((width, height), Image.Resampling.LANCZOS) for image in images] - - -def prepare_text_encoder_prompt( +def prepare_encoder_prompt( prompt: Any, sampling_params_list: Sequence[Any], ) -> Any: - """Build H3's labeled Qwen3-VL presentation for Stage 0. - - The upstream Qwen3-VL multimodal processor expands each image/video - placeholder into the exact number of vision tokens and adds timestamped - video blocks. Audio is represented only by its H3 text label and is not - sent to Qwen3-VL. - """ + prepared = prepare_encoder_inputs(prompt, _diffusion_sampling_params(sampling_params_list)) if isinstance(prompt, str): - return prompt - if not isinstance(prompt, dict): - raise TypeError(f"MiniMax H3 expects a string or dict prompt, got {type(prompt)!r}") - - prompt = copy.copy(prompt) - additional_information = dict(prompt.get("additional_information") or {}) - meta = dict(additional_information.get("meta") or {}) - meta.pop(MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY, None) - additional_information["meta"] = meta - prompt["additional_information"] = additional_information - - text = str(prompt.get("prompt") or "") - if not text: - raise OmniClientError("MiniMax H3 requires a non-empty prompt") - multi_modal_data = prompt.get("multi_modal_data") or {} - if not isinstance(multi_modal_data, Mapping): - raise TypeError("multi_modal_data must be a mapping") - - image_values = _items(multi_modal_data.get("image")) - videos = _items(multi_modal_data.get("video")) - audios = _audio_items(multi_modal_data.get("audio")) - diffusion_sampling = _diffusion_sampling_params(sampling_params_list) - extra_args = diffusion_sampling.extra_args or {} - task = _resolve_task(extra_args, multi_modal_data) - images = _prepare_qwen_images(task, image_values, sampling_params_list) - qwen_video_inputs: list[tuple[np.ndarray, dict[str, Any]]] = [] - condition_labels: list[tuple[str, int]] = [] - - if task == "t2va": - if images or videos or audios: - raise OmniClientError("t2va does not accept image, video, or audio conditions") - elif task == "fl2va": - if not images or videos or audios: - raise OmniClientError("fl2va requires image conditions only") - condition_labels.extend(("image", index) for index in range(1, len(images) + 1)) - elif task == "ref2va": - if not images and not videos: - raise OmniClientError("ref2va requires an image or video condition") - condition_labels.extend(("image", index) for index in range(1, len(images) + 1)) - prepared_videos: list[dict[str, Any]] = [] - artifact_dir: str | None = None - if videos: - artifact_dir = tempfile.mkdtemp(prefix="minimax_h3_ref2va_") - try: - prepared_videos = prepare_reference_videos( - videos, - target_frame_count=_ref2va_target_frame_count(sampling_params_list), - workdir=artifact_dir, - start_time_seconds=extra_args.get("start_time_seconds"), - ) - for item in prepared_videos: - sampled = sample_reference_video_frames(item["prepared_path"]) - frames = np.stack(sampled["frames"]) - frame_count = int(frames.shape[0]) - qwen_video_inputs.append( - ( - frames, - { - "total_num_frames": frame_count, - "fps": MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS, - "duration": frame_count / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS, - "video_backend": "minimax_h3", - "frames_indices": list(range(frame_count)), - "do_sample_frames": False, - }, - ) - ) - except BaseException: - shutil.rmtree(artifact_dir, ignore_errors=True) - raise - audio_index = 0 - for video_index, item in enumerate(prepared_videos, start=1): - if item["input_has_audio"]: - audio_index += 1 - condition_labels.append(("audio", audio_index)) - condition_labels.append(("video", video_index)) - for _ in audios: - audio_index += 1 - condition_labels.append(("audio", audio_index)) - else: - raise OmniClientError(f"unsupported MiniMax H3 task {task!r}") - + prompt = {"prompt": prompt} + text = prepared.prompt + images = prepared.images + qwen_video_inputs = prepared.qwen_videos + condition_labels = prepared.condition_labels + media_input = prepared.media + task = media_input.task transformed = copy.copy(prompt) - if isinstance(prompt.get("additional_information"), Mapping): - transformed["additional_information"] = dict(prompt["additional_information"]) + additional_information = dict(prompt.get("additional_information") or {}) transformed["prompt"] = text - if task == "ref2va" and prepared_videos and artifact_dir is not None: - additional_information = dict(transformed.get("additional_information") or {}) - meta = dict(additional_information.get("meta") or {}) - meta[MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY] = serialize_prepared_reference_videos( - prepared_videos, - artifact_dir, - ) - additional_information["meta"] = meta - transformed["additional_information"] = additional_information - transformed[REQUEST_ARTIFACT_DIRS_KEY] = [artifact_dir] - qwen_mm_data = dict(multi_modal_data) - qwen_mm_data.pop("audio", None) + qwen_mm_data: dict[str, Any] = {} if images: qwen_mm_data["image"] = images if qwen_video_inputs: @@ -271,7 +59,16 @@ def prepare_text_encoder_prompt( mm_processor_kwargs = dict(prompt.get("mm_processor_kwargs") or {}) mm_processor_kwargs[MINIMAX_H3_PRESENTATION_TASK_KEY] = task mm_processor_kwargs[MINIMAX_H3_CONDITION_LABELS_KEY] = condition_labels + media_tensors = media_input.to_mm_tensors() transformed["mm_processor_kwargs"] = mm_processor_kwargs + + hidden_states = dict(additional_information.get("hidden_states") or {}) + hidden_states["layers"] = dict(enumerate(media_tensors)) + additional_information["hidden_states"] = hidden_states + meta = dict(additional_information.get("meta") or {}) + meta[MINIMAX_H3_ENCODER_REQUEST_KEY] = media_input.to_metadata() + additional_information["meta"] = meta + transformed["additional_information"] = additional_information return transformed @@ -295,18 +92,20 @@ def _global_request_id(prompt: Mapping[str, Any]) -> str | None: return str(value) if value is not None else None -def text_encoder2diffusion( +def encoder2diffusion( source_outputs: list[Any], prompt: Any = None, requires_multimodal_data: bool = False, streaming_context: Any | None = None, ) -> dict[str, Any] | None: - """Attach Stage 0 hidden states and token tags to the original request.""" + """Reuse the encoder handoff for all three H3 encoder components.""" del requires_multimodal_data, streaming_context if not source_outputs: return None if len(source_outputs) != 1: - raise RuntimeError(f"MiniMax H3 diffusion requires exactly one text-encoder source, got {len(source_outputs)}") + raise RuntimeError(f"MiniMax H3 DiT requires exactly one encoder source, got {len(source_outputs)}") + if not getattr(source_outputs[0], "finished", True): + return None diffusion_prompt = _original_prompt(prompt) source_output = source_outputs[0] @@ -318,25 +117,32 @@ def text_encoder2diffusion( and str(source_request_id) != expected_request_id ): raise RuntimeError( - "MiniMax H3 text-encoder request ID does not match the diffusion request: " + "MiniMax H3 encoder request ID does not match the diffusion request: " f"source={source_request_id!r}, expected={expected_request_id!r}" ) outputs = getattr(source_output, "outputs", None) if not isinstance(outputs, list) or len(outputs) != 1: output_count = len(outputs) if isinstance(outputs, list) else 0 - raise RuntimeError(f"MiniMax H3 text encoder must return exactly one completion, got {output_count}") - - completion = outputs[0] - payload = completion.multimodal_output + raise RuntimeError(f"MiniMax H3 encoder must return exactly one completion, got {output_count}") + payload = getattr(outputs[0], "multimodal_output", None) if not isinstance(payload, Mapping): - raise RuntimeError("MiniMax H3 text encoder returned no conditioning payload") + raise RuntimeError("MiniMax H3 encoder returned no conditioning payload") try: - conditioning = MiniMaxH3TextConditioning.from_omni_payload(payload) + conditioning = MiniMaxH3EncoderConditioning.from_omni_payload(unflatten_payload(dict(payload))) except ValueError as exc: raise RuntimeError(str(exc)) from exc additional_information = dict(diffusion_prompt.get("additional_information") or {}) - additional_information["text_encoder_output"] = conditioning.to_payload() + additional_information.pop("hidden_states", None) + meta = dict(additional_information.get("meta") or {}) + meta.pop(MINIMAX_H3_ENCODER_REQUEST_KEY, None) + if meta: + additional_information["meta"] = meta + else: + additional_information.pop("meta", None) + additional_information["encoder_output"] = conditioning.to_omni_payload() diffusion_prompt["additional_information"] = additional_information + diffusion_prompt["multi_modal_data"] = None + diffusion_prompt.pop("model_intermediate_buffer", None) return diffusion_prompt