Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8080c18ff4 | ||
|
|
5e3aa88e10 |
@@ -8,7 +8,9 @@ import time
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
import torch
|
import torch
|
||||||
|
from diffusers.utils import numpy_to_pil
|
||||||
|
|
||||||
from vllm_omni.diffusion.data import logger
|
from vllm_omni.diffusion.data import logger
|
||||||
from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
|
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")
|
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:
|
def parse_args() -> argparse.Namespace:
|
||||||
parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
|
parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
@@ -619,6 +632,7 @@ def main():
|
|||||||
|
|
||||||
if not images:
|
if not images:
|
||||||
raise ValueError("No images found in request_output")
|
raise ValueError("No images found in request_output")
|
||||||
|
images = _normalize_images_for_save(images)
|
||||||
|
|
||||||
output_path = Path(args.output)
|
output_path = Path(args.output)
|
||||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
|||||||
Reference in new issue
Block a user