fix: localize NumPy image conversion

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

examples/offline_inference/text_to_image/text_to_image.py

@@ -8,7 +8,9 @@ import time
 from pathlib import Path
 from typing import Any
 
+import numpy as np
 import torch
+from diffusers.utils import numpy_to_pil
 
 from vllm_omni.diffusion.data import logger
 from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
@@ -54,6 +56,17 @@ def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]
 parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
 
 
+def _normalize_images_for_save(images: list[Any]) -> list[Any]:
+    """Convert NumPy diffusion outputs to PIL images before saving."""
+    normalized = []
+    for image in images:
+        if isinstance(image, np.ndarray):
+            normalized.extend(numpy_to_pil(image))
+        else:
+            normalized.append(image)
+    return normalized
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
     parser.add_argument(
@@ -604,10 +617,22 @@ def main():
         raise ValueError("No output generated from omni.generate()")
     logger.info(f"Outputs: {outputs}")
 
-    images = extract_images_from_outputs(outputs)
+    images = None
+    for output in outputs:
+        images = getattr(output, "images", None)
+        if images:
+            break
+        req_out = getattr(output, "request_output", None)
+        images = getattr(req_out, "images", None) if req_out is not None else None
+        if images:
+            break
+
+    if not images:
+        images = extract_images_from_outputs(outputs)
 
     if not images:
         raise ValueError("No images found in request_output")
+    images = _normalize_images_for_save(images)
 
     output_path = Path(args.output)
     output_path.parent.mkdir(parents=True, exist_ok=True)

tests/diffusion/test_image_output.py

@@ -1,37 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import numpy as np
-import pytest
-from PIL import Image
-
-from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
-from vllm_omni.outputs import OmniRequestOutput
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
-
-
-@pytest.mark.parametrize(
-    ("payload", "expected_pixel"),
-    [
-        (np.array([[[-1.0, 0.0, 1.0]]], dtype=np.float32), (0, 127, 255)),
-        (np.array([[[0, 128, 255]]], dtype=np.uint8), (0, 128, 255)),
-    ],
-)
-def test_extract_images_normalizes_numpy_payload(
-    payload: np.ndarray,
-    expected_pixel: tuple[int, int, int],
-    tmp_path,
-) -> None:
-    output = OmniRequestOutput(images=[payload])
-
-    [image] = extract_images_from_outputs(output)
-
-    assert isinstance(image, Image.Image)
-    assert image.mode == "RGB"
-    assert image.size == (1, 1)
-    assert image.getpixel((0, 0)) == expected_pixel
-
-    output_path = tmp_path / "image.png"
-    image.save(output_path)
-    assert output_path.is_file()

vllm_omni/diffusion/utils/image_output.py

@@ -4,7 +4,6 @@
 from collections.abc import Iterable, Mapping, Sequence
 from typing import Any
 
-import numpy as np
 import torch
 from PIL import Image
 from vllm.outputs import RequestOutput
@@ -21,8 +20,6 @@ def extract_images_from_outputs(outputs: ImageOutputs) -> list[Image.Image]:
     Supported inputs:
     - OmniRequestOutput / list[OmniRequestOutput]
     - RequestOutput / list[RequestOutput] (compatibility fallback)
-
-    PIL images, NumPy arrays, and PyTorch tensors are normalized to PIL images.
     """
     for output in _iter_known_outputs(outputs):
         if isinstance(output, OmniRequestOutput):
@@ -61,9 +58,6 @@ def _iter_multimodal_image_payloads(output: OmniRequestOutput | RequestOutput) -
 
 
 def _iter_request_output_payloads(request_output: RequestOutput) -> Iterable[Any]:
-    images = getattr(request_output, "images", None)
-    if images is not None:
-        yield images
     mm = getattr(request_output, "multimodal_output", None)
     if mm is not None:
         yield from _image_values_from_mapping_like(mm)
@@ -88,8 +82,6 @@ def _coerce_images(payload: Any) -> list[Image.Image]:
         return []
     if isinstance(payload, Image.Image):
         return [payload]
-    if isinstance(payload, np.ndarray):
-        return _ndarray_to_images(payload)
     if isinstance(payload, torch.Tensor):
         return _tensor_to_images(payload)
     if isinstance(payload, list | tuple):
@@ -100,37 +92,6 @@ def _coerce_images(payload: Any) -> list[Image.Image]:
     return []
 
 
-def _ndarray_to_images(array: np.ndarray) -> list[Image.Image]:
-    if array.ndim == 4:
-        images: list[Image.Image] = []
-        for single in array:
-            images.extend(_ndarray_to_images(single))
-        return images
-    if array.ndim != 3 or array.size == 0:
-        return []
-
-    image = array
-    if image.shape[-1] not in (1, 3, 4):
-        if image.shape[0] not in (1, 3, 4):
-            return []
-        image = np.moveaxis(image, 0, -1)
-
-    if np.issubdtype(image.dtype, np.floating):
-        image = image.astype(np.float32, copy=False)
-        if image.min() < 0.0:
-            image = image / 2 + 0.5
-        image = np.clip(image, 0.0, 1.0) * 255
-    elif np.issubdtype(image.dtype, np.integer) or np.issubdtype(image.dtype, np.bool_):
-        image = np.clip(image, 0, 255)
-    else:
-        return []
-
-    image = image.astype(np.uint8)
-    if image.shape[-1] == 1:
-        image = image[..., 0]
-    return [Image.fromarray(np.ascontiguousarray(image))]
-
-
 def _tensor_to_images(tensor: torch.Tensor) -> list[Image.Image]:
     img = tensor.detach().to("cpu", dtype=torch.float32)
     if img.ndim == 4: