[cherry-pick] [release/v0.18.0post1] fix: validate layered image layers range (#2334) (#2349)

Signed-off-by: David Chen <530634352@qq.com>

tests/entrypoints/openai_api/test_image_server.py

@@ -615,6 +615,13 @@ def test_parameter_validation():
     with pytest.raises(ValueError):
         ImageGenerationRequest(prompt="test", guidance_scale=21.0)
 
+    # Invalid layers for layered models (must stay within the backend-supported range)
+    with pytest.raises(ValueError):
+        ImageGenerationRequest(prompt="test", layers=2)
+
+    with pytest.raises(ValueError):
+        ImageGenerationRequest(prompt="test", layers=11)
+
 
 # Pass-Through Tests
 
@@ -839,35 +846,36 @@ def test_image_edit_invalid_resolution(async_omni_test_client):
 
 
 def test_image_edit_invalid_layers(async_omni_test_client):
-    """Test that invalid layers values (< 1) are rejected with 400."""
+    """Test that layered image edits reject out-of-range layers with 400."""
     img_bytes = make_test_image_bytes((16, 16))
 
-    # Test layers = 0
+    # Test layers below the supported range
     response = async_omni_test_client.post(
         "/v1/images/edits",
         files=[("image", img_bytes)],
         data={
             "prompt": "test",
-            "layers": 0,
+            "layers": 2,
         },
     )
     assert response.status_code == 400
     detail = response.json()["detail"]
     assert "Invalid layers" in detail
-    assert "layers must be >= 1" in detail
+    assert "layers must be between 3 and 10 inclusive" in detail
 
-    # Test layers = -1
+    # Test layers above the supported range
     response = async_omni_test_client.post(
         "/v1/images/edits",
         files=[("image", img_bytes)],
         data={
             "prompt": "test",
-            "layers": -1,
+            "layers": 11,
         },
     )
     assert response.status_code == 400
     detail = response.json()["detail"]
     assert "Invalid layers" in detail
+    assert "layers must be between 3 and 10 inclusive" in detail
 
 
 def test_image_edit_resolution_and_size_conflict(async_omni_test_client):

vllm_omni/entrypoints/openai/api_server.py

@@ -84,8 +84,10 @@ from vllm.utils.system_utils import decorate_logs
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 from vllm_omni.entrypoints.openai.image_api_utils import (
+    SUPPORTED_LAYERED_RESOLUTIONS,
     encode_image_base64,
     parse_size,
+    validate_layered_layers,
 )
 from vllm_omni.entrypoints.openai.protocol.audio import BatchSpeechRequest, OpenAICreateSpeechRequest
 from vllm_omni.entrypoints.openai.protocol.images import (
@@ -116,8 +118,6 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParam
 logger = init_logger(__name__)
 router = APIRouter()
 
-# Supported resolution buckets for layered models (e.g., Qwen-Image-Layered)
-SUPPORTED_LAYERED_RESOLUTIONS = (640, 1024)
 profiler_router = APIRouter()
 
 
@@ -1314,6 +1314,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
             gen_params, "seed", request.seed if request.seed is not None else random.randint(0, 2**32 - 1)
         )
         _update_if_not_none(gen_params, "generator_device", request.generator_device)
+        _update_if_not_none(gen_params, "layers", request.layers)
 
         request_id = f"img_gen-{random_uuid()}"
 
@@ -1462,11 +1463,13 @@ async def edit_images(
                 detail=f"Invalid resolution {resolution}. Supported resolutions: {SUPPORTED_LAYERED_RESOLUTIONS}.",
             )
         # 3.2.1 Validate layers if provided
-        if layers is not None and layers < 1:
+        try:
+            layers = validate_layered_layers(layers)
+        except ValueError as e:
             raise HTTPException(
                 status_code=HTTPStatus.BAD_REQUEST.value,
-                detail=f"Invalid layers value {layers}. layers must be >= 1.",
-            )
+                detail=str(e),
+            ) from e
         # 3.2.2 Check for conflicting size and resolution parameters
         if resolution is not None and size.lower() != "auto":
             raise HTTPException(

vllm_omni/entrypoints/openai/image_api_utils.py

@@ -13,6 +13,9 @@ import io
 
 import PIL.Image
 
+SUPPORTED_LAYERED_RESOLUTIONS = (640, 1024)
+SUPPORTED_LAYERED_LAYERS_RANGE = range(3, 11)
+
 
 def parse_size(size_str: str) -> tuple[int, int]:
     """Parse size string to width and height tuple.
@@ -63,3 +66,16 @@ def encode_image_base64(image: PIL.Image.Image) -> str:
     image.save(buffer, format="PNG")
     buffer.seek(0)
     return base64.b64encode(buffer.read()).decode("utf-8")
+
+
+def validate_layered_layers(layers: int | None) -> int | None:
+    """Validate the Qwen-Image-Layered ``layers`` parameter."""
+    if layers is None:
+        return None
+    if layers not in SUPPORTED_LAYERED_LAYERS_RANGE:
+        raise ValueError(
+            f"Invalid layers value {layers}. layers must be between "
+            f"{SUPPORTED_LAYERED_LAYERS_RANGE.start} and "
+            f"{SUPPORTED_LAYERED_LAYERS_RANGE.stop - 1} inclusive."
+        )
+    return layers

vllm_omni/entrypoints/openai/protocol/images.py

@@ -12,6 +12,8 @@ from typing import Any
 
 from pydantic import BaseModel, Field, field_validator
 
+from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
+
 
 class ResponseFormat(str, Enum):
     """Image response format"""
@@ -43,6 +45,10 @@ class ImageGenerationRequest(BaseModel):
     )
     response_format: ResponseFormat = Field(default=ResponseFormat.B64_JSON, description="Format of the returned image")
     user: str | None = Field(default=None, description="User identifier for tracking")
+    layers: int | None = Field(
+        default=None,
+        description="Number of output layers for layered image models. Supported range: 3-10.",
+    )
 
     @field_validator("size")
     @classmethod
@@ -67,6 +73,12 @@ class ImageGenerationRequest(BaseModel):
             raise ValueError(f"Only 'b64_json' response format is supported, got: {v}")
         return v
 
+    @field_validator("layers")
+    @classmethod
+    def validate_layers(cls, v):
+        """Validate the layers parameter for layered image models."""
+        return validate_layered_layers(v)
+
     # vllm-omni extensions for diffusion control
     negative_prompt: str | None = Field(default=None, description="Text describing what to avoid in the image")
     num_inference_steps: int | None = Field(

vllm_omni/entrypoints/openai/serving_chat.py

@@ -82,6 +82,7 @@ from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
 from vllm.utils.collection_utils import as_list
 
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
+from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
 from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse
 from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio
 from vllm_omni.entrypoints.openai.utils import parse_lora_request
@@ -2094,6 +2095,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             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)
+
             logger.info(
                 "Diffusion chat request %s: prompt=%r, ref_images=%d, params=%s",
                 request_id,