240 lines
8.4 KiB
Python
240 lines
8.4 KiB
Python
"""
|
|
Definitions running and validating individual tasks, e.g., text to image,
|
|
image to image, and so on. These are called by the core test runner.
|
|
"""
|
|
|
|
import base64
|
|
import io
|
|
from dataclasses import replace
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler, dummy_messages_from_mix_data
|
|
from vllm_omni.entrypoints.omni import Omni
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
|
|
PROMPT = "Dummy prompt"
|
|
IMAGE_DIMS = (512, 512)
|
|
HEIGHT, WIDTH = IMAGE_DIMS
|
|
INPUT_IMAGE = Image.new("RGB", IMAGE_DIMS)
|
|
|
|
# Offline sampling params
|
|
IMAGE_GEN_SAMPLING_PARAMS = OmniDiffusionSamplingParams(
|
|
num_inference_steps=4,
|
|
height=HEIGHT,
|
|
width=WIDTH,
|
|
seed=42,
|
|
)
|
|
|
|
VIDEO_NUM_FRAMES = 9
|
|
VIDEO_GEN_SAMPLING_PARAMS = OmniDiffusionSamplingParams(
|
|
num_inference_steps=4,
|
|
height=HEIGHT,
|
|
width=WIDTH,
|
|
num_frames=VIDEO_NUM_FRAMES,
|
|
seed=42,
|
|
)
|
|
|
|
# Online extra_body for diffusion requests
|
|
IMAGE_GEN_EXTRA_BODY = {
|
|
"height": HEIGHT,
|
|
"width": WIDTH,
|
|
"num_inference_steps": 4,
|
|
"seed": 42,
|
|
}
|
|
|
|
# Online form_data for video generation requests (multipart /v1/videos API)
|
|
VIDEO_GEN_FORM_DATA = {
|
|
"height": HEIGHT,
|
|
"width": WIDTH,
|
|
"num_inference_steps": 4,
|
|
"num_frames": VIDEO_NUM_FRAMES,
|
|
"seed": 42,
|
|
}
|
|
|
|
|
|
### Shared validation
|
|
def _validate_images(images: list[Image.Image], expected_n: int = 1):
|
|
"""Given a set of images, ensure we got the expected count of images,
|
|
and that all provided images match the expected dimensions."""
|
|
assert len(images) == expected_n
|
|
for img in images:
|
|
assert isinstance(img, Image.Image)
|
|
assert img.size == IMAGE_DIMS
|
|
return images
|
|
|
|
|
|
def _validate_image_gen_determinism(images_a: list[Image.Image], images_b: list[Image.Image]):
|
|
"""Ensure that two image sets are valid and that the results match."""
|
|
_validate_images(images_a)
|
|
_validate_images(images_b)
|
|
assert np.array_equal(np.array(images_a[0]), np.array(images_b[0]))
|
|
|
|
|
|
### Output extractor utils for offline / online paths respectively
|
|
def _get_offline_images(outputs: list[OmniRequestOutput]) -> list[Image.Image]:
|
|
"""Extract the images from an Omni .generate() call."""
|
|
assert len(outputs) == 1
|
|
return outputs[0].images
|
|
|
|
|
|
def _get_online_images(responses: list[DiffusionResponse]) -> list[Image.Image]:
|
|
"""Extract the images from a server response."""
|
|
assert len(responses) == 1
|
|
images = responses[0].images
|
|
assert images is not None
|
|
return images
|
|
|
|
|
|
def _validate_video(outputs: list[OmniRequestOutput], expected_n: int = 1):
|
|
"""Given a set of outputs, ensure we got video frames with the expected shape."""
|
|
assert len(outputs) == expected_n
|
|
for output in outputs:
|
|
# Video models return numpy arrays via output.images
|
|
images = output.images
|
|
assert len(images) > 0
|
|
for frame_data in images:
|
|
assert isinstance(frame_data, np.ndarray)
|
|
# (num_outputs, num_frames, H, W, C) or (num_frames, H, W, C)
|
|
assert frame_data.ndim in (4, 5), f"Expected 4D or 5D video array, got shape {frame_data.shape}"
|
|
assert frame_data.shape[-3] == HEIGHT, f"Expected height {HEIGHT}, got {frame_data.shape[-3]}"
|
|
assert frame_data.shape[-2] == WIDTH, f"Expected width {WIDTH}, got {frame_data.shape[-2]}"
|
|
assert frame_data.shape[-1] == 3, f"Expected 3 channels (RGB), got {frame_data.shape[-1]}"
|
|
|
|
|
|
### Offline helpers
|
|
def _run_offline_t2i(omni: Omni, params: OmniDiffusionSamplingParams = IMAGE_GEN_SAMPLING_PARAMS):
|
|
return omni.generate({"prompt": PROMPT}, params)
|
|
|
|
|
|
def _run_offline_t2v(omni: Omni, params: OmniDiffusionSamplingParams = VIDEO_GEN_SAMPLING_PARAMS):
|
|
return omni.generate({"prompt": PROMPT}, params)
|
|
|
|
|
|
def _run_offline_i2i(omni: Omni):
|
|
return omni.generate(
|
|
{"prompt": PROMPT, "multi_modal_data": {"image": INPUT_IMAGE}},
|
|
IMAGE_GEN_SAMPLING_PARAMS,
|
|
)
|
|
|
|
|
|
### Online helpers
|
|
def _build_online_image_data_url() -> str:
|
|
"""Get a valid base 64 encoded data URL corresponding to an image."""
|
|
buf = io.BytesIO()
|
|
INPUT_IMAGE.save(buf, format="PNG")
|
|
b64 = base64.b64encode(buf.getvalue()).decode()
|
|
return f"data:image/png;base64,{b64}"
|
|
|
|
|
|
def _run_online_t2i(
|
|
server: OmniServer, client: OpenAIClientHandler, extra_body: dict | None = None
|
|
) -> list[DiffusionResponse]:
|
|
"""Run a text to image request through the server."""
|
|
messages = dummy_messages_from_mix_data(content_text=PROMPT)
|
|
request_config = {
|
|
"model": server.model,
|
|
"messages": messages,
|
|
"extra_body": extra_body or IMAGE_GEN_EXTRA_BODY,
|
|
}
|
|
return client.send_diffusion_request(request_config)
|
|
|
|
|
|
def _run_online_i2i(server: OmniServer, client: OpenAIClientHandler) -> list[DiffusionResponse]:
|
|
"""Run an image to image request through the server."""
|
|
image_data_url = _build_online_image_data_url()
|
|
messages = dummy_messages_from_mix_data(
|
|
content_text=PROMPT,
|
|
image_data_url=image_data_url,
|
|
)
|
|
request_config = {
|
|
"model": server.model,
|
|
"messages": messages,
|
|
"extra_body": IMAGE_GEN_EXTRA_BODY,
|
|
}
|
|
return client.send_diffusion_request(request_config)
|
|
|
|
|
|
### Offline task runners
|
|
def run_and_validate_text_to_image_request(omni: Omni):
|
|
"""Run and validate a text to image request."""
|
|
_validate_images(_get_offline_images(_run_offline_t2i(omni)))
|
|
|
|
|
|
def run_and_validate_image_to_image_request(omni: Omni):
|
|
"""Run and validate an image to image request."""
|
|
_validate_images(_get_offline_images(_run_offline_i2i(omni)))
|
|
|
|
|
|
def run_and_validate_text_to_video_request(omni: Omni):
|
|
"""Run and validate a text to video request."""
|
|
_validate_video(_run_offline_t2v(omni))
|
|
|
|
|
|
def run_and_validate_text_to_image_determinism(omni: Omni):
|
|
"""Checks for determinism; for now we just keep this for TTI."""
|
|
_validate_image_gen_determinism(
|
|
_get_offline_images(_run_offline_t2i(omni)),
|
|
_get_offline_images(_run_offline_t2i(omni)),
|
|
)
|
|
|
|
|
|
def run_and_validate_text_to_image_multi_output(omni: Omni):
|
|
"""Checks for multi-output; for now we just keep this for TTI."""
|
|
params = replace(IMAGE_GEN_SAMPLING_PARAMS, num_outputs_per_prompt=2)
|
|
_validate_images(_get_offline_images(_run_offline_t2i(omni, params)), expected_n=2)
|
|
|
|
|
|
def _run_online_t2v(
|
|
server: OmniServer, client: OpenAIClientHandler, form_data: dict | None = None
|
|
) -> list[DiffusionResponse]:
|
|
"""Run a text to video request through the server's /v1/videos API."""
|
|
data = dict(form_data or VIDEO_GEN_FORM_DATA)
|
|
data.setdefault("prompt", PROMPT)
|
|
data.setdefault("model", server.model)
|
|
return client.send_video_diffusion_request({"form_data": data})
|
|
|
|
|
|
def _get_online_videos(responses: list[DiffusionResponse]) -> list:
|
|
"""Extract the videos from a server response."""
|
|
assert len(responses) == 1
|
|
videos = responses[0].videos
|
|
assert videos is not None
|
|
assert len(videos) > 0
|
|
return videos
|
|
|
|
|
|
### Online task runners
|
|
def run_and_validate_online_text_to_image_request(server: OmniServer, client: OpenAIClientHandler):
|
|
"""Run and validate a text to image request through the server."""
|
|
_validate_images(_get_online_images(_run_online_t2i(server, client)))
|
|
|
|
|
|
def run_and_validate_online_image_to_image_request(server: OmniServer, client: OpenAIClientHandler):
|
|
"""Run and validate an image to image request through the server."""
|
|
_validate_images(_get_online_images(_run_online_i2i(server, client)))
|
|
|
|
|
|
def run_and_validate_online_text_to_image_determinism(server: OmniServer, client: OpenAIClientHandler):
|
|
"""Checks for determinism through the server; for now we just keep this for TTI."""
|
|
_validate_image_gen_determinism(
|
|
_get_online_images(_run_online_t2i(server, client)),
|
|
_get_online_images(_run_online_t2i(server, client)),
|
|
)
|
|
|
|
|
|
def run_and_validate_online_text_to_image_multi_output(server: OmniServer, client: OpenAIClientHandler):
|
|
"""Checks for multi-output through the server; for now we just keep this for TTI."""
|
|
extra_body = {**IMAGE_GEN_EXTRA_BODY, "num_outputs_per_prompt": 2}
|
|
_validate_images(
|
|
_get_online_images(_run_online_t2i(server, client, extra_body=extra_body)),
|
|
expected_n=2,
|
|
)
|
|
|
|
|
|
def run_and_validate_online_text_to_video_request(server: OmniServer, client: OpenAIClientHandler):
|
|
"""Run and validate a text to video request through the server."""
|
|
_get_online_videos(_run_online_t2v(server, client))
|