# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace from typing import Any import pytest import torch from pydantic import ValidationError from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig, extract_diffusion_stage_config_kwargs from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config from vllm_omni.config.stage_config import build_stage_runtime_overrides from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig from vllm_omni.engine import stage_init_utils from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.entrypoints.cli.serve import OmniServeCommand from vllm_omni.utils.tracking_parser import TrackingArgumentParser pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def _terminal_config(stage_cfg: dict) -> OmniDiffusionConfig: kwargs = extract_diffusion_stage_config_kwargs(stage_cfg["engine_args"], stage_id=stage_cfg["stage_id"]) return OmniDiffusionConfig.from_kwargs(**kwargs) def test_default_stage_config_includes_cache_backend(): """Ensure cache knobs survive the default diffusion-stage builder.""" stage_cfg = StageConfigFactory.create_default_diffusion( { "cache_backend": "cache_dit", "cache_config": '{"Fn_compute_blocks": 2}', "vae_use_slicing": True, "ulysses_degree": 2, "seed": 7, "kv_cache_dtype": "fp8", "diffusion_kv_cache_dtype": "fp8_e4m3", } )[0] engine_args = stage_cfg["engine_args"] assert stage_cfg["stage_type"] == "diffusion" assert engine_args["cache_backend"] == "cache_dit" assert engine_args["cache_config"]["Fn_compute_blocks"] == 2 assert engine_args["vae_use_slicing"] is True assert engine_args["parallel_config"]["ulysses_degree"] == 2 assert engine_args["model_stage"] == "diffusion" assert "seed" not in engine_args assert "kv_cache_dtype" not in engine_args assert engine_args["diffusion_kv_cache_dtype"] == "fp8_e4m3" def test_default_stage_config_preserves_ulysses_a2a_permute() -> None: stage_cfg = StageConfigFactory.create_default_diffusion( { "ulysses_degree": 4, "ulysses_a2a_permute": True, } )[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] assert parallel_config["ulysses_degree"] == 4 assert parallel_config["ulysses_a2a_permute"] is True def test_default_stage_config_preserves_model_extras(): stage_cfg = StageConfigFactory.create_default_diffusion({"extras": {"ltx2_use_conv_vae": True}})[0] assert stage_cfg["engine_args"]["extras"]["ltx2_use_conv_vae"] is True def test_default_stage_config_preserves_and_overrides_promoted_extras(): stage_cfg = StageConfigFactory.create_default_diffusion( { "extras": { "auxiliary_text_encoder": "/models/extras-llama", "default_llama_model_id": "extras/default-llama", }, "auxiliary_text_encoder": None, } )[0] extras = stage_cfg["engine_args"]["extras"] assert extras["auxiliary_text_encoder"] == "/models/extras-llama" assert extras["default_llama_model_id"] == "extras/default-llama" stage_cfg = StageConfigFactory.create_default_diffusion( { "extras": { "auxiliary_text_encoder": "/models/extras-llama", "default_llama_model_id": "extras/default-llama", }, "auxiliary_text_encoder": "/models/top-level-llama", "default_llama_model_id": "top-level/default-llama", } )[0] extras = stage_cfg["engine_args"]["extras"] assert extras["auxiliary_text_encoder"] == "/models/top-level-llama" assert extras["default_llama_model_id"] == "top-level/default-llama" @pytest.mark.parametrize( "stage_overrides", [ {"0": {"extras": {"ltx2_use_conv_vae": True}}}, '{"0":{"extras":{"ltx2_use_conv_vae":true}}}', ], ) def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocker, stage_overrides): """Local/unregistered Diffusers checkpoints still honor stage-0 extras.""" mocker.patch( "vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None, ) mocker.patch( "vllm_omni.config.resolver._resolve_generic_diffusion_model_class", return_value=(True, "LTX2Pipeline"), ) engine = AsyncOmniEngine.__new__(AsyncOmniEngine) _, stage_configs = engine._resolve_stage_configs( "/models/LTX-2.5-Diffusers", {"stage_overrides": stage_overrides}, trust_remote_code=False, ) assert stage_configs[0].diffusion_config.extras["ltx2_use_conv_vae"] is True def test_default_stage_rejects_unknown_nested_parallel_config_key(): unknown_key = "unknown_parallel_field" with pytest.raises(ValidationError, match=unknown_key): StageConfigFactory.create_default_diffusion( {"parallel_config": {unknown_key: 2}}, ) def test_default_stage_routes_ar_profiler_away_before_diffusion_build(mocker): stage_dict = StageConfigFactory.create_default_diffusion({"enable_ar_profiler": True})[0] assert "enable_ar_profiler" not in stage_dict["engine_args"] @pytest.mark.parametrize( ("field_name", "value"), [ ("enable_sleep_mod", None), ("enable_lora", True), ("kv_cache_dtype", "fp8"), ("seed", 7), ], ) def test_legacy_diffusion_stage_rejects_unowned_field(field_name, value): from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.yaml_util import create_config from vllm_omni.engine.stage_init_utils import build_diffusion_config stage_dict = StageConfigFactory.create_default_diffusion({"model": "unused"})[0] stage_dict["engine_args"][field_name] = value stage_cfg = create_config(stage_dict) metadata = SimpleNamespace(stage_id=0, cfg_kv_collect_func=None) with pytest.raises(ValueError, match=rf"stage 0.*{field_name}"): build_diffusion_config("unused", stage_cfg, metadata) @pytest.mark.parametrize( ("field_name", "value"), [ ("enable_sleep_mod", None), ("enable_lora", True), ], ) def test_default_diffusion_factory_rejects_unowned_field(field_name, value): with pytest.raises(ValueError, match=rf"stage 0.*{field_name}"): StageConfigFactory.create_default_diffusion({field_name: value}) def test_legacy_default_stage_build_accepts_engine_adapter_metadata(monkeypatch): from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.yaml_util import create_config from vllm_omni.engine import stage_init_utils stage_dict = StageConfigFactory.create_default_diffusion( { "model": "unused", "api_key": "frontend-owned", } )[0] assert "api_key" not in stage_dict["engine_args"] stage_cfg = create_config(stage_dict) metadata = SimpleNamespace(stage_id=0, cfg_kv_collect_func=None, default_sampling_params=None) monkeypatch.setattr(stage_init_utils.current_omni_platform, "get_device_count", lambda: 1) config = stage_init_utils.build_diffusion_config("unused", stage_cfg, metadata) assert config.model == "unused" def test_default_cache_config_used_when_missing(): """Ensure default cache_config is synthesized when only backend is given.""" stage_cfg = StageConfigFactory.create_default_diffusion( { "cache_backend": "cache_dit", } )[0] cache_config = _terminal_config(stage_cfg).cache_config assert cache_config.Fn_compute_blocks == 1 def test_default_stage_devices_from_sequence_parallel(): """Ensure runtime devices reflect computed diffusion world size.""" stage_cfg = StageConfigFactory.create_default_diffusion( { "ulysses_degree": 2, "ring_degree": 2, } )[0] assert stage_cfg["runtime"]["devices"] == "0,1,2,3" def test_default_stage_devices_and_dp_from_num_gpus(): """Resolve omitted DP before deriving devices from an explicit GPU count.""" stage_cfg = StageConfigFactory.create_default_diffusion( { "num_gpus": 8, "tensor_parallel_size": 2, "ulysses_degree": 2, } )[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] assert parallel_config["data_parallel_size"] == 2 assert stage_cfg["engine_args"]["num_gpus"] == 8 assert stage_cfg["runtime"]["devices"] == "0,1,2,3,4,5,6,7" def test_default_stage_config_uses_parallel_size_kwargs(): """Ensure default diffusion parallel_config uses CLI/API parallel sizes.""" stage_cfg = StageConfigFactory.create_default_diffusion( { "pipeline_parallel_size": 2, "data_parallel_size": 3, "tensor_parallel_size": 4, "enable_expert_parallel": True, } )[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] assert parallel_config["pipeline_parallel_size"] == 2 assert parallel_config["data_parallel_size"] == 3 assert parallel_config["tensor_parallel_size"] == 4 assert parallel_config["enable_expert_parallel"] is True def test_default_stage_config_preserves_omitted_dp_for_runtime_inference(): """Keep omitted DP unresolved until the runtime WORLD size is known.""" stage_cfg = StageConfigFactory.create_default_diffusion( { "pipeline_parallel_size": None, "data_parallel_size": None, "tensor_parallel_size": None, "enable_expert_parallel": None, "enforce_eager": None, "diffusion_compile_granularity": None, "diffusion_compile_dynamic": None, } )[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] assert parallel_config["pipeline_parallel_size"] == 1 assert parallel_config["data_parallel_size"] is None assert parallel_config["tensor_parallel_size"] == 1 assert parallel_config["enable_expert_parallel"] is False terminal_config = _terminal_config(stage_cfg) assert terminal_config.enforce_eager is False assert terminal_config.diffusion_compile_granularity == "regional" assert terminal_config.diffusion_compile_dynamic is True def test_default_stage_config_propagates_ulysses_mode(): """Ensure UAA mode survives default diffusion-stage creation.""" stage_cfg = StageConfigFactory.create_default_diffusion( { "ulysses_degree": 4, "ulysses_mode": "advanced_uaa", } )[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] assert parallel_config["ulysses_degree"] == 4 assert parallel_config["ulysses_mode"] == "advanced_uaa" def test_default_stage_config_includes_default_sampling_params(): """Ensure default sampling params survive the default diffusion-stage builder.""" stage_cfg = StageConfigFactory.create_default_diffusion( { "default_sampling_params": '{"0": {"generator_device":"cpu", "guidance_scale":7.5}}', } )[0] assert stage_cfg["default_sampling_params"] == { "generator_device": "cpu", "guidance_scale": 7.5, } @pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"]) @pytest.mark.parametrize("sampling_defaults", [{"0": {"guidance_scale": 7.5}}, '{"0":{"guidance_scale":7.5}}']) def test_generic_diffusion_sampling_defaults_remain_overridable(typed, sampling_defaults): from vllm_omni.config.yaml_util import create_config from vllm_omni.entrypoints.omni_base import OmniBase from vllm_omni.inputs.data import OmniDiffusionSamplingParams kwargs = {"model_class_name": "QwenImagePipeline", "default_sampling_params": sampling_defaults} if typed: stage = StageConfigFactory.create_typed_default_diffusion("generic-diffusion", kwargs).stage_configs[0] metadata = stage_init_utils.extract_stage_metadata_from_omni_stage_config(stage) else: stage = create_config(StageConfigFactory.create_default_diffusion(kwargs))[0] metadata = stage_init_utils.extract_legacy_stage_metadata(stage) base = OmniBase.__new__(OmniBase) base.engine = SimpleNamespace(num_stages=1, stage_configs=[stage]) base.default_sampling_params_list = [metadata.default_sampling_params] base.sampling_constraints_list = base._get_sampling_constraints_list([stage]) assert base.resolve_sampling_params_list(None)[0].guidance_scale == 7.5 requested = OmniDiffusionSamplingParams(guidance_scale=2.0) assert base.resolve_sampling_params_list(requested)[0].guidance_scale == 2.0 assert requested.guidance_scale == 2.0 assert base.default_sampling_params_list[0].guidance_scale == 7.5 assert base.sampling_constraints_list == [{}] def test_default_stage_config_includes_diffusion_attention_backend(): """Ensure diffusion attention shorthand lands in engine_args.diffusion_attention_config.""" stage_cfg = StageConfigFactory.create_default_diffusion( { "diffusion_attention_backend": "FLASH_ATTN", } )[0] diffusion_attention_config = _terminal_config(stage_cfg).diffusion_attention_config assert isinstance(diffusion_attention_config, AttentionConfig) assert diffusion_attention_config.default is not None assert diffusion_attention_config.default.backend == "FLASH_ATTN" def test_default_stage_config_includes_diffusion_attention_config(): """Ensure structured diffusion attention config survives default stage creation.""" stage_cfg = StageConfigFactory.create_default_diffusion( { "diffusion_attention_config": { "default": {"backend": "FLASH_ATTN"}, "per_role": {"cross": {"backend": "TORCH_SDPA"}}, }, } )[0] diffusion_attention_config = _terminal_config(stage_cfg).diffusion_attention_config assert isinstance(diffusion_attention_config, AttentionConfig) assert diffusion_attention_config.default is not None assert diffusion_attention_config.default.backend == "FLASH_ATTN" assert diffusion_attention_config.per_role["cross"].backend == "TORCH_SDPA" def test_default_stage_config_rejects_conflicting_diffusion_attention_inputs(): """Ensure shorthand and default.backend stay mutually exclusive.""" with pytest.raises(ValueError, match="mutually exclusive"): StageConfigFactory.create_default_diffusion( { "diffusion_attention_backend": "FLASH_ATTN", "diffusion_attention_config": { "default": {"backend": "TORCH_SDPA"}, }, } ) def test_default_stage_config_engine_args(): """Ensure default diffusion-stage builder sets and propagates engine_args.""" stage_cfg = StageConfigFactory.create_default_diffusion( { "distributed_executor_backend": "ray", "boundary_ratio": 0.875, "flow_shift": 5.0, "trust_remote_code": True, } )[0] engine_args = stage_cfg["engine_args"] assert engine_args["distributed_executor_backend"] == "ray" assert engine_args["boundary_ratio"] == 0.875 assert engine_args["flow_shift"] == 5.0 assert engine_args["trust_remote_code"] is True def test_default_stage_config_whitelist_none_fallback(): """DeployConfig / StageDeployConfig whitelist fields with value None fall back to OmniDiffusionConfig dataclass defaults.""" stage_cfg = StageConfigFactory.create_default_diffusion( { # DeployConfig pipeline-wide "trust_remote_code": None, "distributed_executor_backend": None, "dtype": None, # StageDeployConfig "enforce_eager": None, } )[0] terminal = _terminal_config(stage_cfg) assert terminal.trust_remote_code is False assert terminal.distributed_executor_backend is None assert terminal.dtype == torch.bfloat16 assert terminal.enforce_eager is False def test_serve_cli_accepts_ulysses_mode(): """Ensure diffusion serve CLI exposes ulysses_mode and wires it to parallel_config.""" parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) args = parser.parse_args( [ "serve", "Qwen/Qwen-Image", "--omni", "--usp", "4", "--ulysses-mode", "advanced_uaa", ] ) explicit_kwargs = args.get_explicit_kwargs_dict() stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] assert args.ulysses_mode == "advanced_uaa" assert parallel_config["ulysses_degree"] == 4 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 = StageConfigFactory.create_default_diffusion(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") OmniServeCommand().subparser_init(subparsers) args = parser.parse_args( [ "serve", "MiniMaxAI/MiniMax-H3", "--omni", "--task-type", "fl2va", ] ) explicit_kwargs = args.get_explicit_kwargs_dict() stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] assert args.task_type == "fl2va" assert stage_cfg["engine_args"]["task_type"] == "fl2va" def test_serve_cli_accepts_diffusion_pipeline_profiler_flag(): """Ensure diffusion serve CLI exposes the profiler switch.""" parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) args = parser.parse_args( [ "serve", "Wan-AI/Wan2.2-T2V-A14B-Diffusers", "--omni", "--enable-diffusion-pipeline-profiler", ] ) explicit_kwargs = args.get_explicit_kwargs_dict() stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] assert args.enable_diffusion_pipeline_profiler is True assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True def test_serve_cli_forwards_distilled_lora_to_diffusion_stage(): """Ensure startup distilled LoRA options reach the online diffusion stage.""" parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) args = parser.parse_args( [ "serve", "Wan-AI/Wan2.2-T2V-A14B-Diffusers", "--omni", "--lora-backend", "distill", "--lora-path", "/models/high.safetensors", "/models/low.safetensors", ] ) explicit_kwargs = args.get_explicit_kwargs_dict() stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] engine_args = stage_cfg["engine_args"] assert explicit_kwargs["lora_backend"] == "distill" assert engine_args["lora_backend"] == "distill" assert engine_args["lora_path"] == [ "/models/high.safetensors", "/models/low.safetensors", ] def test_serve_cli_forwards_compact_diffusion_offload_config(): """Ensure component-specific layer settings reach the diffusion stage.""" parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) args = parser.parse_args( [ "serve", "MiniMaxAI/MiniMax-H3", "--omni", "--diffusion-offload-config", '{"mode":"layer","components":["dit","text_encoder"],' '"layer_options":{"dit":{"weight_transfer":"rank-local","resident_layers":20},' '"text_encoder":{"weight_transfer":"allgather"}},' '"pin_memory":true}', ] ) explicit_kwargs = args.get_explicit_kwargs_dict() stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] engine_args = stage_cfg["engine_args"] expected = { "mode": "layer", "components": ["dit", "text_encoder"], "layer_options": { "dit": {"weight_transfer": "rank-local", "resident_layers": 20}, "text_encoder": {"weight_transfer": "allgather"}, }, "pin_memory": True, } assert args.diffusion_offload_config == expected assert engine_args["diffusion_offload_config"] == expected def test_invalid_diffusion_offload_config_fails_before_model_loading(monkeypatch, mocker): load_model = mocker.patch("vllm_omni.diffusion.model_loader.diffusers_loader.DiffusersPipelineLoader.load_model") create_client = mocker.patch("vllm_omni.diffusion.stage_diffusion_client.create_diffusion_client") monkeypatch.setattr( stage_init_utils, "build_engine_args_dict", lambda *_args, **_kwargs: { "model": "test", "diffusion_offload_config": { "mode": "layerwise", "components": ["dit"], }, }, ) with pytest.raises(ValueError, match="Unknown diffusion offload mode"): stage_init_utils.initialize_diffusion_stage( stage_id=0, model="test", stage_cfg=object(), metadata=mocker.Mock(), stage_init_timeout=30, ) create_client.assert_not_called() load_model.assert_not_called() def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo(): parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) args = parser.parse_args( [ "serve", "MiniMaxAI/MiniMax-H3", "--omni", "--enable-distributed-layerwise-offload", "--dlo-no-use-allgather", "--host-weight-runtime-mode", "preferred", "--host-weight-runtime-root", "/var/cache/vllm-omni/hwr", "--dlo-host-registration-limit-gib", "80", ] ) explicit_kwargs = args.get_explicit_kwargs_dict() stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] engine_args = stage_cfg["engine_args"] assert explicit_kwargs["host_weight_runtime_mode"] == "preferred" assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr" assert explicit_kwargs["dlo_host_registration_limit_gib"] == 80 assert engine_args["host_weight_runtime_mode"] == "preferred" assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr" assert engine_args["dlo_host_registration_limit_gib"] == 80 def test_serve_cli_accepts_diffusion_compile_controls(): """Ensure both compile controls reach the diffusion stage.""" parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) args = parser.parse_args( [ "serve", "Lightricks/LTX-Video-0.9.8-13B-distilled", "--omni", "--diffusion-compile-granularity", "full", "--no-diffusion-compile-dynamic", ] ) explicit_kwargs = args.get_explicit_kwargs_dict() stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] assert args.diffusion_compile_granularity == "full" assert args.diffusion_compile_dynamic is False assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "full" assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is False def test_serve_cli_accepts_diffusion_attention_backend(): """Ensure diffusion serve CLI exposes the shorthand backend flag.""" parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) args = parser.parse_args( [ "serve", "Qwen/Qwen-Image", "--omni", "--diffusion-attention-backend", "FASTVIDEO_VSA", "--fastvideo-vsa-topk", "96", ] ) explicit_kwargs = args.get_explicit_kwargs_dict() stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"] assert args.diffusion_attention_backend == "FASTVIDEO_VSA" assert args.fastvideo_vsa_topk == 96 assert isinstance(diffusion_attention_config, AttentionConfig) assert diffusion_attention_config.default is not None assert diffusion_attention_config.default.backend == "FASTVIDEO_VSA" assert diffusion_attention_config.default.backend_kwargs() == {"topk": 96} def test_serve_cli_accepts_request_batch_max_wait_ms(): """Ensure diffusion serve CLI forwards request-batch admission wait to stage config.""" parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) args = parser.parse_args( [ "serve", "Qwen/Qwen-Image", "--omni", "--request-batch-max-wait-ms", "250", ] ) explicit_kwargs = args.get_explicit_kwargs_dict() stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] assert args.request_batch_max_wait_ms == 250.0 assert stage_cfg["engine_args"]["request_batch_max_wait_ms"] == 250.0 @pytest.mark.parametrize("bad_wait", ["nan", "inf", "-inf", "-1"]) def test_serve_cli_rejects_invalid_request_batch_max_wait_ms(bad_wait: str): parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) with pytest.raises(SystemExit): parser.parse_args( [ "serve", "Qwen/Qwen-Image", "--omni", "--request-batch-max-wait-ms", bad_wait, ] ) @pytest.mark.parametrize("subcommand_dest", ["command", "subparser"]) def test_serve_cli_accepts_additional_config(subcommand_dest): """Ensure diffusion serve CLI exposes additional_config and forwards it to stage config.""" parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest=subcommand_dest) OmniServeCommand().subparser_init(subparsers) args = parser.parse_args( [ "serve", "Qwen/Qwen-Image", "--omni", "--additional-config", '{"torchair_graph_config":{"enabled":true}}', ] ) stage_cfg = StageConfigFactory.create_default_diffusion(args.get_explicit_kwargs_dict())[0] engine_args = stage_cfg["engine_args"] assert args.additional_config == {"torchair_graph_config": {"enabled": True}} assert engine_args["additional_config"] == {"torchair_graph_config": {"enabled": True}} def test_resolve_stage_configs_delegates_overrides_to_resolver(mocker): """The engine consumes resolver output without a second merge pass.""" additional_config = {"torchair_graph_config": {"enabled": True}} fake_diffusion_stage = SimpleNamespace( stage_type="diffusion", engine_args=SimpleNamespace(additional_config=additional_config), ) resolve_config = mocker.patch( "vllm_omni.engine.omni_engine_base.resolve_omni_config", return_value=OmniConfigResolution( config_path="dummy.yaml", stage_configs=(fake_diffusion_stage,), ), ) engine = AsyncOmniEngine.__new__(AsyncOmniEngine) _, stage_configs = engine._resolve_stage_configs( "dummy-model", { "deploy_config": "dummy.yaml", "additional_config": additional_config, }, trust_remote_code=False, ) assert stage_configs == [fake_diffusion_stage] assert resolve_config.call_args.args == ("dummy-model",) assert resolve_config.call_args.kwargs["deploy_config_path"] == "dummy.yaml" assert resolve_config.call_args.kwargs["cli_overrides"]["additional_config"] is additional_config @pytest.mark.parametrize( ("legacy_arg", "value"), [ ("stage_configs_path", "legacy.yaml"), ("stage_configs", [{"stage_id": 0}]), ], ) def test_resolve_stage_configs_rejects_legacy_config_arguments(legacy_arg, value): engine = AsyncOmniEngine.__new__(AsyncOmniEngine) with pytest.raises(ValueError, match=rf"`{legacy_arg}`.*`deploy_config`"): engine._resolve_stage_configs( "dummy-model", {legacy_arg: value}, trust_remote_code=False, ) def test_default_stage_config_includes_quantization_config(): """Ensure structured quantization_config survives default diffusion-stage creation.""" quantization_config = { "method": "example_quant", "weights": "weights.bin", } stage_cfg = StageConfigFactory.create_default_diffusion({"quantization_config": quantization_config})[0] assert stage_cfg["engine_args"]["quantization_config"] == quantization_config @pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"]) def test_default_diffusion_factory_preserves_engine_quantization(typed, monkeypatch): monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500) monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None) kwargs = {"quantization": "fp8"} if typed: stage = StageConfigFactory.create_typed_default_diffusion("generic-diffusion", kwargs).stage_configs[0] config = stage.diffusion_config config.enrich_config() else: config = _terminal_config(StageConfigFactory.create_default_diffusion(kwargs)[0]) assert config.quantization_config is not None assert config.quantization_config.get_name() == "fp8" assert config.quantization_config_is_auto_detected is False @pytest.mark.parametrize("model_class_name", ["HeliosPipeline", "HunyuanVideo15Pipeline"]) def test_generic_diffusion_uses_canonical_video_output_type(model_class_name): config = StageConfigFactory.create_typed_default_diffusion( "generic-video", {"model_class_name": model_class_name}, ) assert config.stage_configs[0].final_output_type == "video" def test_generic_diffusion_resolves_structured_stage_without_legacy_conversion(mocker): """Generic diffusion reaches runtime as the structured stage itself.""" mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None) mocker.patch( "vllm_omni.config.resolver._resolve_generic_diffusion_model_class", return_value=(True, "FakeDiffusionPipeline"), ) resolved = resolve_omni_config( "generic-diffusion", trust_remote_code=False, deploy_config_path=None, cli_overrides={ "num_gpus": 4, "tensor_parallel_size": 2, "default_sampling_params": '{"0": {"guidance_scale": 7.5}}', }, stage_overrides=None, strategy_config_path=None, ) assert resolved.pipeline_config is not None assert resolved.pipeline_config.model_type == "generic_diffusion" assert len(resolved.stage_configs) == 1 stage = resolved.stage_configs[0] assert isinstance(stage, VllmOmniDiffusionStageConfig) assert stage.model_config.model == "generic-diffusion" assert stage.model_config.default_sampling_params == {"guidance_scale": 7.5} assert stage.parallel_config.tensor_parallel_size == 2 assert stage.parallel_config.data_parallel_size == 2 assert stage.parallel_config.world_size == 4 assert stage.runtime_config.devices == "0,1,2,3" def test_generic_diffusion_structured_stage_reaches_standard_startup(mocker): """Standard runtime resolves and starts the typed stage through the real launcher.""" from vllm_omni.engine import stage_engine_startup as startup_module from vllm_omni.engine import stage_runtime as runtime_module from vllm_omni.engine.stage_runtime import StageRuntime mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None) mocker.patch( "vllm_omni.config.resolver._resolve_generic_diffusion_model_class", return_value=(True, "FakeDiffusionPipeline"), ) resolved = resolve_omni_config( "generic-diffusion", trust_remote_code=False, deploy_config_path=None, cli_overrides={"num_gpus": 1}, stage_overrides=None, strategy_config_path=None, ) stage = resolved.stage_configs[0] launched: dict[str, Any] = {} client = SimpleNamespace(input_address=None, shutdown=mocker.Mock()) mocker.patch.object(runtime_module, "prepare_engine_environment") mocker.patch.object(runtime_module, "load_omni_transfer_config_for_model", return_value=None) mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={}) mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None)) def _initialize_typed_stage(stage_id, model, stage_config, metadata, **kwargs): launched.update( stage_id=stage_id, model=model, stage_config=stage_config, metadata=metadata, **kwargs, ) return client mocker.patch.object(startup_module, "initialize_diffusion_stage", side_effect=_initialize_typed_stage) runtime = StageRuntime( stage_configs=list(resolved.stage_configs), model="generic-diffusion", config_path="", stage_init_timeout=10, async_chunk=False, ) runtime.initialize() assert launched["stage_config"] is stage assert isinstance(launched["stage_config"], VllmOmniDiffusionStageConfig) assert launched["metadata"].stage_type == "diffusion" assert launched["metadata"].model_stage == "diffusion" assert launched["use_inline"] is True assert launched["model"] == "generic-diffusion" assert launched["stage_id"] == 0 assert runtime.stage_pools[0].clients == [client] def test_default_stage_config_preserves_pre_sharded_hsdp_weight_load() -> None: stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"hsdp_weight_load_strategy": "pre_sharded"})[0] assert stage_cfg["engine_args"]["hsdp_weight_load_strategy"] == "pre_sharded" @pytest.mark.parametrize("prefix", ["", "stage_0_"]) def test_pre_sharded_hsdp_weight_load_survives_stage_overrides(prefix: str) -> None: overrides = build_stage_runtime_overrides(0, {f"{prefix}hsdp_weight_load_strategy": "pre_sharded"}) assert overrides["hsdp_weight_load_strategy"] == "pre_sharded" @pytest.mark.parametrize("strategy", ["full", "pre_sharded"]) def test_typed_default_stage_preserves_hsdp_weight_load_strategy(strategy: str) -> None: config = StageConfigFactory.create_typed_default_diffusion( "generic-diffusion", {"hsdp_weight_load_strategy": strategy} ) assert config.stage_configs[0].diffusion_config.hsdp_weight_load_strategy == strategy @pytest.mark.parametrize("strategy", ["memory_limited", "invalid"]) def test_removed_hsdp_loading_strategies_are_rejected(strategy): from vllm_omni.diffusion.data import OmniDiffusionConfig with pytest.raises(ValueError, match="hsdp_weight_load_strategy"): OmniDiffusionConfig(hsdp_weight_load_strategy=strategy)