[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>

tests/entrypoints/openai_api/test_image_server.py

@@ -438,6 +438,90 @@ def async_omni_stage_configs_only_client():
     return TestClient(app)
 
 
+@pytest.fixture
+def mammoth_moda2_test_client():
+    """MammothModa2-style pipeline: AR (llm) + generation-LLM DiT stage.
+
+    The DiT stage is ``stage_type == "llm"`` with ``final_output_type ==
+    "image"`` (no classical diffusion stage) — the legacy topology
+    MammothModa2 served before #7134 migrated it to the shared diffusion
+    runtime. Regression guard for the 503 (stage discovery) and the T2I
+    envelope / max_tokens wiring on that compat path.
+    """
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.entrypoints.openai.api_server import router
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    class FakeAsyncOmniClass(AsyncOmni):
+        def __init__(self):
+            stage_configs = [
+                SimpleNamespace(
+                    stage_type="llm",
+                    is_comprehension=False,
+                    engine_args=SimpleNamespace(model_arch="MammothModa2ForConditionalGeneration"),
+                ),
+                SimpleNamespace(
+                    stage_type="llm",
+                    is_comprehension=False,
+                    final_output=True,
+                    final_output_type="image",
+                    engine_args=SimpleNamespace(model_arch="MammothModa2ForConditionalGeneration"),
+                ),
+            ]
+            default_sampling_params_list = [
+                SamplingParams(temperature=1.0, max_tokens=16),
+                OmniDiffusionSamplingParams(),
+            ]
+            self.engine = SimpleNamespace(
+                stage_configs=stage_configs,
+                default_sampling_params_list=default_sampling_params_list,
+                # Real AsyncOmni engines always carry stage_clients (read by
+                # get_diffusion_od_config); edits-endpoint tests walk through
+                # helpers that touch it before any stage-type gate.
+                stage_clients=[],
+            )
+            self.default_sampling_params_list = default_sampling_params_list
+            self.captured_sampling_params_list = None
+            self.captured_prompt = None
+            self._images = [Image.new("RGB", (64, 64), color="green")]
+
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
+            if sampling_params_list is not None:
+                self.captured_sampling_params_list = sampling_params_list
+            else:
+                self.captured_sampling_params_list = [sampling_params]
+            self.captured_prompt = prompt
+            images = [img.copy() for img in self._images]
+            yield MockGenerationResult(images)
+
+        def __class_getitem__(cls, item):
+            return cls
+
+    app = FastAPI()
+    app.include_router(router)
+
+    engine = FakeAsyncOmniClass()
+    chat_handler = object.__new__(OmniOpenAIServingChat)
+    chat_handler.engine_client = engine
+    chat_handler._diffusion_engine = None
+    app.state.openai_serving_chat = chat_handler
+    app.state.engine_client = engine
+    app.state.stage_configs = [
+        SimpleNamespace(stage_type="llm"),
+        SimpleNamespace(stage_type="llm", final_output=True, final_output_type="image"),
+    ]
+    app.state.openai_serving_models = _DiffusionServingModels(
+        [BaseModelPath(name="Mammoth/MammothModa2-Preview", model_path="Mammoth/MammothModa2-Preview")]
+    )
+    app.state.args = Namespace(
+        default_sampling_params='{"1": {"num_inference_steps":50, "generator_device":"cpu"}}',
+        max_generated_image_size=1024 * 1024,
+    )
+    return TestClient(app)
+
+
 @pytest.fixture
 def streaming_image_edit_client():
     """Create a multi-stage client whose engine yields AR text before image output."""
@@ -2418,3 +2502,190 @@ def test_image_edits_omitted_bot_task_stop_tokens_match_prompt_default(
         f"{ar_params.stop_token_ids} -- this is the full ratio range, meaning "
         "resolve_stop_token_ids disagreed with build_prompt_tokens's default."
     )
+
+
+def test_generate_images_mammoth_moda2_llm_dit_pipeline_accepted(mammoth_moda2_test_client):
+    """Regression: /v1/images/generations must not 503 on generation-LLM DiT stages.
+
+    A generation-LLM DiT stage (``stage_type == "llm"`` with
+    ``final_output_type == "image"``, MammothModa2's topology before #7134):
+    stage discovery previously required a classical diffusion stage and
+    returned 503 (issue #7199).
+    """
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 1,
+            "size": "1024x1024",
+            "seed": 7,
+        },
+    )
+    assert response.status_code == 200, response.text
+    data = response.json()
+    assert len(data["data"]) == 1
+
+
+def test_generate_images_mammoth_moda2_llm_dit_rejects_multiple_outputs(mammoth_moda2_test_client):
+    """Regression: n > 1 must be rejected on the LLM-typed DiT path.
+
+    The LLM engine path for generation-LLM DiT stages produces exactly one
+    image per request; accepting ``n > 1`` used to silently drop the extra
+    outputs (review on PR #7293). The request must fail fast with 400 instead.
+    """
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 4,
+            "size": "1024x1024",
+            "seed": 7,
+        },
+    )
+    assert response.status_code == 400, response.text
+    assert "exactly one image" in response.text
+
+
+def test_generate_images_mammoth_moda2_t2i_envelope_and_max_tokens(mammoth_moda2_test_client):
+    """Regression: the registry T2I builder drives the AR prompt + max_tokens.
+
+    The AR stage must receive the <|image start|>W*H<|image token|> envelope,
+    the structural additional_information, and max_tokens sized from the AR
+    grid (ar_height * (ar_width + 1) + 1) instead of SamplingParams' default.
+    """
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 1,
+            "size": "1024x1024",
+            "seed": 7,
+        },
+    )
+    assert response.status_code == 200, response.text
+
+    engine = mammoth_moda2_test_client.app.state.engine_client
+    captured_prompt = engine.captured_prompt
+    assert captured_prompt["prompt"].endswith("<|image start|>64*64<|image token|>")
+    addi = captured_prompt["additional_information"]
+    assert addi["omni_task"] == ["t2i"]
+    assert addi["ar_width"] == [64]
+    assert addi["ar_height"] == [64]
+
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert len(captured) == 2
+    assert captured[0].max_tokens == 64 * 65 + 1
+    assert captured[0].seed == 7
+
+
+def test_generate_images_mammoth_moda2_extra_args_routed_to_both_stages(mammoth_moda2_test_client):
+    """Regression: declared extra_body params reach AR and DiT extra_args.
+
+    text_guidance_scale / cfg_range / num_inference_steps are declared by the
+    MammothModa2 registry spec; the DiT stage (llm-typed) reads them from
+    sampling_params.extra_args, so the serving layer must route them there.
+    """
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 1,
+            "size": "1024x1024",
+            "seed": 7,
+            "text_guidance_scale": 4.5,
+            "cfg_range": [0.0, 0.8],
+            "num_inference_steps": 30,
+        },
+    )
+    assert response.status_code == 200, response.text
+
+    engine = mammoth_moda2_test_client.app.state.engine_client
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert len(captured) == 2
+    for stage_params in captured:
+        extra_args = getattr(stage_params, "extra_args", None) or {}
+        assert extra_args.get("text_guidance_scale") == 4.5
+        assert extra_args.get("cfg_range") == [0.0, 0.8]
+        assert extra_args.get("num_inference_steps") == 30
+
+
+def test_generate_images_mammoth_moda2_migrated_topology_keeps_t2i_builder(mammoth_moda2_test_client):
+    """Regression: the T2I builder must stay active after #7134's migration.
+
+    Once MammothModa2's DiT stage is migrated to the shared diffusion runtime
+    (``stage_type == "diffusion"``, see #7134), the pipeline contains a
+    classical diffusion stage while stage 0 still needs the AR envelope/grid
+    metadata. An earlier guard skipped the registry T2I builder whenever any
+    diffusion stage was present, which would silently drop the
+    <|image start|>W*H<|image token|> envelope and additional_information
+    after the migration. Dispatch must be by resolved model class, not by
+    pipeline topology.
+    """
+    engine = mammoth_moda2_test_client.app.state.engine_client
+    # Simulate the post-#7134 topology in place: stage 1 becomes a classical
+    # diffusion stage (final_output metadata moves onto it unchanged).
+    engine.engine.stage_configs[1] = SimpleNamespace(
+        stage_type="diffusion",
+        is_comprehension=False,
+        final_output=True,
+        final_output_type="image",
+        engine_args=SimpleNamespace(model_arch="MammothModa2ForConditionalGeneration"),
+    )
+    mammoth_moda2_test_client.app.state.stage_configs = list(engine.engine.stage_configs)
+
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 1,
+            "size": "1024x1024",
+            "seed": 7,
+        },
+    )
+    assert response.status_code == 200, response.text
+
+    captured_prompt = engine.captured_prompt
+    assert captured_prompt["prompt"].endswith("<|image start|>64*64<|image token|>"), (
+        "the registry T2I builder must still wrap the AR prompt when the "
+        "pipeline contains a classical diffusion stage (post-#7134 topology)"
+    )
+    addi = captured_prompt["additional_information"]
+    assert addi["omni_task"] == ["t2i"]
+    assert addi["ar_width"] == [64]
+    assert addi["ar_height"] == [64]
+
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert captured[0].max_tokens == 64 * 65 + 1
+
+
+def test_image_edits_mammoth_moda2_llm_dit_pipeline_rejected(mammoth_moda2_test_client):
+    """Regression: /v1/images/edits stays diffusion-only for now.
+
+    MammothModa2 is admitted to /v1/images/generations via the LLM-typed DiT
+    stage, but image edits (img2img) were never validated for it. Keep the
+    pre-existing diffusion-stage gate on the edits endpoint so the widened
+    generations acceptance does not silently admit it here.
+    """
+    img = make_test_image_bytes((64, 64))
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img)],
+        data={"prompt": "make it neon", "size": "512x512"},
+    )
+    assert response.status_code == 503, response.text
+    assert "No diffusion stage found" in response.json()["detail"]
+
+
+def test_is_image_generation_stage_matrix():
+    """Unit-test the stage classifier used by the image serving endpoints."""
+    from vllm_omni.entrypoints.openai.utils import is_image_generation_stage
+
+    assert is_image_generation_stage(SimpleNamespace(stage_type="diffusion"))
+    assert is_image_generation_stage(SimpleNamespace(stage_type="llm", final_output=True, final_output_type="image"))
+    assert is_image_generation_stage({"stage_type": "llm", "final_output": True, "final_output_type": "images"})
+    assert not is_image_generation_stage(SimpleNamespace(stage_type="llm", final_output=True, final_output_type="text"))
+    assert not is_image_generation_stage(SimpleNamespace(stage_type="llm", final_output=False))
+    assert not is_image_generation_stage(SimpleNamespace(stage_type="llm", final_output=True, final_output_type=None))

tests/examples/offline_inference/test_text_to_image_cli_forwarding.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""
+Offline inference regression tests: explicit CLI parallel-knob forwarding.
+
+Covers the argparse defaults -> ``build_parallel_knob_kwargs`` -> ``main()``'s
+``omni_kwargs`` assembly in examples/offline_inference/text_to_image/text_to_image.py.
+Three invariants:
+
+1. The shipped argparse must default the parallel knobs to ``None`` so unset
+   flags do not enter ``omni_kwargs`` and a deploy YAML value stays in effect.
+2. Explicitly supplied values equal to the old argparse defaults (e.g.
+   ``--ulysses-degree 1``, ``--ulysses-mode strict``) must still be forwarded
+   so they can override the deploy YAML, instead of being silently dropped
+   (review feedback on PR #7293).
+3. The AR grid ``max_tokens`` budget must stay owned by the shared
+   ``ar_grid_max_tokens`` helper used by both this example and the online
+   serving path (review feedback on PR #7293).
+
+Both invariants are asserted against the real script (its ``parse_args()`` and
+its ``build_parallel_knob_kwargs``), not a re-implemented copy. Invariant 3 is
+asserted directly against ``ar_grid_max_tokens`` itself.
+"""
+
+import sys
+from pathlib import Path
+
+import pytest
+
+_script_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "offline_inference" / "text_to_image")
+sys.path.insert(0, _script_dir)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_KNOBS = (
+    "ulysses_degree",
+    "ulysses_mode",
+    "ring_degree",
+    "cfg_parallel_size",
+    "vae_patch_parallel_size",
+)
+
+
+def _parse_real_args(cli_flags: list[str]):
+    """Run the example's real argparse against a synthetic argv."""
+    import text_to_image as t2i
+
+    saved = sys.argv
+    sys.argv = ["text_to_image.py", *cli_flags]
+    try:
+        return t2i.parse_args()
+    finally:
+        sys.argv = saved
+
+
+class TestParallelKnobForwarding:
+    """Drive the example's real argparse + real forwarding helper."""
+
+    def test_unset_knobs_forward_nothing(self):
+        """No CLI flags -> no keys, so the deploy YAML values stay in effect."""
+        import text_to_image as t2i
+
+        args = _parse_real_args([])
+        assert t2i.build_parallel_knob_kwargs(args) == {}
+
+    def test_explicit_default_equivalent_values_are_forwarded(self):
+        """--ulysses-degree 1 / --ulysses-mode strict must override the YAML."""
+        import text_to_image as t2i
+
+        args = _parse_real_args(["--ulysses-degree", "1", "--ulysses-mode", "strict"])
+        assert t2i.build_parallel_knob_kwargs(args) == {
+            "ulysses_degree": 1,
+            "ulysses_mode": "strict",
+        }
+
+        args = _parse_real_args(["--ring-degree", "1", "--cfg-parallel-size", "1", "--vae-patch-parallel-size", "1"])
+        assert t2i.build_parallel_knob_kwargs(args) == {
+            "ring_degree": 1,
+            "cfg_parallel_size": 1,
+            "vae_patch_parallel_size": 1,
+        }
+
+    def test_non_default_values_are_forwarded(self):
+        import text_to_image as t2i
+
+        args = _parse_real_args(
+            [
+                "--ulysses-degree",
+                "2",
+                "--ulysses-mode",
+                "advanced_uaa",
+                "--ring-degree",
+                "2",
+                "--cfg-parallel-size",
+                "2",
+                "--vae-patch-parallel-size",
+                "2",
+            ]
+        )
+        assert t2i.build_parallel_knob_kwargs(args) == {
+            "ulysses_degree": 2,
+            "ulysses_mode": "advanced_uaa",
+            "ring_degree": 2,
+            "cfg_parallel_size": 2,
+            "vae_patch_parallel_size": 2,
+        }
+
+    def test_script_defaults_are_none(self):
+        """The shipped argparse must default to None so unset == defer to YAML."""
+        import text_to_image as t2i
+
+        parser = t2i.build_parser()
+        defaults = {key: parser.get_default(key) for key in _KNOBS}
+        assert defaults == {key: None for key in _KNOBS}
+
+
+class TestARGridMaxTokens:
+    """The AR grid budget formula must stay shared between serving and offline.
+
+    ``ar_grid_max_tokens`` (vllm_omni.diffusion.utils.param_utils) is the
+    single owner of the ``ar_height * (ar_width + 1) + 1`` contract used by
+    both the online serving path (serving_chat.py) and this example script,
+    so a grid-contract change cannot desynchronize one copy (review feedback
+    on PR #7293).
+    """
+
+    def test_grid_budget_matches_formula(self):
+        from vllm_omni.diffusion.utils.param_utils import ar_grid_max_tokens
+
+        assert ar_grid_max_tokens(64, 64) == 64 * 65 + 1
+        assert ar_grid_max_tokens(32, 48) == 48 * 33 + 1
+
+    def test_grid_budget_none_for_absent_or_invalid_grid(self):
+        from vllm_omni.diffusion.utils.param_utils import ar_grid_max_tokens
+
+        # Callers treat None as "keep the existing max_tokens default".
+        assert ar_grid_max_tokens(0, 0) is None
+        assert ar_grid_max_tokens(0, 64) is None
+        assert ar_grid_max_tokens(64, 0) is None