diff --git a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py index dc1085c28..7c91c0323 100644 --- a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py +++ b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py @@ -118,9 +118,6 @@ def parse_args() -> argparse.Namespace: help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', ) - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/helios/end2end.py b/examples/offline_inference/helios/end2end.py index 6cf7dfdcb..88c3b865d 100644 --- a/examples/offline_inference/helios/end2end.py +++ b/examples/offline_inference/helios/end2end.py @@ -196,9 +196,6 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.") parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallelism size.") - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index b857bb22d..7837fe16d 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -378,9 +378,6 @@ def parse_args() -> argparse.Namespace: default=None, help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', ) - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index fea5178d8..e46db4de4 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -211,9 +211,6 @@ def parse_args() -> argparse.Namespace: default=None, help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', ) - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/magi_human/end2end.py b/examples/offline_inference/magi_human/end2end.py index 7ea816138..64f11c465 100644 --- a/examples/offline_inference/magi_human/end2end.py +++ b/examples/offline_inference/magi_human/end2end.py @@ -24,9 +24,6 @@ def parse_args(): parser.add_argument("--width", type=int, default=448, help="Video width.") parser.add_argument("--num-inference-steps", type=int, default=8, help="Number of denoising steps.") parser.add_argument("--seed", type=int, default=52, help="Random seed for generation.") - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index bc18c6859..c0fd337bd 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -9,7 +9,7 @@ from typing import Any import torch -from vllm_omni.diffusion.data import DiffusionParallelConfig, logger +from vllm_omni.diffusion.data import logger from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.lora.request import LoRARequest @@ -342,17 +342,6 @@ def main(): # (e.g., QwenImagePipeline or FluxPipeline) } - # assert args.ring_degree == 1, "Ring attention is not supported yet" - parallel_config = DiffusionParallelConfig( - ulysses_degree=args.ulysses_degree, - ring_degree=args.ring_degree, - ulysses_mode=args.ulysses_mode, - cfg_parallel_size=args.cfg_parallel_size, - tensor_parallel_size=args.tensor_parallel_size, - vae_patch_parallel_size=args.vae_patch_parallel_size, - enable_expert_parallel=args.enable_expert_parallel, - ) - profiler_enabled = args.profiler_config is not None # Prepare LoRA kwargs for Omni initialization @@ -388,7 +377,13 @@ def main(): "cache_backend": args.cache_backend, "cache_config": cache_config, "enable_cache_dit_summary": args.enable_cache_dit_summary, - "parallel_config": parallel_config, + "ulysses_degree": args.ulysses_degree, + "ring_degree": args.ring_degree, + "ulysses_mode": args.ulysses_mode, + "cfg_parallel_size": args.cfg_parallel_size, + "tensor_parallel_size": args.tensor_parallel_size, + "vae_patch_parallel_size": args.vae_patch_parallel_size, + "enable_expert_parallel": args.enable_expert_parallel, "enforce_eager": args.enforce_eager, "enable_cpu_offload": args.enable_cpu_offload, "mode": "text-to-image", diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index a96949fa5..1dba97c20 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -226,9 +226,6 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of HSDP replica groups.", ) - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/vace/vace_video_generation.py b/examples/offline_inference/vace/vace_video_generation.py index 5fad37366..6ca0d74c5 100644 --- a/examples/offline_inference/vace/vace_video_generation.py +++ b/examples/offline_inference/vace/vace_video_generation.py @@ -71,9 +71,6 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.") parser.add_argument("--ring-degree", type=int, default=1, help="Ring attention degree.") parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.") - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 7ed812826..83b465fdb 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -55,6 +55,44 @@ def test_default_stage_devices_from_sequence_parallel(): assert stage_cfg["runtime"]["devices"] == "0,1,2,3" +def test_default_stage_config_uses_parallel_size_kwargs(): + """Ensure default diffusion parallel_config uses CLI/API parallel sizes.""" + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "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_defaults_nullified_parallel_size_kwargs(): + """Ensure nullified diffusion parallel-size kwargs fall back to defaults.""" + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "pipeline_parallel_size": None, + "data_parallel_size": None, + "tensor_parallel_size": None, + "enable_expert_parallel": None, + "enforce_eager": None, + } + )[0] + + parallel_config = stage_cfg["engine_args"]["parallel_config"] + assert parallel_config.pipeline_parallel_size == 1 + assert parallel_config.data_parallel_size == 1 + assert parallel_config.tensor_parallel_size == 1 + assert parallel_config.enable_expert_parallel is False + assert stage_cfg["engine_args"]["enforce_eager"] is False + + def test_default_stage_config_propagates_ulysses_mode(): """Ensure UAA mode survives default diffusion-stage creation.""" stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( diff --git a/tests/examples/offline_inference/test_text_to_image.py b/tests/examples/offline_inference/test_text_to_image.py index 9b8a25184..a63af7cfb 100644 --- a/tests/examples/offline_inference/test_text_to_image.py +++ b/tests/examples/offline_inference/test_text_to_image.py @@ -26,27 +26,9 @@ def _skip_readme_snippet(language: str, code: str, h2_title: str) -> tuple[bool, README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH, skipif=_skip_readme_snippet) -# Track: https://github.com/vllm-project/vllm-omni/issues/3123 — CLI example subprocess exit 1 in CI; re-enable when fixed. -_SKIP_T2I_SNIPPET_IDS = frozenset( - { - "quick_start_002", - "more_cli_examples_001", - "more_cli_examples_002", - "more_cli_examples_003", - "more_cli_examples_007", - "more_cli_examples_008", - } -) -_SKIP_T2I_REASON = ( - "Skipped: text_to_image README examples failing in CI (subprocess exit 1). " - "See https://github.com/vllm-project/vllm-omni/issues/3123" -) - @pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id) def test_text_to_image(snippet: ReadmeSnippet, example_runner: ExampleRunner): - if snippet.test_id in _SKIP_T2I_SNIPPET_IDS: - pytest.skip(_SKIP_T2I_REASON) should_skip, reason = snippet.skip if should_skip: pytest.skip(reason) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 831d41092..54c9d32d9 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1294,9 +1294,12 @@ class AsyncOmniEngine: ring_degree = normalized_kwargs.get("ring_degree") or 1 ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict" sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size") + pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 + data_parallel_size = normalized_kwargs.get("data_parallel_size") or 1 tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1 cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1 vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1 + enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False use_hsdp = normalized_kwargs.get("use_hsdp", False) hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1) hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size", 1) @@ -1304,9 +1307,10 @@ class AsyncOmniEngine: sequence_parallel_size = ulysses_degree * ring_degree parallel_config = DiffusionParallelConfig( - pipeline_parallel_size=1, - data_parallel_size=1, + pipeline_parallel_size=pipeline_parallel_size, + data_parallel_size=data_parallel_size, tensor_parallel_size=tensor_parallel_size, + enable_expert_parallel=enable_expert_parallel, sequence_parallel_size=sequence_parallel_size, ulysses_degree=ulysses_degree, ring_degree=ring_degree, @@ -1333,7 +1337,7 @@ class AsyncOmniEngine: "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), "enable_cpu_offload": kwargs.get("enable_cpu_offload", False), "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), - "enforce_eager": kwargs.get("enforce_eager", False), + "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), "boundary_ratio": kwargs.get("boundary_ratio", None), "flow_shift": kwargs.get("flow_shift", None), "diffusion_load_format": kwargs.get("diffusion_load_format", "default"),