[Bugfix] fix GLM-Image multi-stage online generation fail (#3084)

Signed-off-by: Lancer <maruixiang6688@gmail.com>

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})

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

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

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)