From b6a4baa659238e5e13a2f295430820a3ae36dde5 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Fri, 28 Aug 2026 14:27:05 +0800 Subject: [PATCH] [Feature] Support text encoder Disaggregation for H3 (#5885) Signed-off-by: Canlin Guo Signed-off-by: lishunyang12 Signed-off-by: yuanwu Co-authored-by: lishunyang12 Co-authored-by: yuanwu --- recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md | 82 +++ tests/config/test_config_factory.py | 127 +++++ tests/config/test_omni_config.py | 14 +- tests/config/test_pipeline_registry.py | 7 + .../minimax_h3/test_minimax_h3_contract.py | 125 ++++- .../test_minimax_h3_quantization.py | 15 - .../test_minimax_h3_quantization_quality.py | 1 - .../test_minimax_h3_qwen3vl_swiglu.py | 9 +- .../test_diffusion_output_formatter.py | 6 +- .../test_diffusion_streaming_output.py | 4 +- .../test_inline_stage_diffusion_client.py | 5 + .../comfyui/test_comfyui_integration.py | 10 +- tests/engine/test_async_omni_engine_input.py | 33 ++ .../test_async_omni_engine_stage_init.py | 41 +- tests/engine/test_data_entry_keys.py | 28 ++ .../test_orchestrator_error_handling.py | 18 +- tests/engine/test_single_stage_mode.py | 42 ++ tests/engine/test_stage_engine_args.py | 6 +- .../test_serving_video_output_stream.py | 24 +- .../test_video_pipeline_capability.py | 66 +++ .../openai_api/test_video_server.py | 20 +- .../test_async_omni_diffusion_config.py | 53 ++ .../models/test_minimax_h3_token_tags.py | 62 +++ .../stage_input_processors/test_minimax_h3.py | 387 ++++++++++++++ vllm_omni/config/config_factory.py | 6 +- vllm_omni/config/omni_config.py | 7 + vllm_omni/config/pipeline_registry.py | 2 + vllm_omni/config/stage_config.py | 55 ++ vllm_omni/data_entry_keys.py | 7 + .../deploy/minimax_h3_disaggregated.yaml | 28 ++ .../minimax_h3_disaggregated_turbo.yaml | 31 ++ vllm_omni/diffusion/data.py | 9 +- .../inline_stage_diffusion_client.py | 2 + vllm_omni/diffusion/io_support.py | 10 + vllm_omni/diffusion/model_metadata.py | 55 +- .../diffusion/models/minimax_h3/__init__.py | 3 + .../models/minimax_h3/pipeline_minimax_h3.py | 472 +++++++++--------- vllm_omni/diffusion/output_formatter.py | 12 +- vllm_omni/diffusion/stage_diffusion_client.py | 1 + vllm_omni/engine/arg_utils.py | 1 - vllm_omni/engine/async_omni_engine.py | 142 +++++- vllm_omni/engine/messages.py | 1 + vllm_omni/engine/orchestrator.py | 11 + vllm_omni/engine/stage_client.py | 1 + vllm_omni/engine/stage_engine_core_client.py | 1 + vllm_omni/engine/stage_init_utils.py | 30 ++ vllm_omni/engine/stage_runtime.py | 29 +- vllm_omni/entrypoints/cli/serve.py | 5 +- vllm_omni/entrypoints/openai/serving_video.py | 39 +- .../openai/serving_video_output_stream.py | 14 +- vllm_omni/entrypoints/openai/utils.py | 21 + .../models/minimax_h3/__init__.py | 1 + .../models/minimax_h3/checkpoint.py | 62 +++ .../models/minimax_h3/conditioning.py | 51 ++ .../models/minimax_h3/pipeline.py | 53 ++ .../models/minimax_h3/preprocessing.py} | 234 ++++++++- .../models/minimax_h3/reference_video.py | 53 +- .../models/minimax_h3/text_encoder.py | 311 ++++++++++++ vllm_omni/model_executor/models/registry.py | 5 + .../stage_input_processors/minimax_h3.py | 364 ++++++++++++++ vllm_omni/platforms/npu/models/minimax_h3.py | 4 +- 61 files changed, 2952 insertions(+), 366 deletions(-) create mode 100644 recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md create mode 100644 tests/entrypoints/openai_api/test_video_pipeline_capability.py create mode 100644 tests/model_executor/models/test_minimax_h3_token_tags.py create mode 100644 tests/model_executor/stage_input_processors/test_minimax_h3.py create mode 100644 vllm_omni/deploy/minimax_h3_disaggregated.yaml create mode 100644 vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml create mode 100644 vllm_omni/model_executor/models/minimax_h3/__init__.py create mode 100644 vllm_omni/model_executor/models/minimax_h3/checkpoint.py create mode 100644 vllm_omni/model_executor/models/minimax_h3/conditioning.py create mode 100644 vllm_omni/model_executor/models/minimax_h3/pipeline.py rename vllm_omni/{diffusion/models/minimax_h3/presentation.py => model_executor/models/minimax_h3/preprocessing.py} (53%) rename vllm_omni/{diffusion => model_executor}/models/minimax_h3/reference_video.py (92%) create mode 100644 vllm_omni/model_executor/models/minimax_h3/text_encoder.py create mode 100644 vllm_omni/model_executor/stage_input_processors/minimax_h3.py diff --git a/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md b/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md new file mode 100644 index 000000000..debc89d89 --- /dev/null +++ b/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md @@ -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. diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index f1a7f33df..9ba96dd89 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -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))) diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index c216206f4..1a9efcbed 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -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", diff --git a/tests/config/test_pipeline_registry.py b/tests/config/test_pipeline_registry.py index 4bc19b116..9f4976c06 100644 --- a/tests/config/test_pipeline_registry.py +++ b/tests/config/test_pipeline_registry.py @@ -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" diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 76adf7471..336be1503 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -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, ) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py index cc14d92e3..92bd09ac2 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py @@ -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 diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py index 1ba382703..5d8c0fa25 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py @@ -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, } diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py index 2664990b4..3ed7649c6 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py @@ -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") diff --git a/tests/diffusion/test_diffusion_output_formatter.py b/tests/diffusion/test_diffusion_output_formatter.py index dc21351ff..28714ac97 100644 --- a/tests/diffusion/test_diffusion_output_formatter.py +++ b/tests/diffusion/test_diffusion_output_formatter.py @@ -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}, diff --git a/tests/diffusion/test_diffusion_streaming_output.py b/tests/diffusion/test_diffusion_streaming_output.py index 6c5828902..2e1d24540 100644 --- a/tests/diffusion/test_diffusion_streaming_output.py +++ b/tests/diffusion/test_diffusion_streaming_output.py @@ -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", diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py index 824c1483a..25d4f1feb 100644 --- a/tests/diffusion/test_inline_stage_diffusion_client.py +++ b/tests/diffusion/test_inline_stage_diffusion_client.py @@ -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()]) diff --git a/tests/e2e/features/comfyui/test_comfyui_integration.py b/tests/e2e/features/comfyui/test_comfyui_integration.py index b58e33cb8..61cf76dcd 100644 --- a/tests/e2e/features/comfyui/test_comfyui_integration.py +++ b/tests/e2e/features/comfyui/test_comfyui_integration.py @@ -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), diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index ee6d322ca..ee603a57e 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/tests/engine/test_async_omni_engine_input.py @@ -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) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 6203ed726..23e56d95b 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -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 diff --git a/tests/engine/test_data_entry_keys.py b/tests/engine/test_data_entry_keys.py index d7b4d1b20..df0d7b2ed 100644 --- a/tests/engine/test_data_entry_keys.py +++ b/tests/engine/test_data_entry_keys.py @@ -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 diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index 3f1674133..e1ec355ce 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -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] diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index b576e72d4..334001a65 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -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 diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 632b3891e..6afde7469 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -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)}" diff --git a/tests/entrypoints/openai_api/test_serving_video_output_stream.py b/tests/entrypoints/openai_api/test_serving_video_output_stream.py index 27f1bed1e..e57e3e279 100644 --- a/tests/entrypoints/openai_api/test_serving_video_output_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_output_stream.py @@ -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()], diff --git a/tests/entrypoints/openai_api/test_video_pipeline_capability.py b/tests/entrypoints/openai_api/test_video_pipeline_capability.py new file mode 100644 index 000000000..cfa11f96d --- /dev/null +++ b/tests/entrypoints/openai_api/test_video_pipeline_capability.py @@ -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" diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 1d2bfe204..b48029c30 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -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 diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 1127e136d..641729601 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -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( diff --git a/tests/model_executor/models/test_minimax_h3_token_tags.py b/tests/model_executor/models/test_minimax_h3_token_tags.py new file mode 100644 index 000000000..69ee7cb84 --- /dev/null +++ b/tests/model_executor/models/test_minimax_h3_token_tags.py @@ -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 diff --git a/tests/model_executor/stage_input_processors/test_minimax_h3.py b/tests/model_executor/stage_input_processors/test_minimax_h3.py new file mode 100644 index 000000000..599be81f2 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_minimax_h3.py @@ -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) diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py index 48bd0b476..b2eb64bca 100644 --- a/vllm_omni/config/config_factory.py +++ b/vllm_omni/config/config_factory.py @@ -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] diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 18befa4fc..3043a1309 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -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, diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index e6c1982c6..373a84739 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -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, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index beeb5bdc3..661c0bf6a 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -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: diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index acedc8a86..685558a00 100644 --- a/vllm_omni/data_entry_keys.py +++ b/vllm_omni/data_entry_keys.py @@ -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): diff --git a/vllm_omni/deploy/minimax_h3_disaggregated.yaml b/vllm_omni/deploy/minimax_h3_disaggregated.yaml new file mode 100644 index 000000000..1c1ee6c40 --- /dev/null +++ b/vllm_omni/deploy/minimax_h3_disaggregated.yaml @@ -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 diff --git a/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml b/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml new file mode 100644 index 000000000..d7e84db42 --- /dev/null +++ b/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml @@ -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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 120f74000..e89825e0e 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -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 diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index e5456f0cc..b8983ab20 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -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 diff --git a/vllm_omni/diffusion/io_support.py b/vllm_omni/diffusion/io_support.py index 8fecc50df..f4cf5bb36 100644 --- a/vllm_omni/diffusion/io_support.py +++ b/vllm_omni/diffusion/io_support.py @@ -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.""" diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index e11a7541e..1665d6a16 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -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() diff --git a/vllm_omni/diffusion/models/minimax_h3/__init__.py b/vllm_omni/diffusion/models/minimax_h3/__init__.py index 53052ad42..02937c058 100644 --- a/vllm_omni/diffusion/models/minimax_h3/__init__.py +++ b/vllm_omni/diffusion/models/minimax_h3/__init__.py @@ -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", ] diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 33c3cd3b0..4f42e7734 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -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)) diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py index 4e68bfd25..62c2330d5 100644 --- a/vllm_omni/diffusion/output_formatter.py +++ b/vllm_omni/diffusion/output_formatter.py @@ -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, diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 96140ef91..2bd4f2888 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -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) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index ef611886f..12dbacefe 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -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 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 3c66290a9..0a5065022 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -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: diff --git a/vllm_omni/engine/messages.py b/vllm_omni/engine/messages.py index 65ec53b10..0b6630ba7 100644 --- a/vllm_omni/engine/messages.py +++ b/vllm_omni/engine/messages.py @@ -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): diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index ab052b494..dc09c2968 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -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 diff --git a/vllm_omni/engine/stage_client.py b/vllm_omni/engine/stage_client.py index 74f65d8c2..06110e005 100644 --- a/vllm_omni/engine/stage_client.py +++ b/vllm_omni/engine/stage_client.py @@ -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 diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 5fbbbc46d..cef4dbf62 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -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 diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index f24eda0a1..6895bdad8 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -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 diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index dea5242d1..add223ea5 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -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") diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 87f4f8834..b3a1e91af 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -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( diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index cdb5114e2..e8cbcbb48 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -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) diff --git a/vllm_omni/entrypoints/openai/serving_video_output_stream.py b/vllm_omni/entrypoints/openai/serving_video_output_stream.py index 770890faf..a9abad146 100644 --- a/vllm_omni/entrypoints/openai/serving_video_output_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_output_stream.py @@ -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), diff --git a/vllm_omni/entrypoints/openai/utils.py b/vllm_omni/entrypoints/openai/utils.py index 5cca49129..4aac073d4 100644 --- a/vllm_omni/entrypoints/openai/utils.py +++ b/vllm_omni/entrypoints/openai/utils.py @@ -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. diff --git a/vllm_omni/model_executor/models/minimax_h3/__init__.py b/vllm_omni/model_executor/models/minimax_h3/__init__.py new file mode 100644 index 000000000..bd977b2ec --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_h3/__init__.py @@ -0,0 +1 @@ +"""MiniMax H3 model components.""" diff --git a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py new file mode 100644 index 000000000..e8350f41a --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py @@ -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"] diff --git a/vllm_omni/model_executor/models/minimax_h3/conditioning.py b/vllm_omni/model_executor/models/minimax_h3/conditioning.py new file mode 100644 index 000000000..8a22c7737 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_h3/conditioning.py @@ -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, + } diff --git a/vllm_omni/model_executor/models/minimax_h3/pipeline.py b/vllm_omni/model_executor/models/minimax_h3/pipeline.py new file mode 100644 index 000000000..5ac94b5ee --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_h3/pipeline.py @@ -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, + ), + ), +) diff --git a/vllm_omni/diffusion/models/minimax_h3/presentation.py b/vllm_omni/model_executor/models/minimax_h3/preprocessing.py similarity index 53% rename from vllm_omni/diffusion/models/minimax_h3/presentation.py rename to vllm_omni/model_executor/models/minimax_h3/preprocessing.py index a8e142189..36a1e20d6 100644 --- a/vllm_omni/diffusion/models/minimax_h3/presentation.py +++ b/vllm_omni/model_executor/models/minimax_h3/preprocessing.py @@ -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: ': ' 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", ] diff --git a/vllm_omni/diffusion/models/minimax_h3/reference_video.py b/vllm_omni/model_executor/models/minimax_h3/reference_video.py similarity index 92% rename from vllm_omni/diffusion/models/minimax_h3/reference_video.py rename to vllm_omni/model_executor/models/minimax_h3/reference_video.py index 18c8de5b4..4f98d9e13 100644 --- a/vllm_omni/diffusion/models/minimax_h3/reference_video.py +++ b/vllm_omni/model_executor/models/minimax_h3/reference_video.py @@ -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, ), } ) diff --git a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py new file mode 100644 index 000000000..85583f243 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py @@ -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, + ) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 0064f1c4a..4438d2200 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -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", diff --git a/vllm_omni/model_executor/stage_input_processors/minimax_h3.py b/vllm_omni/model_executor/stage_input_processors/minimax_h3.py new file mode 100644 index 000000000..372d675b6 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/minimax_h3.py @@ -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 diff --git a/vllm_omni/platforms/npu/models/minimax_h3.py b/vllm_omni/platforms/npu/models/minimax_h3.py index b70f90e30..c8889f18f 100644 --- a/vllm_omni/platforms/npu/models/minimax_h3.py +++ b/vllm_omni/platforms/npu/models/minimax_h3.py @@ -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