[feat] Extend disaggregated MiniMax-H3 stage with VAE encoders (#6939)

Signed-off-by: asukaqaq-s <1311722138@qq.com>
Signed-off-by: Jader <yjader@foxmail.com>
Co-authored-by: Jader <yjader@foxmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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 |
 

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.

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.

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}

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

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

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

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,
     }
 

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)

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"

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

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

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

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

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

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"],

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",
+]

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",
+]

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"]

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

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,

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

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

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": (

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