diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index c91c5a5c7..4b38692da 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -1165,3 +1165,91 @@ def test_image_edit_with_seed_zero_single_stage(test_client): f"Expected seed=0, but got seed={captured_sampling_params.seed}. " "This indicates the bug where seed=0 is treated as falsy." ) + + +def test_normalize_image(): + """Test _normalize_image with various input types""" + import numpy as np + + from vllm_omni.entrypoints.openai.api_server import _normalize_image + + # Test PIL Image input + img = Image.new("RGB", (64, 64), color="red") + result = _normalize_image(img) + assert isinstance(result, Image.Image) + assert result.size == (64, 64) + + # Test uint8 numpy array + arr = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8) + result = _normalize_image(arr) + assert isinstance(result, Image.Image) + assert result.size == (64, 64) + + # Test float [0, 1] numpy array + arr = np.random.rand(64, 64, 3).astype(np.float32) + result = _normalize_image(arr) + assert isinstance(result, Image.Image) + assert result.size == (64, 64) + + # Test float [-1, 1] numpy array + arr = np.random.rand(64, 64, 3).astype(np.float32) * 2 - 1 + result = _normalize_image(arr) + assert isinstance(result, Image.Image) + assert result.size == (64, 64) + + # Test batch dimensions (1, 1, H, W, C) + arr = np.random.randint(0, 255, (1, 1, 64, 64, 3), dtype=np.uint8) + result = _normalize_image(arr) + assert isinstance(result, Image.Image) + assert result.size == (64, 64) + + +def test_extract_images_from_result(): + """Test _extract_images_from_result with various result formats""" + import numpy as np + + from vllm_omni.entrypoints.openai.api_server import _extract_images_from_result + + # Test empty result + class EmptyResult: + pass + + result = EmptyResult() + images = _extract_images_from_result(result) + assert images == [] + + # Test nested batch: [np.array(shape=(3, 64, 64, 3))] + batch = np.random.randint(0, 255, (3, 1, 64, 64, 3), dtype=np.uint8) + + class BatchResult: + def __init__(self): + self.images = [batch] + + result = BatchResult() + images = _extract_images_from_result(result) + assert len(images) == 3 + assert all(isinstance(img, Image.Image) for img in images) + assert all(img.size == (64, 64) for img in images) + + # Test dict path: result.request_output["images"] + class DictRequestOutput: + def __init__(self): + self.request_output = {"images": [np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)]} + + result = DictRequestOutput() + images = _extract_images_from_result(result) + assert len(images) == 1 + assert isinstance(images[0], Image.Image) + + # Test attribute path: result.request_output.images + class AttrRequestOutput: + def __init__(self): + self.request_output = type( + "obj", (), {"images": [np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)]} + )() + + result = AttrRequestOutput() + images = _extract_images_from_result(result) + assert len(images) == 1 + assert isinstance(images[0], Image.Image) + assert images[0].size == (32, 32) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 6a65f4433..d847a96db 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -18,6 +18,7 @@ from http import HTTPStatus from typing import Annotated, Any, Literal, cast import httpx +import numpy as np import vllm.envs as envs from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile, WebSocket from fastapi.responses import FileResponse, JSONResponse, Response, StreamingResponse @@ -1767,6 +1768,34 @@ def _update_if_not_none(object: Any, key: str, val: Any) -> None: setattr(object, key, val) +def _normalize_image(image: Any) -> Any: + """Normalize a single image output to a PIL-compatible format.""" + if isinstance(image, Image.Image): + return image + if not isinstance(image, np.ndarray): + raise ValueError(f"Unsupported image type: {type(image)}") + if not np.issubdtype(image.dtype, np.integer) and not np.issubdtype(image.dtype, np.floating): + raise ValueError(f"Unsupported dtype: {image.dtype}") + if isinstance(image, np.ndarray): + while image.ndim > 3: + image = image[0] + if image.min() < 0: + if image.min() < -1.01 or image.max() > 1.01: + logger.warning( + f"Image float range [{image.min():.2f}, {image.max():.2f}] outside expected [-1, 1]. " + f"Clipping to [-1, 1] before normalization." + ) + image = np.clip(image, -1.0, 1.0) * 0.5 + 0.5 + elif image.max() > 1.01: + logger.warning( + f"Image float range [{image.min():.2f}, {image.max():.2f}] outside expected [0, 1]. " + f"Clipping to [0, 1] before normalization." + ) + image = (np.clip(image, 0.0, 1.0) * 255).astype(np.uint8) + image = Image.fromarray(image) + return image + + def _extract_images_from_result(result: Any) -> list[Any]: images = [] if hasattr(result, "images") and result.images: @@ -1777,6 +1806,10 @@ def _extract_images_from_result(result: Any) -> list[Any]: images = request_output["images"] elif hasattr(request_output, "images") and request_output.images: images = request_output.images + # Handle when generate more than one image + if images and isinstance(images[0], np.ndarray) and images[0].shape[0] > 1 and images[0].ndim == 5: + # Unwrap batch: (N, T, H, W, C) -> [img1, img2, ...] + images = list(images[0]) # Flatten nested lists (e.g., from layered models like Qwen-Image-Layered). # Note: This only flattens one level deep. Deeper nesting is not supported. flattened = [] @@ -1785,7 +1818,7 @@ def _extract_images_from_result(result: Any) -> list[Any]: flattened.extend(img) else: flattened.append(img) - return flattened + return [_normalize_image(img) for img in flattened] async def _load_input_images(