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 81acebe00..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( @@ -619,6 +632,7 @@ def main(): 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)