fcf1085395
Signed-off-by: Boao Shi <aoibosh@connect.hku.hk> Signed-off-by: bowieshi <aoibosh@connect.hku.hk> Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: R2-Y <ruiruyang2@gmail.com> Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com> Signed-off-by: Shi Boao Bowie <aoibosh@connect.hku.hk> Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: ruirui(rein) yang <73573651+R2-Y@users.noreply.github.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
118 lines
3.9 KiB
Python
118 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from collections.abc import Iterable, Mapping, Sequence
|
|
from typing import Any
|
|
|
|
import torch
|
|
from PIL import Image
|
|
from vllm.outputs import RequestOutput
|
|
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
|
|
ImageOutput = OmniRequestOutput | RequestOutput
|
|
ImageOutputs = ImageOutput | Sequence[ImageOutput] | None
|
|
|
|
|
|
def extract_images_from_outputs(outputs: ImageOutputs) -> list[Image.Image]:
|
|
"""Extract PIL images from known Omni output types.
|
|
|
|
Supported inputs:
|
|
- OmniRequestOutput / list[OmniRequestOutput]
|
|
- RequestOutput / list[RequestOutput] (compatibility fallback)
|
|
"""
|
|
for output in _iter_known_outputs(outputs):
|
|
if isinstance(output, OmniRequestOutput):
|
|
images = _coerce_images(output.images)
|
|
if images:
|
|
return images
|
|
|
|
for payload in _iter_multimodal_image_payloads(output):
|
|
images = _coerce_images(payload)
|
|
if images:
|
|
return images
|
|
|
|
return []
|
|
|
|
|
|
def _iter_known_outputs(value: ImageOutputs) -> Iterable[ImageOutput]:
|
|
if value is None:
|
|
return
|
|
if isinstance(value, OmniRequestOutput | RequestOutput):
|
|
yield value
|
|
return
|
|
if isinstance(value, Sequence):
|
|
for item in value:
|
|
if isinstance(item, OmniRequestOutput | RequestOutput):
|
|
yield item
|
|
|
|
|
|
def _iter_multimodal_image_payloads(output: OmniRequestOutput | RequestOutput) -> Iterable[Any]:
|
|
if isinstance(output, OmniRequestOutput):
|
|
if output.multimodal_output is not None:
|
|
yield from _image_values_from_mapping_like(output.multimodal_output)
|
|
# OmniRequestOutput IS the RequestOutput now; iterate completion
|
|
# outputs directly for multimodal payloads that the property
|
|
# (which returns only the first match) may have missed.
|
|
for completion in getattr(output, "outputs", []):
|
|
mm = getattr(completion, "multimodal_output", None)
|
|
if mm is not None:
|
|
yield from _image_values_from_mapping_like(mm)
|
|
else:
|
|
yield from _iter_request_output_payloads(output)
|
|
|
|
|
|
def _iter_request_output_payloads(request_output: RequestOutput) -> Iterable[Any]:
|
|
mm = getattr(request_output, "multimodal_output", None)
|
|
if mm is not None:
|
|
yield from _image_values_from_mapping_like(mm)
|
|
for completion in getattr(request_output, "outputs", []):
|
|
mm = getattr(completion, "multimodal_output", None)
|
|
if mm is not None:
|
|
yield from _image_values_from_mapping_like(mm)
|
|
|
|
|
|
def _image_values_from_mapping_like(value: Any) -> Iterable[Any]:
|
|
if hasattr(value, "to_dict") and callable(value.to_dict):
|
|
value = value.to_dict()
|
|
if not isinstance(value, Mapping):
|
|
return
|
|
for key in ("image", "images", "model_outputs"):
|
|
if key in value:
|
|
yield value[key]
|
|
|
|
|
|
def _coerce_images(payload: Any) -> list[Image.Image]:
|
|
if payload is None:
|
|
return []
|
|
if isinstance(payload, Image.Image):
|
|
return [payload]
|
|
if isinstance(payload, torch.Tensor):
|
|
return _tensor_to_images(payload)
|
|
if isinstance(payload, list | tuple):
|
|
images: list[Image.Image] = []
|
|
for item in payload:
|
|
images.extend(_coerce_images(item))
|
|
return images
|
|
return []
|
|
|
|
|
|
def _tensor_to_images(tensor: torch.Tensor) -> list[Image.Image]:
|
|
img = tensor.detach().to("cpu", dtype=torch.float32)
|
|
if img.ndim == 4:
|
|
images: list[Image.Image] = []
|
|
for single in img:
|
|
images.extend(_tensor_to_images(single))
|
|
return images
|
|
if img.ndim != 3:
|
|
return []
|
|
|
|
if img.shape[0] in (1, 3, 4):
|
|
img = img.permute(1, 2, 0)
|
|
|
|
# Some diffusion models emit image tensors normalized in [-1, 1].
|
|
if img.min().item() < 0.0:
|
|
img = img / 2 + 0.5
|
|
img = img.clamp(0, 1).mul(255).to(torch.uint8).contiguous().numpy()
|
|
return [Image.fromarray(img)]
|