[BugFix] Fix nullify regressions in pure diffusion offline examples (#3181)

Signed-off-by: xiaohajiayou <923390377@qq.com>

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

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

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

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

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

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

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

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

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(

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)

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