[Bugfix] Handle numpy array outputs when generate image (#1680)

Signed-off-by: rongfu.leng <lenronfu@gmail.com>

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)

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(