From 0c88f8a0b662ae497a6485db6f3b91b79620608c Mon Sep 17 00:00:00 2001 From: TaffyOfficial <2587297563@qq.com> Date: Thu, 30 Jul 2026 15:08:12 +0800 Subject: [PATCH] [Refactor] Normalize diffusion request extra arguments in Chat serving (#5171) Signed-off-by: TaffyOfficial <2324465096@qq.com> Co-authored-by: TaffyOfficial <2324465096@qq.com> --- .../models/bagel/test_negative_prompt.py | 78 ++++ .../comfyui/test_comfyui_integration.py | 4 +- .../test_diffusion_request_utils.py | 130 +++++++ .../test_serving_chat_sampling_params.py | 350 ++++++++++++------ .../diffusion/models/bagel/pipeline_bagel.py | 3 +- .../openai/diffusion_request_utils.py | 105 ++++++ vllm_omni/entrypoints/openai/serving_chat.py | 175 ++++++--- 7 files changed, 669 insertions(+), 176 deletions(-) create mode 100644 tests/diffusion/models/bagel/test_negative_prompt.py create mode 100644 tests/entrypoints/openai_api/test_diffusion_request_utils.py create mode 100644 vllm_omni/entrypoints/openai/diffusion_request_utils.py diff --git a/tests/diffusion/models/bagel/test_negative_prompt.py b/tests/diffusion/models/bagel/test_negative_prompt.py new file mode 100644 index 000000000..37715600c --- /dev/null +++ b/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] diff --git a/tests/e2e/features/comfyui/test_comfyui_integration.py b/tests/e2e/features/comfyui/test_comfyui_integration.py index 2024cce8b..bfb26e897 100644 --- a/tests/e2e/features/comfyui/test_comfyui_integration.py +++ b/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 diff --git a/tests/entrypoints/openai_api/test_diffusion_request_utils.py b/tests/entrypoints/openai_api/test_diffusion_request_utils.py new file mode 100644 index 000000000..0b3091f24 --- /dev/null +++ b/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"}}) diff --git a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py index a6c7baffe..de3158d43 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py +++ b/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 diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 9786e2b47..6ca306c60 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/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"], diff --git a/vllm_omni/entrypoints/openai/diffusion_request_utils.py b/vllm_omni/entrypoints/openai/diffusion_request_utils.py new file mode 100644 index 000000000..84fec8a3a --- /dev/null +++ b/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, + } diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index a0bce7894..7c807cfc9 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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]