[Tests] Tiny Diffusion Model Pattern (#4524)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/model_tests/diffusion/task_runners.py

@@ -0,0 +1,172 @@
+"""
+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,
+)
+
+# Online extra_body for diffusion requests
+IMAGE_GEN_EXTRA_BODY = {
+    "height": HEIGHT,
+    "width": WIDTH,
+    "num_inference_steps": 4,
+    "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
+
+
+### Offline helpers
+def _run_offline_t2i(omni: Omni, params: OmniDiffusionSamplingParams = IMAGE_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_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)
+
+
+### 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,
+    )