[Refactor] Normalize diffusion request extra arguments in Chat serving (#5171)

Signed-off-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: TaffyOfficial <2324465096@qq.com>

tests/diffusion/models/bagel/test_negative_prompt.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+from pytest_mock import MockerFixture
+
+from vllm_omni.diffusion.models.bagel.pipeline_bagel import BagelPipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _NegativePromptPreparedError(Exception):
+    pass
+
+
+@pytest.mark.parametrize(
+    ("prompt_negative", "extra_negative", "expected"),
+    [
+        ("prompt negative", "legacy negative", "prompt negative"),
+        (None, "legacy negative", "legacy negative"),
+    ],
+    ids=["prompt-owner", "legacy-extra-args"],
+)
+def test_forward_reads_negative_prompt_from_prompt_with_legacy_fallback(
+    mocker: MockerFixture,
+    prompt_negative: str | None,
+    extra_negative: str,
+    expected: str,
+) -> None:
+    pipeline = object.__new__(BagelPipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.od_config = mocker.Mock(dtype=torch.float32)
+    pipeline.tokenizer = mocker.sentinel.tokenizer
+    pipeline.new_token_ids = {}
+    pipeline.language_model = mocker.Mock(vocab_size=10)
+    pipeline.image_processor = None
+    pipeline.vae = None
+
+    bagel = mocker.MagicMock()
+    bagel.max_latent_size = 32
+    bagel.latent_downsample = 8
+    bagel.config.llm_config.num_hidden_layers = 1
+    prepared_prompts: list[str] = []
+
+    def prepare_prompts(*, prompts, **_):
+        prepared_prompts.append(prompts[0])
+        if len(prepared_prompts) == 1:
+            return {"packed_text_ids": torch.tensor([1])}, [1], [1]
+        raise _NegativePromptPreparedError
+
+    bagel.prepare_prompts.side_effect = prepare_prompts
+    pipeline.bagel = bagel
+
+    prompt = {
+        "prompt": "positive",
+        "negative_prompt": prompt_negative,
+    }
+    request = DiffusionRequestBatch(
+        requests=[
+            OmniDiffusionRequest(
+                prompt=prompt,
+                sampling_params=OmniDiffusionSamplingParams(
+                    extra_args={"negative_prompt": extra_negative},
+                ),
+                request_id="test",
+            )
+        ]
+    )
+
+    with pytest.raises(_NegativePromptPreparedError):
+        pipeline.forward(request)
+
+    assert prepared_prompts == ["positive", expected]

tests/e2e/features/comfyui/test_comfyui_integration.py

@@ -35,7 +35,7 @@ from vllm.outputs import CompletionOutput, RequestOutput
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand
-from vllm_omni.inputs.data import OmniSamplingParams
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
@@ -417,7 +417,7 @@ def mock_async_omni(
     mock_instance.stage_configs = server_case.stage_configs
     mock_instance.output_modalities = _build_output_modalities(server_case.stage_configs)
     mock_instance.default_sampling_params_list = [
-        SamplingParams() if _stage_type(stage) != "diffusion" else mocker.MagicMock()
+        OmniDiffusionSamplingParams() if _stage_type(stage) == "diffusion" else SamplingParams()
         for stage in server_case.stage_configs
     ]
     mock_instance.errored = False

tests/entrypoints/openai_api/test_diffusion_request_utils.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.entrypoints.openai.diffusion_request_utils import (
+    apply_normalized_diffusion_request_extra_args,
+    normalize_diffusion_request_args,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+REQUEST_SOURCES = (
+    "root.declared",
+    "root.extra_args",
+    "root.extra_params",
+    "nested.declared",
+    "nested.extra_args",
+    "nested.extra_params",
+)
+
+
+def _source_args(source: str, value: object) -> dict[str, object]:
+    scope, container = source.split(".")
+    payload = {"solver": value} if container == "declared" else {container: {"solver": value}}
+    kwargs: dict[str, object] = {scope: payload}
+    if container == "declared":
+        kwargs["registered_extra_fields"] = {"solver"}
+    return kwargs
+
+
+@pytest.mark.parametrize("source", REQUEST_SOURCES)
+@pytest.mark.parametrize(("value", "expected"), [("ddim", "ddim"), (None, "euler")])
+def test_normalizer_routes_values_without_erasing_defaults(
+    source: str,
+    value: object,
+    expected: str,
+) -> None:
+    normalized, _ = normalize_diffusion_request_args(**_source_args(source, value))
+    sampling_params = OmniDiffusionSamplingParams(
+        extra_args={"solver": "euler", "stage_default": True},
+    )
+    apply_normalized_diffusion_request_extra_args(sampling_params, normalized)
+
+    assert sampling_params.extra_args == {"solver": expected, "stage_default": True}
+
+
+def test_non_overlapping_legacy_and_canonical_values_are_merged(mocker) -> None:
+    warning_once = mocker.patch("vllm_omni.entrypoints.openai.diffusion_request_utils.logger.warning_once")
+
+    normalized, _ = normalize_diffusion_request_args(
+        root={"extra_args": {"cfg_text_scale": 7.0}},
+        nested={"extra_params": {"sample_solver": "euler"}},
+    )
+
+    assert normalized == {"cfg_text_scale": 7.0, "sample_solver": "euler"}
+    warning_once.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    "kwargs, expected_sources",
+    [
+        (
+            {
+                "root": {"seed": 1, "extra_args": {"seed": 2}},
+                "serving_root_fields": {"seed"},
+            },
+            "request.seed, request.extra_args.seed",
+        ),
+        (
+            {
+                "root": {"cfg_scale": 7.0},
+                "nested": {"true_cfg_scale": 2.0},
+                "serving_root_fields": {"true_cfg_scale"},
+                "root_field_aliases": {"cfg_scale": "true_cfg_scale"},
+            },
+            "request.cfg_scale, request.extra_body.true_cfg_scale",
+        ),
+        (
+            {
+                "root": {
+                    "extra_args": {"solver": "euler"},
+                    "extra_params": {"solver": "ddim"},
+                }
+            },
+            "request.extra_args.solver, request.extra_params.solver",
+        ),
+        (
+            {
+                "root": {"cfg_text_scale": 7.0},
+                "nested": {"extra_args": {"cfg_text_scale": 8.0}},
+                "registered_extra_fields": {"cfg_text_scale"},
+            },
+            "request.cfg_text_scale, request.extra_body.extra_args.cfg_text_scale",
+        ),
+    ],
+)
+def test_duplicate_sources_are_rejected(
+    kwargs: dict[str, object],
+    expected_sources: str,
+    mocker,
+) -> None:
+    warning_once = mocker.patch("vllm_omni.entrypoints.openai.diffusion_request_utils.logger.warning_once")
+
+    with pytest.raises(ValueError, match=expected_sources):
+        normalize_diffusion_request_args(**kwargs)
+
+    warning_once.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("scope", "name"),
+    (
+        ("root", "extra_args"),
+        ("root", "extra_params"),
+        ("nested", "extra_args"),
+        ("nested", "extra_params"),
+    ),
+)
+def test_public_containers_must_be_objects(scope: str, name: str) -> None:
+    with pytest.raises(ValueError, match="must be a JSON object"):
+        normalize_diffusion_request_args(**{scope: {name: ["not", "an", "object"]}})
+
+
+def test_container_keys_must_be_strings() -> None:
+    with pytest.raises(ValueError, match="must be a JSON object with string keys"):
+        normalize_diffusion_request_args(root={"extra_args": {1: "invalid"}})

tests/entrypoints/openai_api/test_serving_chat_sampling_params.py

@@ -6,10 +6,12 @@ Tests that standard OpenAI API parameters (max_tokens, temperature, etc.)
 are correctly applied to the comprehension stage while preserving YAML defaults.
 """
 
+import asyncio
 from types import SimpleNamespace
 
 import pytest
 from pytest_mock import MockerFixture
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.sampling_params import SamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -86,6 +88,233 @@ def serving_chat(mock_engine_client):
     return instance
 
 
+def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_chat):
+    serving_chat._diffusion_extra_body_params = frozenset({"cfg_text_scale", "negative_prompt"})
+    request = ChatCompletionRequest(
+        model="test",
+        messages=[],
+        modalities=["image"],
+        num_inference_steps=7,
+        size="768x512",
+        negative_prompt="avoid blur",
+        lora={"name": "adapter"},
+        cfg_text_scale=7.0,
+        extra_body={"extra_args": {"sample_solver": "euler"}},
+    )
+
+    normalized_extra_args, diffusion_request_args = serving_chat._normalize_diffusion_request_args(request)
+
+    assert normalized_extra_args == {
+        "cfg_text_scale": 7.0,
+        "sample_solver": "euler",
+    }
+    assert diffusion_request_args["cfg_text_scale"] == 7.0
+    assert diffusion_request_args["num_inference_steps"] == 7
+    assert diffusion_request_args["size"] == "768x512"
+    assert diffusion_request_args["negative_prompt"] == "avoid blur"
+    assert diffusion_request_args["lora"] == {"name": "adapter"}
+    assert diffusion_request_args["modalities"] == ["image"]
+
+
+def test_unknown_root_extra_does_not_claim_canonical_extra(serving_chat):
+    serving_chat._diffusion_extra_body_params = frozenset()
+    request = ChatCompletionRequest(
+        model="test",
+        messages=[],
+        pipeline_option="ignored-root-value",
+        extra_args={"pipeline_option": "canonical-value"},
+    )
+
+    normalized_extra_args, diffusion_request_args = serving_chat._normalize_diffusion_request_args(request)
+
+    assert normalized_extra_args == {
+        "pipeline_option": "canonical-value",
+    }
+    assert "pipeline_option" not in diffusion_request_args
+
+
+def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(serving_chat):
+    serving_chat._diffusion_extra_body_params = frozenset()
+    request = ChatCompletionRequest(
+        model="test",
+        messages=[],
+        cfg_scale=7.0,
+    )
+
+    normalized_extra_args, diffusion_request_args = serving_chat._normalize_diffusion_request_args(request)
+
+    assert normalized_extra_args == {}
+    assert diffusion_request_args == {"true_cfg_scale": 7.0}
+
+
+@pytest.mark.parametrize(
+    ("registered", "request_kwargs"),
+    [
+        ({"cfg_scale"}, {"cfg_scale": 7.0, "extra_args": {"cfg_scale": 8.0}}),
+        (set(), {"cfg_scale": 7.0, "true_cfg_scale": 2.0}),
+    ],
+    ids=["registered-model-extra", "unregistered-common-alias"],
+)
+def test_cfg_scale_owner_conflicts_are_rejected(serving_chat, registered, request_kwargs):
+    serving_chat._diffusion_extra_body_params = frozenset(registered)
+    request = ChatCompletionRequest(model="test", messages=[], **request_kwargs)
+
+    with pytest.raises(ValueError, match="provided more than once"):
+        serving_chat._normalize_diffusion_request_args(request)
+
+
+@pytest.mark.parametrize("diffusion_mode", [True, False], ids=["pure", "mixed"])
+def test_duplicate_extras_return_the_same_bad_request_before_dispatch(
+    serving_chat,
+    mocker: MockerFixture,
+    diffusion_mode: bool,
+):
+    serving_chat._diffusion_mode = diffusion_mode
+    serving_chat._diffusion_extra_body_params = frozenset({"sample_solver"})
+    serving_chat.engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
+    check_model = mocker.patch.object(serving_chat, "_check_model", new=mocker.AsyncMock())
+    pure_dispatch = mocker.patch.object(
+        serving_chat,
+        "_create_diffusion_chat_completion",
+        new=mocker.AsyncMock(),
+    )
+    request = ChatCompletionRequest(
+        model="test",
+        messages=[],
+        sample_solver="euler",
+        extra_args={"sample_solver": "ddim"},
+    )
+
+    response = asyncio.run(serving_chat._create_chat_completion(request))
+
+    assert response.error.code == 400
+    assert response.error.message == (
+        'Diffusion request parameters were provided more than once: "sample_solver": '
+        "request.sample_solver, request.extra_args.sample_solver."
+    )
+    check_model.assert_not_awaited()
+    pure_dispatch.assert_not_awaited()
+
+
+def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat, mocker: MockerFixture):
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    captured: dict[str, object] = {}
+
+    async def generate(**kwargs):
+        captured.update(kwargs)
+        if False:
+            yield None
+
+    serving_chat._diffusion_model_name = "test"
+    serving_chat._diffusion_engine = SimpleNamespace(
+        stage_configs=[SimpleNamespace(stage_type="diffusion")],
+        default_sampling_params_list=[
+            OmniDiffusionSamplingParams(extra_args={"solver": "euler", "stage_default": True}),
+        ],
+        generate=generate,
+    )
+    serving_chat._diffusion_mode = True
+    serving_chat._diffusion_extra_body_params = frozenset({"cfg_scale"})
+    serving_chat._extract_diffusion_prompt_and_media = mocker.Mock(return_value=("prompt", [], [], []))
+    request = ChatCompletionRequest(
+        model="test",
+        messages=[],
+        modalities=["image"],
+        num_inference_steps=7,
+        size="768x512",
+        negative_prompt="avoid blur",
+        cfg_scale=7.0,
+        true_cfg_scale=2.0,
+        extra_body={"extra_args": {"solver": "ddim"}},
+    )
+
+    response = asyncio.run(serving_chat._create_chat_completion(request))
+
+    (sampling_params,) = captured["sampling_params_list"]
+    assert sampling_params.extra_args == {
+        "cfg_scale": 7.0,
+        "solver": "ddim",
+        "stage_default": True,
+    }
+    assert sampling_params.true_cfg_scale == 2.0
+    assert sampling_params.num_inference_steps == 7
+    assert (sampling_params.height, sampling_params.width) == (512, 768)
+    assert captured["prompt"]["negative_prompt"] == "avoid blur"
+    assert response.error.message == "No output generated from AsyncOmni"
+
+
+def test_mixed_consumer_keeps_root_common_args_with_nested_extras(serving_chat, mocker: MockerFixture):
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    serving_chat.engine_client.stage_configs = [
+        SimpleNamespace(stage_type="llm", is_comprehension=True),
+        SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+    ]
+    serving_chat.engine_client.default_sampling_params_list = [
+        SamplingParams(),
+        OmniDiffusionSamplingParams(),
+    ]
+    serving_chat.engine_client.output_modalities = ["image"]
+    serving_chat.engine_client.errored = False
+
+    captured: dict[str, object] = {}
+
+    async def results():
+        if False:
+            yield None
+
+    def generate(**kwargs):
+        captured.update(kwargs)
+        return results()
+
+    serving_chat.engine_client.generate = generate
+    serving_chat.__dict__.update(
+        _diffusion_mode=False,
+        _diffusion_extra_body_params=frozenset(),
+        models=SimpleNamespace(model_name=lambda _: "test"),
+        renderer=SimpleNamespace(get_tokenizer=lambda: object()),
+        online_renderer=SimpleNamespace(validate_chat_template=lambda **_: None),
+        parser_cls=None,
+        use_harmony=False,
+        enable_auto_tools=False,
+        exclude_tools_when_tool_choice_none=False,
+        trust_request_chat_template=True,
+        chat_template=None,
+        chat_template_content_format="auto",
+    )
+    mocker.patch.multiple(
+        serving_chat,
+        _check_model=mocker.AsyncMock(return_value=None),
+        _maybe_get_adapters=mocker.Mock(return_value=None),
+        _effective_chat_template_kwargs=mocker.Mock(return_value={}),
+        _preprocess_chat=mocker.AsyncMock(return_value=([], [{"prompt": "raw"}])),
+        _base_request_id=mocker.Mock(return_value="test"),
+        _extract_diffusion_prompt_and_images_from_messages=mocker.Mock(return_value=("prompt", [])),
+        _log_inputs=mocker.Mock(),
+        chat_completion_full_generator=mocker.AsyncMock(return_value="done"),
+    )
+    request = ChatCompletionRequest(
+        model="test",
+        messages=[{"role": "user", "content": "prompt"}],
+        modalities=["image"],
+        num_inference_steps=7,
+        size="768x512",
+        negative_prompt="avoid blur",
+        extra_body={"extra_args": {"solver": "euler"}},
+    )
+
+    assert asyncio.run(serving_chat._create_chat_completion(request)) == "done"
+
+    sampling_params_list = captured["sampling_params_list"]
+    assert sampling_params_list[0].extra_args == {"target_h": 512, "target_w": 768}
+    diffusion_params = sampling_params_list[1]
+    assert diffusion_params.num_inference_steps == 7
+    assert (diffusion_params.height, diffusion_params.width) == (512, 768)
+    assert diffusion_params.extra_args == {"solver": "euler"}
+    assert captured["prompt"]["negative_prompt"] == "avoid blur"
+
+
 @pytest.fixture
 def mock_request(mocker: MockerFixture):
     """Create a mock request with all OpenAI sampling params set to None."""
@@ -616,124 +845,3 @@ class TestResolveHeightWidth:
         h, w = OmniOpenAIServingChat._resolve_height_width_from_extra_body({"size": "invalid"})
         assert h is None
         assert w is None
-
-
-# =============================================================================
-# Tests for _apply_request_overrides with GLM-Image (target_h/w injection)
-# =============================================================================
-
-
-class TestApplyRequestOverridesGLMImage:
-    """Test target_h/w injection for GLM-Image AR stage.
-
-    max_tokens is NOT computed dynamically — it comes from the deploy YAML
-    default (e.g. 4353). _apply_request_overrides only injects target_h/w
-    into extra_args so the model can build M-RoPE position grids.
-    """
-
-    @pytest.fixture
-    def glm_serving_chat(self, mock_engine_client, mocker: MockerFixture):
-        from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
-        instance = object.__new__(OmniOpenAIServingChat)
-        instance.engine_client = mock_engine_client
-        instance._extract_diffusion_prompt_and_images_from_messages = mocker.MagicMock(return_value=("a cat", []))
-        return instance
-
-    @pytest.fixture
-    def glm_request(self, mocker: MockerFixture):
-        req = mocker.MagicMock()
-        req.temperature = None
-        req.top_p = None
-        req.top_k = None
-        req.max_tokens = None
-        req.min_tokens = None
-        req.seed = None
-        req.ignore_eos = None
-        req.stop = None
-        req.stop_token_ids = None
-        req.frequency_penalty = None
-        req.presence_penalty = None
-        req.extra_body = {"height": 1024, "width": 1024}
-        req.model_fields_set = set()
-        return req
-
-    def test_t2i_injects_target_h_w(self, glm_serving_chat, glm_request, default_comprehension_params):
-        """t2i mode: target_h/w injected into extra_args, max_tokens unchanged."""
-        result = glm_serving_chat._apply_request_overrides(default_comprehension_params, glm_request)
-        assert result.extra_args["target_h"] == 1024
-        assert result.extra_args["target_w"] == 1024
-        # max_tokens stays at YAML default (not dynamically computed)
-        assert result.max_tokens == 4353
-
-    def test_i2i_injects_target_h_w(
-        self, glm_serving_chat, glm_request, default_comprehension_params, mocker: MockerFixture
-    ):
-        """i2i mode: target_h/w injected, max_tokens unchanged."""
-        glm_serving_chat._extract_diffusion_prompt_and_images_from_messages = mocker.MagicMock(
-            return_value=("edit this", ["fake_image"])
-        )
-        result = glm_serving_chat._apply_request_overrides(default_comprehension_params, glm_request)
-        assert result.extra_args["target_h"] == 1024
-        assert result.extra_args["target_w"] == 1024
-        # max_tokens stays at YAML default regardless of t2i/i2i
-        assert result.max_tokens == 4353
-
-    def test_user_max_tokens_preserved(self, glm_serving_chat, glm_request, default_comprehension_params):
-        """User-provided max_tokens is respected (not overridden by dynamic computation)."""
-        glm_request.max_tokens = 500
-        glm_request.model_fields_set = {"max_tokens"}
-
-        result = glm_serving_chat._apply_request_overrides(default_comprehension_params, glm_request)
-        assert result.max_tokens == 500
-        assert result.extra_args["target_h"] == 1024
-        assert result.extra_args["target_w"] == 1024
-
-    def test_no_height_width_preserves_default(
-        self, glm_serving_chat, mocker: MockerFixture, default_comprehension_params
-    ):
-        """When no height/width in extra_body, keep YAML default max_tokens, no target_h/w."""
-        req = mocker.MagicMock()
-        req.temperature = None
-        req.top_p = None
-        req.top_k = None
-        req.max_tokens = None
-        req.min_tokens = None
-        req.seed = None
-        req.ignore_eos = None
-        req.stop = None
-        req.stop_token_ids = None
-        req.frequency_penalty = None
-        req.presence_penalty = None
-        req.extra_body = {}
-        req.model_fields_set = set()
-
-        result = glm_serving_chat._apply_request_overrides(default_comprehension_params, req)
-        assert result.max_tokens == 4353  # YAML default
-        # No target_h/w injected when dimensions not provided
-        assert not result.extra_args or "target_h" not in (result.extra_args or {})
-
-    def test_size_string_parsed_for_glm_image(
-        self, glm_serving_chat, mocker: MockerFixture, default_comprehension_params
-    ):
-        """'size' in extra_body is parsed as fallback for height/width."""
-        req = mocker.MagicMock()
-        req.temperature = None
-        req.top_p = None
-        req.top_k = None
-        req.max_tokens = None
-        req.min_tokens = None
-        req.seed = None
-        req.ignore_eos = None
-        req.stop = None
-        req.stop_token_ids = None
-        req.frequency_penalty = None
-        req.presence_penalty = None
-        req.extra_body = {"size": "512x512"}
-        req.model_fields_set = set()
-
-        result = glm_serving_chat._apply_request_overrides(default_comprehension_params, req)
-        assert result.extra_args["target_h"] == 512
-        assert result.extra_args["target_w"] == 512
-        # max_tokens stays at YAML default (not dynamically computed)
-        assert result.max_tokens == 4353

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -610,7 +610,8 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
             # cfg_text_context: update with negative prompt (no text condition).
             # When empty, keep cfg_text_context as-is (kv_lens=0) to match
             # original BAGEL.
-            neg_prompt = extra_args.get("negative_prompt", "")
+            prompt_negative = first_prompt.get("negative_prompt") if isinstance(first_prompt, dict) else None
+            neg_prompt = prompt_negative if prompt_negative is not None else extra_args.get("negative_prompt", "")
             if neg_prompt:
                 neg_input, neg_newlens, neg_rope = self.bagel.prepare_prompts(
                     curr_kvlens=cfg_text_context["kv_lens"],

vllm_omni/entrypoints/openai/diffusion_request_utils.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from collections.abc import Collection, Mapping
+
+from vllm.logger import init_logger
+
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+logger = init_logger(__name__)
+
+
+def _request_mapping(name: str, value: object | None) -> dict[str, object]:
+    if value is None:
+        return {}
+    if not isinstance(value, Mapping):
+        raise ValueError(f"{name} must be a JSON object.")
+    if any(not isinstance(key, str) for key in value):
+        raise ValueError(f"{name} must be a JSON object with string keys.")
+    return dict(value)
+
+
+def normalize_diffusion_request_args(
+    *,
+    root: object | None = None,
+    nested: object | None = None,
+    explicit_root_args: object | None = None,
+    serving_root_fields: Collection[str] = (),
+    registered_extra_fields: Collection[str] = (),
+    root_field_aliases: Mapping[str, str] | None = None,
+) -> tuple[dict[str, object], dict[str, object]]:
+    """Validate request sources and project diffusion consumer arguments."""
+    root_args = _request_mapping("request", root)
+    root_args.update(_request_mapping("explicit_root_args", explicit_root_args))
+    nested_args = _request_mapping("request.extra_body", nested)
+
+    registered = set(registered_extra_fields)
+    aliases = {alias: canonical for alias, canonical in (root_field_aliases or {}).items() if alias not in registered}
+    serving_fields = set(serving_root_fields) | aliases.keys()
+    declared_fields = registered - serving_fields
+    consumer_fields = serving_fields | declared_fields
+
+    root_declared = {key: root_args[key] for key in declared_fields if key in root_args}
+    nested_declared = {key: nested_args[key] for key in declared_fields if key in nested_args}
+    canonical = _request_mapping("extra_args", root_args.get("extra_args"))
+    legacy = _request_mapping("extra_params", root_args.get("extra_params"))
+    nested_canonical = _request_mapping("extra_body.extra_args", nested_args.get("extra_args"))
+    nested_legacy = _request_mapping("extra_body.extra_params", nested_args.get("extra_params"))
+
+    sources_by_key: dict[str, list[str]] = {}
+    sources = (
+        ((root_args.keys() & serving_fields) | root_declared.keys(), "request"),
+        ((nested_args.keys() & serving_fields) | nested_declared.keys(), "request.extra_body"),
+        (canonical, "request.extra_args"),
+        (legacy, "request.extra_params"),
+        (nested_canonical, "request.extra_body.extra_args"),
+        (nested_legacy, "request.extra_body.extra_params"),
+    )
+    for keys, prefix in sources:
+        for key in sorted(keys):
+            sources_by_key.setdefault(aliases.get(key, key), []).append(f"{prefix}.{key}")
+
+    conflicts = {key: list(dict.fromkeys(paths)) for key, paths in sources_by_key.items() if len(set(paths)) > 1}
+    if conflicts:
+        details = "; ".join(f'"{key}": {", ".join(conflicts[key])}' for key in sorted(conflicts))
+        raise ValueError(f"Diffusion request parameters were provided more than once: {details}.")
+
+    if root_args.get("extra_params") is not None or nested_args.get("extra_params") is not None:
+        logger.warning_once(
+            "extra_params is deprecated; use extra_args for model-specific diffusion request parameters."
+        )
+
+    # Duplicate request keys were rejected above, so ordering is not precedence.
+    normalized_extra_args = {
+        **nested_declared,
+        **root_declared,
+        **legacy,
+        **nested_legacy,
+        **canonical,
+        **nested_canonical,
+    }
+    normalized_extra_args = {key: value for key, value in normalized_extra_args.items() if value is not None}
+
+    request_args = {
+        **{key: nested_args[key] for key in consumer_fields if key in nested_args},
+        **{key: root_args[key] for key in consumer_fields if key in root_args},
+    }
+    for alias, canonical_name in aliases.items():
+        if alias in request_args:
+            request_args[canonical_name] = request_args.pop(alias)
+    return normalized_extra_args, request_args
+
+
+def apply_normalized_diffusion_request_extra_args(
+    sampling_params: OmniDiffusionSamplingParams,
+    normalized_extra_args: Mapping[str, object],
+) -> None:
+    """Overlay request extras without discarding stage defaults."""
+    if normalized_extra_args:
+        sampling_params.extra_args = {
+            **(sampling_params.extra_args or {}),
+            **normalized_extra_args,
+        }

vllm_omni/entrypoints/openai/serving_chat.py

@@ -25,6 +25,10 @@ from vllm.entrypoints.chat_utils import (
 
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.openai.diffusion_request_utils import (
+    apply_normalized_diffusion_request_extra_args,
+    normalize_diffusion_request_args,
+)
 from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse
 from vllm_omni.entrypoints.utils import coerce_param_message_types
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
@@ -33,7 +37,10 @@ from vllm_omni.metrics.modality import (
     observe_audio_first_packet,
     observe_audio_streaming_finalize,
 )
-from vllm_omni.model_extras import get_extra_body_params, get_extra_output_params
+from vllm_omni.model_extras import (
+    get_extra_body_params,
+    get_extra_output_params,
+)
 
 try:
     import soundfile
@@ -95,7 +102,10 @@ from vllm.utils.collection_utils import as_list
 from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
-from vllm_omni.entrypoints.openai.image_api_utils import encode_image_base64_with_compression, validate_layered_layers
+from vllm_omni.entrypoints.openai.image_api_utils import (
+    encode_image_base64_with_compression,
+    validate_layered_layers,
+)
 from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse
 from vllm_omni.entrypoints.openai.protocol.audio import (
     DEFAULT_AUDIO_FORMAT,
@@ -152,6 +162,33 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
     _supported_speakers: set[str] | None = None
     _diffusion_extra_body_params: frozenset[str] | None = None
     _diffusion_extra_output_params: frozenset[str] | None = None
+    # Common fields stay at the request root, but still participate in
+    # conflicts with model-specific compatibility forms.
+    _diffusion_common_root_fields = frozenset(
+        {
+            "height",
+            "width",
+            "num_outputs_per_prompt",
+            "seed",
+            "num_inference_steps",
+            "guidance_scale",
+            "true_cfg_scale",
+            "num_frames",
+            "guidance_scale_2",
+            "layers",
+            "resolution",
+        }
+    )
+    _diffusion_existing_control_fields = frozenset(
+        {
+            "size",
+            "negative_prompt",
+            "lora",
+            "modalities",
+        }
+    )
+    _diffusion_root_field_aliases = {"cfg_scale": "true_cfg_scale"}
+    _diffusion_serving_root_fields = _diffusion_common_root_fields | _diffusion_existing_control_fields
 
     # Harmony flag (always False for vllm-omni models)
     use_harmony: bool = False
@@ -291,6 +328,30 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         self._diffusion_extra_body_params = params
         return params
 
+    def _normalize_diffusion_request_args(
+        self,
+        request: ChatCompletionRequest,
+    ) -> tuple[dict[str, object], dict[str, Any]]:
+        """Normalize extras and preserve validated common root fields."""
+        root = dict(request.model_extra or {})
+        nested = root.pop("extra_body", None)
+        if nested is None:
+            nested = getattr(request, "extra_body", None)
+        explicit = getattr(request, "model_fields_set", None)
+        if explicit is None:
+            explicit = getattr(request, "__fields_set__", ())
+        explicit_root_args = {
+            key: getattr(request, key) for key in explicit if key != "extra_body" and hasattr(request, key)
+        }
+        return normalize_diffusion_request_args(
+            root=root,
+            nested=nested,
+            explicit_root_args=explicit_root_args,
+            serving_root_fields=self._diffusion_serving_root_fields,
+            registered_extra_fields=self._get_diffusion_extra_body_params(),
+            root_field_aliases=self._diffusion_root_field_aliases,
+        )
+
     def _get_diffusion_extra_output_params(
         self,
         output: object,
@@ -386,9 +447,26 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         request: ChatCompletionRequest,
         raw_request: Request | None = None,
     ) -> AsyncGenerator[str, None] | ChatCompletionResponse | ErrorResponse:
+        stage_configs = getattr(self.engine_client, "stage_configs", ()) or ()
+        serves_diffusion = self._diffusion_mode or any(
+            get_stage_type(stage_config) == "diffusion" for stage_config in stage_configs
+        )
+        normalized_extra_args: dict[str, object] = {}
+        diffusion_request_args: dict[str, Any] = {}
+        if serves_diffusion:
+            try:
+                normalized_extra_args, diffusion_request_args = self._normalize_diffusion_request_args(request)
+            except ValueError as exc:
+                return self._create_error_response(str(exc), status_code=400)
+
         # Handle diffusion mode
         if self._diffusion_mode:
-            return await self._create_diffusion_chat_completion(request, raw_request)
+            return await self._create_diffusion_chat_completion(
+                request,
+                normalized_extra_args,
+                diffusion_request_args,
+                raw_request,
+            )
 
         request_timestamp = time.time()
         if raw_request is not None:
@@ -530,7 +608,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 return audio_format_check
 
         num_inference_steps = None
-        extra_body: dict[str, Any] = {}
+        extra_body = diffusion_request_args
         # Omni multistage image generation: Stage-0 (AR) should receive a clean
         # text prompt (and optional conditioning image/size) so the model's own
         # processor can construct the correct inputs.
@@ -544,12 +622,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 if not extracted_prompt:
                     return self.create_error_response("No text prompt found in messages")
 
-                # [NOTE] When sending request via openai client Python library,
-                #   `extra_body` is flattented and merged into the payload's root.
-                #   These extra fields are accessible via `model_extra` property (from Pydantic base class).
-                #   When sending raw request with curl, no flattening happens. Directly read the `extra_body` dict.
-                extra_body = getattr(request, "extra_body", None) or request.model_extra or {}
-
                 height, width = self._resolve_height_width_from_extra_body(extra_body)
 
                 num_inference_steps = extra_body.get("num_inference_steps")
@@ -668,15 +740,32 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         output_modalities,
                     )
 
+                comprehension_idx = (
+                    self._get_comprehension_stage_index()
+                    if _image_gen_height is not None and _image_gen_width is not None
+                    else None
+                )
                 # Apply user-specified overrides to diffusion stage(s) for image generation
                 for idx, sp in enumerate(sampling_params_list):
+                    if idx == comprehension_idx:
+                        extra_args = dict(getattr(sp, "extra_args", {}) or {})
+                        extra_args["target_h"] = int(_image_gen_height)
+                        extra_args["target_w"] = int(_image_gen_width)
+                        sp.extra_args = extra_args
                     if hasattr(sp, "height") and _image_gen_height is not None:
                         sp.height = _image_gen_height
                     if hasattr(sp, "width") and _image_gen_width is not None:
                         sp.width = _image_gen_width
                     if hasattr(sp, "num_inference_steps") and num_inference_steps is not None:
                         sp.num_inference_steps = num_inference_steps
-                    apply_declared_extra_args(sp, self._get_diffusion_extra_body_params(), extra_body)
+                    if isinstance(sp, OmniDiffusionSamplingParams):
+                        apply_normalized_diffusion_request_extra_args(sp, normalized_extra_args)
+                    else:
+                        apply_declared_extra_args(
+                            sp,
+                            self._get_diffusion_extra_body_params(),
+                            extra_body,
+                        )
 
                 self._log_inputs(
                     request_id,
@@ -1166,8 +1255,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         Starts with YAML defaults and only overrides fields that the user
         explicitly provided (non-None values) in the request.
 
-        For models needing spatial metadata (e.g. GLM-Image), target_h/w is
-        injected into extra_args so the runner can build M-RoPE position grids.
         max_tokens is NOT computed dynamically — it uses the deploy YAML default.
 
         Args:
@@ -1193,19 +1280,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             if (value is not None and not isinstance(value, list)) or (isinstance(value, list) and len(value) > 0):
                 setattr(params, field_name, value)
 
-        # For GLM-Image: compute max_tokens from height/width with mode-aware
-        # budgeting (t2i vs i2i).
-        extra_body = getattr(request, "extra_body", {}) or {}
-        height, width = self._resolve_height_width_from_extra_body(extra_body)
-
-        if height is not None and width is not None:
-            # Keep target size in stage-0 sampling params so runner/model can
-            # build deterministic M-RoPE grids for t2i (no MM features).
-            extra_args = dict(getattr(params, "extra_args", {}) or {})
-            extra_args["target_h"] = int(height)
-            extra_args["target_w"] = int(width)
-            params.extra_args = extra_args
-
         return params
 
     @staticmethod
@@ -3366,12 +3440,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
     async def _create_diffusion_chat_completion(
         self,
         request: ChatCompletionRequest,
+        normalized_extra_args: dict[str, object],
+        diffusion_request_args: dict[str, Any],
         raw_request: Request | None = None,
     ) -> ChatCompletionResponse | ErrorResponse:
         """Generate images via chat completion interface for diffusion models.
 
         Args:
             request: Chat completion request
+            normalized_extra_args: Validated model-specific request arguments
+            diffusion_request_args: Compatibility arguments plus validated common root fields
             raw_request: Raw FastAPI request object
 
         Returns:
@@ -3396,15 +3474,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 messages
             )
 
-            # Extract generation parameters from extra_body (preferred)
             # Reference: text_to_image.py and text_to_video.py for supported parameters
-            # [NOTE] When sending request via openai client Python library,
-            #   `extra_body` is flattented and merged into the payload's root.
-            #   These extra fields are accessible via `model_extra` property (from Pydantic base class).
-            #   When sending raw request with curl, no flattening happens. Directly read the `extra_body` dict.
-            extra_body = getattr(request, "extra_body", None)
-            if not extra_body:
-                extra_body = request.model_extra or {}
+            extra_body = diffusion_request_args
 
             # Parse size if provided (supports "1024x1024" format)
             height, width = self._resolve_height_width_from_extra_body(extra_body)
@@ -3415,7 +3486,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             # not provide a value.
             num_inference_steps = extra_body.get("num_inference_steps")
             guidance_scale = extra_body.get("guidance_scale")
-            true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale")
+            true_cfg_scale = extra_body.get("true_cfg_scale")
             seed = extra_body.get("seed")
             if seed is None:
                 seed = getattr(request, "seed", None)
@@ -3430,18 +3501,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             # Qwen-Image-Layered parameters
             layers = extra_body.get("layers")
             resolution = extra_body.get("resolution")
-
             try:
                 layers = validate_layered_layers(layers)
-            except ValueError as e:
-                return self._create_error_response(str(e), status_code=400)
+            except ValueError as exc:
+                return self._create_error_response(str(exc), status_code=400)
 
             logger.info(
                 "Diffusion chat request %s: prompt=%r, ref_images=%d, params=%s",
                 request_id,
                 prompt[:50] + "..." if len(prompt) > 50 else prompt,
                 len(reference_images),
-                {k: v for k, v in extra_body.items() if v is not None},
+                {key: value for key, value in extra_body.items() if value is not None},
             )
 
             # Decode reference images if provided
@@ -3473,7 +3543,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 gen_params.guidance_scale = guidance_scale
             if true_cfg_scale is not None:
                 gen_params.true_cfg_scale = true_cfg_scale
-            apply_declared_extra_args(gen_params, self._get_diffusion_extra_body_params(), extra_body)
+            apply_normalized_diffusion_request_extra_args(gen_params, normalized_extra_args)
             if num_frames is not None:
                 gen_params.num_frames = num_frames
             if guidance_scale_2 is not None:
@@ -3483,14 +3553,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             if resolution is not None:
                 gen_params.resolution = resolution
 
-            # Pipeline-agnostic escape hatch (mirrors ``extra_params`` on the /v1/videos
-            # endpoint in ``serving_video.py``): a single reserved ``extra_args`` dict in
-            # ``extra_body`` flows straight into ``gen_params.extra_args``, with no keys
-            # hardcoded here.
-            extra_args_body = extra_body.get("extra_args")
-            if isinstance(extra_args_body, dict):
-                gen_params.extra_args.update(extra_args_body)
-
             # Parse per-request LoRA.
             if lora_body and isinstance(lora_body, dict):
                 try:
@@ -3499,8 +3561,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         gen_params.lora_request = lora_req
                         if lora_scale is not None:
                             gen_params.lora_scale = lora_scale
-                except Exception as e:  # pragma: no cover - safeguard
-                    logger.warning("Failed to parse LoRA request: %s", e)
+                except Exception as exc:  # pragma: no cover - safeguard
+                    logger.warning("Failed to parse LoRA request: %s", exc)
 
             # Route text modality for single-stage diffusion (img2text / text2text)
             requested_modalities = extra_body.get("modalities") or []
@@ -3539,13 +3601,22 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             # Generate image or audio (e.g. AudioX) via AsyncOmni
             diffusion_engine = cast(AsyncOmni, self._diffusion_engine)
             stage_configs = list(getattr(diffusion_engine, "stage_configs", []) or [])
+            default_sampling_params_list = get_default_sampling_params_list(diffusion_engine)
             sampling_params_list = build_stage_sampling_params_list(
                 stage_configs,
-                get_default_sampling_params_list(diffusion_engine),
+                default_sampling_params_list,
                 diffusion_params=gen_params,
                 replace_diffusion_params=True,
             )
-
+            for index, (stage_config, sampling_params) in enumerate(zip(stage_configs, sampling_params_list)):
+                if get_stage_type(stage_config) != "diffusion" or index >= len(default_sampling_params_list):
+                    continue
+                default_extra_args = getattr(default_sampling_params_list[index], "extra_args", None)
+                if default_extra_args:
+                    sampling_params.extra_args = {
+                        **default_extra_args,
+                        **(getattr(sampling_params, "extra_args", None) or {}),
+                    }
             if not sampling_params_list:
                 sampling_params_list = [gen_params]