[Frontend][Model] Support batch request with refined OmniDiffusionReq… (#797)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>

tests/entrypoints/openai_api/test_image_server.py

@@ -14,11 +14,13 @@ from unittest.mock import AsyncMock, Mock
 import pytest
 from fastapi.testclient import TestClient
 from PIL import Image
+from vllm import SamplingParams
 
 from vllm_omni.entrypoints.openai.image_api_utils import (
     encode_image_base64,
     parse_size,
 )
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 # Unit Tests
 
@@ -111,7 +113,7 @@ class FakeAsyncOmni:
 
     def __init__(self):
         self.stage_list = ["llm", "diffusion"]
-        self.default_sampling_params_list = [{"temperature": 0.1}, {"top_p": 0.9}]
+        self.default_sampling_params_list = [SamplingParams(temperature=0.1), OmniDiffusionSamplingParams()]
         self.captured_sampling_params_list = None
 
     async def generate(self, prompt, request_id, sampling_params_list):
@@ -129,7 +131,8 @@ def mock_async_diffusion():
 
     async def generate(**kwargs):
         # Return n PIL images wrapped in result object
-        n = kwargs.get("num_outputs_per_prompt", 1)
+        print("!!!!!!!!!!!!!!!!!!!!! kwargs", kwargs)
+        n = kwargs["sampling_params_list"][0].num_outputs_per_prompt
         images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)]
         return MockGenerationResult(images)
 
@@ -268,11 +271,11 @@ def test_generate_images_async_omni_sampling_params(async_omni_test_client):
     captured = engine.captured_sampling_params_list
     assert captured is not None
     assert len(captured) == 2
-    assert captured[0] == {"temperature": 0.1}
-    assert captured[1]["num_outputs_per_prompt"] == 2
-    assert captured[1]["height"] == 256
-    assert captured[1]["width"] == 256
-    assert captured[1]["seed"] == 7
+    assert captured[0].temperature == 0.1
+    assert captured[1].num_outputs_per_prompt == 2
+    assert captured[1].height == 256
+    assert captured[1].width == 256
+    assert captured[1].seed == 7
 
 
 def test_generate_multiple_images(test_client):
@@ -505,30 +508,11 @@ def test_parameters_passed_through(test_client, mock_async_diffusion):
 
     # Ensure generate() was called exactly once
     mock_async_diffusion.generate.assert_awaited_once()
-    call_kwargs = mock_async_diffusion.generate.call_args[1]
-    assert call_kwargs["num_inference_steps"] == 100
-    assert call_kwargs["guidance_scale"] == 7.5
-    assert call_kwargs["true_cfg_scale"] == 3.0
-    assert call_kwargs["seed"] == 42
-
-
-def test_optional_parameters_omitted(test_client, mock_async_diffusion):
-    """Verify optional parameters not passed when omitted"""
-    response = test_client.post(
-        "/v1/images/generations",
-        json={
-            "prompt": "test",
-            "size": "512x512",
-        },
-    )
-    assert response.status_code == 200
-
-    # Ensure generate() was called exactly once
-    mock_async_diffusion.generate.assert_awaited_once()
-    call_kwargs = mock_async_diffusion.generate.call_args[1]
-    assert "num_inference_steps" not in call_kwargs
-    assert "guidance_scale" not in call_kwargs
-    assert "true_cfg_scale" not in call_kwargs
+    call_kwargs = mock_async_diffusion.generate.call_args[1]["sampling_params_list"][0]
+    assert call_kwargs.num_inference_steps == 100
+    assert call_kwargs.guidance_scale == 7.5
+    assert call_kwargs.true_cfg_scale == 3.0
+    assert call_kwargs.seed == 42
 
 
 def test_model_field_omitted_works(test_client):

tests/entrypoints/test_omni_llm.py

@@ -5,6 +5,7 @@ from typing import Any
 from unittest.mock import MagicMock
 
 import pytest
+from vllm import SamplingParams
 
 from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK
 
@@ -82,8 +83,9 @@ class _FakeStage:
         self.stage_id = getattr(config, "stage_id", 0)
         self.engine_args = config.engine_args
         self.model_stage = getattr(config.engine_args, "model_stage", None)
+        self.stage_type = "llm"
         # set default sampling params
-        self.default_sampling_params = {"temperature": 1.0}
+        self.default_sampling_params = SamplingParams(temperature=1.0)
         # Allow configuring final_output and final_output_type
         self.final_output = config.final_output if hasattr(config, "final_output") else False
         self.final_output_type = getattr(config, "final_output_type", None)
@@ -637,8 +639,10 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
         }
     )
 
-    # Use dicts instead of object() for serializable sampling params
-    sampling_params_list = [{"temperature": 0.7}, {"temperature": 0.8}]
+    sampling_params_list = [
+        SamplingParams(temperature=0.7),
+        SamplingParams(temperature=0.8),
+    ]
     prompts = ["hi"]
     outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
 
@@ -722,8 +726,13 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
         }
     )
 
-    # Use dicts instead of object() for serializable sampling params
-    outputs = omni.generate(prompts=["p"], sampling_params_list=[{"temperature": 0.7}, {"temperature": 0.8}])
+    outputs = omni.generate(
+        prompts=["p"],
+        sampling_params_list=[
+            SamplingParams(temperature=0.7),
+            SamplingParams(temperature=0.8),
+        ],
+    )
     assert outputs == []
 
 
@@ -922,8 +931,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     )
 
     # Generate should handle error gracefully (log but continue)
-    # Use dict instead of object() for serializable sampling params
-    sampling_params_list = [{"temperature": 0.7}]
+    sampling_params_list = [SamplingParams(temperature=0.7)]
     outputs = omni.generate(prompts=["hi"], sampling_params_list=sampling_params_list)
     # Should return final output (error was logged but didn't stop processing)
     assert isinstance(outputs, list)