Benchmark and Bugfix for GLM-Image (#3024)

Signed-off-by: JaredforReal <w13431838023@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/entrypoints/openai_api/test_image_server.py

@@ -28,7 +28,6 @@ from vllm_omni.entrypoints.openai.image_api_utils import (
 )
 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]
 
@@ -216,7 +215,11 @@ def async_omni_test_client():
             ]
             default_sampling_params_list = [
                 SamplingParams(temperature=0.1),
-                OmniDiffusionSamplingParams(),
+                OmniDiffusionSamplingParams(
+                    num_inference_steps=4,
+                    guidance_scale=7.5,
+                    generator_device="cpu",
+                ),
             ]
             self.engine = SimpleNamespace(
                 stage_configs=stage_configs,
@@ -576,7 +579,11 @@ def test_multistage_images_async_omni_construction(async_omni_test_client):
 
 
 def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens():
-    """GLM-Image multistage requests must compute stage-0 max_tokens from size."""
+    """GLM-Image multistage: stage-0 gets target_h/w from requested size.
+
+    max_tokens comes from the deploy YAML default (upper-bound ceiling),
+    NOT computed dynamically from height/width.
+    """
 
     class FakeAsyncOmniClass(AsyncOmni):
         def __init__(self):
@@ -584,8 +591,9 @@ def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens():
                 SimpleNamespace(stage_type="llm", is_comprehension=True, model_arch="GlmImageForConditionalGeneration"),
                 SimpleNamespace(stage_type="diffusion", is_comprehension=False, model_arch="GlmImagePipeline"),
             ]
+            # YAML default max_tokens for GLM-Image AR stage (upper bound for 2048x2048 t2i)
             default_sampling_params_list = [
-                SamplingParams(temperature=0.1, seed=42),
+                SamplingParams(temperature=0.1, seed=42, max_tokens=4353),
                 OmniDiffusionSamplingParams(height=1024, width=1024),
             ]
             self.engine = SimpleNamespace(
@@ -646,7 +654,8 @@ def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens():
     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)
+    # max_tokens comes from YAML default, not computed dynamically
+    assert captured[0].max_tokens == 4353
     assert captured[0].extra_args["target_h"] == 1024
     assert captured[0].extra_args["target_w"] == 1024
     assert captured[1].height == 1024
@@ -1161,9 +1170,10 @@ def test_image_edit_ignores_mock_like_multimodal_limit(async_omni_test_client):
     assert response.status_code == 200
     captured_prompt = engine.captured_prompt
     assert captured_prompt is not None
-    processed_images = captured_prompt["multi_modal_data"]["image"]
-    assert len(processed_images) == 1
-    assert processed_images[0].size == (16, 16)
+    # Multi-stage path uses "img2img" key for single reference image
+    processed_images = captured_prompt["multi_modal_data"]["img2img"]
+    assert isinstance(processed_images, Image.Image)
+    assert processed_images.size == (16, 16)
 
 
 def test_image_edit_parameter_pass(async_omni_test_client):

tests/entrypoints/openai_api/test_serving_chat_sampling_params.py

@@ -38,7 +38,7 @@ def default_comprehension_params():
         temperature=0.4,
         top_p=0.9,
         top_k=1,
-        max_tokens=2048,
+        max_tokens=4353,
         seed=42,
         repetition_penalty=1.05,
     )
@@ -146,7 +146,7 @@ def test_preserves_yaml_defaults_when_no_request_params(serving_chat, mock_reque
     assert comprehension_params.temperature == 0.4
     assert comprehension_params.top_p == 0.9
     assert comprehension_params.top_k == 1  # YAML custom param preserved
-    assert comprehension_params.max_tokens == 2048
+    assert comprehension_params.max_tokens == 4353
     assert comprehension_params.seed == 42
     assert comprehension_params.repetition_penalty == 1.05  # YAML custom param preserved
 
@@ -190,7 +190,7 @@ def test_max_tokens_uses_yaml_default_when_not_specified(serving_chat, mock_requ
     """Test that max_tokens falls back to YAML default when not in request."""
     result = serving_chat._build_sampling_params_list_from_request(mock_request)
 
-    assert result[0].max_tokens == 2048
+    assert result[0].max_tokens == 4353
 
 
 def test_request_seed_overrides_yaml_default(serving_chat, mock_request):
@@ -588,12 +588,17 @@ class TestResolveHeightWidth:
 
 
 # =============================================================================
-# Tests for _apply_request_overrides with GLM-Image (max_tokens computation)
+# Tests for _apply_request_overrides with GLM-Image (target_h/w injection)
 # =============================================================================
 
 
 class TestApplyRequestOverridesGLMImage:
-    """Test dynamic max_tokens computation for GLM-Image AR stage."""
+    """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):
@@ -601,7 +606,6 @@ class TestApplyRequestOverridesGLMImage:
 
         instance = object.__new__(OmniOpenAIServingChat)
         instance.engine_client = mock_engine_client
-        # Mock the image extraction to return no reference images (t2i by default)
         instance._extract_diffusion_prompt_and_images_from_messages = mocker.MagicMock(return_value=("a cat", []))
         return instance
 
@@ -623,40 +627,41 @@ class TestApplyRequestOverridesGLMImage:
         req.model_fields_set = set()
         return req
 
-    def test_t2i_computes_max_tokens(self, glm_serving_chat, glm_request, default_comprehension_params):
-        """t2i mode: max_tokens computed from height/width, no reference images."""
+    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)
-        # t2i 1024x1024 = 256 + 1024 + 1 = 1281
-        assert result.max_tokens == 1281
         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_computes_fewer_tokens(
+    def test_i2i_injects_target_h_w(
         self, glm_serving_chat, glm_request, default_comprehension_params, mocker: MockerFixture
     ):
-        """i2i mode: max_tokens should be smaller than t2i for same dimensions."""
-        # Make it detect reference images
+        """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)
-        # i2i 1024x1024 = 1024 + 1 = 1025
-        assert result.max_tokens == 1025
+        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_dynamic_max_tokens_overrides_user_value(self, glm_serving_chat, glm_request, default_comprehension_params):
-        """When height/width are provided, dynamic computation overrides user max_tokens."""
+    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)
-        # Dynamic computation from height/width always wins when present
-        assert result.max_tokens == 1281
+        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."""
+        """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
@@ -673,7 +678,9 @@ class TestApplyRequestOverridesGLMImage:
         req.model_fields_set = set()
 
         result = glm_serving_chat._apply_request_overrides(default_comprehension_params, req)
-        assert result.max_tokens == 2048  # YAML default
+        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
@@ -695,5 +702,7 @@ class TestApplyRequestOverridesGLMImage:
         req.model_fields_set = set()
 
         result = glm_serving_chat._apply_request_overrides(default_comprehension_params, req)
-        # 512x512 t2i = 256 + 256 + 1 = 513
-        assert result.max_tokens == 513
+        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