From 0fe76857acc0447419829710d50d8a482266be93 Mon Sep 17 00:00:00 2001 From: Lancer Date: Fri, 24 Apr 2026 20:00:14 +0800 Subject: [PATCH] [Bugfix] fix GLM-Image multi-stage online generation fail (#3084) Signed-off-by: Lancer --- .../test_orchestrator_kv_sender_info.py | 46 ++++++++++ .../openai_api/test_image_server.py | 84 +++++++++++++++++++ vllm_omni/engine/orchestrator.py | 27 ++++++ vllm_omni/entrypoints/openai/serving_chat.py | 37 ++++++++ 4 files changed, 194 insertions(+) diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index 7e3fe0906..ec3a42e35 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -183,6 +183,52 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info(): } +def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs(): + orchestrator = object.__new__(Orchestrator) + sender_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151}) + diffusion_stage = _DummyDiffusionStage(engine_input_source=[0]) + diffusion_stage.custom_process_input_func = lambda *_args, **_kwargs: [] + + class _AsyncQueue: + def __init__(self): + self.items = [] + + async def put(self, item): + self.items.append(item) + + orchestrator.num_stages = 2 + orchestrator.stage_clients = [sender_stage, diffusion_stage] + orchestrator._cfg_tracker = CfgCompanionTracker() + orchestrator.stage_vllm_configs = [None, None] + orchestrator.output_processors = [None, None] + orchestrator.output_async_queue = _AsyncQueue() + orchestrator.request_states = {} + orchestrator._pd_kv_params = {} + + params = OmniDiffusionSamplingParams() + req_state = OrchestratorRequestState( + request_id="req-empty", + prompt={"prompt": "hello"}, + sampling_params_list=[SamplingParams(max_tokens=4), params], + final_stage_id=1, + ) + orchestrator.request_states["req-empty"] = req_state + + output = SimpleNamespace(request_id="req-empty", finished=True) + asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-empty", 0, output, req_state)) + + assert sender_stage.engine_outputs == [output] + assert diffusion_stage.calls == [] + assert len(orchestrator.output_async_queue.items) == 1 + terminal_msg = orchestrator.output_async_queue.items[0] + assert terminal_msg["type"] == "output" + assert terminal_msg["request_id"] == "req-empty" + assert terminal_msg["stage_id"] == 1 + assert terminal_msg["finished"] is True + assert "produced no valid inputs" in terminal_msg["engine_outputs"].error + assert "req-empty" not in orchestrator.request_states + + def test_prewarm_diffusion_attaches_kv_sender_info(): orchestrator = object.__new__(Orchestrator) sender_stage = _DummySenderStage({"host": "10.0.0.3", "zmq_port": 50151}) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 4cf46f18c..7fcec2a32 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -13,16 +13,22 @@ from argparse import Namespace from types import SimpleNamespace import pytest +from fastapi import FastAPI from fastapi.testclient import TestClient from PIL import Image from pytest_mock import MockerFixture from vllm import SamplingParams +from vllm.entrypoints.openai.models.protocol import BaseModelPath +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.openai.api_server import _DiffusionServingModels, router from vllm_omni.entrypoints.openai.image_api_utils import ( encode_image_base64, parse_size, ) +from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_executor.stage_input_processors.glm_image import compute_max_tokens pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -569,6 +575,84 @@ def test_multistage_images_async_omni_construction(async_omni_test_client): assert captured[1].guidance_scale == 6.5 +def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens(): + """GLM-Image multistage requests must compute stage-0 max_tokens from size.""" + + class FakeAsyncOmniClass(AsyncOmni): + def __init__(self): + stage_configs = [ + SimpleNamespace(stage_type="llm", is_comprehension=True, model_arch="GlmImageForConditionalGeneration"), + SimpleNamespace(stage_type="diffusion", is_comprehension=False, model_arch="GlmImagePipeline"), + ] + default_sampling_params_list = [ + SamplingParams(temperature=0.1, seed=42), + OmniDiffusionSamplingParams(height=1024, width=1024), + ] + self.engine = SimpleNamespace( + stage_configs=stage_configs, + default_sampling_params_list=default_sampling_params_list, + ) + 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")] + self.od_config = SimpleNamespace(supports_multimodal_inputs=True) + + async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None): + self.captured_sampling_params_list = ( + sampling_params_list if sampling_params_list is not None else [sampling_params] + ) + self.captured_prompt = prompt + yield MockGenerationResult([img.copy() for img in self._images]) + + def __class_getitem__(cls, item): + return cls + + def get_diffusion_od_config(self): + return self.od_config + + 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", model_arch="GlmImageForConditionalGeneration"), + SimpleNamespace(stage_type="diffusion", model_arch="GlmImagePipeline"), + ] + app.state.openai_serving_models = _DiffusionServingModels( + [BaseModelPath(name="THUDM/GLM-4.5V", model_path="THUDM/GLM-4.5V")] + ) + app.state.args = Namespace( + default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}', + max_generated_image_size=1048576, + ) + client = TestClient(app) + + response = client.post( + "/v1/images/generations", + json={ + "prompt": "a coral reef", + "n": 1, + "size": "1024x1024", + "seed": 7, + }, + ) + assert response.status_code == 200 + + captured = engine.captured_sampling_params_list + assert captured is not None + assert len(captured) == 2 + assert captured[0].max_tokens == compute_max_tokens(1024, 1024, is_i2i=False) + assert captured[0].extra_args["target_h"] == 1024 + assert captured[0].extra_args["target_w"] == 1024 + assert captured[1].height == 1024 + assert captured[1].width == 1024 + + def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only_client): """Regression: image edits accepts refactored AsyncOmni without stage_list.""" img_bytes = make_test_image_bytes((16, 16)) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index ff5117934..a7edf4339 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -33,6 +33,7 @@ from vllm_omni.engine.serialization import serialize_additional_information from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics from vllm_omni.metrics.stats import StageStats from vllm_omni.metrics.utils import count_tokens_from_outputs +from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) @@ -689,6 +690,32 @@ class Orchestrator: next_stage_id, ) if isinstance(diffusion_prompt, list): + if not diffusion_prompt: + error_output = OmniRequestOutput.from_error( + req_id, + f"Stage-{stage_id} produced no valid inputs for diffusion stage-{next_stage_id}", + ) + logger.warning( + "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; " + "routing terminal error output", + req_id, + stage_id, + next_stage_id, + ) + await self.output_async_queue.put( + { + "type": "output", + "request_id": req_id, + "stage_id": next_stage_id, + "engine_outputs": error_output, + "metrics": None, + "finished": True, + } + ) + self._pd_kv_params.pop(req_id, None) + self._cfg_tracker.cleanup_parent(req_id) + self.request_states.pop(req_id, None) + return diffusion_prompt = diffusion_prompt[0] else: diffusion_prompt = req_state.prompt diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index fb96b397e..0e4919f82 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -2192,6 +2192,24 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): comprehension_idx = idx break + is_glm_image_multistage = False + for stage_cfg in stage_configs: + if isinstance(stage_cfg, dict): + engine_args = stage_cfg.get("engine_args") + model_arch = stage_cfg.get("model_arch") + else: + engine_args = getattr(stage_cfg, "engine_args", None) + model_arch = getattr(stage_cfg, "model_arch", None) + + if model_arch is None and engine_args is not None: + if isinstance(engine_args, dict): + model_arch = engine_args.get("model_arch") + else: + model_arch = getattr(engine_args, "model_arch", None) + if model_arch in {"GlmImageForConditionalGeneration", "GlmImagePipeline"}: + is_glm_image_multistage = True + break + sampling_params_list: list[Any] = [] for idx, stage_cfg in enumerate(stage_configs): stage_type = get_stage_type(stage_cfg) @@ -2240,6 +2258,25 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): default_stage_params.lora_scale = lora_scale except Exception as e: # pragma: no cover - safeguard logger.warning("Failed to parse LoRA request: %s", e) + elif ( + is_glm_image_multistage + and comprehension_idx is not None + and idx == comprehension_idx + and isinstance(default_stage_params, SamplingParams) + and height is not None + and width is not None + ): + from vllm_omni.model_executor.stage_input_processors.glm_image import compute_max_tokens + + default_stage_params.max_tokens = compute_max_tokens( + int(height), + int(width), + is_i2i=bool(reference_images), + ) + extra_args = dict(getattr(default_stage_params, "extra_args", {}) or {}) + extra_args["target_h"] = int(height) + extra_args["target_w"] = int(width) + default_stage_params.extra_args = extra_args sampling_params_list.append(default_stage_params)