diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index a09a0d5fd..03bdf7080 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/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) diff --git a/tests/diffusion/test_image_output.py b/tests/diffusion/test_image_output.py deleted file mode 100644 index f9cebde32..000000000 --- a/tests/diffusion/test_image_output.py +++ /dev/null @@ -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() diff --git a/vllm_omni/diffusion/utils/image_output.py b/vllm_omni/diffusion/utils/image_output.py index 5b5e02720..5f8c9903a 100644 --- a/vllm_omni/diffusion/utils/image_output.py +++ b/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: