[Bugfix][Examples][MammothModa2] Serve T2I via LLM-typed DiT stage (#7199) + benchmark harness and profiler wiring (#7293)

Signed-off-by: quan-yi-ai <82917002+quan-yi-ai@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

examples/offline_inference/text_to_image/text_to_image.py

@@ -14,7 +14,7 @@ from diffusers.utils import numpy_to_pil
 
 from vllm_omni.diffusion.data import logger
 from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
-from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
+from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args, ar_grid_max_tokens
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -72,6 +72,24 @@ def build_text_to_image_prompt(prompt: str, negative_prompt: str | None) -> dict
     return result
 
 
+def build_parallel_knob_kwargs(args: argparse.Namespace) -> dict[str, Any]:
+    """Forward explicitly-set parallel knobs so they can override a deploy YAML.
+
+    The argparse defaults are ``None``; only explicitly supplied values enter
+    the returned dict, letting a deploy YAML's parallel_config stay in effect
+    when the flags are omitted (and letting explicit values equal to the old
+    argparse defaults, e.g. ``--ulysses-degree 1``, still override the YAML).
+    """
+    knob_names = (
+        "ulysses_degree",
+        "ring_degree",
+        "ulysses_mode",
+        "cfg_parallel_size",
+        "vae_patch_parallel_size",
+    )
+    return {name: getattr(args, name) for name in knob_names if getattr(args, name) is not None}
+
+
 def _normalize_images_for_save(images: list[Any]) -> list[Any]:
     """Convert NumPy diffusion outputs to PIL images before saving."""
     normalized = []
@@ -83,7 +101,8 @@ def _normalize_images_for_save(images: list[Any]) -> list[Any]:
     return normalized
 
 
-def parse_args() -> argparse.Namespace:
+def build_parser() -> argparse.ArgumentParser:
+    """Build the example's CLI parser (kept separate for testability)."""
     parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
     parser.add_argument(
         "--model",
@@ -166,28 +185,31 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--ulysses-degree",
         type=int,
-        default=1,
-        help="Number of GPUs used for ulysses sequence parallelism.",
+        default=None,
+        help="Number of GPUs used for ulysses sequence parallelism. "
+        "Explicitly passing 1 overrides a deploy YAML value.",
     )
     parser.add_argument(
         "--ulysses-mode",
         type=str,
-        default="strict",
+        default=None,
         choices=["strict", "advanced_uaa"],
-        help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).",
+        help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA). "
+        "Explicitly passing 'strict' overrides a deploy YAML value.",
     )
     parser.add_argument(
         "--ring-degree",
         type=int,
-        default=1,
-        help="Number of GPUs used for ring sequence parallelism.",
+        default=None,
+        help="Number of GPUs used for ring sequence parallelism. Explicitly passing 1 overrides a deploy YAML value.",
     )
     parser.add_argument(
         "--cfg-parallel-size",
         type=int,
-        default=1,
+        default=None,
         choices=[1, 2],
-        help="Number of GPUs used for classifier free guidance parallel size.",
+        help="Number of GPUs used for classifier free guidance parallel size. "
+        "Explicitly passing 1 overrides a deploy YAML value.",
     )
     parser.add_argument(
         "--enforce-eager",
@@ -296,8 +318,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--vae-patch-parallel-size",
         type=int,
-        default=1,
-        help="Number of ranks used for VAE patch/tile parallelism (decode/encode).",
+        default=None,
+        help="Number of ranks used for VAE patch/tile parallelism (decode/encode). "
+        "Explicitly passing 1 overrides a deploy YAML value.",
     )
     # NextStep-1.1 specific arguments
     parser.add_argument(
@@ -413,7 +436,11 @@ def parse_args() -> argparse.Namespace:
         ),
     )
     current_omni_platform.pre_register_and_update(parser)
-    return parser.parse_args()
+    return parser
+
+
+def parse_args() -> argparse.Namespace:
+    return build_parser().parse_args()
 
 
 def _apply_ar_stage_inputs(
@@ -556,21 +583,11 @@ def main():
 
     omni_kwargs = {
         "model": args.model,
-        "enable_layerwise_offload": args.enable_layerwise_offload,
-        "vae_use_slicing": args.vae_use_slicing,
-        "vae_use_tiling": args.vae_use_tiling,
         "cache_backend": args.cache_backend,
         "cache_config": cache_config,
-        "enable_cache_dit_summary": args.enable_cache_dit_summary,
-        "ulysses_degree": args.ulysses_degree,
-        "ring_degree": args.ring_degree,
-        "ulysses_mode": args.ulysses_mode,
-        "cfg_parallel_size": args.cfg_parallel_size,
-        "vae_patch_parallel_size": args.vae_patch_parallel_size,
         "enable_expert_parallel": args.enable_expert_parallel,
         "enable_cpu_offload": args.enable_cpu_offload,
         "log_stats": args.log_stats,
-        "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
         "profiler_config": args.profiler_config,
         "init_timeout": args.init_timeout,
         "stage_init_timeout": args.stage_init_timeout,
@@ -578,6 +595,8 @@ def main():
         **lora_args,
         **quant_kwargs,
     }
+    # Diffusion-only engine args: pass only when explicitly set, so pipelines
+    # without a DIFFUSION stage pass ownership validation.
     if args.tensor_parallel_size is not None:
         omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
     if args.enforce_eager is not None:
@@ -597,6 +616,18 @@ def main():
     # gate is an engine-level config (offline analog of the server's --no-guardrails).
     if args.extra_body and "guardrails" in args.extra_body:
         omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
+    if args.enable_layerwise_offload:
+        omni_kwargs["enable_layerwise_offload"] = True
+    if args.vae_use_slicing:
+        omni_kwargs["vae_use_slicing"] = True
+    if args.vae_use_tiling:
+        omni_kwargs["vae_use_tiling"] = True
+    if args.enable_cache_dit_summary:
+        omni_kwargs["enable_cache_dit_summary"] = True
+    if args.enable_diffusion_pipeline_profiler:
+        omni_kwargs["enable_diffusion_pipeline_profiler"] = True
+    omni_kwargs.update(build_parallel_knob_kwargs(args))
+
     omni = Omni(**omni_kwargs)
     model_class_name = get_model_class_name(omni)
     declared_extra_body_params = get_extra_body_params(model_class_name)
@@ -615,11 +646,17 @@ def main():
     if ignored_layers:
         print(f"  Ignored layers: {ignored_layers}")
     tp_display = args.tensor_parallel_size if args.tensor_parallel_size is not None else "deploy/default"
+
+    def _knob_display(value: Any, default: Any) -> Any:
+        return value if value is not None else f"deploy/default ({default})"
+
     print(
         f"  Parallel configuration: tensor_parallel_size={tp_display}, "
-        f"ulysses_degree={args.ulysses_degree}, ulysses_mode={args.ulysses_mode}, "
-        f"ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
-        f"vae_patch_parallel_size={args.vae_patch_parallel_size}, "
+        f"ulysses_degree={_knob_display(args.ulysses_degree, 1)}, "
+        f"ulysses_mode={_knob_display(args.ulysses_mode, 'strict')}, "
+        f"ring_degree={_knob_display(args.ring_degree, 1)}, "
+        f"cfg_parallel_size={_knob_display(args.cfg_parallel_size, 1)}, "
+        f"vae_patch_parallel_size={_knob_display(args.vae_patch_parallel_size, 1)}, "
         f"enable_expert_parallel={args.enable_expert_parallel}."
     )
     print(f"  CPU offload: {args.enable_cpu_offload}; CPU Layerwise Offload: {args.enable_layerwise_offload}")
@@ -730,8 +767,9 @@ def main():
             if idx == 0 and prompt_info.get("omni_task") == ["t2i"]:
                 ar_width = int(prompt_info.get("ar_width", [0])[0])
                 ar_height = int(prompt_info.get("ar_height", [0])[0])
-                if ar_width > 0 and ar_height > 0:
-                    params.max_tokens = ar_height * (ar_width + 1) + 1
+                ar_grid_budget = ar_grid_max_tokens(ar_width, ar_height)
+                if ar_grid_budget is not None:
+                    params.max_tokens = ar_grid_budget
 
     if not diffusion_replaced and len(sampling_params_list) == 1:
         sampling_params_list = [diffusion_params]