[Feature] Support text encoder Disaggregation for H3 (#5885)

Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: yuanwu <yuan.wu@intel.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: yuanwu <yuan.wu@intel.com>
This commit is contained in:
authored and GitHub committed 2026-08-28 14:27:05 +08:00
1 parent ed837e350e
commit b6a4baa659
61 files changed
+2952 -366

No files matched your search

@@ -0,0 +1,82 @@
# MiniMax H3 disaggregated text encoder
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.
## Start the server
Choose the topology explicitly and load its deployment defaults:
```bash
vllm-omni serve MiniMaxAI/MiniMax-H3 \
--omni \
--deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml
```
The default deployment assigns stage 0 to GPUs 0-1 with tensor parallel size
2 and `max_num_seqs: 1`. Stage 1 uses GPUs 2-5 with diffusion tensor parallel
size 1, Ulysses degree 4, and VAE patch parallel size 4. Adjust the
`devices`, `tensor_parallel_size`, and stage 1 `parallel_config` values in a
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).
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:
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3,4 \
vllm-omni serve MiniMaxAI/MiniMax-H3 \
--omni \
--deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml \
--stage-overrides '{"0":{"devices":"0","tensor_parallel_size":1},"1":{"devices":"1,2,3,4","tensor_parallel_size":1,"ulysses_degree":4,"vae_patch_parallel_size":4}}'
```
For memory-constrained deployments, start from the CPU-offload or distributed
layerwise-offload profiles in [MiniMax-H3.md](MiniMax-H3.md). Apply memory and
quantization options only to Stage 1 with `--stage-overrides`; retain the
encoder's BF16 configuration and the video/audio VAEs' FP32 precision. Select
one offload strategy per deployment:
```bash
# Stage 1 model-level CPU offload.
--stage-overrides '{"1":{"enable_cpu_offload":true}}'
# Stage 1 distributed layerwise offload. Tune resident layers for available RAM.
--stage-overrides '{"1":{"enable_distributed_layerwise_offload":true,"dlo_use_allgather":false,"dlo_resident_layers":20}}'
# Stage 1 online FP8 quantization of the DiT only.
--stage-overrides '{"1":{"diffusion_quantization_config":"{\"transformer\":{\"method\":\"fp8\"}}"}}'
```
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
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.
The `/v1/videos` request schema and `extra_params.task` values (`t2va`,
`fl2va`, and `ref2va`) are unchanged from the single-stage recipe.
## Turbo LoRA
MiniMax-H3 Turbo uses five sigma points, `flow_shift=6`, and
`audio_flow_shift=3`. Start Turbo deployments with the dedicated defaults so
requests that omit sampling controls do not inherit the 50-step base schedule:
```bash
vllm-omni serve MiniMaxAI/MiniMax-H3 \
--omni \
--lora-path /path/to/MiniMax-H3-Turbo \
--deploy-config vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml
```
The base deployment intentionally retains 50 inference steps for non-LoRA
quality. Turbo LoRA supports T2VA and FL2VA requests, not Ref2VA.
+127
View File
@@ -34,6 +34,7 @@ from vllm_omni.config.stage_config import (
build_stage_runtime_overrides,
load_deploy_config,
merge_pipeline_deploy,
normalize_pipeline_cli_overrides,
pipeline_cfg_resolver,
)
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
@@ -429,6 +430,20 @@ class TestStageConfigFactory:
assert cfg["final_output"] is True
assert cfg["final_output_type"] == "image"
@pytest.mark.parametrize(
"model_class_name",
[
"MiniMaxH3Pipeline",
"WanPipeline",
"LTX2Pipeline",
"Cosmos3OmniDiffusersPipeline",
],
)
def test_default_video_diffusion_output_type(self, model_class_name):
configs = StageConfigFactory.create_default_diffusion({"model_class_name": model_class_name})
assert configs[0]["final_output_type"] == "video"
def test_default_diffusion_with_parallel_config(self):
"""Test diffusion config calculates devices from parallel_config."""
@@ -1324,6 +1339,13 @@ stages:
("step_audio2_ci.yaml", "step_audio_2_asr", 1, "text", "step_audio_2_asr"),
("hunyuan_video_15.yaml", "hunyuan_video_15", 1, "video", None),
("wan2_2_ti2v.yaml", "wan2_2_ti2v", 1, "video", None),
(
"minimax_h3_disaggregated.yaml",
"minimax_h3_disaggregated",
2,
"video",
"minimax_h3_disaggregated",
),
],
)
def test_load_new_registry_backed_deploy_configs(
@@ -1344,9 +1366,91 @@ stages:
assert len(stages) == stage_count
assert stages[-1].final_output is True
assert stages[-1].final_output_type == final_output_type
if pipeline_name == "minimax_h3_disaggregated":
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")
)
assert "model_path_resolver" not in stages[0].yaml_extras
assert "model_path_resolver" not in stages[1].yaml_extras
if deploy.trust_remote_code is not None:
assert {s.yaml_engine_args.get("trust_remote_code") for s in stages} == {deploy.trust_remote_code}
def test_minimax_h3_resolvers_reach_production_engine_args(self, tmp_path: Path, monkeypatch):
from vllm_omni.engine import stage_init_utils
model_root = tmp_path / "MiniMax-H3"
(model_root / "FL2VA" / "text_encoder").mkdir(parents=True)
(model_root / "Ref2VA" / "text_encoder").mkdir(parents=True)
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}
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]
assert resolved[0]["model"] == str(model_root / "FL2VA" / "text_encoder")
assert resolved[1]["model"] == str(model_root / "FL2VA")
assert all("model_path_resolver" not in engine_args for engine_args in resolved)
def test_minimax_h3_text_encoder_tp_reaches_only_stage_zero(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
stages, _ = StageConfigFactory._create_legacy_from_registry(
pipeline,
{"text_encoder_tp_size": 4},
)
assert stages[0].runtime_overrides["tensor_parallel_size"] == 4
assert "text_encoder_tp_size" not in stages[0].runtime_overrides
assert "text_encoder_tp_size" not in stages[1].runtime_overrides
assert stages[1].yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1
def test_minimax_h3_disaggregated_defaults_match_validated_topology(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
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["max_num_seqs"] == 1
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"]
assert parallel["tensor_parallel_size"] == 1
assert parallel["ulysses_degree"] == 4
assert parallel["vae_patch_parallel_size"] == 4
assert stages[1].yaml_engine_args["inline_diffusion"] is True
assert stages[1].yaml_extras["default_sampling_params"]["num_inference_steps"] == 50
turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml")))
turbo_stages = merge_pipeline_deploy(pipeline, turbo)
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}
def test_minimax_h3_text_encoder_tp_alias_targets_stage_zero(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
assert normalize_pipeline_cli_overrides(pipeline, {}) == {}
assert normalize_pipeline_cli_overrides(pipeline, {"text_encoder_tp_size": 4}) == {
"stage_0_tensor_parallel_size": 4
}
def test_minimax_h3_stage_zero_tp_override_wins_alias_conflict(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
with pytest.warns(UserWarning, match="stage_0_tensor_parallel_size=2 takes precedence"):
normalized = normalize_pipeline_cli_overrides(
pipeline,
{"text_encoder_tp_size": 4, "stage_0_tensor_parallel_size": 2},
)
assert normalized == {"stage_0_tensor_parallel_size": 2}
def test_minimax_h3_rejects_text_encoder_tp_on_diffusion_stage(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
with pytest.raises(ValueError, match="stage_1_text_encoder_tp_size cannot be set"):
normalize_pipeline_cli_overrides(pipeline, {"stage_1_text_encoder_tp_size": 4})
@pytest.mark.parametrize(
("config_json", "model_index", "expected_pipeline"),
[
@@ -1382,6 +1486,29 @@ stages:
assert pipeline is not None
assert pipeline.model_type == expected_pipeline
def test_minimax_h3_disaggregation_is_not_auto_discovered(self):
def get_model_file(filename, _model, revision=None):
del revision
if filename == "config.json":
return None
if filename == "model_index.json":
return {"_class_name": "MiniMaxH3Pipeline"}
return None
with (
patch.object(StageConfigFactory, "get_hf_config", return_value=None),
patch(
"vllm_omni.config.config_factory.get_hf_file_to_dict",
side_effect=get_model_file,
),
):
pipeline = StageConfigFactory.get_pipeline_config(
model="/models/unrelated-checkpoint-name",
trust_remote_code=False,
)
assert pipeline is None
@pytest.mark.parametrize("deploy_name", ["step_audio_2.yaml", "step_audio_2_async_chunk.yaml"])
def test_step_audio2_deploy_configs_fit_two_gpus(self, deploy_name):
deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
+13 -1
View File
@@ -92,6 +92,18 @@ def _from_pipeline_key(
)
def test_minimax_h3_text_encoder_tp_targets_only_structured_stage_zero() -> None:
config = _from_pipeline_key(
"minimax_h3_disaggregated",
cli_overrides={"text_encoder_tp_size": 4},
)
assert config.stage_by_id(0).parallel_config.tensor_parallel_size == 4
diffusion = config.stage_by_id(1)
assert diffusion.parallel_config.tensor_parallel_size == 1
assert diffusion.parallel_config.text_encoder_tp_size == 1
def test_duplex_session_capacity_propagates_to_every_structured_stage() -> None:
omni_config = _from_pipeline_key("personaplex")
@@ -643,8 +655,8 @@ def test_sub_config_fields_match_structured_scopes():
"ulysses_mode",
"cfg_parallel_size",
"vae_patch_parallel_size",
"text_encoder_tp_size",
"vae_parallel_mode",
"text_encoder_tp_size",
"use_hsdp",
"mask_sp_padding",
"hsdp_shard_size",
+7
View File
@@ -66,3 +66,10 @@ def test_register_resolver_requires_model_type(custom_resolver, clean_pipeline_r
"""Ensure that registering a custom resolver to OMNI_PIPELINES requires an explicit model_type."""
with pytest.raises(ValueError):
register_pipeline(custom_resolver)
def test_minimax_h3_disaggregation_is_explicit_opt_in():
assert "minimax_h3" not in OMNI_PIPELINES
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
assert isinstance(pipeline, PipelineConfig)
assert pipeline.model_type == "minimax_h3_disaggregated"
@@ -38,6 +38,7 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
pipeline.default_audio_shift = 3.0
pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace()
pipeline.text_encoder = object()
cache_spec = object()
quality_plan = SimpleNamespace(cache_dit=cache_spec)
pipeline._quality_policy = Mock()
@@ -110,6 +111,40 @@ def test_pipeline_import_registry_and_component_discovery():
assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
def test_encoder_free_stage_skips_text_encoder_during_dlo_discovery():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.transformer = torch.nn.Linear(1, 1)
pipeline.video_vae = torch.nn.Linear(1, 1)
pipeline.audio_vae = torch.nn.Linear(1, 1)
pipeline.text_encoder = None
pipeline._dit_modules = ["transformer"]
pipeline._encoder_modules = []
pipeline._vae_modules = ["video_vae", "audio_vae"]
discovered = ModuleDiscovery.discover(pipeline)
assert discovered.dit_names == ["transformer"]
assert discovered.encoder_names == []
assert discovered.vae_names == ["video_vae", "audio_vae"]
def test_encoder_free_stage_skips_missing_component_during_dlo_registration():
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
cache = Mock()
video_vae = Mock()
audio_vae = Mock()
pipeline_module._register_dlo_component_cache(cache, None, video_vae, audio_vae)
video_vae.set_omni_component_cache.assert_called_once_with(cache)
audio_vae.set_omni_component_cache.assert_called_once_with(cache)
def _write_partition_index(path, *, partition, tasks):
path.mkdir(parents=True)
(path / "model_index.json").write_text(
@@ -575,6 +610,7 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition):
pipeline.default_audio_shift = 3.0
pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace()
pipeline.text_encoder = Mock()
pipeline._base_schedule_by_partition = schedules
pipeline._quality_policy = Mock()
pipeline._quality_policy.resolve.return_value = SimpleNamespace(cache_dit=None)
@@ -998,7 +1034,7 @@ def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch):
config.hidden_size,
bias=False,
)
attn_call = {}
attn_call: dict[str, object] = {}
def fake_attention(query, key, value):
attn_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape)
@@ -1020,7 +1056,7 @@ def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch):
def test_text_attention_sdpa_helper_preserves_expanded_kv_fallback(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
sdpa_call = {}
sdpa_call: dict[str, object] = {}
def fake_sdpa(query, key, value, **kwargs):
sdpa_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape, kwargs=kwargs)
@@ -1095,7 +1131,7 @@ def test_encoder_forward_forwards_video_inputs():
def test_reference_video_shape_uses_h3_adapt_shape_policy():
from vllm_omni.diffusion.models.minimax_h3.reference_video import (
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
_reference_video_shape,
)
@@ -1681,6 +1717,85 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti
assert transcode_calls[0][1]["target_frame_count"] == 209
def test_ref2va_transcode_zero_frame_count_keeps_video_stream(monkeypatch, tmp_path):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
commands: list[list[str]] = []
monkeypatch.setattr(
reference_video_module.subprocess,
"run",
lambda command, **_kwargs: commands.append(command),
)
reference_video_module._transcode_reference_video(
"input.mp4",
target_width=768,
target_height=448,
target_frame_count=0,
workdir=str(tmp_path),
duration_seconds=4.0,
)
assert "-frames:v" not in commands[0]
def test_prepared_reference_video_descriptor_round_trip(tmp_path):
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
deserialize_prepared_reference_videos,
serialize_prepared_reference_videos,
)
item = {
"original_path": str(tmp_path / "source.mp4"),
"prepared_path": str(tmp_path / "prepared.mp4"),
"input_has_audio": True,
"width": 1344,
"height": 768,
"start_time_seconds": 0.0,
"duration_seconds": 5.2,
"audio_duration_seconds": 5.2,
}
descriptor = serialize_prepared_reference_videos([item], str(tmp_path))
assert deserialize_prepared_reference_videos(descriptor) == (str(tmp_path), [item])
def test_prepared_reference_video_descriptor_rejects_unknown_fields():
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
deserialize_prepared_reference_videos,
)
with pytest.raises(ValueError, match="invalid MiniMax H3 prepared-reference-video descriptor"):
deserialize_prepared_reference_videos('{"artifact_dir":"/tmp","videos":[{"path":"x"}]}')
def test_stage_one_reuses_existing_prepared_reference_video(tmp_path):
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_reuse_prepared_reference_videos,
)
prepared_path = tmp_path / "prepared.mp4"
prepared_path.touch()
prepared = [{"prepared_path": str(prepared_path)}]
assert _reuse_prepared_reference_videos(prepared, expected_count=1) is prepared
def test_stage_one_rejects_invalid_prepared_reference_video(tmp_path):
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_reuse_prepared_reference_videos,
)
with pytest.raises(ValueError, match="count does not match"):
_reuse_prepared_reference_videos([], expected_count=1)
with pytest.raises(ValueError, match="prepared reference video is unavailable"):
_reuse_prepared_reference_videos(
[{"prepared_path": str(tmp_path / "missing.mp4")}],
expected_count=1,
)
def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
@@ -2120,7 +2235,7 @@ def test_g4_reference_image_file_format_and_size_contract(tmp_path):
def test_g4_standalone_audio_duration_and_total_duration_contract():
from vllm_omni.diffusion.models.minimax_h3.reference_video import (
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
validate_reference_audio_waveforms,
)
@@ -2195,7 +2310,7 @@ def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration():
],
)
def test_g4_reference_video_metadata_validation(field, value, message, tmp_path):
from vllm_omni.diffusion.models.minimax_h3.reference_video import (
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
_validate_reference_video_metadata,
)
@@ -254,18 +254,3 @@ def test_pipeline_resolves_transformer_component_quant_config():
assert transformer_config.ignored_layers == ignored_layers
assert _resolve_component_quant_config(component_config, "transformer") is transformer_config
assert _resolve_component_quant_config(transformer_config, "transformer") is transformer_config
def test_pipeline_strips_prequantized_text_encoder_config():
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _resolve_minimax_h3_text_encoder_quant_config
from vllm_omni.quantization import ComponentQuantizationConfig
prequantized = Mock()
prequantized.get_name.return_value = "modelopt"
online_fp8 = Mock()
online_fp8.get_name.return_value = "fp8"
component_config = ComponentQuantizationConfig({"text_encoder": prequantized})
assert _resolve_minimax_h3_text_encoder_quant_config(online_fp8) is online_fp8
assert _resolve_minimax_h3_text_encoder_quant_config(prequantized) is None
assert _resolve_minimax_h3_text_encoder_quant_config(component_config) is None
@@ -187,7 +187,6 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
"model": model_ref,
"enforce_eager": True,
"tensor_parallel_size": 2,
"text_encoder_tp_size": 2,
"vae_use_tiling": True,
}
@@ -1,13 +1,18 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
import torch
import torch.nn.functional as F
from tests.helpers.mark import hardware_marks
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
pytestmark = [
pytest.mark.core_model,
pytest.mark.diffusion,
*hardware_marks(res={"npu": "A3"}, num_cards=1),
]
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -76,7 +76,7 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
result = results[0]
assert result.images == ["frame-0"]
assert result.prompt == "prompt-0"
assert result.final_output_type == "image"
assert result.final_output_type == "video"
assert result.multimodal_output == {
"metadata": {
"audio": {"sample_rate": 48000},
@@ -315,7 +315,7 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
assert result.images == ["frame-0"]
assert result.prompt == "watch and listen"
assert result.final_output_type == "image"
assert result.final_output_type == "video"
assert result.multimodal_output == {
"metadata": {
"audio": {"sample_rate": 16000},
@@ -311,8 +311,8 @@ class TestPipelineStreamingOutputToEntrypoint:
lambda *, output_format, fps, video_codec_options=None: FakeStreamingVideoEncoder(),
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
return_value="diffusion",
"vllm_omni.entrypoints.openai.serving_video_output_stream.is_video_generation_pipeline",
return_value=True,
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
@@ -52,6 +52,11 @@ def client(mock_engine):
c.shutdown()
def test_inline_client_implements_shared_prompt_hook_protocol(client):
assert client.prompt_transform_func is None
assert client.prompt_expand_func is None
@pytest.mark.asyncio
async def test_inline_dispatch_request_success(client, mock_engine):
mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()])
@@ -537,7 +537,7 @@ def api_server(server_case: ServerCase, mock_async_omni, tmp_path):
cmd = OmniServeCommand()
cmd.subparser_init(subparsers)
port = get_open_port()
port = get_open_port(host="0.0.0.0")
args = parser.parse_args(["serve", server_case.served_model, "--omni", "--port", str(port)])
def run_server():
@@ -814,7 +814,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
ServerCase(
served_model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
stage_list=["diffusion"],
stage_configs=[{"stage_type": "diffusion"}],
stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
outputs=[_build_diffusion_video_output()],
),
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
@@ -825,7 +825,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
ServerCase(
served_model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
stage_list=["diffusion"],
stage_configs=[{"stage_type": "diffusion"}],
stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
outputs=[_build_diffusion_video_output()],
),
"Wan-AI/Wan2.2-I2V-A14B-Diffusers",
@@ -886,7 +886,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
ServerCase(
served_model="MiniMaxAI/MiniMax-H3",
stage_list=["diffusion"],
stage_configs=[{"stage_type": "diffusion"}],
stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
outputs=[_build_diffusion_video_output()],
),
SamplingCase(kind=SamplingKind.VIDEO_REF2VA_IMAGE_AUDIO, sampling_params=None),
@@ -897,7 +897,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
ServerCase(
served_model="MiniMaxAI/MiniMax-H3",
stage_list=["diffusion"],
stage_configs=[{"stage_type": "diffusion"}],
stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
outputs=[_build_diffusion_video_output()],
),
SamplingCase(kind=SamplingKind.VIDEO_REF2VA_MULTI_VIDEO, sampling_params=None),
@@ -67,6 +67,39 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
output_processor.add_request.assert_not_called()
def test_build_add_request_message_injects_global_id_before_prompt_transform(mocker: MockerFixture):
engine = object.__new__(AsyncOmniEngine)
params = SamplingParams(max_tokens=8)
engine.default_sampling_params_list = [params]
engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
engine.supported_tasks = ("speech",)
input_processor = mocker.Mock()
input_processor.process_inputs.return_value = _make_engine_core_request()
engine.input_processor = input_processor
engine.output_processors = [mocker.Mock()]
seen_global_ids: list[list[str]] = []
def transform(prompt, _sampling_params):
seen_global_ids.append(prompt["additional_information"]["global_request_id"])
return prompt
engine.prompt_transform_func = transform
prompt = {"prompt_token_ids": [1, 1, 1]}
message = engine._build_add_request_message(
request_id="req-1",
prompt=prompt,
sampling_params_list=[params],
final_stage_id=0,
arrival_time=0.0,
)
assert seen_global_ids == [["req-1"]]
assert message.original_prompt["additional_information"]["global_request_id"] == ["req-1"]
def test_build_add_request_message_preserves_model_intermediate_buffer(mocker: MockerFixture):
engine = object.__new__(AsyncOmniEngine)
params = SamplingParams(max_tokens=8)
@@ -433,11 +433,35 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline
"stage_id": 0,
"stage_init_timeout": 302,
"batch_size": 4,
"use_inline": True,
"use_inline": False,
"omni_master_server": None,
}
def test_initialize_local_diffusion_replica_uses_explicit_inline_opt_in(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
from vllm_omni.engine.stage_engine_startup import StageReplicaResources
runtime = _make_stage_runtime()
plan = _make_diffusion_plan(0, stage_id=1).replicas[0]
plan.stage_cfg.engine_args["inline_diffusion"] = True
captured: dict[str, object] = {}
monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None)
monkeypatch.setattr(
runtime_mod,
"launch_diffusion_stage_replica",
lambda **kwargs: (
captured.update(use_inline=kwargs["use_inline"]),
(types.SimpleNamespace(), StageReplicaResources()),
)[1],
)
runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1)
assert captured["use_inline"] is True
def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
@@ -1140,6 +1164,21 @@ def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path):
assert engine_args["tokenizer"] == os.path.join(str(tmp_path), "tokenizer")
def test_model_path_resolver_is_generic_and_model_owned(tmp_path):
from vllm_omni.engine.stage_init_utils import _resolve_model_path
engine_args = {
"model_path_resolver": ("vllm_omni.model_executor.models.minimax_h3.checkpoint.resolve_minimax_h3_model_root"),
"revision": None,
"task_type": "ref2va",
}
resolved = _resolve_model_path(str(tmp_path), engine_args)
assert resolved == str(tmp_path / "Ref2VA" / "text_encoder")
assert "model_path_resolver" not in engine_args
def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch):
import vllm_omni.engine.stage_init_utils as init_mod
+28
View File
@@ -74,6 +74,25 @@ class TestOmniPayloadStruct:
for top, sub in original.items():
assert sorted(d[top].keys()) == sorted(sub.keys())
def test_text_conditioning_uses_shared_payload_fields(self):
hidden = torch.zeros(4, 5120)
token_role_ids = torch.tensor([[1], [1], [0], [0]])
payload = to_dict(
to_struct(
{
"hidden_states": {"output": hidden},
"meta": {"token_role_ids": token_role_ids},
}
)
)
assert torch.equal(payload["hidden_states"]["output"], hidden)
assert torch.equal(payload["meta"]["token_role_ids"], token_role_ids)
for legacy_key in ("encoder_hidden_states", "token_tags"):
with pytest.raises(msgspec.ValidationError, match="unknown field"):
to_struct({legacy_key: torch.zeros(1)})
def test_to_dict_drops_unset_fields(self):
s = OmniPayloadStruct(meta=MetaStruct(left_context_size=10))
d = to_dict(s)
@@ -397,3 +416,12 @@ class TestSerializeDeserializePayload:
original: OmniPayload = {"meta": {"finished": None}}
wire = serialize_payload(original)
assert wire is None
def test_minimax_h3_prepared_reference_video_descriptor_round_trip(self):
descriptor = '{"artifact_dir":"/tmp/h3","videos":[]}'
wire = serialize_payload({"meta": {"minimax_h3_prepared_reference_videos": descriptor}})
assert wire is not None
restored = deserialize_payload(wire)
assert restored["meta"]["minimax_h3_prepared_reference_videos"] == descriptor
@@ -27,7 +27,11 @@ from vllm_omni.engine.messages import (
ShutdownRequestMessage,
StageSubmissionMessage,
)
from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
from vllm_omni.engine.orchestrator import (
Orchestrator,
OrchestratorRequestState,
cleanup_request_artifact_dirs,
)
from vllm_omni.engine.stage_pool import StageUnavailableError
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
@@ -44,6 +48,18 @@ from .test_orchestrator import (
_wait_for,
)
def test_request_artifact_cleanup_is_idempotent(tmp_path):
artifact_dir = tmp_path / "request-artifacts"
artifact_dir.mkdir()
(artifact_dir / "prepared.mp4").touch()
cleanup_request_artifact_dirs([str(artifact_dir)])
cleanup_request_artifact_dirs([str(artifact_dir)])
assert not artifact_dir.exists()
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+42
View File
@@ -1207,6 +1207,48 @@ class TestSingleStageReplicaInitialization:
)
assert mock_from_addresses.call_args.kwargs["proc_manager"].proc is proc
def test_initialize_local_custom_diffusion_replica_uses_inline_client(self, mocker: MockerFixture):
import vllm_omni.engine.stage_runtime as runtime_mod
from vllm_omni.platforms import current_omni_platform
runtime = DistStageRuntime(
stage_configs=[],
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=4,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
omni_master_port=26000,
)
runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer)
runtime._coordinator_runtime = None
plan = _make_diffusion_plan(0, stage_id=0, launch_mode="local").replicas[0]
plan.stage_cfg.engine_args = {"custom_pipeline_args": {"pipeline_class": "test.CustomPipeline"}}
sentinel_client = SimpleNamespace()
mock_launch = mocker.patch.object(
runtime_mod,
"launch_diffusion_stage_replica",
return_value=(sentinel_client, StageReplicaResources()),
)
mocker.patch.object(runtime_mod, "inject_kv_stage_info")
device_env_var = current_omni_platform.device_control_env_var
prev_device_env = os.environ.get(device_env_var)
os.environ[device_env_var] = "0"
runtime._init_visible_devices_baseline = "0"
try:
result = runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=60)
finally:
if prev_device_env is None:
os.environ.pop(device_env_var, None)
else:
os.environ[device_env_var] = prev_device_env
assert result is sentinel_client
assert mock_launch.call_args.kwargs["use_inline"] is True
def test_initialize_local_diffusion_replica_failure_terminates_proc(self, mocker: MockerFixture):
import vllm_omni.engine.stage_runtime as runtime_mod
from vllm_omni.platforms import current_omni_platform
+5 -1
View File
@@ -72,6 +72,7 @@ def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str,
_LLM_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniEngineArgs))
_DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusionConfig))
_TOPOLOGY_ONLY_ENGINE_ARGS = frozenset({"inline_diffusion"})
_OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
{
"active_stream_window",
@@ -368,7 +369,9 @@ def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
)
typed_args_by_stage[stage_id] = typed_args
assert {name: typed_args[name] for name in legacy_args} == legacy_args
expected_args = {name: value for name, value in legacy_args.items() if name not in _TOPOLOGY_ONLY_ENGINE_ARGS}
assert {name: typed_args[name] for name in expected_args} == expected_args
assert _TOPOLOGY_ONLY_ENGINE_ARGS.isdisjoint(typed_args)
thinker_args = typed_args_by_stage[0]
inherited_vllm_fields = {
@@ -658,6 +661,7 @@ def test_typed_engine_args_match_current_registry_backend_semantics(model_type):
backend_fields = _LLM_BACKEND_FIELDS
backend_config_cls = OmniEngineArgs
backend_fields -= _TOPOLOGY_ONLY_ENGINE_ARGS
missing_fields = (legacy_args.keys() & backend_fields) - typed_args.keys()
assert not missing_fields, f"{model_type} stage {stage_id} lost backend fields: {sorted(missing_fields)}"
@@ -47,7 +47,13 @@ def _build_test_app(
engine_client.abort = mocker.AsyncMock()
engine_client.submit_interaction_async = mocker.AsyncMock()
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
engine_client.stage_configs = [
SimpleNamespace(
stage_type="diffusion",
final_output=True,
final_output_type="video",
)
]
engine_client.generate = mock_generate
class FakeStreamingVideoEncoder:
@@ -72,10 +78,6 @@ def _build_test_app(
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
encoder_factory,
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
return_value="diffusion",
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
return_value=[OmniDiffusionSamplingParams()],
@@ -398,7 +400,13 @@ def _build_handler_for_async_tests(
engine_client.abort = mocker.AsyncMock()
engine_client.submit_interaction_async = mocker.AsyncMock()
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
engine_client.stage_configs = [
SimpleNamespace(
stage_type="diffusion",
final_output=True,
final_output_type="video",
)
]
engine_client.generate = mock_generate
chunks = streaming_chunks or [(b"mp4-chunk-0", True)]
@@ -421,10 +429,6 @@ def _build_handler_for_async_tests(
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
encoder_factory,
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
return_value="diffusion",
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
return_value=[OmniDiffusionSamplingParams()],
@@ -0,0 +1,66 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
import pytest
from vllm_omni.diffusion.io_support import get_diffusion_output_type
from vllm_omni.entrypoints.openai.utils import is_video_generation_pipeline
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_video_pipeline_requires_declared_final_video_stage():
assert is_video_generation_pipeline(
[
SimpleNamespace(
stage_type="llm",
final_output=False,
final_output_type=None,
),
SimpleNamespace(
stage_type="diffusion",
final_output=True,
final_output_type="video",
),
]
)
@pytest.mark.parametrize(
"stage_configs",
[
[SimpleNamespace(stage_type="diffusion")],
[
SimpleNamespace(
stage_type="diffusion",
final_output=True,
final_output_type="image",
)
],
[
{
"stage_type": "diffusion",
"final_output": False,
"final_output_type": "video",
}
],
],
)
def test_video_pipeline_rejects_non_video_final_outputs(stage_configs):
assert not is_video_generation_pipeline(stage_configs)
@pytest.mark.parametrize(
"model_class_name",
[
"LTX2TwoStagePipeline",
"LTX2DistilledOneStagePipeline",
"WanDMDPipeline",
"LingBotWorldCausalDMDPipeline",
"LongCatVideoAvatarPipeline",
],
)
def test_registered_video_aliases_declare_video_output(model_class_name):
assert get_diffusion_output_type(model_class_name) == "video"
@@ -64,7 +64,13 @@ class MockVideoResult:
class FakeAsyncOmni:
def __init__(self):
self.stage_configs = [SimpleNamespace(stage_type="diffusion")]
self.stage_configs = [
SimpleNamespace(
stage_type="diffusion",
final_output=True,
final_output_type="video",
)
]
self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
self.model_class_name = "WanPipeline"
self.captured_prompt = None
@@ -828,6 +834,18 @@ def test_multi_video_generation_preserves_uploaded_files_until_generation(
assert engine.captured_sampling_params_list[0].extra_args["duration"] == 15.0
def test_mixed_reference_capability_uses_model_metadata_when_config_defaults_false(test_client):
handler = test_client.app.state.openai_serving_video
handler._engine_client.model_class_name = None
handler._engine_client.stage_configs = [
SimpleNamespace(engine_args={"model_class_name": "MiniMaxH3TextEncoder"}),
SimpleNamespace(engine_args={"model_class_name": "MiniMaxH3Pipeline"}),
]
handler._stage_configs = handler._engine_client.stage_configs
assert handler.supports_mixed_reference_inputs
def test_decode_video_bytes_can_keep_first_frames():
from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
@@ -315,6 +315,29 @@ def test_serve_cli_accepts_ulysses_mode():
assert parallel_config.ulysses_mode == "advanced_uaa"
def test_serve_cli_accepts_text_encoder_tp_size():
parser = TrackingArgumentParser()
subparsers = parser.add_subparsers(dest="command")
OmniServeCommand().subparser_init(subparsers)
args = parser.parse_args(
[
"serve",
"MiniMaxAI/MiniMax-H3",
"--omni",
"--text-encoder-tp-size",
"4",
]
)
explicit_kwargs = args.get_explicit_kwargs_dict()
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
parallel_config = stage_cfg["engine_args"]["parallel_config"]
assert args.text_encoder_tp_size == 4
assert parallel_config.text_encoder_tp_size == 4
def test_serve_cli_forwards_model_defined_task_type_to_diffusion_stage():
parser = TrackingArgumentParser()
subparsers = parser.add_subparsers(dest="command")
@@ -574,6 +597,36 @@ def test_serve_cli_accepts_additional_config():
assert engine_args["additional_config"] == {"torchair_graph_config": {"enabled": True}}
def test_default_stage_resolves_video_output_from_checkpoint(mocker):
captured = {}
def resolve_with_default(*args, default_stage_cfg_factory, **kwargs):
del args, kwargs
captured.update(default_stage_cfg_factory()[0])
stage = SimpleNamespace(stage_type="diffusion", engine_args=SimpleNamespace())
return ("", [stage], None)
resolver = mocker.patch(
"vllm_omni.engine.async_omni_engine.resolve_model_class_name",
return_value="MiniMaxH3Pipeline",
)
mocker.patch(
"vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
side_effect=resolve_with_default,
)
engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
engine._strip_single_engine_args = lambda kwargs: kwargs
engine._resolve_stage_configs(
"/models/MiniMax-H3/FL2VA",
{},
trust_remote_code=False,
)
resolver.assert_called_once_with("/models/MiniMax-H3/FL2VA", "default")
assert captured["final_output_type"] == "video"
def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mocker):
"""Ensure YAML/deploy stage resolution forwards top-level additional_config."""
fake_diffusion_stage = SimpleNamespace(
@@ -0,0 +1,62 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
from unittest.mock import patch
import pytest
import torch
from vllm.model_executor.models.qwen3_vl import Qwen3VLForConditionalGeneration
from vllm_omni.model_executor.models.minimax_h3.text_encoder import MiniMaxH3TextEncoder
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _model_without_weights() -> MiniMaxH3TextEncoder:
model = MiniMaxH3TextEncoder.__new__(MiniMaxH3TextEncoder)
model.config = SimpleNamespace(
vision_start_token_id=1,
vision_end_token_id=2,
image_token_id=3,
video_token_id=4,
)
model._token_tags = None
return model
def test_embed_failure_clears_token_tags():
model = _model_without_weights()
model._token_tags = torch.tensor([0])
with (
patch.object(Qwen3VLForConditionalGeneration, "embed_input_ids", side_effect=RuntimeError("failed")),
pytest.raises(RuntimeError, match="failed"),
):
model.embed_input_ids(torch.tensor([5, 3]))
assert model._token_tags is None
def test_invalid_output_length_consumes_token_tags():
model = _model_without_weights()
model._token_tags = torch.tensor([1, 0])
with pytest.raises(RuntimeError, match="token tags exceed"):
model.make_omni_output(torch.zeros(1, 8))
assert model._token_tags is None
def test_make_omni_output_uses_shared_stage_payload_fields():
model = _model_without_weights()
model._token_tags = torch.tensor([1, 0])
hidden = torch.zeros(2, 5120)
output = model.make_omni_output(hidden)
assert output.multimodal_outputs is not None
assert torch.equal(output.multimodal_outputs["hidden_states"]["output"], hidden)
assert torch.equal(output.multimodal_outputs["meta"]["token_role_ids"], torch.tensor([[1], [0]]))
assert "encoder_hidden_states" not in output.multimodal_outputs
assert "token_tags" not in output.multimodal_outputs
@@ -0,0 +1,387 @@
# SPDX-License-Identifier: Apache-2.0
# 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
import numpy as np
import pytest
import torch
from PIL import Image
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_load_audio,
resolve_minimax_h3_diffusion_model_path,
)
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_partition,
)
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
MINIMAX_H3_CONDITION_LABELS_KEY,
MINIMAX_H3_PRESENTATION_TASK_KEY,
)
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,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _SegmentTokenizer:
_special_ids = {
"<|vision_start|>": 1,
"<|vision_end|>": 2,
"<|image_pad|>": 3,
"<|video_pad|>": 4,
}
def __call__(self, text, *, add_special_tokens):
assert add_special_tokens is False
return {"input_ids": [100 + len(text), 1000 + sum(text.encode())]}
def convert_tokens_to_ids(self, token):
return self._special_ids[token]
def test_h3_processor_reprocesses_media_instead_of_using_partial_sender_cache(mocker):
processor = object.__new__(MiniMaxH3MultiModalProcessor)
sentinel = ([1, 2, 3], object(), True)
apply_processor = mocker.patch.object(
MiniMaxH3MultiModalProcessor,
"_apply_hf_processor",
return_value=sentinel,
)
inputs = object()
timing_ctx = object()
result = processor._cached_apply_hf_processor(inputs, timing_ctx)
assert result is sentinel
apply_processor.assert_called_once_with(inputs, timing_ctx)
@pytest.mark.parametrize(
("value", "expected_count"),
[
((torch.zeros(16), 16_000), 1),
([np.zeros(16), 16_000], 1),
([(torch.zeros(16), 16_000), (torch.ones(16), 24_000)], 2),
(["first.wav", "second.wav"], 2),
],
)
def test_audio_items_preserves_waveform_pairs(value, expected_count):
assert len(_audio_items(value)) == expected_count
def test_h3_selects_the_single_explicit_diffusion_stage_params() -> None:
stage_zero = SimpleNamespace(extra_args={"task": "t2va"})
diffusion = OmniDiffusionSamplingParams(extra_args={"task": "ref2va"})
assert _diffusion_sampling_params([stage_zero, diffusion]) is diffusion
def test_h3_rejects_missing_or_ambiguous_diffusion_stage_params() -> None:
with pytest.raises(RuntimeError, match="exactly one OmniDiffusionSamplingParams"):
_diffusion_sampling_params([SimpleNamespace(extra_args={"task": "t2va"})])
with pytest.raises(RuntimeError, match="got 2"):
_diffusion_sampling_params([OmniDiffusionSamplingParams(), OmniDiffusionSamplingParams()])
def test_fused_audio_loader_accepts_list_waveform_pair():
waveform, sample_rate = _load_audio([[0.0, 0.5, -0.5], 16_000])
assert sample_rate == 16_000
torch.testing.assert_close(waveform, torch.tensor([0.0, 0.5, -0.5]))
def test_prepare_ref2va_keeps_original_text_and_exact_condition_order():
prompt = {
"prompt": "hello",
"multi_modal_data": {
"image": Image.new("RGB", (256, 256)),
"audio": [np.zeros(16), 16_000],
},
}
sampling = OmniDiffusionSamplingParams(
height=256,
width=448,
extra_args={"task": "ref2va"},
)
transformed = prepare_text_encoder_prompt(prompt, [sampling])
assert transformed["prompt"] == "hello"
assert len(transformed["multi_modal_data"]["image"]) == 1
assert "audio" not in transformed["multi_modal_data"]
assert transformed["mm_processor_kwargs"][MINIMAX_H3_PRESENTATION_TASK_KEY] == "ref2va"
assert transformed["mm_processor_kwargs"][MINIMAX_H3_CONDITION_LABELS_KEY] == [
("image", 1),
("audio", 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"})
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)
monkeypatch.setattr(
"vllm_omni.model_executor.stage_input_processors.minimax_h3.sample_reference_video_frames",
sample_frames,
)
prompt = {
"prompt": "hello",
"multi_modal_data": {"video": "/tmp/input.mp4"},
}
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"}
assert transformed["mm_processor_kwargs"][MINIMAX_H3_CONDITION_LABELS_KEY] == [
("audio", 1),
("video", 1),
]
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 test_stage_wire_rejects_request_id_mismatch():
source = SimpleNamespace(request_id="other", outputs=[SimpleNamespace()])
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 test_text_encoder2diffusion_reads_hidden_states_output_and_token_role_ids() -> None:
hidden = torch.randn(4, 5120)
token_role_ids = torch.tensor([[1], [1], [0], [0]])
result = text_encoder2diffusion(
[
_source_output(
{
"hidden_states": {"output": hidden},
"meta": {"token_role_ids": token_role_ids},
}
)
],
_diffusion_prompt(),
)
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))
@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())
def test_ref2va_one_image_tokens_and_tags_match_fused_presentation():
tokenizer = _SegmentTokenizer()
labels = [("image", 1), ("audio", 1)]
image_grid = torch.tensor([[1, 4, 4]])
actual = _build_minimax_h3_presentation(
tokenizer,
prompt="hello",
task="ref2va",
condition_labels=labels,
image_grid_thw=image_grid,
video_grid_thw=None,
video_timestamps=None,
merge_size=2,
)
expected = minimax_h3_ref2va_presentation(
tokenizer,
prompt="hello",
condition_labels=labels,
image_token_count=[4],
)
torch.testing.assert_close(actual[0], expected[0])
torch.testing.assert_close(actual[1], expected[1])
def test_ref2va_video_tokens_and_tags_match_fused_without_outer_markers():
tokenizer = _SegmentTokenizer()
labels = [("audio", 1), ("video", 1)]
video_grid = torch.tensor([[2, 4, 4]])
timestamps = [[0.2, 0.4]]
actual = _build_minimax_h3_presentation(
tokenizer,
prompt="hello",
task="ref2va",
condition_labels=labels,
image_grid_thw=None,
video_grid_thw=video_grid,
video_timestamps=timestamps,
merge_size=2,
)
expected = minimax_h3_ref2va_video_presentation(
tokenizer,
prompt="hello",
condition_labels=labels,
image_token_count=None,
video_block_token_counts=[[4, 4]],
video_block_timestamps=timestamps,
)
torch.testing.assert_close(actual[0], expected[0])
torch.testing.assert_close(actual[1], expected[1])
assert int((actual[0] == tokenizer._special_ids["<|vision_start|>"]).sum()) == 2
assert int((actual[0] == tokenizer._special_ids["<|vision_end|>"]).sum()) == 2
def test_checkpoint_resolver_selects_local_partition(tmp_path):
root = tmp_path / "MiniMax-H3"
(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")
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")
def test_partition_resolver_preserves_consumer_auto_default(tmp_path):
root = tmp_path / "MiniMax-H3"
ref2va = root / "Ref2VA"
ref2va.mkdir(parents=True)
assert resolve_minimax_h3_partition(str(root), "auto", auto_partition="fl2va") == "fl2va"
assert resolve_minimax_h3_partition(str(root), "auto", auto_partition="combined") == "combined"
assert resolve_minimax_h3_partition(str(ref2va), "auto", auto_partition="combined") == "ref2va"
def test_diffusion_resolver_selects_startup_partition(tmp_path):
root = tmp_path / "MiniMax-H3"
fl2va = root / "FL2VA"
ref2va = root / "Ref2VA"
fl2va.mkdir(parents=True)
ref2va.mkdir()
(fl2va / "model_index.json").write_text("{}")
(ref2va / "model_index.json").write_text("{}")
assert resolve_minimax_h3_diffusion_model_path(str(root), None, "fl2va") == str(fl2va)
assert resolve_minimax_h3_diffusion_model_path(str(root), None, "ref2va") == str(ref2va)
assert resolve_minimax_h3_diffusion_model_path(str(root), None, None) == str(fl2va)
assert resolve_minimax_h3_diffusion_model_path(str(root), None, "combined") == str(root)
assert resolve_minimax_h3_diffusion_model_path(str(ref2va), None, None) == str(ref2va)
def test_diffusion_resolver_normalizes_partial_partition_directory(tmp_path):
root = tmp_path / "MiniMax-H3"
ref2va = root / "Ref2VA"
(ref2va / "text_encoder").mkdir(parents=True)
assert resolve_minimax_h3_diffusion_model_path(str(ref2va), None, "ref2va") == str(ref2va)
+5 -1
View File
@@ -29,8 +29,10 @@ from vllm_omni.config.stage_config import (
build_stage_runtime_overrides,
load_deploy_config,
merge_pipeline_deploy,
normalize_pipeline_cli_overrides,
)
from vllm_omni.config.yaml_util import create_config
from vllm_omni.diffusion.io_support import get_diffusion_output_type
from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
logger = init_logger(__name__)
@@ -436,6 +438,7 @@ class StageConfigFactory:
load-balancer policy (``None`` when no strategy set one) travels with the
stages instead of through a mutable out-param.
"""
cli_overrides = normalize_pipeline_cli_overrides(pipeline_cfg, cli_overrides)
deploy_cfg: DeployConfig | None
if user_deploy_config is not None:
deploy_cfg = user_deploy_config
@@ -628,6 +631,7 @@ class StageConfigFactory:
engine_args["dtype"] = str(engine_args["dtype"])
engine_args.setdefault("max_num_seqs", 1)
model_class_name = engine_args.get("model_class_name")
config_dict: dict[str, Any] = {
"stage_id": 0,
@@ -638,7 +642,7 @@ class StageConfigFactory:
},
"engine_args": create_config(engine_args),
"final_output": True,
"final_output_type": "image",
"final_output_type": get_diffusion_output_type(model_class_name),
}
return [config_dict]
+7
View File
@@ -46,6 +46,7 @@ from vllm_omni.config.stage_config import (
_select_processor_funcs,
build_stage_runtime_overrides,
load_deploy_config,
normalize_pipeline_cli_overrides,
)
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
@@ -730,6 +731,7 @@ class _DiffusionConfigProjection:
default_factory=lambda: {
"transformer": True,
"vae": True,
"text_encoder": True,
}
)
override_transformer_cls_name: str | None = None
@@ -1373,6 +1375,10 @@ class BaseVllmOmniStageConfig:
def prompt_expand_func(self) -> str | None:
return self.stage_pipeline_config.prompt_expand_func
@property
def prompt_transform_func(self) -> str | None:
return self.stage_pipeline_config.prompt_transform_func
@property
def cfg_kv_collect_func(self) -> str | None:
return self.stage_pipeline_config.cfg_kv_collect_func
@@ -1836,6 +1842,7 @@ class VllmOmniConfig:
"""Create a structured config from a resolved pipeline and deploy YAML."""
if cli_overrides is None:
cli_overrides = {}
cli_overrides = normalize_pipeline_cli_overrides(pipeline_cfg, cli_overrides)
deploy, loaded_deploy_config_path = _get_deploy_config(
pipeline_cfg,
+2
View File
@@ -85,6 +85,7 @@ from vllm_omni.model_executor.models.ming_tts.pipeline import (
MING_TTS_PIPELINE,
)
from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE
from vllm_omni.model_executor.models.minimax_h3.pipeline import MINIMAX_H3_PIPELINE
from vllm_omni.model_executor.models.minimax_music3.pipeline import MINIMAX_MUSIC3_PIPELINE
from vllm_omni.model_executor.models.moss_tts.pipeline import (
MOSS_TTS_LOCAL_PIPELINE,
@@ -171,6 +172,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
"ming_flash_omni_thinker_only": MING_FLASH_OMNI_THINKER_ONLY_PIPELINE,
"ming_flash_omni_image": MING_FLASH_OMNI_IMAGE_PIPELINE,
"moss_tts_nano": MOSS_TTS_NANO_PIPELINE,
"minimax_h3_disaggregated": MINIMAX_H3_PIPELINE,
"omnivoice": OMNIVOICE_PIPELINE,
"mammoth_moda2": MAMMOTH_MODA2_PIPELINE,
"mammoth_moda2_ar": MAMMOTH_MODA2_AR_PIPELINE,
+55
View File
@@ -89,6 +89,43 @@ def build_stage_runtime_overrides(
return result
def normalize_pipeline_cli_overrides(
pipeline: PipelineConfig,
cli_overrides: dict[str, Any],
) -> dict[str, Any]:
"""Translate pipeline-owned global CLI aliases into stage-scoped overrides."""
normalized = dict(cli_overrides)
for source, (stage_id, target) in pipeline.stage_cli_aliases.items():
invalid_stage_keys = [
key
for key, value in normalized.items()
if value is not None
and (match := _STAGE_OVERRIDE_PATTERN.match(key)) is not None
and match.group(2) == source
and int(match.group(1)) != stage_id
]
if invalid_stage_keys:
invalid = ", ".join(sorted(invalid_stage_keys))
raise ValueError(
f"{invalid} cannot be set for pipeline {pipeline.model_type!r}; "
f"{source} belongs to stage {stage_id} as {target}."
)
value = normalized.pop(source, None)
if value is None:
continue
stage_key = f"stage_{stage_id}_{target}"
stage_value = normalized.get(stage_key)
if stage_value is not None and stage_value != value:
warnings.warn(
f"Ignoring {source}={value!r} because {stage_key}={stage_value!r} takes precedence.",
UserWarning,
stacklevel=2,
)
continue
normalized[stage_key] = value
return normalized
def _apply_diffusion_parallel_runtime_overrides(
engine_args: dict[str, Any],
runtime_overrides: dict[str, Any],
@@ -191,6 +228,10 @@ class StagePipelineConfig:
# Alternates picked by ``merge_pipeline_deploy`` based on ``deploy.async_chunk``.
async_chunk_process_next_stage_input_func: str | None = None
sync_process_input_func: str | None = None
# Rewrites the Stage-0 view of a raw prompt before vLLM input processing.
# The callable receives ``(prompt, sampling_params_list)``; downstream
# stages continue to receive the original prompt.
prompt_transform_func: str | None = None
prompt_expand_func: str | None = None
cfg_kv_collect_func: str | None = None
omni_kv_config: dict[str, Any] | None = None
@@ -201,6 +242,13 @@ class StagePipelineConfig:
# by ``stage_init_utils._resolve_model_tokenizer_paths``.
model_subdir: str | None = None
tokenizer_subdir: str | None = None
# Model-owned hook that resolves a pipeline root to this stage's checkpoint.
# Consumed and removed before backend engine args are constructed.
model_path_resolver: str | None = None
# Keep a single-replica diffusion stage in the orchestrator process.
# Disabled by default so existing multi-stage pipelines retain subprocess
# isolation unless their topology explicitly opts in.
inline_diffusion: bool = False
# Whether the non-async path waits for a complete upstream payload from
# the model-runner connector before scheduling this stage.
requires_full_payload_input: bool = False
@@ -251,6 +299,8 @@ class PipelineConfig:
# Bundled deploy defaults for this concrete pipeline topology. The file is
# loaded from vllm_omni/deploy; None uses DeployConfig defaults.
default_deploy_config_name: str | None = None
# Global CLI spelling -> (stage id, stage-local spelling).
stage_cli_aliases: dict[str, tuple[int, str]] = field(default_factory=dict)
def get_stage(self, stage_id: int) -> StagePipelineConfig | None:
"""Look up a stage by its ID."""
@@ -834,6 +884,9 @@ def _build_engine_args(
engine_args["model_subdir"] = ps.model_subdir
if ps.tokenizer_subdir:
engine_args["tokenizer_subdir"] = ps.tokenizer_subdir
if ps.model_path_resolver:
engine_args["model_path_resolver"] = ps.model_path_resolver
engine_args["inline_diffusion"] = ps.inline_diffusion
# Pipeline-wide top-level DeployConfig settings, applied to every stage.
for name in _PIPELINE_WIDE_ENGINE_FIELDS:
@@ -881,6 +934,8 @@ def _build_extras(
extras["output_connectors"] = dict(ds.output_connectors)
if ds is not None and ds.input_connectors:
extras["input_connectors"] = dict(ds.input_connectors)
if ps.prompt_transform_func:
extras["prompt_transform_func"] = ps.prompt_transform_func
if ps.prompt_expand_func:
extras["prompt_expand_func"] = ps.prompt_expand_func
if ps.cfg_kv_collect_func:
+7
View File
@@ -19,6 +19,9 @@ import msgspec
import numpy as np
import torch
REQUEST_ARTIFACT_DIRS_KEY = "_omni_request_artifact_dirs"
TRANSFORM_OWNED_META_KEYS = frozenset({"minimax_h3_prepared_reference_videos"})
if TYPE_CHECKING:
from vllm_omni.engine import AdditionalInformationEntry, AdditionalInformationPayload
@@ -98,6 +101,8 @@ class OmniPayloadMeta(TypedDict, total=False):
# reproducible, and the producing stage's SamplingParams do not travel
# with the payload.
audio_seed: int
token_role_ids: torch.Tensor
minimax_h3_prepared_reference_videos: str
class OmniPayload(TypedDict, total=False):
@@ -200,6 +205,8 @@ class MetaStruct(_StructBase):
code_flat_numel: int | None = None
omni_final_stage_id: int | None = None
audio_seed: int | None = None
token_role_ids: torch.Tensor | None = None
minimax_h3_prepared_reference_videos: str | None = None
class OmniPayloadStruct(_StructBase):
@@ -0,0 +1,28 @@
# MiniMax H3: vLLM Qwen3-VL text encoder + diffusion video/audio generation.
pipeline: minimax_h3_disaggregated
async_chunk: false
trust_remote_code: true
distributed_executor_backend: mp
stages:
- stage_id: 0
devices: "0,1"
max_num_seqs: 1
tensor_parallel_size: 2
default_sampling_params:
max_tokens: 1
temperature: 0.0
detokenize: false
- stage_id: 1
devices: "2,3,4,5"
max_num_seqs: 1
model_loaded:
text_encoder: false
parallel_config:
tensor_parallel_size: 1
ulysses_degree: 4
vae_patch_parallel_size: 4
default_sampling_params:
seed: 42
num_inference_steps: 50
@@ -0,0 +1,31 @@
# MiniMax H3 Turbo: five-sigma defaults for the disaggregated topology.
pipeline: minimax_h3_disaggregated
async_chunk: false
trust_remote_code: true
distributed_executor_backend: mp
stages:
- stage_id: 0
devices: "0,1"
max_num_seqs: 1
tensor_parallel_size: 2
default_sampling_params:
max_tokens: 1
temperature: 0.0
detokenize: false
- stage_id: 1
devices: "2,3,4,5"
max_num_seqs: 1
model_loaded:
text_encoder: false
parallel_config:
tensor_parallel_size: 1
ulysses_degree: 4
vae_patch_parallel_size: 4
default_sampling_params:
seed: 42
num_inference_steps: 5
extra_args:
flow_shift: 6.0
audio_flow_shift: 3.0
+7 -2
View File
@@ -632,8 +632,12 @@ def resolve_model_class_name(
return None
is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"}
# Diffusers models: read _class_name from the pipeline index.
model_index = get_diffusion_model_index(model, revision=revision)
# Diffusers models: read _class_name from the pipeline index. Missing
# local paths can otherwise be interpreted as invalid Hub repo IDs.
try:
model_index = get_diffusion_model_index(model, revision=revision)
except Exception:
model_index = None
if model_index is not None:
return model_index.get("_class_name")
if diffusion_load_format == "diffusers":
@@ -881,6 +885,7 @@ class OmniDiffusionConfig:
default_factory=lambda: {
"transformer": True,
"vae": True,
"text_encoder": True,
}
)
override_transformer_cls_name: str | None = None
@@ -60,6 +60,8 @@ class InlineStageDiffusionClient(StageClientBase):
self.default_sampling_params = metadata.default_sampling_params
self.requires_multimodal_data = metadata.requires_multimodal_data
self.custom_process_input_func = metadata.custom_process_input_func
self.prompt_transform_func = None
self.prompt_expand_func = None
self.engine_input_source = metadata.engine_input_source
self.batch_size = batch_size
+10
View File
@@ -45,6 +45,16 @@ def supports_audio_output(model_class_name: str) -> bool:
return bool(getattr(model_cls, "support_audio_output", False))
def get_diffusion_output_type(model_class_name: str | None) -> str:
"""Return the declared final modality for a diffusion model."""
declared_output_type = get_diffusion_model_metadata(model_class_name).final_output_type
if declared_output_type is not None:
return declared_output_type
if model_class_name and supports_audio_output(model_class_name):
return "audio"
return "image"
def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool) -> int:
"""Get num_frames for the dummy warmup run. Returns 0 to skip warmup."""
+49 -6
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from dataclasses import dataclass
@@ -12,6 +12,7 @@ class DiffusionModelMetadata:
max_multimodal_image_inputs: int | None = None
supports_mixed_reference_inputs: bool = False
attention_mask_free: bool = False
final_output_type: str | None = None
QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4
@@ -41,6 +42,7 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
supports_multimodal_inputs=True,
max_multimodal_image_inputs=9,
supports_mixed_reference_inputs=True,
final_output_type="video",
# H3 represents alignment padding as a second packed sequence. The
# packed TRTLLM backend consumes cu_seqlens and isolates that padding.
attention_mask_free=True,
@@ -51,17 +53,54 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
supports_multimodal_inputs=True,
max_multimodal_image_inputs=9,
supports_mixed_reference_inputs=True,
final_output_type="video",
),
"WanPipeline": DiffusionModelMetadata(attention_mask_free=True),
"WanImageToVideoPipeline": DiffusionModelMetadata(attention_mask_free=True),
"WanVACEPipeline": DiffusionModelMetadata(attention_mask_free=True),
"WanS2VPipeline": DiffusionModelMetadata(attention_mask_free=True),
"WanPipeline": DiffusionModelMetadata(
attention_mask_free=True,
final_output_type="video",
),
"WanImageToVideoPipeline": DiffusionModelMetadata(
attention_mask_free=True,
final_output_type="video",
),
"WanVACEPipeline": DiffusionModelMetadata(
attention_mask_free=True,
final_output_type="video",
),
"WanS2VPipeline": DiffusionModelMetadata(
attention_mask_free=True,
final_output_type="video",
),
"WanT2VDMD2Pipeline": DiffusionModelMetadata(final_output_type="video"),
"WanI2VDMD2Pipeline": DiffusionModelMetadata(final_output_type="video"),
"LTX2Pipeline": DiffusionModelMetadata(final_output_type="video"),
"LTX2DistilledPipeline": DiffusionModelMetadata(final_output_type="video"),
"LTX2T2VDMD2Pipeline": DiffusionModelMetadata(final_output_type="video"),
"LTX2I2VDMD2Pipeline": DiffusionModelMetadata(final_output_type="video"),
"HeliosPipeline": DiffusionModelMetadata(final_output_type="video"),
"HeliosPyramidPipeline": DiffusionModelMetadata(final_output_type="video"),
"HunyuanVideo15Pipeline": DiffusionModelMetadata(final_output_type="video"),
"HunyuanVideo15ImageToVideoPipeline": DiffusionModelMetadata(final_output_type="video"),
"LingBotVideoPipeline": DiffusionModelMetadata(final_output_type="video"),
"LongCatVideoAvatarPipeline": DiffusionModelMetadata(final_output_type="video"),
"MagiHumanPipeline": DiffusionModelMetadata(final_output_type="video"),
"DreamIDOmniPipeline": DiffusionModelMetadata(final_output_type="video"),
"Cosmos3OmniDiffusersPipeline": DiffusionModelMetadata(final_output_type="video"),
"Cosmos3OmniPipeline": DiffusionModelMetadata(final_output_type="video"),
"SanaVideoPipeline": DiffusionModelMetadata(final_output_type="video"),
"SanaWmPipeline": DiffusionModelMetadata(
supports_multimodal_inputs=True,
max_multimodal_image_inputs=1,
),
}
_DIFFUSION_MODEL_METADATA_ALIASES = {
"WanDMDPipeline": "WanPipeline",
"LTX2TwoStagePipeline": "LTX2Pipeline",
"LTX2DistilledOneStagePipeline": "LTX2DistilledPipeline",
"LingBotWorldCausalDMDPipeline": "LingBotVideoPipeline",
}
def get_diffusion_model_metadata(model_class_name: str | None) -> DiffusionModelMetadata:
# Unknown models fall back to "no special multimodal capabilities" so new
@@ -71,6 +110,9 @@ def get_diffusion_model_metadata(model_class_name: str | None) -> DiffusionModel
metadata = _DIFFUSION_MODEL_METADATA.get(model_class_name)
if metadata is not None:
return metadata
canonical_name = _DIFFUSION_MODEL_METADATA_ALIASES.get(model_class_name)
if canonical_name is not None:
return _DIFFUSION_MODEL_METADATA[canonical_name]
# Some checkpoints report the HF architecture name diff from internal pipeline class name
# (e.g. HunyuanImage3ForCausalMM, WanVACEPipeline, OmniVoice ...).
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
@@ -84,5 +126,6 @@ def get_diffusion_model_metadata(model_class_name: str | None) -> DiffusionModel
# key spaces. Aliases whose pipeline class has no metadata entry (e.g.
# Wan22VACEPipeline, OmniVoicePipeline) still fall back to the defaults;
# that is not a regression, it just means no capability override.
return _DIFFUSION_MODEL_METADATA.get(pipeline_cls_name, DiffusionModelMetadata())
canonical_name = _DIFFUSION_MODEL_METADATA_ALIASES.get(pipeline_cls_name, pipeline_cls_name)
return _DIFFUSION_MODEL_METADATA.get(canonical_name, DiffusionModelMetadata())
return DiffusionModelMetadata()
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
"""MiniMax H3 audio-video diffusion support."""
from vllm_omni.model_executor.models.minimax_h3 import reference_video
from .pipeline_minimax_h3 import (
MiniMaxH3Pipeline,
get_minimax_h3_post_process_func,
@@ -9,4 +11,5 @@ from .pipeline_minimax_h3 import (
__all__ = [
"MiniMaxH3Pipeline",
"get_minimax_h3_post_process_func",
"reference_video",
]
@@ -62,6 +62,41 @@ from vllm_omni.errors import OmniClientError, client_error_from_metadata
from vllm_omni.model_executor.model_loader.weight_utils import (
download_weights_from_hf_specific,
)
from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
resolve_minimax_h3_partition,
)
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
MiniMaxH3TextConditioning,
)
from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
load_minimax_h3_images as _load_images,
)
from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
minimax_h3_multi_image_presentation,
minimax_h3_ref2va_presentation,
minimax_h3_ref2va_video_presentation,
minimax_h3_text_only_ids,
)
from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
resolve_minimax_h3_aspect_ratio as _resolve_minimax_h3_aspect_ratio,
)
from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
resolve_minimax_h3_output_canvas as _resolve_output_canvas,
)
from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
resolve_minimax_h3_reference_image_shape as _reference_image_shape,
)
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY,
deserialize_prepared_reference_videos,
load_audio_file,
load_video_audio,
load_video_frames,
prepare_reference_videos,
sample_reference_video_frames,
validate_reference_audio_files,
validate_reference_audio_waveforms,
)
from vllm_omni.platforms import current_omni_platform
from vllm_omni.quantization import (
resolve_component_quant_config as _resolve_component_quant_config,
@@ -97,23 +132,7 @@ from .packed_tokens import (
minimax_h3_unpack_audio_tokens,
minimax_h3_unpatchify_video_tokens,
)
from .presentation import (
minimax_h3_multi_image_presentation_ids,
minimax_h3_multi_image_presentation_token_tags,
minimax_h3_ref2va_presentation,
minimax_h3_ref2va_video_presentation,
minimax_h3_text_only_ids,
)
from .quality_policy import MINIMAX_H3_GENERIC_CACHE_KEY, MiniMaxH3QualityPolicy
from .reference_video import (
load_audio_file,
load_video_audio,
load_video_frames,
prepare_reference_videos,
sample_reference_video_frames,
validate_reference_audio_files,
validate_reference_audio_waveforms,
)
from .scheduling_minimax_h3_euler_ancestral import (
minimax_h3_euler_eta0_step,
minimax_h3_rf_v_to_x0,
@@ -169,6 +188,28 @@ MINIMAX_H3_TASK_DOWNLOAD_PATTERNS = {
"fl2va": ["FL2VA/**"],
"ref2va": ["Ref2VA/**"],
}
MINIMAX_H3_DIFFUSION_DOWNLOAD_PATTERNS = {
"fl2va": [
"FL2VA/model_index.json",
"FL2VA/transformer/**",
"FL2VA/video_vae/**",
"FL2VA/audio_vae/**",
],
"ref2va": [
"Ref2VA/model_index.json",
"Ref2VA/transformer/**",
"Ref2VA/video_vae/**",
"Ref2VA/audio_vae/**",
],
"combined": [
"FL2VA/model_index.json",
"FL2VA/transformer/**",
"FL2VA/video_vae/**",
"FL2VA/audio_vae/**",
"Ref2VA/model_index.json",
"Ref2VA/transformer/**",
],
}
def _resolve_minimax_h3_text_encoder_quant_config(
@@ -182,33 +223,27 @@ def _minimax_h3_partition_for_task(
task_type: str | None,
model: str | None = None,
) -> str:
task = str(task_type or "auto").lower()
if task == "auto" and model is not None:
path = Path(model)
if path.is_dir() and path.name in {"FL2VA", "Ref2VA"} and (path / "model_index.json").is_file():
return path.name.lower()
if task in {"auto", "combined"}:
return "combined"
if task in {"t2va", "fl2va"}:
return "fl2va"
if task == "ref2va":
return "ref2va"
raise ValueError(f"MiniMax-H3 task_type must be one of auto, t2va, fl2va, or ref2va; got {task_type!r}")
return resolve_minimax_h3_partition(model or "", task_type, auto_partition="combined")
def _resolve_minimax_h3_model_root(
model: str,
revision: str | None,
partition: str,
*,
load_text_encoder: bool,
) -> Path:
path = Path(model)
if path.is_dir():
if path.name in {"FL2VA", "Ref2VA"} and (path / "model_index.json").is_file():
if path.name in {"FL2VA", "Ref2VA"}:
return path.parent
return path
allow_patterns = (
MINIMAX_H3_DOWNLOAD_PATTERNS if partition == "combined" else MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[partition]
)
if load_text_encoder:
allow_patterns = (
MINIMAX_H3_DOWNLOAD_PATTERNS if partition == "combined" else MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[partition]
)
else:
allow_patterns = MINIMAX_H3_DIFFUSION_DOWNLOAD_PATTERNS[partition]
return Path(
download_weights_from_hf_specific(
model_name_or_path=model,
@@ -287,6 +322,29 @@ def _read_base_schedule(release: Mapping[str, Any]) -> DMD2SigmaSchedule | None:
return DMD2SigmaSchedule.from_metadata(release)
def resolve_minimax_h3_diffusion_model_path(
model: str,
revision: str | None,
task_type: str | None,
) -> str:
"""Resolve a repository root or Hub ID to its startup partition."""
partition = (
"combined"
if str(task_type or "").lower() == "combined"
else resolve_minimax_h3_partition(model, task_type, auto_partition="fl2va")
)
model_root = _resolve_minimax_h3_model_root(
model,
revision,
partition,
load_text_encoder=False,
)
if partition == "combined":
return str(model_root)
subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
return str(model_root / subdir)
def _minimax_h3_post_process(output, output_type: str = "np"):
"""Convert the joint video/audio output without capturing worker state.
@@ -310,6 +368,12 @@ def _minimax_h3_post_process(output, output_type: str = "np"):
}
def _register_dlo_component_cache(cache: BoundedAllocatorCache, *components: Any) -> None:
for component in components:
if component is not None:
component.set_omni_component_cache(cache)
def get_minimax_h3_post_process_func(
od_config: OmniDiffusionConfig,
):
@@ -328,39 +392,6 @@ def _load_image(value: Any) -> Image.Image:
return images[0]
def _load_images(value: Any) -> list[Image.Image]:
if isinstance(value, (list, tuple)):
if not value:
raise OmniClientError("MiniMax H3 image input must not be empty")
return [_load_image(item) for item in value]
if isinstance(value, (str, os.PathLike)):
file_size = os.path.getsize(value)
if file_size > MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES:
raise OmniClientError("MiniMax H3 reference image exceeds the 30 MiB size limit")
with Image.open(value) as image:
image_format = str(image.format or "").lower()
if image_format and image_format not in MINIMAX_H3_REFERENCE_IMAGE_FORMATS:
raise OmniClientError(
f"MiniMax H3 reference image must use JPG, JPEG, PNG, WEBP, HEIC, or HEIF, got {image.format}"
)
return [image.convert("RGB")]
if isinstance(value, Image.Image):
return [value.convert("RGB")]
if isinstance(value, torch.Tensor):
tensor = value.detach().float().cpu()
if tensor.ndim == 4 and tensor.shape[0] == 1:
tensor = tensor[0]
if tensor.ndim != 3:
raise OmniClientError(f"image tensor must be [C,H,W], got {tuple(tensor.shape)}")
if tensor.shape[0] in (1, 3, 4):
tensor = tensor.permute(1, 2, 0)
array = tensor.numpy()
if array.max(initial=0) <= 1.0:
array = array * 255.0
return [Image.fromarray(array.clip(0, 255).astype(np.uint8)).convert("RGB")]
raise OmniClientError(f"unsupported MiniMax H3 image input {type(value)!r}")
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)
@@ -369,7 +400,7 @@ def _load_audio(value: Any) -> tuple[torch.Tensor, int]:
return audios[0]
if isinstance(value, (str, os.PathLike)):
return load_audio_file(str(value))
if isinstance(value, tuple) and len(value) == 2:
if isinstance(value, (list, tuple)) and len(value) == 2:
waveform, sample_rate = value
waveform = torch.as_tensor(waveform).float()
return waveform, int(sample_rate)
@@ -423,6 +454,21 @@ def _resolve_fl2va_keyframe_indices(extra: Mapping[str, Any], image_count: int)
return raw_indices
def _reuse_prepared_reference_videos(
prepared: list[dict[str, Any]] | None,
*,
expected_count: int,
) -> list[dict[str, Any]] | None:
if prepared is None:
return None
if len(prepared) != expected_count:
raise OmniClientError("MiniMax H3 prepared-reference-video count does not match the request")
for item in prepared:
if not os.path.isfile(item["prepared_path"]):
raise OmniClientError(f"MiniMax H3 prepared reference video is unavailable: {item['prepared_path']}")
return prepared
def _validate_ref2va_reference_counts(
image_count: int,
video_count: int,
@@ -443,55 +489,6 @@ def _validate_ref2va_reference_counts(
raise OmniClientError("ref2va accepts at most 12 total references")
def _resolve_minimax_h3_aspect_ratio(
task: str,
value: Any,
image: Image.Image | None,
) -> float:
"""Resolve H3's task-specific ratio policy.
T2VA must name one of the official ratios. FL2VA always follows the
first input image, even when a generic client sends ``aspect_ratio``.
Ref2VA defaults to 16:9; ``adaptive``/``auto`` are retained as aliases
for that default for compatibility with existing clients.
"""
if task == "fl2va":
if image is None:
raise OmniClientError("fl2va requires an input image to resolve its aspect ratio")
return float(image.width) / float(image.height)
if value is None:
if task == "t2va":
raise OmniClientError("t2va requires an explicit aspect_ratio")
return MINIMAX_H3_SUPPORTED_ASPECT_RATIOS["16:9"]
if isinstance(value, str):
normalized = value.strip().lower()
if normalized in {"adaptive", "auto"}:
if task == "t2va":
raise OmniClientError("t2va requires an explicit named aspect_ratio, not adaptive")
return MINIMAX_H3_SUPPORTED_ASPECT_RATIOS["16:9"]
if normalized in MINIMAX_H3_SUPPORTED_ASPECT_RATIOS:
return MINIMAX_H3_SUPPORTED_ASPECT_RATIOS[normalized]
try:
numeric_value = float(normalized)
except (TypeError, ValueError) as exc:
supported = ", ".join(MINIMAX_H3_SUPPORTED_ASPECT_RATIOS)
raise OmniClientError(f"MiniMax H3 aspect_ratio must be one of {supported}, got {value!r}") from exc
elif isinstance(value, (int, float, np.integer, np.floating)) and not isinstance(value, bool):
numeric_value = float(value)
else:
raise OmniClientError(f"MiniMax H3 aspect_ratio must be a string ratio, got {value!r}")
if not math.isfinite(numeric_value) or not any(
math.isclose(numeric_value, ratio, rel_tol=0.0, abs_tol=1e-6)
for ratio in MINIMAX_H3_SUPPORTED_ASPECT_RATIOS.values()
):
supported = ", ".join(MINIMAX_H3_SUPPORTED_ASPECT_RATIOS)
raise OmniClientError(f"MiniMax H3 aspect_ratio must be one of {supported}, got {value!r}")
return numeric_value
def _resolve_minimax_h3_num_outputs(value: Any) -> int:
if value is None:
return 1
@@ -609,49 +606,6 @@ def _broadcast_tensor(
return output
def _reference_image_shape(image: Image.Image) -> tuple[int, int]:
width, height = image.size
ratio = width / height
if not 0.4 <= ratio <= 2.5:
raise OmniClientError(f"reference image aspect ratio must be in [0.4, 2.5], got {width}x{height}")
if min(width, height) < 256 or max(width, height) > 5760:
raise OmniClientError(f"reference image dimensions must be in [256, 5760] pixels, got {width}x{height}")
scale = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE / min(width, height)
return (
_align_multiple(
width * scale,
MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE,
),
_align_multiple(
height * scale,
MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE,
),
)
def _resolve_output_canvas(aspect_ratio: float, short_edge: int) -> tuple[int, int]:
"""Resolve the official H3 ratio/area policy to a 32-pixel canvas."""
if not math.isfinite(float(aspect_ratio)) or float(aspect_ratio) <= 0:
raise OmniClientError(f"MiniMax H3 canvas aspect ratio must be positive, got {aspect_ratio!r}")
if short_edge != MINIMAX_H3_OUTPUT_SHORT_EDGE:
raise OmniClientError(f"MiniMax H3 target.short_edge must be {MINIMAX_H3_OUTPUT_SHORT_EDGE}, got {short_edge}")
if aspect_ratio >= 1.0:
width = float(short_edge) * aspect_ratio
height = float(short_edge)
else:
width = float(short_edge)
height = float(short_edge) / aspect_ratio
area = width * height
if area > MINIMAX_H3_OUTPUT_MAX_PIXELS:
scale = (MINIMAX_H3_OUTPUT_MAX_PIXELS / area) ** 0.5
width *= scale
height *= scale
return (
_align_multiple(height, 32),
_align_multiple(width, 32),
)
class _SingleRankEncoderGroup:
"""Lightweight encoder group for ``text_encoder_tp_size == 1``.
@@ -806,6 +760,7 @@ class MiniMaxH3Pipeline(
if int(self.parallel_config.cfg_parallel_size) != 1:
raise ValueError("MiniMax-H3 is CFG-distilled and has no negative branch; cfg_parallel_size must be 1")
self.device = get_local_device()
self.load_text_encoder = od_config.model_loaded["text_encoder"]
self.partition = _minimax_h3_partition_for_task(
getattr(od_config, "task_type", None),
str(od_config.model),
@@ -815,6 +770,7 @@ class MiniMaxH3Pipeline(
str(od_config.model),
od_config.revision,
self.partition,
load_text_encoder=self.load_text_encoder,
)
model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA")
model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8"))
@@ -886,54 +842,62 @@ class MiniMaxH3Pipeline(
quant_config=transformer_quant_config,
)
self.tokenizer = Qwen2TokenizerFast.from_pretrained(
str(model_path),
subfolder="tokenizer",
local_files_only=os.path.isdir(model_path),
)
self.processor = Qwen3VLProcessor.from_pretrained(
str(model_path),
subfolder="processor",
local_files_only=os.path.isdir(model_path),
)
if self.load_text_encoder:
self.tokenizer = Qwen2TokenizerFast.from_pretrained(
str(model_path),
subfolder="tokenizer",
local_files_only=os.path.isdir(model_path),
)
self.processor = Qwen3VLProcessor.from_pretrained(
str(model_path),
subfolder="processor",
local_files_only=os.path.isdir(model_path),
)
else:
self.tokenizer = None
self.processor = None
_, rank, dit_world = _dit_rank_world()
self._dit_rank = rank
text_encoder_tp_size = int(getattr(self.parallel_config, "text_encoder_tp_size", 1))
if text_encoder_tp_size < 1:
raise ValueError(f"text_encoder_tp_size must be >= 1, got {text_encoder_tp_size}")
if text_encoder_tp_size > dit_world:
raise ValueError(
f"text_encoder_tp_size must not exceed the DiT group size ({dit_world}), got {text_encoder_tp_size}"
)
# The Qwen3-VL text model uses 64 attention heads / 8 KV heads; the
# encoder shards them across the encoder TP ranks.
if 64 % text_encoder_tp_size or 8 % text_encoder_tp_size:
raise ValueError(
"text_encoder_tp_size must divide both Qwen3-VL "
f"num_attention_heads (64) and num_key_value_heads (8), "
f"got {text_encoder_tp_size}"
)
self.text_encoder_tp_size = text_encoder_tp_size
self.text_encoder_group = self._build_text_encoder_group(text_encoder_tp_size)
load_text_encoder = rank < text_encoder_tp_size
self.text_encoder = MiniMaxH3Qwen3VLEncoder(
os.path.join(model_path, "text_encoder"),
device=self.device,
load_model=load_text_encoder,
encoder_group=self.text_encoder_group,
quant_config=_resolve_minimax_h3_text_encoder_quant_config(od_config.quantization_config),
)
if load_text_encoder:
self.weights_sources.append(
DiffusersPipelineLoader.ComponentSource(
model_or_path=str(model_path),
subfolder="text_encoder",
revision=od_config.revision,
prefix="text_encoder.",
fall_back_to_pt=False,
if self.load_text_encoder:
text_encoder_tp_size = int(getattr(self.parallel_config, "text_encoder_tp_size", 1))
if text_encoder_tp_size < 1:
raise ValueError(f"text_encoder_tp_size must be >= 1, got {text_encoder_tp_size}")
if text_encoder_tp_size > dit_world:
raise ValueError(
f"text_encoder_tp_size must not exceed the DiT group size ({dit_world}), got {text_encoder_tp_size}"
)
# The Qwen3-VL text model uses 64 attention heads / 8 KV heads.
if 64 % text_encoder_tp_size or 8 % text_encoder_tp_size:
raise ValueError(
"text_encoder_tp_size must divide both Qwen3-VL "
f"num_attention_heads (64) and num_key_value_heads (8), "
f"got {text_encoder_tp_size}"
)
self.text_encoder_tp_size = text_encoder_tp_size
self.text_encoder_group = self._build_text_encoder_group(text_encoder_tp_size)
self.text_encoder = MiniMaxH3Qwen3VLEncoder(
os.path.join(model_path, "text_encoder"),
device=self.device,
load_model=rank < text_encoder_tp_size,
encoder_group=self.text_encoder_group,
quant_config=_resolve_minimax_h3_text_encoder_quant_config(od_config.quantization_config),
)
if rank < text_encoder_tp_size:
self.weights_sources.append(
DiffusersPipelineLoader.ComponentSource(
model_or_path=str(model_path),
subfolder="text_encoder",
revision=od_config.revision,
prefix="text_encoder.",
fall_back_to_pt=False,
)
)
else:
self.text_encoder_tp_size = 0
self.text_encoder_group = None
self.text_encoder = None
self._encoder_modules = []
stage_components = bool(
od_config.enable_layerwise_offload or getattr(od_config, "enable_distributed_layerwise_offload", False)
)
@@ -954,8 +918,12 @@ class MiniMaxH3Pipeline(
self._dlo_component_cache = None
if getattr(od_config, "enable_distributed_layerwise_offload", False):
self._dlo_component_cache = BoundedAllocatorCache(self.device)
for component in (self.text_encoder, self.video_vae, self.audio_vae):
component.set_omni_component_cache(self._dlo_component_cache)
_register_dlo_component_cache(
self._dlo_component_cache,
self.text_encoder,
self.video_vae,
self.audio_vae,
)
self._quality_policy = MiniMaxH3QualityPolicy(od_config)
self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self)
@@ -992,6 +960,8 @@ class MiniMaxH3Pipeline(
# path as the DiT.
for component_name in ("video_vae", "audio_vae"):
component = getattr(self, component_name)
if component is None:
continue
loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters())
return loaded_with_prefix
@@ -1191,12 +1161,7 @@ class MiniMaxH3Pipeline(
if task == "fl2va":
if prepared_videos:
raise OmniClientError("fl2va does not accept video conditions")
ids = minimax_h3_multi_image_presentation_ids(
self.tokenizer,
prompt=prompt,
image_token_counts=image_token_counts,
)
tags = minimax_h3_multi_image_presentation_token_tags(
ids, tags = minimax_h3_multi_image_presentation(
self.tokenizer,
prompt=prompt,
image_token_counts=image_token_counts,
@@ -1985,12 +1950,46 @@ class MiniMaxH3Pipeline(
raise OmniClientError("MiniMax H3 requires a non-empty prompt")
return prompt, multi_modal_data
@staticmethod
def _extract_text_conditioning(raw_prompt: Any) -> MiniMaxH3TextConditioning | None:
if isinstance(raw_prompt, str):
return None
additional_information = raw_prompt.get("additional_information") or {}
text_encoder_output = additional_information.get("text_encoder_output")
if text_encoder_output is None:
return None
if not isinstance(text_encoder_output, Mapping):
raise OmniClientError("text_encoder_output must be a mapping")
try:
return MiniMaxH3TextConditioning.from_payload(text_encoder_output)
except ValueError as exc:
raise OmniClientError(str(exc)) from exc
@staticmethod
def _extract_prepared_reference_videos(raw_prompt: Any) -> list[dict[str, Any]] | None:
if isinstance(raw_prompt, str):
return None
additional_information = raw_prompt.get("additional_information") or {}
meta = additional_information.get("meta") or {}
descriptor = meta.get(MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY)
if descriptor is None:
return None
if not isinstance(descriptor, str):
raise OmniClientError("MiniMax H3 prepared-reference-video descriptor must be a string")
try:
_, videos = deserialize_prepared_reference_videos(descriptor)
except ValueError as exc:
raise OmniClientError(str(exc)) from exc
return videos
def _prepare_request_inputs(
self,
*,
prompt: str,
multi_modal_data: dict[str, Any],
sampling: Any,
text_conditioning: MiniMaxH3TextConditioning | None = None,
prepared_reference_videos: list[dict[str, Any]] | None = None,
) -> dict[str, Any]:
"""Resolve the task and output shape, then run every request-level encode.
@@ -2083,12 +2082,20 @@ class MiniMaxH3Pipeline(
# before starting any subsequent collective.
prep_error: Exception | None = None
try:
prepared_videos = self._prepare_reference_videos(
raw_videos,
target_frame_count=num_frames,
workdir=workdir,
start_time_seconds=extra.get("start_time_seconds"),
)
_, rank, _ = _dit_rank_world()
if prepared_reference_videos is not None:
if rank == 0:
prepared_videos = _reuse_prepared_reference_videos(
prepared_reference_videos,
expected_count=video_count,
)
else:
prepared_videos = self._prepare_reference_videos(
raw_videos,
target_frame_count=num_frames,
workdir=workdir,
start_time_seconds=extra.get("start_time_seconds"),
)
except Exception as exc:
prep_error = exc
_broadcast_rank0_exception(prep_error)
@@ -2129,13 +2136,27 @@ class MiniMaxH3Pipeline(
audio_index += 1
condition_labels.append(("audio", audio_index))
text_embeddings, text_tags = self.encode_prompt(
task=task,
prompt=prompt,
images=prepared_images,
prepared_videos=prepared_videos,
condition_labels=condition_labels if task == "ref2va" else None,
)
if text_conditioning is not None:
text_embeddings = text_conditioning.hidden_states.to(
device=self.device,
dtype=torch.bfloat16,
)
text_tags = text_conditioning.token_tags.to(
device=self.device,
dtype=torch.long,
)
elif getattr(self, "text_encoder", None) is not None:
text_embeddings, text_tags = self.encode_prompt(
task=task,
prompt=prompt,
images=prepared_images,
prepared_videos=prepared_videos,
condition_labels=condition_labels if task == "ref2va" else None,
)
else:
raise OmniClientError(
"MiniMax H3 diffusion stage requires text_encoder_output when text_encoder is not loaded"
)
# ``prepared_videos`` is intentionally ``None`` on non-zero DiT
# ranks; the distributed video encoder broadcasts the prepared
@@ -2264,11 +2285,14 @@ class MiniMaxH3Pipeline(
def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput:
if len(request.prompts) != 1:
raise OmniClientError("MiniMax H3 supports one request at a time")
prompt, multi_modal_data = self._extract_prompt(request.prompts[0])
raw_prompt = request.prompts[0]
prompt, multi_modal_data = self._extract_prompt(raw_prompt)
context = self._prepare_request_inputs(
prompt=prompt,
multi_modal_data=multi_modal_data,
sampling=request.sampling_params,
text_conditioning=self._extract_text_conditioning(raw_prompt),
prepared_reference_videos=self._extract_prepared_reference_videos(raw_prompt),
)
denoise_kwargs = self._denoise_kwargs(context)
num_outputs = context["num_outputs"]
@@ -2357,6 +2381,8 @@ class MiniMaxH3Pipeline(
prompt=prompt,
multi_modal_data=multi_modal_data,
sampling=state.sampling,
text_conditioning=self._extract_text_conditioning(state.prompt),
prepared_reference_videos=self._extract_prepared_reference_videos(state.prompt),
)
inputs = self._build_denoise_inputs(**self._denoise_kwargs(context))
+10 -2
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -8,7 +8,7 @@ from dataclasses import dataclass, field
from typing import TypeGuard, cast
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.io_support import supports_audio_output
from vllm_omni.diffusion.io_support import get_diffusion_output_type, supports_audio_output
from vllm_omni.diffusion.registry import DiffusionModelRegistry
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.inputs.data import OmniPromptType
@@ -146,6 +146,11 @@ def format_diffusion_outputs(
is_text_output = postprocess_output.primary_key == "text"
is_audio_output = supports_audio_output(od_config.model_class_name)
final_output_type = (
postprocess_output.primary_key
if postprocess_output.primary_key in {"image", "video"}
else get_diffusion_output_type(od_config.model_class_name)
)
audio_sample_rate = _metadata_audio_sample_rate(postprocess_output.metadata)
if is_audio_output and audio_sample_rate is None:
model_cls = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name)
@@ -167,6 +172,7 @@ def format_diffusion_outputs(
postprocess_output=postprocess_output,
is_text_output=is_text_output,
is_audio_output=is_audio_output,
final_output_type=final_output_type,
audio_sample_rate=audio_sample_rate,
finished=diffusion_output.finished,
)
@@ -272,6 +278,7 @@ def _format_single_prompt_output(
postprocess_output: DiffusionPostprocessOutput,
is_text_output: bool,
is_audio_output: bool,
final_output_type: str,
audio_sample_rate: int | None,
finished: bool = True,
) -> list[OmniRequestOutput]:
@@ -339,6 +346,7 @@ def _format_single_prompt_output(
trajectory_log_probs=trajectory_log_probs,
trajectory_decoded=trajectory_decoded,
multimodal_output=mm_output,
final_output_type=final_output_type,
stage_durations=diffusion_output.stage_durations,
peak_memory_mb=diffusion_output.peak_memory_mb,
finished=finished,
@@ -158,6 +158,7 @@ class StageDiffusionClient(StageClientBase):
self.final_output_type = metadata.final_output_type
self.model_stage = metadata.model_stage
self.default_sampling_params = metadata.default_sampling_params
self.prompt_transform_func = metadata.prompt_transform_func
self.prompt_expand_func = metadata.prompt_expand_func
self.requires_multimodal_data = getattr(metadata, "requires_multimodal_data", False)
self.custom_process_input_func = getattr(metadata, "custom_process_input_func", None)
-1
View File
@@ -583,7 +583,6 @@ class OrchestratorArgs:
cfg_parallel_size: int = 1
vae_patch_parallel_size: int = 1
vae_parallel_mode: str = "tile"
text_encoder_tp_size: int = 1
default_sampling_params: str | None = None
max_generated_image_size: int | None = None
tts_max_instructions_length: int | None = None
+115 -27
View File
@@ -12,8 +12,10 @@ from __future__ import annotations
import asyncio
import concurrent.futures
import copy
import json
import queue
import shutil
import threading
import time
import uuid
@@ -36,8 +38,14 @@ from vllm_omni.config.stage_config import (
DuplexSessionRuntimeConfig,
load_deploy_config,
)
from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config
from vllm_omni.diffusion.diffusion_engine import supports_audio_output
from vllm_omni.data_entry_keys import REQUEST_ARTIFACT_DIRS_KEY, TRANSFORM_OWNED_META_KEYS
from vllm_omni.diffusion.data import (
DiffusionParallelConfig,
parse_attention_config,
resolve_model_class_name,
)
from vllm_omni.diffusion.io_support import get_diffusion_output_type
from vllm_omni.engine import OmniEngineCoreRequest
from vllm_omni.engine.async_engine_utils import (
SHUTDOWN_ENQUEUE_TIMEOUT_S,
SHUTDOWN_JOIN_TIMEOUT_S,
@@ -61,6 +69,7 @@ from vllm_omni.engine.messages import (
)
from vllm_omni.engine.orchestrator import Orchestrator
from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient
from vllm_omni.engine.serialization import deserialize_additional_information
from vllm_omni.engine.stage_client import StageClient
from vllm_omni.engine.stage_init_utils import build_stage0_input_processor
from vllm_omni.engine.stage_pool import StagePool
@@ -224,6 +233,7 @@ class AsyncOmniEngine:
self.stage_pools: list[StagePool] = []
self.stage_clients: list[StageClient] = [] # logical-stage view for external readers
self.input_processor: InputProcessor | None = None
self.prompt_transform_func: Any | None = None
self.prompt_expand_func: Any | None = None
self.supported_tasks: tuple[str, ...] = ("generate",)
self.default_sampling_params_list: list[OmniSamplingParams] = []
@@ -332,6 +342,9 @@ class AsyncOmniEngine:
if self.stage_vllm_configs and self.stage_vllm_configs[0] is not None
else None
)
self.prompt_transform_func = (
getattr(self.stage_clients[0], "prompt_transform_func", None) if self.stage_clients else None
)
self.prompt_expand_func = next(
(
getattr(client, "prompt_expand_func", None)
@@ -679,14 +692,45 @@ class AsyncOmniEngine:
# Keep the original prompt for downstream stages (they need the raw
# dict, e.g. for multi_modal_data).
if isinstance(prompt, dict):
raw_info = prompt.get("additional_information")
if isinstance(raw_info, dict):
raw_meta = raw_info.get("meta")
if isinstance(raw_meta, dict):
for key in TRANSFORM_OWNED_META_KEYS:
raw_meta.pop(key, None)
original_prompt = prompt
preselected_stage0_replica: int | None = None
request_artifact_dirs: list[str] = []
stage_type = self.stage_metadata[0].stage_type
output_prompt_text: Any = None
_preprocess_ms = 0.0
if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest):
# Inject global_request_id into the raw prompt.
# Stage transforms and downstream stages must share the same
# request identity, including when the transform replaces the
# prompt object.
if isinstance(prompt, dict):
inject_global_id(prompt, request_id)
elif isinstance(prompt, list):
for item in prompt:
inject_global_id(item, request_id)
prompt_transform_func = getattr(self, "prompt_transform_func", None)
if prompt_transform_func is not None:
if isinstance(prompt, dict):
prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None)
prompt = prompt_transform_func(
copy.copy(prompt),
effective_sampling_params_list,
)
if isinstance(prompt, dict):
raw_dirs = prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None)
if isinstance(raw_dirs, list) and all(isinstance(path, str) for path in raw_dirs):
request_artifact_dirs = list(raw_dirs)
if isinstance(original_prompt, dict):
original_prompt[REQUEST_ARTIFACT_DIRS_KEY] = request_artifact_dirs
if isinstance(prompt, dict):
inject_global_id(prompt, request_id)
elif isinstance(prompt, list):
@@ -718,12 +762,25 @@ class AsyncOmniEngine:
except Exception:
if preselected_stage0_replica is not None and self.stage_pools:
self.stage_pools[0].release_binding(request_id)
for artifact_dir in request_artifact_dirs:
shutil.rmtree(artifact_dir, ignore_errors=True)
raise
_preprocess_ms = (time.perf_counter() - _t_preprocess) * 1000.0
# TODO (Peiqi): add this for Qwen3-TTS only. Other models don't have
# additional_information field in the prompt.
request = upgrade_to_omni_request(request, prompt)
if isinstance(request, OmniEngineCoreRequest) and request.additional_information is not None:
processed_info = deserialize_additional_information(request.additional_information)
processed_meta = processed_info.get("meta")
if isinstance(processed_meta, dict):
if isinstance(original_prompt, dict):
original_info = dict(original_prompt.get("additional_information") or {})
original_meta = dict(original_info.get("meta") or {})
original_meta.update(processed_meta)
original_info["meta"] = original_meta
original_prompt["additional_information"] = original_info
if reasoning_ended is not None:
request.reasoning_ended = reasoning_ended
@@ -743,6 +800,8 @@ class AsyncOmniEngine:
if output_prompt_text is None and isinstance(original_prompt, dict):
output_prompt_text = original_prompt.get("prompt")
prompt = request
else:
request_artifact_dirs = []
return StageSubmissionMessage(
type=message_type,
@@ -756,6 +815,7 @@ class AsyncOmniEngine:
preprocess_ms=_preprocess_ms,
request_timestamp=request_timestamp,
enqueue_ts=time.perf_counter(),
request_artifact_dirs=request_artifact_dirs or None,
)
def _build_cfg_companions(
@@ -1029,7 +1089,7 @@ class AsyncOmniEngine:
num_devices = max(1, int(parallel_config.world_size))
devices = ",".join(str(i) for i in range(num_devices))
model_class_name = kwargs.get("model_class_name", None)
final_output_type = "audio" if model_class_name and supports_audio_output(model_class_name) else "image"
final_output_type = get_diffusion_output_type(model_class_name)
attention_config = None
if (
@@ -1222,11 +1282,22 @@ class AsyncOmniEngine:
},
}
def create_default_stage_config() -> list:
fallback_kwargs = dict(default_stage_kwargs)
if not fallback_kwargs.get("model_class_name"):
model_class_name = resolve_model_class_name(
model,
fallback_kwargs.get("diffusion_load_format", "default"),
)
if model_class_name is not None:
fallback_kwargs["model_class_name"] = model_class_name
return self._create_default_diffusion_stage_cfg(fallback_kwargs)
config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs(
model,
kwargs,
trust_remote_code=trust_remote_code,
default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(default_stage_kwargs),
default_stage_cfg_factory=create_default_stage_config,
deploy_config_path=deploy_config_path,
stage_overrides=stage_overrides,
strategy_config_path=strategy_config_path,
@@ -1358,22 +1429,29 @@ class AsyncOmniEngine:
a queue + coroutine-switch round-trip. The Orchestrator receives a
ready-to-submit OmniEngineCoreRequest.
"""
msg = self._build_add_request_message(
request_id=request_id,
prompt=prompt,
prompt_text=prompt_text,
sampling_params_list=sampling_params_list,
final_stage_id=final_stage_id,
final_output_stage_ids=final_output_stage_ids,
arrival_time=arrival_time,
lora_request=lora_request,
tokenization_kwargs=tokenization_kwargs,
trace_headers=trace_headers,
priority=priority,
data_parallel_rank=data_parallel_rank,
reasoning_ended=reasoning_ended,
resumable=resumable,
)
try:
msg = self._build_add_request_message(
request_id=request_id,
prompt=prompt,
prompt_text=prompt_text,
sampling_params_list=sampling_params_list,
final_stage_id=final_stage_id,
final_output_stage_ids=final_output_stage_ids,
arrival_time=arrival_time,
lora_request=lora_request,
tokenization_kwargs=tokenization_kwargs,
trace_headers=trace_headers,
priority=priority,
data_parallel_rank=data_parallel_rank,
reasoning_ended=reasoning_ended,
resumable=resumable,
)
except BaseException:
if isinstance(prompt, dict):
for artifact_dir in prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None) or ():
if isinstance(artifact_dir, str):
shutil.rmtree(artifact_dir, ignore_errors=True)
raise
# CFG companions are built before the parent is admitted, so the group
# is all-or-nothing: a build failure raises here, nothing is enqueued,
# and the caller sees the error. Admitting the parent first would leave
@@ -1381,13 +1459,23 @@ class AsyncOmniEngine:
# because a model whose guidance is mandatory cannot decode a request
# whose companion never arrived.
companions: list[AddCompanionRequestMessage] = []
if self.prompt_expand_func is not None and final_stage_id > 0:
effective_spl = msg.sampling_params_list
stage0_params = effective_spl[0] if effective_spl else None
if stage0_params is not None:
companions = self._build_cfg_companions(request_id, msg.original_prompt, stage0_params, effective_spl)
try:
if self.prompt_expand_func is not None and final_stage_id > 0:
effective_spl = msg.sampling_params_list
stage0_params = effective_spl[0] if effective_spl else None
if stage0_params is not None:
companions = self._build_cfg_companions(
request_id, msg.original_prompt, stage0_params, effective_spl
)
self.request_queue.sync_q.put(msg)
self.request_queue.sync_q.put(msg)
except BaseException:
for artifact_dir in msg.request_artifact_dirs or ():
shutil.rmtree(artifact_dir, ignore_errors=True)
raise
finally:
if isinstance(msg.original_prompt, dict):
msg.original_prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None)
for companion in companions:
self.request_queue.sync_q.put(companion)
if companions:
+1
View File
@@ -27,6 +27,7 @@ class StageSubmissionMessage(EngineQueueMessage, kw_only=True):
request_timestamp: float
enqueue_ts: float
final_output_stage_ids: list[int] | None = None
request_artifact_dirs: list[str] | None = None
class AddCompanionRequestMessage(EngineQueueMessage, kw_only=True):
+11
View File
@@ -16,6 +16,7 @@ from __future__ import annotations
import asyncio
import os
import shutil
import time as _time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
@@ -66,6 +67,12 @@ from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
def cleanup_request_artifact_dirs(artifact_dirs: set[str] | list[str]) -> None:
for artifact_dir in artifact_dirs:
shutil.rmtree(artifact_dir, ignore_errors=True)
# VLLM_OMNI_EVENT_DRIVEN_ORCH=1 switches the orchestration loop (and the
# serving-side final-output drain in entrypoints/async_omni.py) from the legacy
# 1 ms poll cadence to event-driven wakeups: one reader task per live LLM stage
@@ -220,6 +227,7 @@ class OrchestratorRequestState:
duplex_stage_fences: dict[int, DuplexFence] = field(default_factory=dict)
duplex_config_generation: int = -1
running_counter_registered: bool = False
request_artifact_dirs: set[str] = field(default_factory=set)
@dataclass
@@ -723,6 +731,7 @@ class Orchestrator:
final_output_stage_ids=final_output_stage_ids,
request_timestamp=float(msg.request_timestamp or _time.time()),
mm_features=getattr(prompt, "mm_features", None),
request_artifact_dirs=set(msg.request_artifact_dirs or ()),
)
self.request_states[request_id] = req_state
self._register_running_request(req_state)
@@ -1720,6 +1729,8 @@ class Orchestrator:
for request_id in cleanup_ids:
self._pd_kv_params.pop(request_id, None)
req_state = self.request_states.pop(request_id, None)
if req_state is not None:
cleanup_request_artifact_dirs(getattr(req_state, "request_artifact_dirs", ()))
if req_state is not None and req_state.running_counter_registered and self._running_counter is not None:
self._running_counter.decrement()
req_state.running_counter_registered = False
+1
View File
@@ -35,6 +35,7 @@ class StageClient(Protocol):
final_output: bool
final_output_type: FinalOutputModalityType | None
default_sampling_params: OmniSamplingParams
prompt_transform_func: Any | None
prompt_expand_func: Any | None
requires_multimodal_data: bool
custom_process_input_func: Any | None
@@ -143,6 +143,7 @@ class StageEngineCoreClientBase(StageClientBase):
self.final_output = metadata.final_output
self.final_output_type = metadata.final_output_type
self.default_sampling_params = metadata.default_sampling_params
self.prompt_transform_func = metadata.prompt_transform_func
self.prompt_expand_func = metadata.prompt_expand_func
self.custom_process_input_func = metadata.custom_process_input_func
+30
View File
@@ -125,6 +125,23 @@ def _resolve_model_tokenizer_paths(model: str, engine_args: dict[str, Any]) -> s
return model
def _resolve_model_path(model: str, engine_args: dict[str, Any]) -> str:
"""Apply a model-owned path resolver from stage engine args."""
resolver_path = engine_args.pop("model_path_resolver", None)
if resolver_path is None:
return model
resolver = _resolve_omni_metadata_hook(str(resolver_path))
if resolver is None:
return model
return str(
resolver(
model,
engine_args.get("revision"),
engine_args.get("task_type"),
)
)
def apply_cli_tokenizer(
engine_args: dict[str, Any],
*,
@@ -358,6 +375,7 @@ class StageMetadata:
custom_process_input_func: Callable | None
model_stage: str | None
runtime_cfg: Any
prompt_transform_func: Callable | None = None
prompt_expand_func: Callable | None = None
cfg_kv_collect_func: Callable | None = None
# Multi-replica: replica_id distinguishes replicas of the same stage.
@@ -427,6 +445,12 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata:
mod_path, fn_name = _cpif_path.rsplit(".", 1)
custom_process_input_func = getattr(importlib.import_module(mod_path), fn_name)
prompt_transform_func: Callable | None = None
_ptf_path = _get_attr_or_item(stage_config, "prompt_transform_func")
if _ptf_path:
_mod, _fn = _ptf_path.rsplit(".", 1)
prompt_transform_func = getattr(importlib.import_module(_mod), _fn)
prompt_expand_func: Callable | None = None
_pef_path = _get_attr_or_item(stage_config, "prompt_expand_func")
if _pef_path:
@@ -455,6 +479,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata:
custom_process_input_func=custom_process_input_func,
model_stage=model_stage,
runtime_cfg=runtime_cfg,
prompt_transform_func=prompt_transform_func,
cfg_kv_collect_func=cfg_kv_collect_func,
)
@@ -475,6 +500,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata:
custom_process_input_func=custom_process_input_func,
model_stage=model_stage,
runtime_cfg=runtime_cfg,
prompt_transform_func=prompt_transform_func,
prompt_expand_func=prompt_expand_func,
)
@@ -521,6 +547,7 @@ def extract_stage_metadata_from_omni_stage_config(
custom_process_input_func=custom_process_input_func,
model_stage=stage_config.model_stage,
runtime_cfg=stage_config.runtime_config,
prompt_transform_func=_resolve_omni_metadata_hook(stage_config.prompt_transform_func),
cfg_kv_collect_func=_resolve_omni_metadata_hook(stage_config.cfg_kv_collect_func),
)
@@ -537,6 +564,7 @@ def extract_stage_metadata_from_omni_stage_config(
custom_process_input_func=custom_process_input_func,
model_stage=stage_config.model_stage,
runtime_cfg=stage_config.runtime_config,
prompt_transform_func=_resolve_omni_metadata_hook(stage_config.prompt_transform_func),
prompt_expand_func=_resolve_omni_metadata_hook(stage_config.prompt_expand_func),
)
@@ -920,6 +948,7 @@ def _project_omni_stage_engine_args(
"hf_config_name": stage_config.hf_config_name,
"engine_output_type": stage_config.engine_output_type,
"custom_process_next_stage_input_func": stage_config.custom_process_next_stage_input_func,
"model_path_resolver": topology.model_path_resolver,
"retains_state_across_chunks": topology.retains_state_across_chunks,
}
engine_args.update(
@@ -986,6 +1015,7 @@ def _finalize_engine_args_dict(
stage_defines_tokenizer = (
engine_args_dict.get("tokenizer") is not None or engine_args_dict.get("tokenizer_subdir") is not None
)
model = _resolve_model_path(model, engine_args_dict)
audex_stage = str(engine_args_dict.get("model_stage") or "")
if audex_stage == "audex_xcodec":
# TTA stage 1 decodes with the external XCodec1 checkpoint, not a
+21 -8
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Stage Runtime implementations for single-node and distributed omni stages."""
@@ -12,7 +12,7 @@ import threading
from collections.abc import Callable, Iterator, Mapping, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any
from typing import Any, cast
import janus
from omegaconf import OmegaConf
@@ -545,7 +545,7 @@ class StageRuntime:
self,
plan: ReplicaInitPlan,
stage_init_timeout: int,
) -> StageEngineCoreClientBase:
) -> StagePoolClient:
"""Initialize one local LLM replica using vLLM's launch/attach pattern."""
resources: StageReplicaResources | None = None
stage_client = None
@@ -608,7 +608,7 @@ class StageRuntime:
)
logger.info("[StageRuntime] Stage %s initialized", plan.metadata.stage_id)
return stage_client
return cast(StagePoolClient, stage_client)
except Exception:
if stage_client is not None:
try:
@@ -641,7 +641,7 @@ class StageRuntime:
self,
plan: ReplicaInitPlan,
stage_init_timeout: int,
) -> Any:
) -> StagePoolClient:
"""Initialize one local diffusion replica end-to-end."""
client = None
resources = None
@@ -652,15 +652,28 @@ class StageRuntime:
):
omni_conn_cfg, omni_from, omni_to = plan.omni_kv_connector
if omni_conn_cfg:
if omni_from is None or omni_to is None:
raise RuntimeError("Omni KV connector requires source and destination stages")
inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to)
inject_kv_stage_info(plan.stage_cfg, plan.metadata.stage_id, self._stage_configs)
engine_args = getattr(plan.stage_cfg, "engine_args", {})
inline_diffusion = (
engine_args.get("inline_diffusion", False)
if hasattr(engine_args, "get")
else getattr(engine_args, "inline_diffusion", False)
)
custom_pipeline_args = (
engine_args.get("custom_pipeline_args")
if hasattr(engine_args, "get")
else getattr(engine_args, "custom_pipeline_args", None)
)
client, resources = launch_diffusion_stage_replica(
model=self._model,
stage_config=plan.stage_cfg,
metadata=plan.metadata,
stage_init_timeout=stage_init_timeout,
batch_size=self._diffusion_batch_size,
use_inline=self._num_stages == 1 and plan.num_replicas == 1,
use_inline=plan.num_replicas == 1 and bool(inline_diffusion or custom_pipeline_args),
replica_id=plan.replica_id,
omni_master_server=self._get_omni_master_server(),
omni_coordinator_address=self._get_coordinator_address(),
@@ -672,7 +685,7 @@ class StageRuntime:
plan.replica_id,
self._diffusion_batch_size,
)
return client
return cast(StagePoolClient, client)
except Exception:
if client is not None:
try:
@@ -1038,7 +1051,7 @@ class DistStageRuntime(StageRuntime):
stage_id,
replica_id,
)
return client
return cast(StagePoolClient, client)
if ctx.vllm_config is None:
raise RuntimeError(f"Remote LLM stage {stage_id} is missing vllm_config")
+2 -3
View File
@@ -821,10 +821,9 @@ class OmniServeCommand(CLISubcommand):
omni_config_group.add_argument(
"--text-encoder-tp-size",
type=int,
default=1,
default=None,
help="Tensor-parallel degree for the diffusion text encoder. "
"Shards the Qwen3-VL encoder across the first N DiT ranks, "
"removing the rank-0 encoder memory hotspot in no-offload runs. "
"Shards the encoder across the first N DiT ranks. "
"Equivalent to setting DiffusionParallelConfig.text_encoder_tp_size.",
)
omni_config_group.add_argument(
+26 -13
View File
@@ -29,7 +29,7 @@ from vllm_omni.entrypoints.openai.stage_params import (
build_stage_sampling_params_list,
get_default_sampling_params_list,
)
from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
from vllm_omni.entrypoints.openai.utils import is_video_generation_pipeline, parse_lora_request
from vllm_omni.entrypoints.openai.video_api_utils import (
_encode_video_bytes,
encode_video_base64,
@@ -137,11 +137,25 @@ class OmniOpenAIServingVideo:
return False
capability = getattr(od_config, "supports_mixed_reference_inputs", None)
if isinstance(capability, bool):
return capability
model_class_name = getattr(od_config, "model_class_name", None)
return get_diffusion_model_metadata(model_class_name).supports_mixed_reference_inputs
model_archs = [model_class_name]
for stage_config in self.stage_configs or ():
stage_get = (
stage_config.get if isinstance(stage_config, Mapping) else lambda key: getattr(stage_config, key, None)
)
engine_args = stage_get("engine_args") or {}
model_archs.extend(
(
stage_get("model_arch"),
engine_args.get("model_class_name")
if isinstance(engine_args, Mapping)
else getattr(engine_args, "model_class_name", None),
)
)
metadata_capability = any(
get_diffusion_model_metadata(model_arch).supports_mixed_reference_inputs for model_arch in model_archs
)
return capability is True or metadata_capability
@classmethod
def for_diffusion(
@@ -478,14 +492,13 @@ class OmniOpenAIServingVideo:
detail="Stage configs not found. Start server with an omni diffusion model.",
)
# Video generation endpoint only supports diffusion stages.
for stage in stage_configs:
stage_type = get_stage_type(stage)
if stage_type != "diffusion":
raise HTTPException(
status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
detail=f"Video generation only supports diffusion stages, found '{stage_type}' stage.",
)
# Video pipelines may use preparatory AR stages (for example, a
# vLLM-hosted text encoder) before the diffusion stage.
if not is_video_generation_pipeline(stage_configs):
raise HTTPException(
status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
detail="No final video output stage found in video generation pipeline.",
)
# Common generation logic for both paths
engine_client = cast(AsyncOmni, self._engine_client)
@@ -51,7 +51,7 @@ from vllm_omni.entrypoints.openai.stage_params import (
build_stage_sampling_params_list,
get_default_sampling_params_list,
)
from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
from vllm_omni.entrypoints.openai.utils import is_video_generation_pipeline, parse_lora_request
from vllm_omni.entrypoints.openai.video_api_utils import (
StreamingVideoFormat,
create_streaming_video_encoder,
@@ -616,13 +616,11 @@ class OmniStreamingVideoOutputHandler:
detail="Stage configs not found. Start server with an omni diffusion model.",
)
for stage in stage_configs:
stage_type = get_stage_type(stage)
if stage_type != "diffusion":
raise HTTPException(
status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
detail=f"Video generation only supports diffusion stages, found '{stage_type}' stage.",
)
if not is_video_generation_pipeline(stage_configs):
raise HTTPException(
status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
detail="No final video output stage found in video generation pipeline.",
)
sampling_params_list = build_stage_sampling_params_list(
list(stage_configs),
+21
View File
@@ -20,6 +20,27 @@ def get_stage_type(stage_cfg: Any) -> str:
return getattr(stage_cfg, "stage_type", "llm")
def is_video_generation_pipeline(stage_configs: list[Any] | None) -> bool:
"""Return whether a pipeline declares a final video output stage."""
for stage in stage_configs or ():
if isinstance(stage, dict):
final_output = stage.get("final_output", False)
final_output_type = stage.get("final_output_type")
elif hasattr(stage, "get"):
try:
final_output = stage.get("final_output", False)
final_output_type = stage.get("final_output_type")
except Exception:
final_output = getattr(stage, "final_output", False)
final_output_type = getattr(stage, "final_output_type", None)
else:
final_output = getattr(stage, "final_output", False)
final_output_type = getattr(stage, "final_output_type", None)
if final_output and final_output_type in {"video", "videos"}:
return True
return False
def parse_lora_request(lora_body: Any) -> tuple[LoRARequest | None, float | None]:
"""Parse a request-level LoRA object into a LoRARequest and optional scale.
@@ -0,0 +1 @@
"""MiniMax H3 model components."""
@@ -0,0 +1,62 @@
# 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/**"],
}
def resolve_minimax_h3_partition(
model: str,
task_type: str | None,
*,
auto_partition: str,
) -> str:
task = str(task_type or "auto").lower()
if task not in {"auto", "combined", "t2va", "fl2va", "ref2va"}:
raise ValueError(
f"MiniMax-H3 task_type must be one of auto, combined, t2va, fl2va, or ref2va; got {task_type!r}"
)
path = Path(model)
if task == "auto" and path.is_dir() and path.name in {"FL2VA", "Ref2VA"}:
return path.name.lower()
if task in {"auto", "combined"}:
return auto_partition
return "ref2va" if task == "ref2va" else "fl2va"
def resolve_minimax_h3_model_root(
model: str,
revision: str | None,
task_type: str | None,
) -> str:
path = Path(model)
partition = resolve_minimax_h3_partition(model, task_type, auto_partition="fl2va")
if path.is_dir():
if path.name == "text_encoder" and (path / "config.json").is_file():
return str(path)
if path.name in {"FL2VA", "Ref2VA"}:
path = path.parent
subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
return str(path / subdir / "text_encoder")
snapshot = download_weights_from_hf_specific(
model_name_or_path=model,
cache_dir=None,
allow_patterns=MINIMAX_H3_TEXT_ENCODER_DOWNLOAD_PATTERNS[partition],
revision=revision,
require_all=True,
)
subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
return str(Path(snapshot) / subdir / "text_encoder")
__all__ = ["resolve_minimax_h3_model_root", "resolve_minimax_h3_partition"]
@@ -0,0 +1,51 @@
# SPDX-License-Identifier: Apache-2.0
"""MiniMax H3 text-conditioning contract."""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any
import torch
MINIMAX_H3_TEXT_HIDDEN_SIZE = 5120
MINIMAX_H3_PRESENTATION_TASK_KEY = "_minimax_h3_presentation_task"
MINIMAX_H3_CONDITION_LABELS_KEY = "_minimax_h3_condition_labels"
@dataclass(frozen=True)
class MiniMaxH3TextConditioning:
"""Layer-50 Qwen3-VL hidden states and their H3 token roles."""
hidden_states: torch.Tensor
token_tags: torch.Tensor
@classmethod
def from_payload(
cls,
payload: Mapping[str, Any],
) -> MiniMaxH3TextConditioning:
hidden_states = payload.get("hidden_states")
token_tags = payload.get("token_tags")
if not isinstance(hidden_states, torch.Tensor) or not isinstance(token_tags, torch.Tensor):
raise ValueError("MiniMax H3 conditioning requires hidden_states and token_tags tensors")
if hidden_states.ndim != 2 or hidden_states.shape[-1] != MINIMAX_H3_TEXT_HIDDEN_SIZE:
raise ValueError(
"MiniMax H3 hidden_states must have shape "
f"[tokens, {MINIMAX_H3_TEXT_HIDDEN_SIZE}], got {tuple(hidden_states.shape)}"
)
if token_tags.ndim != 1 or token_tags.shape[0] != hidden_states.shape[0]:
raise ValueError(
"MiniMax H3 token_tags must align with hidden_states, got "
f"token_tags={tuple(token_tags.shape)} and hidden_states={tuple(hidden_states.shape)}"
)
if not torch.all((token_tags == 0) | (token_tags == 1)):
raise ValueError("MiniMax H3 text-encoder token_tags must contain only 0 and 1")
return cls(hidden_states=hidden_states, token_tags=token_tags)
def to_payload(self) -> dict[str, torch.Tensor]:
return {
"hidden_states": self.hidden_states,
"token_tags": self.token_tags,
}
@@ -0,0 +1,53 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Two-stage MiniMax H3 topology."""
from vllm_omni.config.stage_config import (
PipelineConfig,
StageExecutionType,
StagePipelineConfig,
)
_PROCESSOR = "vllm_omni.model_executor.stage_input_processors.minimax_h3"
_CHECKPOINT = "vllm_omni.model_executor.models.minimax_h3.checkpoint"
_DIFFUSION_PIPELINE = "vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3"
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",
stages=(
StagePipelineConfig(
stage_id=0,
model_stage="text_encoder",
execution_type=StageExecutionType.LLM_AR,
input_sources=(),
owns_tokenizer=True,
requires_multimodal_data=True,
model_arch="MiniMaxH3TextEncoder",
engine_output_type="latent",
prompt_transform_func=f"{_PROCESSOR}.prepare_text_encoder_prompt",
sampling_constraints={
"max_tokens": 1,
"temperature": 0.0,
"detokenize": False,
},
model_path_resolver=f"{_CHECKPOINT}.resolve_minimax_h3_model_root",
),
StagePipelineConfig(
stage_id=1,
model_stage="dit",
execution_type=StageExecutionType.DIFFUSION,
input_sources=(0,),
final_output=True,
final_output_type="video",
requires_multimodal_data=True,
model_arch="MiniMaxH3Pipeline",
custom_process_input_func=f"{_PROCESSOR}.text_encoder2diffusion",
omni_kv_config={"need_recv_cache": False},
model_path_resolver=f"{_DIFFUSION_PIPELINE}.resolve_minimax_h3_diffusion_model_path",
inline_diffusion=True,
),
),
)
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
"""MiniMax H3 Qwen presentation building.
"""Shared MiniMax H3 media normalization and Qwen presentation building.
Builds the positive presentation token stream:
- fl2va: '<Picture 1>: ' label + vision block (<|vision_start|> +
@@ -13,10 +13,16 @@ accumulator so ids and AdaLN token tags cannot drift apart.
from __future__ import annotations
import math
import os
from collections.abc import Sequence
from typing import Any
import numpy as np
import torch
from PIL import Image
from vllm_omni.errors import OmniClientError
VISION_START = "<|vision_start|>"
VISION_END = "<|vision_end|>"
@@ -26,6 +32,142 @@ VIDEO_PAD = "<|video_pad|>"
_TEXT_TAG = 1
_VIDEO_TAG = 0
MINIMAX_H3_OUTPUT_SHORT_EDGE = 768
MINIMAX_H3_OUTPUT_MAX_PIXELS = 768 * 1344
MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE = 2048
MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE = 32
MINIMAX_H3_SUPPORTED_ASPECT_RATIOS = {
"21:9": 21.0 / 9.0,
"16:9": 16.0 / 9.0,
"4:3": 4.0 / 3.0,
"1:1": 1.0,
"3:4": 3.0 / 4.0,
"9:16": 9.0 / 16.0,
}
MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES = 30 * 1024 * 1024
MINIMAX_H3_REFERENCE_IMAGE_FORMATS = frozenset({"jpeg", "png", "webp", "heic", "heif"})
def _align_multiple(value: float, multiple: int = 32) -> int:
return max(multiple, int(round(float(value) / multiple)) * multiple)
def load_minimax_h3_images(value: Any) -> list[Image.Image]:
"""Normalize one or more H3 image inputs to RGB PIL images."""
if isinstance(value, (list, tuple)):
if not value:
raise OmniClientError("MiniMax H3 image input must not be empty")
images: list[Image.Image] = []
for item in value:
loaded = load_minimax_h3_images(item)
if len(loaded) != 1:
raise OmniClientError(f"MiniMax H3 expected one image, got {len(loaded)}")
images.extend(loaded)
return images
if isinstance(value, (str, os.PathLike)):
if os.path.getsize(value) > MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES:
raise OmniClientError("MiniMax H3 reference image exceeds the 30 MiB size limit")
with Image.open(value) as image:
image_format = str(image.format or "").lower()
if image_format and image_format not in MINIMAX_H3_REFERENCE_IMAGE_FORMATS:
raise OmniClientError(
f"MiniMax H3 reference image must use JPG, JPEG, PNG, WEBP, HEIC, or HEIF, got {image.format}"
)
return [image.convert("RGB")]
if isinstance(value, Image.Image):
return [value.convert("RGB")]
if isinstance(value, torch.Tensor):
tensor = value.detach().float().cpu()
if tensor.ndim == 4 and tensor.shape[0] == 1:
tensor = tensor[0]
if tensor.ndim != 3:
raise OmniClientError(f"image tensor must be [C,H,W], got {tuple(tensor.shape)}")
if tensor.shape[0] in (1, 3, 4):
tensor = tensor.permute(1, 2, 0)
array = tensor.numpy()
if array.max(initial=0) <= 1.0:
array = array * 255.0
return [Image.fromarray(array.clip(0, 255).astype(np.uint8)).convert("RGB")]
raise OmniClientError(f"unsupported MiniMax H3 image input {type(value)!r}")
def resolve_minimax_h3_aspect_ratio(
task: str,
value: Any,
image: Image.Image | None,
) -> float:
"""Resolve H3's task-specific output aspect-ratio policy."""
if task == "fl2va":
if image is None:
raise OmniClientError("fl2va requires an input image to resolve its aspect ratio")
return float(image.width) / float(image.height)
if value is None:
if task == "t2va":
raise OmniClientError("t2va requires an explicit aspect_ratio")
return MINIMAX_H3_SUPPORTED_ASPECT_RATIOS["16:9"]
if isinstance(value, str):
normalized = value.strip().lower()
if normalized in {"adaptive", "auto"}:
if task == "t2va":
raise OmniClientError("t2va requires an explicit named aspect_ratio, not adaptive")
return MINIMAX_H3_SUPPORTED_ASPECT_RATIOS["16:9"]
if normalized in MINIMAX_H3_SUPPORTED_ASPECT_RATIOS:
return MINIMAX_H3_SUPPORTED_ASPECT_RATIOS[normalized]
try:
numeric_value = float(normalized)
except (TypeError, ValueError) as exc:
supported = ", ".join(MINIMAX_H3_SUPPORTED_ASPECT_RATIOS)
raise OmniClientError(f"MiniMax H3 aspect_ratio must be one of {supported}, got {value!r}") from exc
elif isinstance(value, (int, float, np.integer, np.floating)) and not isinstance(value, bool):
numeric_value = float(value)
else:
raise OmniClientError(f"MiniMax H3 aspect_ratio must be a string ratio, got {value!r}")
if not math.isfinite(numeric_value) or not any(
math.isclose(numeric_value, ratio, rel_tol=0.0, abs_tol=1e-6)
for ratio in MINIMAX_H3_SUPPORTED_ASPECT_RATIOS.values()
):
supported = ", ".join(MINIMAX_H3_SUPPORTED_ASPECT_RATIOS)
raise OmniClientError(f"MiniMax H3 aspect_ratio must be one of {supported}, got {value!r}")
return numeric_value
def resolve_minimax_h3_reference_image_shape(image: Image.Image) -> tuple[int, int]:
"""Resize an H3 reference image to the official 2048-short-edge canvas."""
width, height = image.size
ratio = width / height
if not 0.4 <= ratio <= 2.5:
raise OmniClientError(f"reference image aspect ratio must be in [0.4, 2.5], got {width}x{height}")
if min(width, height) < 256 or max(width, height) > 5760:
raise OmniClientError(f"reference image dimensions must be in [256, 5760] pixels, got {width}x{height}")
scale = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE / min(width, height)
return (
_align_multiple(width * scale, MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE),
_align_multiple(height * scale, MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE),
)
def resolve_minimax_h3_output_canvas(aspect_ratio: float, short_edge: int) -> tuple[int, int]:
"""Resolve the official H3 ratio/area policy to a 32-pixel canvas."""
if not math.isfinite(float(aspect_ratio)) or float(aspect_ratio) <= 0:
raise OmniClientError(f"MiniMax H3 canvas aspect ratio must be positive, got {aspect_ratio!r}")
if short_edge != MINIMAX_H3_OUTPUT_SHORT_EDGE:
raise OmniClientError(f"MiniMax H3 target.short_edge must be {MINIMAX_H3_OUTPUT_SHORT_EDGE}, got {short_edge}")
if aspect_ratio >= 1.0:
width = float(short_edge) * aspect_ratio
height = float(short_edge)
else:
width = float(short_edge)
height = float(short_edge) / aspect_ratio
area = width * height
if area > MINIMAX_H3_OUTPUT_MAX_PIXELS:
scale = (MINIMAX_H3_OUTPUT_MAX_PIXELS / area) ** 0.5
width *= scale
height *= scale
return _align_multiple(height, 32), _align_multiple(width, 32)
def _text_ids(tokenizer: Any, text: str) -> list[int]:
return list(tokenizer(text, add_special_tokens=False)["input_ids"])
@@ -107,26 +249,18 @@ def _multi_image_presentation(
return presentation.build()
def minimax_h3_multi_image_presentation_ids(
def minimax_h3_multi_image_presentation(
tokenizer: Any,
*,
prompt: str,
image_token_counts: list[int],
) -> torch.Tensor:
"""Positive fl2va presentation with one or more labeled image slots."""
ids, _ = _multi_image_presentation(tokenizer, prompt=prompt, image_token_counts=image_token_counts)
return ids
def minimax_h3_multi_image_presentation_token_tags(
tokenizer: Any,
*,
prompt: str,
image_token_counts: list[int],
) -> torch.Tensor:
"""AdaLN token tags aligned with minimax_h3_multi_image_presentation_ids."""
_, tags = _multi_image_presentation(tokenizer, prompt=prompt, image_token_counts=image_token_counts)
return tags
) -> tuple[torch.Tensor, torch.Tensor]:
"""Return aligned FL2VA presentation token IDs and role IDs."""
return _multi_image_presentation(
tokenizer,
prompt=prompt,
image_token_counts=image_token_counts,
)
def minimax_h3_ref2va_presentation(
@@ -288,10 +422,72 @@ def minimax_h3_ref2va_video_presentation(
return presentation.build()
def build_minimax_h3_presentation(
tokenizer: Any,
*,
prompt: str,
task: str,
condition_labels: list[tuple[str, int]],
image_grid_thw: torch.Tensor | None,
video_grid_thw: torch.Tensor | None,
video_timestamps: Sequence[Sequence[float]] | None,
merge_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Build the token IDs and role IDs shared by fused and split H3."""
if task == "t2va":
ids = minimax_h3_text_only_ids(tokenizer, prompt)
return ids, torch.ones_like(ids)
merge_length = int(merge_size) ** 2
image_counts = (
[int(grid.prod().item()) // merge_length for grid in image_grid_thw] if image_grid_thw is not None else []
)
if task == "fl2va":
return minimax_h3_multi_image_presentation(
tokenizer,
prompt=prompt,
image_token_counts=image_counts,
)
video_counts: list[list[int]] = []
if video_grid_thw is not None:
for grid in video_grid_thw:
block_count = int(grid[0].item())
tokens_per_block = int(grid[1:].prod().item()) // merge_length
video_counts.append([tokens_per_block] * block_count)
timestamps = (
[[float(value) for value in group] for group in video_timestamps] if video_timestamps is not None else []
)
if video_counts:
return minimax_h3_ref2va_video_presentation(
tokenizer,
prompt=prompt,
condition_labels=condition_labels,
image_token_count=image_counts or None,
video_block_token_counts=video_counts,
video_block_timestamps=timestamps,
)
return minimax_h3_ref2va_presentation(
tokenizer,
prompt=prompt,
condition_labels=condition_labels,
image_token_count=image_counts or None,
)
__all__ = [
"minimax_h3_multi_image_presentation_ids",
"minimax_h3_multi_image_presentation_token_tags",
"IMAGE_PAD",
"MINIMAX_H3_OUTPUT_SHORT_EDGE",
"VIDEO_PAD",
"VISION_END",
"VISION_START",
"build_minimax_h3_presentation",
"load_minimax_h3_images",
"minimax_h3_multi_image_presentation",
"minimax_h3_ref2va_presentation",
"minimax_h3_ref2va_video_presentation",
"minimax_h3_text_only_ids",
"resolve_minimax_h3_aspect_ratio",
"resolve_minimax_h3_output_canvas",
"resolve_minimax_h3_reference_image_shape",
]
@@ -15,6 +15,53 @@ import torch
from vllm_omni.errors import OmniClientError
MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY = "minimax_h3_prepared_reference_videos"
def serialize_prepared_reference_videos(items: list[dict[str, Any]], artifact_dir: str) -> str:
fields = (
"original_path",
"prepared_path",
"input_has_audio",
"width",
"height",
"start_time_seconds",
"duration_seconds",
"audio_duration_seconds",
)
return json.dumps(
{
"artifact_dir": artifact_dir,
"videos": [{field: item[field] for field in fields} for item in items],
},
separators=(",", ":"),
)
def deserialize_prepared_reference_videos(value: str) -> tuple[str, list[dict[str, Any]]]:
try:
payload = json.loads(value)
artifact_dir = payload["artifact_dir"]
videos = payload["videos"]
except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
raise ValueError("invalid MiniMax H3 prepared-reference-video descriptor") from exc
if not isinstance(artifact_dir, str) or not isinstance(videos, list):
raise ValueError("invalid MiniMax H3 prepared-reference-video descriptor")
required = {
"original_path",
"prepared_path",
"input_has_audio",
"width",
"height",
"start_time_seconds",
"duration_seconds",
"audio_duration_seconds",
}
if any(not isinstance(item, dict) or set(item) != required for item in videos):
raise ValueError("invalid MiniMax H3 prepared-reference-video descriptor")
return artifact_dir, videos
MINIMAX_H3_FPS = 24.0
MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS = 2.0
MINIMAX_H3_QWEN_TEMPORAL_PATCH = 2
@@ -280,6 +327,7 @@ def _transcode_reference_video(
) -> str:
output = str(Path(workdir) / "prepared.mp4")
duration_args = ["-t", f"{float(duration_seconds):.6f}"] if duration_seconds is not None else []
frame_count_args = ["-frames:v", str(int(target_frame_count))] if target_frame_count > 0 else []
subprocess.run(
[
"ffmpeg",
@@ -296,8 +344,7 @@ def _transcode_reference_video(
"-vf",
(f"fps={MINIMAX_H3_FPS:g},scale={target_width}:{target_height}:flags=lanczos,setsar=1"),
*duration_args,
"-frames:v",
str(int(target_frame_count)),
*frame_count_args,
"-metadata:s:v:0",
"rotate=0",
# Keep the transformed RGB pixels lossless. The same prepared
@@ -394,7 +441,7 @@ def prepare_reference_videos(
# bounded duration separately for audio extraction.
"audio_duration_seconds": min(
duration,
float(target_frame_count) / MINIMAX_H3_FPS,
float(target_frame_count) / MINIMAX_H3_FPS if target_frame_count > 0 else duration,
),
}
)
@@ -0,0 +1,311 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""MiniMax H3 text encoder on the vLLM Qwen3-VL model runner."""
from __future__ import annotations
import copy
from collections.abc import Iterable, Mapping, Sequence
from typing import Any
import regex as re
import torch
import torch.nn as nn
from vllm.config import VllmConfig
from vllm.model_executor.models.qwen3_vl import (
Qwen3VLDummyInputsBuilder,
Qwen3VLForConditionalGeneration,
Qwen3VLMultiModalProcessor,
Qwen3VLProcessingInfo,
)
from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
MINIMAX_H3_CONDITION_LABELS_KEY,
MINIMAX_H3_PRESENTATION_TASK_KEY,
)
from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
IMAGE_PAD,
VIDEO_PAD,
VISION_END,
VISION_START,
build_minimax_h3_presentation,
)
from vllm_omni.model_executor.models.output_templates import OmniOutput
def _build_minimax_h3_presentation(
tokenizer: Any,
*,
prompt: str,
task: str,
condition_labels: list[tuple[str, int]],
image_grid_thw: torch.Tensor | None,
video_grid_thw: torch.Tensor | None,
video_timestamps: Sequence[Sequence[float]] | None,
merge_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Build the same token stream consumed by the fused H3 encoder."""
return build_minimax_h3_presentation(
tokenizer,
prompt=prompt,
task=task,
condition_labels=condition_labels,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
video_timestamps=video_timestamps,
merge_size=merge_size,
)
class MiniMaxH3MultiModalProcessor(Qwen3VLMultiModalProcessor):
"""Qwen3-VL media processing with H3's exact segmented presentation."""
@staticmethod
def _base_processor_kwargs(
kwargs: Mapping[str, object],
) -> dict[str, object]:
return {
key: value
for key, value in kwargs.items()
if key
not in {
MINIMAX_H3_PRESENTATION_TASK_KEY,
MINIMAX_H3_CONDITION_LABELS_KEY,
}
}
@staticmethod
def _condition_labels(value: object) -> list[tuple[str, int]]:
if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
return []
labels: list[tuple[str, int]] = []
for item in value:
if not isinstance(item, Sequence) or isinstance(item, (str, bytes)) or len(item) != 2:
raise ValueError("MiniMax H3 condition labels must be (type, index) pairs")
labels.append((str(item[0]), int(item[1])))
return labels
def _cached_apply_hf_processor(self, inputs: Any, timing_ctx: Any):
# H3 presentation IDs depend on the processed grids for every media
# item. On a sender-cache hit, the generic processor passes only cache
# misses to _apply_hf_processor_main, so the full presentation cannot
# be reconstructed there. Reprocess H3 media to preserve exact IDs.
return self._apply_hf_processor(inputs, timing_ctx)
def _apply_hf_processor_main(
self,
prompt: str | list[int],
mm_items: Any,
hf_processor_mm_kwargs: Mapping[str, object],
tokenization_kwargs: Mapping[str, object],
*,
enable_hf_prompt_update: bool,
):
task_value = hf_processor_mm_kwargs.get(MINIMAX_H3_PRESENTATION_TASK_KEY)
if task_value is None:
return super()._apply_hf_processor_main(
prompt,
mm_items,
hf_processor_mm_kwargs,
tokenization_kwargs,
enable_hf_prompt_update=enable_hf_prompt_update,
)
if not isinstance(prompt, str):
raise ValueError("MiniMax H3 presentation requires the original prompt text")
task = str(task_value)
condition_labels = self._condition_labels(hf_processor_mm_kwargs.get(MINIMAX_H3_CONDITION_LABELS_KEY))
base_kwargs = self._base_processor_kwargs(hf_processor_mm_kwargs)
image_count = sum(kind == "image" for kind, _ in condition_labels)
video_count = sum(kind == "video" for kind, _ in condition_labels)
image_slot = f"{VISION_START}{IMAGE_PAD}{VISION_END}"
hf_processor = self.info.get_hf_processor(**base_kwargs)
video_slot = (
VIDEO_PAD if self._expands_only_video_token(hf_processor) else f"{VISION_START}{VIDEO_PAD}{VISION_END}"
)
processing_prompt = image_slot * image_count + video_slot * video_count
if not processing_prompt:
processing_prompt = prompt
_, processed_data, _ = super()._apply_hf_processor_main(
processing_prompt,
mm_items,
base_kwargs,
tokenization_kwargs,
enable_hf_prompt_update=enable_hf_prompt_update,
)
image_processor = self.info.get_image_processor(**base_kwargs)
ids, _ = _build_minimax_h3_presentation(
self.info.get_tokenizer(),
prompt=prompt,
task=task,
condition_labels=condition_labels,
image_grid_thw=processed_data.get("image_grid_thw"),
video_grid_thw=processed_data.get("video_grid_thw"),
video_timestamps=processed_data.get("timestamps"),
merge_size=int(image_processor.merge_size),
)
return ids.tolist(), processed_data, True
def _get_prompt_updates(
self,
mm_items: Any,
hf_processor_mm_kwargs: Mapping[str, Any],
out_mm_kwargs: Any,
):
return super()._get_prompt_updates(
mm_items,
self._base_processor_kwargs(hf_processor_mm_kwargs),
out_mm_kwargs,
)
class _ResidualMerge(nn.Module):
"""Return the decoder residual stream without the final RMSNorm."""
def forward(
self,
hidden_states: torch.Tensor,
residual: torch.Tensor | None,
) -> tuple[torch.Tensor, None]:
if residual is not None:
hidden_states = hidden_states + residual
return hidden_states, None
@MULTIMODAL_REGISTRY.register_processor(
MiniMaxH3MultiModalProcessor,
info=Qwen3VLProcessingInfo,
dummy_inputs=Qwen3VLDummyInputsBuilder,
)
class MiniMaxH3TextEncoder(Qwen3VLForConditionalGeneration):
"""Qwen3-VL encoder used by MiniMax H3.
MiniMax H3 consumes the residual stream after decoder layer 50. The
upstream vLLM Qwen3-VL implementation supplies the vision tower,
multimodal processor, M-RoPE, DeepStack injection, paged attention, and
tensor-parallel decoder layers. This class changes only the two parts of
the contract that differ from causal generation: decoder depth and output
normalization.
"""
have_multimodal_outputs = True
omni_pooler_payload_include_hidden = False
requires_full_prefix_cached_hidden_states = False
num_encoder_layers = 50
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
if vllm_config.parallel_config.pipeline_parallel_size != 1:
raise ValueError("MiniMax H3 text encoding supports tensor parallelism, not pipeline parallelism")
hf_config = copy.deepcopy(vllm_config.model_config.hf_config)
hf_config.text_config.num_hidden_layers = self.num_encoder_layers
# Reuse the token embedding as the unused LM head, avoiding a second
# vocab-sized allocation in the upstream constructor. Set the flag
# on both levels because VllmConfig.with_hf_config propagates the
# multimodal parent value into its nested text config.
hf_config.tie_word_embeddings = True
hf_config.text_config.tie_word_embeddings = True
encoder_config = vllm_config.with_hf_config(hf_config)
super().__init__(vllm_config=encoder_config, prefix=prefix)
self.language_model.model.norm = _ResidualMerge()
self._token_tags: torch.Tensor | None = None
def embed_input_ids(
self,
input_ids: torch.Tensor,
multimodal_embeddings=None,
*,
is_multimodal: torch.Tensor | None = None,
) -> torch.Tensor:
# Qwen3-VL expands every visual block to exactly
# vision_start + image/video tokens + vision_end. Classifying each
# token directly makes tags independent of prefill chunk boundaries.
visual_token = (
(input_ids == self.config.vision_start_token_id)
| (input_ids == self.config.vision_end_token_id)
| (input_ids == self.config.image_token_id)
| (input_ids == self.config.video_token_id)
)
token_tags = (~visual_token).long()
self._token_tags = None
embeddings = super().embed_input_ids(
input_ids,
multimodal_embeddings,
is_multimodal=is_multimodal,
)
self._token_tags = token_tags
return embeddings
def make_omni_output(
self,
model_outputs: torch.Tensor | OmniOutput,
**_: object,
) -> OmniOutput:
if isinstance(model_outputs, OmniOutput):
return model_outputs
if self._token_tags is None:
raise RuntimeError("token tags were not prepared before MiniMax H3 text encoding")
token_tags = self._token_tags
self._token_tags = None
padding = model_outputs.shape[0] - token_tags.shape[0]
if padding < 0:
raise RuntimeError(
"MiniMax H3 token tags exceed the model output length: "
f"tags={token_tags.shape[0]}, outputs={model_outputs.shape[0]}"
)
if padding:
# Multimodal embedding runs on scheduled tokens, while the AR
# forward may use a CUDA-graph-padded token buffer. Match that
# buffer so both conditioning tensors are recognized as
# per-token prefix-cache entries; runner slicing removes these
# padding rows before the stage payload is emitted.
token_tags = torch.cat((token_tags, token_tags.new_ones(padding)))
# Keep the stage-wire representation two-dimensional. The Omni
# runner treats tensors with shape [tokens, features] as per-token
# data, so they are concatenated across chunked-prefill steps and
# reconstructed on prefix-cache hits. The Stage 1 adapter restores
# H3's semantic [tokens] representation.
token_tags = token_tags.unsqueeze(-1)
return OmniOutput(
text_hidden_states=model_outputs,
multimodal_outputs={
"hidden_states": {"output": model_outputs},
"meta": {"token_role_ids": token_tags},
},
)
def compute_logits(
self,
hidden_states: torch.Tensor,
**_: object,
) -> torch.Tensor:
"""Emit EOS immediately; the useful result is the prefill hidden state."""
logits = hidden_states.new_full(
(hidden_states.shape[0], self.config.text_config.vocab_size),
torch.finfo(hidden_states.dtype).min,
)
logits[:, int(self._tokenizer.eos_token_id)] = 0
return logits
def load_weights(
self,
weights: Iterable[tuple[str, torch.Tensor]],
) -> set[str]:
def encoder_weights():
for name, weight in weights:
layer_match = re.match(r"model\.language_model\.layers\.(\d+)\.", name)
if layer_match and int(layer_match.group(1)) >= self.num_encoder_layers:
continue
if name == "model.language_model.norm.weight" or name.startswith("lm_head."):
continue
yield name, weight
loader = AutoWeightsLoader(self)
return loader.load_weights(
encoder_weights(),
mapper=self.hf_to_vllm_mapper,
)
@@ -457,6 +457,11 @@ _OMNI_MODELS = {
"qwen3_vl",
"AuraQwen3VLForConditionalGeneration",
),
"MiniMaxH3TextEncoder": (
"minimax_h3",
"text_encoder",
"MiniMaxH3TextEncoder",
),
## MiniMax-Music3 (text-to-music; AR talker -> flow-matching acoustic decoder)
"MiniMaxMusic3TalkerForConditionalGeneration": (
"minimax_music3",
@@ -0,0 +1,364 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""MiniMax H3 text-encoder stage input and output adapters."""
from __future__ import annotations
import copy
import shutil
import tempfile
from collections.abc import Mapping, Sequence
from typing import Any
import numpy as np
import torch
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.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
MINIMAX_H3_CONDITION_LABELS_KEY,
MINIMAX_H3_PRESENTATION_TASK_KEY,
MiniMaxH3TextConditioning,
)
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"
def _diffusion_sampling_params(sampling_params_list: Sequence[Any]) -> Any:
diffusion_params = [
sampling_params
for sampling_params in sampling_params_list
if isinstance(sampling_params, OmniDiffusionSamplingParams)
]
if len(diffusion_params) != 1:
raise RuntimeError(
"MiniMax H3 text 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(
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.
"""
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}")
transformed = copy.copy(prompt)
if isinstance(prompt.get("additional_information"), Mapping):
transformed["additional_information"] = dict(prompt["additional_information"])
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)
if images:
qwen_mm_data["image"] = images
if qwen_video_inputs:
qwen_mm_data["video"] = qwen_video_inputs
transformed["multi_modal_data"] = qwen_mm_data or None
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
transformed["mm_processor_kwargs"] = mm_processor_kwargs
return transformed
def _original_prompt(prompt: Any) -> dict[str, Any]:
if isinstance(prompt, list):
prompt = prompt[0] if prompt else {}
if isinstance(prompt, dict):
return copy.copy(prompt)
if isinstance(prompt, str):
return {"prompt": prompt}
raise TypeError(f"invalid MiniMax H3 prompt type {type(prompt)!r}")
def _global_request_id(prompt: Mapping[str, Any]) -> str | None:
additional_information = prompt.get("additional_information")
if not isinstance(additional_information, Mapping):
return None
value = additional_information.get("global_request_id")
if isinstance(value, (list, tuple)):
value = value[0] if value else None
return str(value) if value is not None else None
def text_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."""
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)}")
diffusion_prompt = _original_prompt(prompt)
source_output = source_outputs[0]
source_request_id = getattr(source_output, "request_id", None)
expected_request_id = _global_request_id(diffusion_prompt)
if (
source_request_id is not None
and expected_request_id is not None
and str(source_request_id) != expected_request_id
):
raise RuntimeError(
"MiniMax H3 text-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
if not isinstance(payload, Mapping):
raise RuntimeError("MiniMax H3 text encoder returned no conditioning payload")
hidden_states = payload.get("hidden_states")
if not isinstance(hidden_states, Mapping):
raise RuntimeError("MiniMax H3 text encoder returned no hidden_states payload")
hidden = hidden_states.get("output")
if not isinstance(hidden, torch.Tensor):
raise RuntimeError("MiniMax H3 text encoder returned no hidden_states.output tensor")
meta = payload.get("meta")
if not isinstance(meta, Mapping):
raise RuntimeError("MiniMax H3 text encoder returned no conditioning metadata")
token_tags = meta.get("token_role_ids")
if not isinstance(token_tags, torch.Tensor):
raise RuntimeError("MiniMax H3 text encoder returned no token_role_ids tensor")
if token_tags.ndim != 2 or token_tags.shape[-1] != 1:
raise RuntimeError(
f"MiniMax H3 stage-wire token_tags must have shape [tokens, 1], got {tuple(token_tags.shape)}"
)
try:
conditioning = MiniMaxH3TextConditioning.from_payload(
{
"hidden_states": hidden,
"token_tags": token_tags.squeeze(-1),
}
)
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()
diffusion_prompt["additional_information"] = additional_information
return diffusion_prompt
+2 -2
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""NPU patches for the MiniMax H3 Qwen3-VL text encoder."""
@@ -112,5 +112,5 @@ def apply_minimax_h3_qwen3vl_swiglu_patch() -> None:
from vllm_omni.diffusion.models.minimax_h3 import encoder
encoder.MiniMaxH3Qwen3VLTextMLP.forward = _forward_minimax_h3_qwen3vl_text_mlp_npu
encoder.MiniMaxH3Qwen3VLTextMLP.forward = _forward_minimax_h3_qwen3vl_text_mlp_npu # type: ignore[method-assign]
_SWIGLU_PATCHED = True