[Tests] Tiny Diffusion Model Pattern (#4524)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
This commit is contained in:
Alex Brooks authored and GitHub committed
1 parent
a5db2d839a
commit
690864f277
20 files changed
+962
-263
No files matched your search
@@ -627,11 +627,10 @@ steps:
|
||||
- label: "Diffusion · FLUX.2 Klein Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_flux2_klein.py
|
||||
- tests/e2e/online_serving/test_flux2_klein.py
|
||||
- vllm_omni/diffusion/models/flux2_klein/
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py tests/e2e/online_serving/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
|
||||
- pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
|
||||
@@ -598,6 +598,38 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]"
|
||||
key: nightly-tiny-model-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: >-
|
||||
build.env("NIGHTLY") == "1" ||
|
||||
build.pull_request.labels includes "nightly-test"
|
||||
steps:
|
||||
- label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations)"
|
||||
timeout_in_minutes: 30
|
||||
# NOTE: Single-GPU (core_model) tiny model tests run on the ready label,
|
||||
# so nightly only adds the multi-GPU parallelism configurations.
|
||||
#
|
||||
# This is a bit confusing, but `full_model` is used to mark the multigpu tests, and
|
||||
# --run-level `core_model` is used to indicate that we should use tiny model weights
|
||||
# for the CI, since these tests should pass irrespective of whether or not we use
|
||||
# the full weights or not.
|
||||
commands:
|
||||
- timeout 25m pytest -sv tests/model_tests/diffusion -m 'full_model and cuda' --run-level "core_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
# Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below.
|
||||
- group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
|
||||
key: nightly-diffusion-x2iat-group
|
||||
|
||||
@@ -157,6 +157,27 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
|
||||
- group: ":card_index_dividers: Tiny Model Tests [Single GPU]"
|
||||
depends_on: upload-ready-pipeline
|
||||
steps:
|
||||
- label: "Tiny Model Tests · Diffusion (base)"
|
||||
commands:
|
||||
- timeout 15m pytest -sv tests/model_tests/diffusion -m 'core_model and cuda' --run-level "core_model"
|
||||
agents:
|
||||
# Ready label only runs the cases that don't need parallelism for tiny
|
||||
# model tests, which is usually just a smoke test for TTI and I2I
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Engine&Model Executor Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
|
||||
@@ -41,6 +41,7 @@ dev = [
|
||||
"pytest-cov>=4.0.0",
|
||||
"pytest-mock>=3.10.0",
|
||||
"websockets>=15.0.1",
|
||||
"pytest-subtests>=0.11.0",
|
||||
"datasets>=2.14.0",
|
||||
"mypy==1.11.1",
|
||||
"pre-commit==4.0.1",
|
||||
|
||||
@@ -18,6 +18,8 @@ from tests.helpers.runtime import DiffusionResponse, OmniRunnerHandler
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
# NOTE: Most of the flux2klein tests have migrated to the common diffusion tests.
|
||||
# This model will also use the tiny model builder when run at the core model level.
|
||||
MODEL = "black-forest-labs/FLUX.2-klein-4B"
|
||||
|
||||
_OMNI_RUNNER_PARAM = (MODEL, None)
|
||||
@@ -70,41 +72,6 @@ def _send_inpaint_with_generator(
|
||||
)
|
||||
|
||||
|
||||
# Regression test for https://github.com/vllm-project/vllm-omni/issues/3097
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
def test_flux2_klein_can_accept_text_inputs(omni_runner_handler: OmniRunnerHandler):
|
||||
omni_runner_handler.send_diffusion_request(
|
||||
{
|
||||
"model": MODEL,
|
||||
"prompt": "a cup of coffee on the table",
|
||||
"sampling_params": OmniDiffusionSamplingParams(num_inference_steps=2, seed=42),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_flux2_klein_inpaint_basic(omni_runner_handler: OmniRunnerHandler):
|
||||
input_image, mask_image = _create_test_inputs()
|
||||
omni_runner_handler.send_diffusion_request(
|
||||
{
|
||||
"model": MODEL,
|
||||
"prompt": "Fill in the masked area with a beautiful garden",
|
||||
"multi_modal_data": {"image": input_image, "mask_image": mask_image},
|
||||
"sampling_params": OmniDiffusionSamplingParams(
|
||||
height=_HEIGHT,
|
||||
width=_WIDTH,
|
||||
num_inference_steps=_NUM_INFERENCE_STEPS,
|
||||
guidance_scale=0.0,
|
||||
generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
|
||||
num_outputs_per_prompt=1,
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
|
||||
@@ -1,62 +0,0 @@
|
||||
"""
|
||||
Tests for Flux2 Klein; currently Dev is implemented separately,
|
||||
but ideally these models will fold together in the future.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import (
|
||||
OmniServer,
|
||||
OmniServerParams,
|
||||
OpenAIClientHandler,
|
||||
dummy_messages_from_mix_data,
|
||||
)
|
||||
|
||||
FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4)
|
||||
POSITIVE_PROMPT = "A cat sitting on a windowsill"
|
||||
NEGATIVE_PROMPT = "blurry, low quality"
|
||||
|
||||
|
||||
# For now, we only test a good configuration for this model to keep the CI
|
||||
# light, but in the near future may explore using random models to check
|
||||
# cross-feature compatibility more generally.
|
||||
def _get_diffusion_feature_cases(model: str):
|
||||
return [
|
||||
# FP8 / Hybrid sequence parallelism
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
),
|
||||
id="default",
|
||||
marks=FOUR_CARD_FEATURE_MARKS,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases(
|
||||
model="black-forest-labs/FLUX.2-klein-4B",
|
||||
),
|
||||
indirect=True,
|
||||
)
|
||||
def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT)
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"extra_body": {
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 2,
|
||||
"negative_prompt": NEGATIVE_PROMPT,
|
||||
"true_cfg_scale": 4.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
@@ -1,163 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
End-to-end tests for Flux2 Klein inpainting in online serving mode.
|
||||
|
||||
Uses /v1/images/edits endpoint which is the correct API for image inpainting.
|
||||
"""
|
||||
|
||||
import base64
|
||||
from io import BytesIO
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams
|
||||
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
|
||||
|
||||
MODEL = "black-forest-labs/FLUX.2-klein-4B"
|
||||
|
||||
_HEIGHT = 512
|
||||
_WIDTH = 512
|
||||
_NUM_INFERENCE_STEPS = 4
|
||||
|
||||
|
||||
def _get_diffusion_feature_cases(model: str):
|
||||
return [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=["--tensor-parallel-size", "2", "--enable-cpu-offload"],
|
||||
),
|
||||
id="tp2_cpu_offload",
|
||||
),
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=["--tensor-parallel-size", "2"],
|
||||
),
|
||||
id="tp2_basic",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _image_to_base64_jpeg(image: Image.Image) -> str:
|
||||
buffer = BytesIO()
|
||||
image.save(buffer, format="JPEG")
|
||||
buffer.seek(0)
|
||||
return base64.b64encode(buffer.read()).decode("utf-8")
|
||||
|
||||
|
||||
def _create_test_mask_base64(width: int = _WIDTH, height: int = _HEIGHT) -> str:
|
||||
mask = Image.new("L", (width, height), 0)
|
||||
draw = ImageDraw.Draw(mask)
|
||||
draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255)
|
||||
return _image_to_base64_jpeg(mask)
|
||||
|
||||
|
||||
def _compare_images(img1: Image.Image, img2: Image.Image) -> bool:
|
||||
return list(img1.getdata()) == list(img2.getdata())
|
||||
|
||||
|
||||
def _send_edit_request(host: str, port: int, model: str, image_b64: str, mask_b64: str, prompt: str, **kwargs):
|
||||
url = f"http://{host}:{port}/v1/images/edits"
|
||||
files = {
|
||||
"image": ("image.jpg", base64.b64decode(image_b64), "image/jpeg"),
|
||||
"mask_image": ("mask.jpg", base64.b64decode(mask_b64), "image/jpeg"),
|
||||
}
|
||||
data = {"prompt": prompt, "model": model, **kwargs}
|
||||
with httpx.Client(timeout=60.0) as client:
|
||||
response = client.post(url, files=files, data=data)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
|
||||
def test_flux2_klein_inpaint_basic(omni_server: OmniServer):
|
||||
input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
|
||||
mask_b64 = _create_test_mask_base64()
|
||||
|
||||
result = _send_edit_request(
|
||||
host=omni_server.host,
|
||||
port=omni_server.port,
|
||||
model=MODEL,
|
||||
image_b64=input_image_b64,
|
||||
mask_b64=mask_b64,
|
||||
prompt="Fill in the masked area with a beautiful garden",
|
||||
guidance_scale=1.0,
|
||||
num_inference_steps=_NUM_INFERENCE_STEPS,
|
||||
n=1,
|
||||
seed=42,
|
||||
)
|
||||
|
||||
assert "data" in result and len(result["data"]) == 1
|
||||
img_data = result["data"][0].get("b64_json") or result["data"][0].get("url", "").split(",")[-1]
|
||||
img = Image.open(BytesIO(base64.b64decode(img_data)))
|
||||
assert img.size == (_WIDTH, _HEIGHT)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
|
||||
def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer):
|
||||
input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
|
||||
mask_b64 = _create_test_mask_base64()
|
||||
prompt = "A red flower in a field"
|
||||
|
||||
result1 = _send_edit_request(
|
||||
host=omni_server.host,
|
||||
port=omni_server.port,
|
||||
model=MODEL,
|
||||
image_b64=input_image_b64,
|
||||
mask_b64=mask_b64,
|
||||
prompt=prompt,
|
||||
guidance_scale=1.0,
|
||||
num_inference_steps=_NUM_INFERENCE_STEPS,
|
||||
n=1,
|
||||
seed=12345,
|
||||
)
|
||||
|
||||
result2 = _send_edit_request(
|
||||
host=omni_server.host,
|
||||
port=omni_server.port,
|
||||
model=MODEL,
|
||||
image_b64=input_image_b64,
|
||||
mask_b64=mask_b64,
|
||||
prompt=prompt,
|
||||
guidance_scale=1.0,
|
||||
num_inference_steps=_NUM_INFERENCE_STEPS,
|
||||
n=1,
|
||||
seed=12345,
|
||||
)
|
||||
|
||||
img1_data = result1["data"][0].get("b64_json") or result1["data"][0].get("url", "").split(",")[-1]
|
||||
img2_data = result2["data"][0].get("b64_json") or result2["data"][0].get("url", "").split(",")[-1]
|
||||
|
||||
img1 = Image.open(BytesIO(base64.b64decode(img1_data)))
|
||||
img2 = Image.open(BytesIO(base64.b64decode(img2_data)))
|
||||
|
||||
assert _compare_images(img1, img2), (
|
||||
"Same input with same seed should produce identical output. This is critical for offline/online consistency."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
|
||||
def test_flux2_klein_inpaint_multiple_outputs(omni_server: OmniServer):
|
||||
input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
|
||||
mask_b64 = _create_test_mask_base64()
|
||||
|
||||
result = _send_edit_request(
|
||||
host=omni_server.host,
|
||||
port=omni_server.port,
|
||||
model=MODEL,
|
||||
image_b64=input_image_b64,
|
||||
mask_b64=mask_b64,
|
||||
prompt="A beautiful landscape",
|
||||
guidance_scale=1.0,
|
||||
num_inference_steps=_NUM_INFERENCE_STEPS,
|
||||
n=2,
|
||||
seed=42,
|
||||
)
|
||||
|
||||
assert "data" in result and len(result["data"]) == 2
|
||||
@@ -47,6 +47,7 @@ from tests.helpers.media import (
|
||||
_merge_base64_audio_to_segment,
|
||||
decode_b64_image,
|
||||
)
|
||||
from tests.model_tests.diffusion.utils import resolve_tiny_model_path
|
||||
from vllm_omni.config.stage_config import resolve_deploy_yaml
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
@@ -2815,11 +2816,24 @@ def iter_omni_server(
|
||||
|
||||
with omni_fixture_lock:
|
||||
params: OmniServerParams = request.param
|
||||
model = model_prefix + params.model
|
||||
# For now, when a tiny model is substituted, we preserve the original model
|
||||
# name via --served-model-name (so that the server still accepts requests with
|
||||
# the original name). We also do the same for server.model so that tests reading
|
||||
# server.model send the correct name in requests.
|
||||
#
|
||||
# TODO: core models on this path currently do not clean up tiny models, although
|
||||
# tiny model paths are deterministic, so it's not a huge footprint. Still, it would
|
||||
# be ideal to cleanup consistently everywhere.
|
||||
original_model = model_prefix + params.model
|
||||
model = original_model
|
||||
if run_level == "core_model":
|
||||
model = resolve_tiny_model_path(model)
|
||||
port = params.port
|
||||
stage_config_path = stage_config_path_for_run_level(params.stage_config_path, run_level)
|
||||
|
||||
server_args = params.server_args or []
|
||||
if model != original_model:
|
||||
server_args = [*server_args, "--served-model-name", original_model]
|
||||
if params.use_omni and params.stage_init_timeout is not None:
|
||||
server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)]
|
||||
else:
|
||||
@@ -2845,6 +2859,8 @@ def iter_omni_server(
|
||||
port=port,
|
||||
env_dict=params.env_dict,
|
||||
) as server:
|
||||
if model != original_model:
|
||||
server.model = original_model
|
||||
print("OmniServer started successfully")
|
||||
yield server
|
||||
print("OmniServer stopping...")
|
||||
@@ -2868,6 +2884,8 @@ def iter_omni_server(
|
||||
use_omni=params.use_omni,
|
||||
)
|
||||
) as server:
|
||||
if model != original_model:
|
||||
server.model = original_model
|
||||
print("OmniServer started successfully")
|
||||
yield server
|
||||
print("OmniServer stopping...")
|
||||
@@ -2899,6 +2917,8 @@ def iter_omni_runner(
|
||||
extra_omni_kwargs = dict(extra) if extra is not None else {}
|
||||
stage_config_path = stage_config_path_for_run_level(stage_config_path, run_level)
|
||||
model = model_prefix + model
|
||||
if run_level == "core_model":
|
||||
model = resolve_tiny_model_path(model)
|
||||
with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner:
|
||||
print("OmniRunner started successfully")
|
||||
yield runner
|
||||
|
||||
Whitespace-only changes.
@@ -0,0 +1,41 @@
|
||||
from shutil import rmtree
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def tiny_model_paths(request, run_level):
|
||||
"""Build or download the tiny models for the selected tests.
|
||||
|
||||
At core_model level, builds tiny models via the builder function.
|
||||
At advanced_model / full_model level, uses the real HF model path.
|
||||
|
||||
NOTE: this is session scoped to avoid churn in tiny model creation,
|
||||
but will ensure all the tiny models you need are created for the selected tests
|
||||
before it starts to execute them."""
|
||||
model_paths = {}
|
||||
built_paths = []
|
||||
print("Initializing tiny models...")
|
||||
for item in request.session.items:
|
||||
if not hasattr(item, "callspec"):
|
||||
continue
|
||||
model_name = item.callspec.params.get("model_name")
|
||||
if model_name is None:
|
||||
continue
|
||||
if model_name not in model_paths:
|
||||
settings = DIFFUSION_TEST_SETTINGS[model_name]
|
||||
if run_level == "core_model":
|
||||
print(f"Calling tiny model builder for: {model_name}")
|
||||
path = settings.builder()
|
||||
built_paths.append(path)
|
||||
else:
|
||||
print(f"Run level is {run_level}; {model_name} will use the full model")
|
||||
path = settings.model
|
||||
model_paths[model_name] = path
|
||||
|
||||
yield model_paths
|
||||
for path in built_paths:
|
||||
print(f"Removing tiny model: {path}")
|
||||
rmtree(path, ignore_errors=True)
|
||||
Whitespace-only changes.
@@ -0,0 +1,104 @@
|
||||
"""
|
||||
Analogous to: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/vlm_utils/case_filtering.py
|
||||
|
||||
This follows a similar pattern so that we can use `get_parametrized_options` to expand our test
|
||||
configurations out, similar to the way in which vLLM expands supported input type tests. For now,
|
||||
this mostly just means parametrizing over the test groups, which may be unique per model, but
|
||||
doing so in this way lets us cleanly define test groups that may be heterogeneous + avoid having
|
||||
excessive skip marks for features that may not be supported in every pipeline.
|
||||
"""
|
||||
|
||||
# Engine initialization is currently quite expensive. For efficiency,
|
||||
# when we run the tests, we should try to collapse compatible accelerations,
|
||||
# because even with tiny models, the overhead of spinning up a new instance
|
||||
# stacks up fast.
|
||||
import itertools
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.model_tests.diffusion.config_types import DiffusionAccs, DiffusionModelTestOpts, get_required_device_count
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
# These tests are intended to run on L4 GPUs since they test parallelism (i.e., need GPUs),
|
||||
# but are intended to be smoke tests that we can properly validate without real model
|
||||
# weights in the CI.
|
||||
#
|
||||
# As such, no DiffusionAccs may require more than 4 GPUs, since we would otherwise skip
|
||||
# them in the CI.
|
||||
MAX_CI_DEVICES = 4
|
||||
|
||||
|
||||
def get_test_group_marks(model_name: str, test_group: list[DiffusionAccs] | None, model_marks: list | None) -> list:
|
||||
"""Build the full set of pytest marks for a test group.
|
||||
|
||||
For now, single device groups default to core model, and multi device groups
|
||||
default to advanced_model marks. We also validate and throw if a test group exceeds
|
||||
the max number of GPUs for the CI to prevent accidentally adding configurations that
|
||||
would be skipped in the CI."""
|
||||
marks = list(model_marks) if model_marks is not None else []
|
||||
|
||||
required_devices = get_required_device_count(test_group)
|
||||
if required_devices > MAX_CI_DEVICES:
|
||||
raise ValueError(
|
||||
f"Test group {test_group} for {model_name} requires {required_devices} devices, "
|
||||
f"but the max CI device count is {MAX_CI_DEVICES}. "
|
||||
)
|
||||
|
||||
marks.extend(hardware_marks(res={"cuda": "L4"}, num_cards=required_devices))
|
||||
# hardware_marks only adds skipif for num_cards > 1, so we currently handle the single-device case
|
||||
# directly here. This should probably be handled in a more common way later on.
|
||||
assert current_omni_platform is not None and current_omni_platform.device_count is not None
|
||||
device_count = current_omni_platform.device_count()
|
||||
if current_omni_platform.is_cuda() and device_count < required_devices:
|
||||
marks.append(pytest.mark.skip(reason=f"Need {required_devices} devices, got {device_count}"))
|
||||
if required_devices > 1:
|
||||
marks.append(pytest.mark.full_model)
|
||||
else:
|
||||
marks.append(pytest.mark.core_model)
|
||||
return marks
|
||||
|
||||
|
||||
def get_model_parametrization(model_name: str, test_info: DiffusionModelTestOpts, online: bool):
|
||||
"""Given a model & its corresponding test options, build the list of pytest params
|
||||
to be run for this model. The base case (no accelerations) is always included. Extra
|
||||
test groups are always appended for offline tests, but are only added to the online
|
||||
tests if online_base_only is overridden to False to avoid redundant testing.
|
||||
"""
|
||||
test_groups: list[list[DiffusionAccs] | None] = [None]
|
||||
if not (online and test_info.online_base_only) and test_info.extra_test_groups:
|
||||
test_groups.extend(test_info.extra_test_groups)
|
||||
|
||||
return [
|
||||
pytest.param(
|
||||
model_name,
|
||||
test_group,
|
||||
test_info.supported_tasks,
|
||||
test_info.check_multi_output and test_group is None,
|
||||
test_info.check_determinism and test_group is None,
|
||||
id=f"{model_name}[{'+'.join(test_group)}]" if test_group else model_name,
|
||||
marks=get_test_group_marks(model_name, test_group, test_info.marks),
|
||||
)
|
||||
for test_group in test_groups
|
||||
]
|
||||
|
||||
|
||||
def get_parametrized_options(
|
||||
test_settings: dict[str, DiffusionModelTestOpts],
|
||||
online: bool = False,
|
||||
):
|
||||
"""Converts all the DiffusionModelTestOpts into an expanded list of parameters
|
||||
based on which accelerations are available.
|
||||
|
||||
When online=True and a model has online_base_only=True, only the base case
|
||||
(no accelerations) is included for that model's parametrization to minimize
|
||||
redundancy between offline / online tests.
|
||||
"""
|
||||
# Get a list per model type, where each entry contains a tuple of all of
|
||||
# that model type's cases, then flatten them into the top level so that
|
||||
# we can consume them in one mark.parametrize call.
|
||||
parametrization = [
|
||||
get_model_parametrization(model_name, test_info, online=online)
|
||||
for model_name, test_info in test_settings.items()
|
||||
]
|
||||
return list(itertools.chain(*parametrization))
|
||||
@@ -0,0 +1,187 @@
|
||||
"""
|
||||
Common definitions for controlling what tests run where.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from enum import StrEnum, auto
|
||||
from typing import NamedTuple, TypeAlias
|
||||
|
||||
from pytest import MarkDecorator
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
# All builder funcs take no params and return a path
|
||||
TinyDiffusionBuilder: TypeAlias = Callable[[], str]
|
||||
|
||||
|
||||
class DiffusionAccs(StrEnum):
|
||||
"""Supported acceleration types / test settings for Diffusion Models."""
|
||||
|
||||
HSDP = auto()
|
||||
TEA_CACHE = auto()
|
||||
CACHE_DIT = auto()
|
||||
SEQUENCE_PARALLEL = auto()
|
||||
CFG_PARALLEL = auto()
|
||||
TENSOR_PARALLEL = auto()
|
||||
CPU_OFFLOAD = auto()
|
||||
LAYERWISE_OFFLOAD = auto()
|
||||
VAE_PATCH_PARALLEL = auto()
|
||||
|
||||
|
||||
class DiffusionTasks(StrEnum):
|
||||
"""Supported tasks for Diffusion Models."""
|
||||
|
||||
TEXT_TO_IMAGE = auto()
|
||||
IMAGE_TO_IMAGE = auto()
|
||||
# Text to video, text to audio, etc should be added here as needed
|
||||
|
||||
|
||||
class DiffusionModelTestOpts(NamedTuple):
|
||||
"""Configuration for one Diffusion model's tests."""
|
||||
|
||||
# HF model name for real-weight tests (advanced_model / full_model level).
|
||||
# For now, whether we use the real weights vs tiny weights in the common tests
|
||||
# depends on the run level.
|
||||
model: str
|
||||
|
||||
# Creates a tiny model for the given architecture. We should always use tiny
|
||||
# model weights for tests that do not require us to check the model quality.
|
||||
builder: TinyDiffusionBuilder
|
||||
|
||||
# Actual tasks which controls the tests actually run
|
||||
supported_tasks: list[DiffusionTasks]
|
||||
|
||||
# Additional acceleration groups to run beyond the base case (no acceleration).
|
||||
# The base case is always run for every model in the test settings. None means
|
||||
# we only run the base case.
|
||||
extra_test_groups: list[list[DiffusionAccs]] | None = None
|
||||
|
||||
# Pytest Marks for this model. This may be useful for selecting which models
|
||||
# we want to run where, similar to the way vLLM's multimodal tests mark some
|
||||
# as core models to always run in the CI.
|
||||
# Example: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/test_common.py#L131
|
||||
marks: list[MarkDecorator] | None = None
|
||||
|
||||
# When True (default), online tests only run the base case (no accelerations).
|
||||
# When False, online tests run all test_groups, same as offline. This should be
|
||||
# True unless there is a good reason for it not to be, because the execution of
|
||||
# the acceleration should be the same on both codepaths, and CLI parsing etc should
|
||||
# be tested by tests adding the acceleration, and not per model.
|
||||
online_base_only: bool = True
|
||||
|
||||
# Additional checks to run for the base case.
|
||||
check_multi_output: bool = True # Runs multiple generations in one request
|
||||
check_determinism: bool = True # Runs 2 generations with the same seed and check determinism
|
||||
|
||||
|
||||
### Mappings & utils for building offline Omni() instances given a list of enabled accelerations
|
||||
ACC_OMNI_KWARGS = {
|
||||
DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_use_tiling": True},
|
||||
DiffusionAccs.CPU_OFFLOAD: {"enable_cpu_offload": True},
|
||||
DiffusionAccs.LAYERWISE_OFFLOAD: {"enable_layerwise_offload": True},
|
||||
DiffusionAccs.CACHE_DIT: {"cache_backend": "cache_dit"},
|
||||
DiffusionAccs.TEA_CACHE: {"cache_backend": "tea_cache"},
|
||||
}
|
||||
|
||||
ACC_PARALLEL_KWARGS = {
|
||||
DiffusionAccs.HSDP: {"use_hsdp": True, "hsdp_shard_size": 2},
|
||||
DiffusionAccs.TENSOR_PARALLEL: {"tensor_parallel_size": 2},
|
||||
DiffusionAccs.CFG_PARALLEL: {"cfg_parallel_size": 2},
|
||||
DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_patch_parallel_size": 2},
|
||||
# For SP, we don't run ring here to conserve devices, since it is easy
|
||||
# to blow up the number of needed devices fast. Compatibility for ring
|
||||
# and ulysses together should be tested generically outside of these tests.
|
||||
DiffusionAccs.SEQUENCE_PARALLEL: {"ulysses_degree": 2},
|
||||
}
|
||||
|
||||
ACC_DEVICE_COUNT_KEYS = {
|
||||
DiffusionAccs.TENSOR_PARALLEL: "tensor_parallel_size",
|
||||
DiffusionAccs.CFG_PARALLEL: "cfg_parallel_size",
|
||||
DiffusionAccs.SEQUENCE_PARALLEL: "ulysses_degree",
|
||||
DiffusionAccs.VAE_PATCH_PARALLEL: "vae_patch_parallel_size",
|
||||
DiffusionAccs.HSDP: "hsdp_shard_size",
|
||||
}
|
||||
|
||||
### CLI args for launching an OmniServer subprocess per acceleration.
|
||||
# Follows the same pattern as e2e online serving tests, which hand-code
|
||||
# their server_args lists in OmniServerParams.
|
||||
ACC_SERVER_ARGS: dict[DiffusionAccs, list[str]] = {
|
||||
DiffusionAccs.HSDP: ["--use-hsdp", "--hsdp-shard-size", "2"],
|
||||
DiffusionAccs.TEA_CACHE: ["--cache-backend", "tea_cache"],
|
||||
DiffusionAccs.CACHE_DIT: ["--cache-backend", "cache_dit"],
|
||||
DiffusionAccs.SEQUENCE_PARALLEL: ["--usp", "2"],
|
||||
DiffusionAccs.CFG_PARALLEL: ["--cfg-parallel-size", "2"],
|
||||
DiffusionAccs.TENSOR_PARALLEL: ["--tensor-parallel-size", "2"],
|
||||
DiffusionAccs.CPU_OFFLOAD: ["--enable-cpu-offload"],
|
||||
DiffusionAccs.LAYERWISE_OFFLOAD: ["--enable-layerwise-offload"],
|
||||
DiffusionAccs.VAE_PATCH_PARALLEL: ["--vae-use-tiling", "--vae-patch-parallel-size", "2"],
|
||||
}
|
||||
|
||||
## TODO ^ These should be in the same object, it's getting messy.
|
||||
|
||||
|
||||
def get_required_device_count(accelerations: list[DiffusionAccs] | None) -> int:
|
||||
"""Compute the minimum number of devices needed for a set of accelerations.
|
||||
The total is the product of all parallel dimensions (defaulting to 1).
|
||||
|
||||
If not enough devices are available for a test group's accelerations,
|
||||
that test will be skipped."""
|
||||
count = 1
|
||||
if accelerations is None:
|
||||
return count
|
||||
|
||||
for acc in accelerations:
|
||||
key = ACC_DEVICE_COUNT_KEYS.get(acc)
|
||||
if key is not None:
|
||||
count *= ACC_PARALLEL_KWARGS[acc][key]
|
||||
return count
|
||||
|
||||
|
||||
def build_parallel_config_from_diff_accelerations(accelerations: list[DiffusionAccs]) -> DiffusionParallelConfig | None:
|
||||
"""Given a list of accelerations pertaining to the current test group,
|
||||
build the parallel config needed for the Omni() object (if any)."""
|
||||
config_kwargs = {}
|
||||
for acc in accelerations:
|
||||
update_dict = ACC_PARALLEL_KWARGS.get(acc, {})
|
||||
config_kwargs.update(update_dict)
|
||||
if config_kwargs:
|
||||
return DiffusionParallelConfig(**config_kwargs)
|
||||
return None
|
||||
|
||||
|
||||
### Offline Omni() object builder
|
||||
def build_omni_from_diff_accelerations(accelerations: list[DiffusionAccs] | None, **kwargs) -> Omni:
|
||||
"""Given one or more acceleration types, build the corresponding Omni() object."""
|
||||
# Coerce to a list and build the parallel config, since that depends on the accelerations
|
||||
if accelerations is None:
|
||||
accelerations = []
|
||||
parallel_config = build_parallel_config_from_diff_accelerations(accelerations)
|
||||
|
||||
# Then add anything else that's a top-level kwarg
|
||||
acc_kwargs = {}
|
||||
if parallel_config is not None:
|
||||
acc_kwargs["parallel_config"] = parallel_config
|
||||
for acc in accelerations:
|
||||
update_dict = ACC_OMNI_KWARGS.get(acc, {})
|
||||
acc_kwargs.update(update_dict)
|
||||
|
||||
# Keys passed through should mostly be things like enforce_eager;
|
||||
# if there's overlap, it's probably due to a misconfiguration
|
||||
shared_keys = acc_kwargs.keys() & kwargs.keys()
|
||||
if shared_keys:
|
||||
raise ValueError(f"Explicit Omni kwargs and inferred Omni kwargs for accelerations overlap: {shared_keys}")
|
||||
omni_kwargs = {**acc_kwargs, **kwargs}
|
||||
return Omni(**omni_kwargs)
|
||||
|
||||
|
||||
### Online server flag builder
|
||||
def build_server_args_from_diff_accelerations(accelerations: list[DiffusionAccs] | None) -> list[str]:
|
||||
"""Given one or more acceleration types, build the corresponding CLI args
|
||||
for launching an OmniServer subprocess."""
|
||||
if accelerations is None:
|
||||
return []
|
||||
args = []
|
||||
for acc in accelerations:
|
||||
args.extend(ACC_SERVER_ARGS.get(acc, []))
|
||||
return args
|
||||
@@ -0,0 +1,63 @@
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import torch
|
||||
from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
|
||||
from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
|
||||
from diffusers.pipelines.flux2.pipeline_flux2_klein import Flux2KleinPipeline
|
||||
from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
|
||||
from transformers import AutoTokenizer, Qwen3Config, Qwen3ForCausalLM
|
||||
|
||||
TINY_MODEL_DIR = os.path.join(tempfile.gettempdir(), "vllm-omni-tiny-models")
|
||||
|
||||
|
||||
def _get_tiny_model_path(name: str) -> str:
|
||||
path = os.path.join(TINY_MODEL_DIR, name)
|
||||
os.makedirs(path, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def tiny_flux2_klein_builder() -> str:
|
||||
"""Build a tiny Flux2Klein model."""
|
||||
model_dir = _get_tiny_model_path("Flux2KleinPipeline")
|
||||
|
||||
pipe = Flux2KleinPipeline(
|
||||
scheduler=FlowMatchEulerDiscreteScheduler(),
|
||||
vae=AutoencoderKLFlux2(
|
||||
down_block_types=("DownEncoderBlock2D",),
|
||||
up_block_types=("UpDecoderBlock2D",),
|
||||
block_out_channels=(32,),
|
||||
layers_per_block=1,
|
||||
latent_channels=16,
|
||||
norm_num_groups=16,
|
||||
),
|
||||
# NOTE: For now we need 28 layers because of hardcoded stuff in the model :(
|
||||
text_encoder=Qwen3ForCausalLM(
|
||||
Qwen3Config(
|
||||
hidden_size=32,
|
||||
intermediate_size=64,
|
||||
num_hidden_layers=28,
|
||||
num_attention_heads=2,
|
||||
num_key_value_heads=2,
|
||||
head_dim=16,
|
||||
vocab_size=151936,
|
||||
max_position_embeddings=512,
|
||||
)
|
||||
),
|
||||
tokenizer=AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B"),
|
||||
# NOTE: For now we need at least 2 layers for the transformer
|
||||
# due to hardcoded hacks in CacheDiT for Flux2Klein specifically.
|
||||
transformer=Flux2Transformer2DModel(
|
||||
in_channels=64,
|
||||
num_layers=2,
|
||||
num_single_layers=2,
|
||||
attention_head_dim=16,
|
||||
num_attention_heads=2,
|
||||
joint_attention_dim=96,
|
||||
timestep_guidance_channels=32,
|
||||
axes_dims_rope=(4, 4, 4, 4),
|
||||
),
|
||||
)
|
||||
# Need dtypes to be consistent; for now we just put it on bfloat16
|
||||
pipe.to(torch.bfloat16).save_pretrained(model_dir)
|
||||
return model_dir
|
||||
@@ -0,0 +1,42 @@
|
||||
from tests.model_tests.diffusion import diff_model_builders
|
||||
from tests.model_tests.diffusion.config_types import (
|
||||
DiffusionAccs,
|
||||
DiffusionModelTestOpts,
|
||||
DiffusionTasks,
|
||||
)
|
||||
|
||||
# This object defines the (tiny) model configurations for common tests.
|
||||
#
|
||||
# TO ADD A NEW MODEL:
|
||||
# Map the pipeline name (i.e., the name of the pipeline in _DIFFUSION_MODELS)
|
||||
# to a DiffusionModelTestOpts and define the supported tasks. A base case with no
|
||||
# accelerations is always run. Use extra_test_groups to define additional
|
||||
# acceleration configurations to run a basic smoke test for (i.e., simple inference).
|
||||
#
|
||||
# The model will only be loaded once per test group, and will execute each of its
|
||||
# supported tasks as a pytest subtest. The base case also runs additional checks
|
||||
# that are independent of accelerations, e.g., for determinism and compatibility with
|
||||
# multi-output, unless explicitly disabled in the test settings.
|
||||
#
|
||||
# TO RUN A SUBSET OF THE TESTS:
|
||||
# These tests should be fast, but if you'd like to run a subset of the tests, you can do so
|
||||
# with `-k`; the IDs of the tests are the name of the pipeline joined by +.
|
||||
#
|
||||
# Example:
|
||||
# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline]
|
||||
# ^ Runs only the base (no accelerations) case for Flux2KleinPipeline
|
||||
#
|
||||
# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline
|
||||
# ^ Runs all test groups for Flux2KleinPipeline only
|
||||
DIFFUSION_TEST_SETTINGS = {
|
||||
"Flux2KleinPipeline": DiffusionModelTestOpts(
|
||||
model="black-forest-labs/FLUX.2-klein-4B",
|
||||
builder=diff_model_builders.tiny_flux2_klein_builder,
|
||||
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE],
|
||||
extra_test_groups=[
|
||||
[DiffusionAccs.HSDP, DiffusionAccs.TEA_CACHE],
|
||||
[DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
|
||||
[DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
|
||||
],
|
||||
)
|
||||
}
|
||||
@@ -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,
|
||||
)
|
||||
@@ -0,0 +1,98 @@
|
||||
"""
|
||||
Tests for alignment between _DIFFUSION_MODELS and DIFFUSION_TEST_SETTINGS; if
|
||||
tests in this file are failing, you are probably adding a new model, and need
|
||||
to account for it in DIFFUSION_TEST_SETTINGS
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
|
||||
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
|
||||
# These diffusion models currently do not have any added (tiny) model tests.
|
||||
# In general, pipelines should only be in the list below if they have not been
|
||||
# migrated yet or there is a compelling reason that they should not be (which for
|
||||
# now may just be CI pressure etc due to the volume of models and the cost of
|
||||
# initializing the server).
|
||||
#
|
||||
# The tests here validate that all keys in the diffusion registry are accounted for;
|
||||
# if you are adding a new model and see the tests below fail, please follow the pattern
|
||||
# for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
|
||||
EXCLUDED_MODELS = [
|
||||
"QwenImagePipeline",
|
||||
"QwenImageEditPipeline",
|
||||
"QwenImageEditPlusPipeline",
|
||||
"QwenImageLayeredPipeline",
|
||||
"GlmImagePipeline",
|
||||
"ZImagePipeline",
|
||||
"OvisImagePipeline",
|
||||
"WanPipeline",
|
||||
"WanVACEPipeline",
|
||||
"LTX2Pipeline",
|
||||
"LTX2ImageToVideoPipeline",
|
||||
"LTX2TwoStagesPipeline",
|
||||
"LTX2ImageToVideoTwoStagesPipeline",
|
||||
"LTX2T2VDMD2Pipeline",
|
||||
"LTX2I2VDMD2Pipeline",
|
||||
"LTX23Pipeline",
|
||||
"LTX23ImageToVideoPipeline",
|
||||
"StableAudioPipeline",
|
||||
"WanImageToVideoPipeline",
|
||||
"WanS2VPipeline",
|
||||
"WanT2VDMD2Pipeline",
|
||||
"WanI2VDMD2Pipeline",
|
||||
"LongCatImagePipeline",
|
||||
"BagelPipeline",
|
||||
"LancePipeline",
|
||||
"MingImagePipeline",
|
||||
"InternVLAA1Pipeline",
|
||||
"LongCatImageEditPipeline",
|
||||
"StableDiffusion3Pipeline",
|
||||
"FluxKontextPipeline",
|
||||
"HunyuanImage3ForCausalMM",
|
||||
"ErnieImagePipeline",
|
||||
"NextStep11Pipeline",
|
||||
"FluxPipeline",
|
||||
"FluxDMD2Pipeline",
|
||||
"QwenImageDMD2Pipeline",
|
||||
"OmniGen2Pipeline",
|
||||
"HeliosPipeline",
|
||||
"HeliosPyramidPipeline",
|
||||
"Flux2Pipeline",
|
||||
"DreamIDOmniPipeline",
|
||||
"SenseNovaU1Pipeline",
|
||||
"AudioXPipeline",
|
||||
"HunyuanVideo15Pipeline",
|
||||
"HunyuanVideo15ImageToVideoPipeline",
|
||||
"MagiHumanPipeline",
|
||||
"OmniVoicePipeline",
|
||||
"OmniVoice",
|
||||
"Cosmos3OmniDiffusersPipeline",
|
||||
"DiffusersAdapterPipeline",
|
||||
"HiDreamImagePipeline",
|
||||
"DreamZeroPipeline",
|
||||
"StableDiffusionXLPipeline",
|
||||
"Gr00tN1d7Pipeline",
|
||||
"SoulXSingerPipeline",
|
||||
"SoulXSingerSVCPipeline",
|
||||
]
|
||||
|
||||
|
||||
def test_all_non_excluded_pipelines_have_tests():
|
||||
"""Ensure that every pipeline in the diffusion registry either has a test
|
||||
configuration in DIFFUSION_TEST_SETTINGS, or is explicitly excluded for a
|
||||
known reason."""
|
||||
non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS)
|
||||
missing_tests = non_excluded_diff_pipes - DIFFUSION_TEST_SETTINGS.keys()
|
||||
assert len(missing_tests) == 0, f"Pipelines missing test settings: {missing_tests}"
|
||||
|
||||
|
||||
def test_no_excluded_models_have_test_settings():
|
||||
"""Ensure that no models in the exclude list have a test configuration in
|
||||
DIFFUSION_TEST_SETTINGS (i.e., if they have test settings, we should not
|
||||
exclude them)."""
|
||||
non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS)
|
||||
missing_pipes = DIFFUSION_TEST_SETTINGS.keys() - set(non_excluded_diff_pipes)
|
||||
assert len(missing_pipes) == 0
|
||||
@@ -0,0 +1,69 @@
|
||||
import pytest
|
||||
|
||||
from tests.model_tests.diffusion.case_filtering import get_parametrized_options
|
||||
from tests.model_tests.diffusion.config_types import (
|
||||
DiffusionAccs,
|
||||
DiffusionTasks,
|
||||
build_omni_from_diff_accelerations,
|
||||
)
|
||||
from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
|
||||
from tests.model_tests.diffusion.task_runners import (
|
||||
run_and_validate_image_to_image_request,
|
||||
run_and_validate_text_to_image_determinism,
|
||||
run_and_validate_text_to_image_multi_output,
|
||||
run_and_validate_text_to_image_request,
|
||||
)
|
||||
|
||||
# NOTE : Hardware marks are added dynamically based on test requirements
|
||||
pytestmark = [pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
|
||||
get_parametrized_options(DIFFUSION_TEST_SETTINGS),
|
||||
)
|
||||
def test_pipeline_on_supported_tasks(
|
||||
model_name,
|
||||
accelerations: list[DiffusionAccs] | None,
|
||||
supported_tasks: list[DiffusionTasks],
|
||||
check_multioutput: bool,
|
||||
check_determinism: bool,
|
||||
tiny_model_paths: dict[str, str],
|
||||
subtests,
|
||||
):
|
||||
"""Run a smoke test on all of the pipelines supported tasks using a set of enabled accelerations."""
|
||||
assert len(supported_tasks) > 0
|
||||
# We initialize the Omni object before running the tasks, then run each task as a pytest subtask.
|
||||
# This lets us init the model once, but display separate failures in pytest, and avoid halting the
|
||||
# checks on other tasks if one fails.
|
||||
#
|
||||
# This allows us to have some degree of test isolation without the cost of redundant initialization,
|
||||
# since starting the server can take 10+ seconds, even for tiny models.
|
||||
#
|
||||
# NOTE: Be sure to install pytest-subtests if you're running on pytest < 9
|
||||
omni = build_omni_from_diff_accelerations(
|
||||
accelerations=accelerations,
|
||||
model=tiny_model_paths[model_name],
|
||||
enforce_eager=True,
|
||||
)
|
||||
try:
|
||||
for task_type in supported_tasks:
|
||||
with subtests.test(msg=task_type):
|
||||
if task_type == DiffusionTasks.TEXT_TO_IMAGE:
|
||||
run_and_validate_text_to_image_request(omni)
|
||||
elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
|
||||
run_and_validate_image_to_image_request(omni)
|
||||
else:
|
||||
raise ValueError(f"Task type {task_type} is not yet supported")
|
||||
|
||||
# NOTE: For now, we only check determinism + multi output for the base case,
|
||||
# since checking it on every extra acceleration configuration is redundant
|
||||
# (see case_filtering).
|
||||
if check_determinism:
|
||||
with subtests.test(msg="determinism"):
|
||||
run_and_validate_text_to_image_determinism(omni)
|
||||
if check_multioutput:
|
||||
with subtests.test(msg="multi_output"):
|
||||
run_and_validate_text_to_image_multi_output(omni)
|
||||
finally:
|
||||
omni.close()
|
||||
@@ -0,0 +1,75 @@
|
||||
"""
|
||||
Analogous to test_common_offline, but for server tests. Validates the full
|
||||
online serving stack (CLI arg parsing, subprocess, API routing, response
|
||||
encoding) using tiny models.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.runtime import OmniServer, OpenAIClientHandler
|
||||
from tests.model_tests.diffusion.case_filtering import get_parametrized_options
|
||||
from tests.model_tests.diffusion.config_types import (
|
||||
DiffusionAccs,
|
||||
DiffusionTasks,
|
||||
build_server_args_from_diff_accelerations,
|
||||
)
|
||||
from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
|
||||
from tests.model_tests.diffusion.task_runners import (
|
||||
run_and_validate_online_image_to_image_request,
|
||||
run_and_validate_online_text_to_image_determinism,
|
||||
run_and_validate_online_text_to_image_multi_output,
|
||||
run_and_validate_online_text_to_image_request,
|
||||
)
|
||||
|
||||
# NOTE : Hardware marks are added dynamically based on test requirements
|
||||
pytestmark = [pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
|
||||
get_parametrized_options(DIFFUSION_TEST_SETTINGS, online=True),
|
||||
)
|
||||
def test_online_on_supported_tasks(
|
||||
model_name: str,
|
||||
accelerations: list[DiffusionAccs] | None,
|
||||
supported_tasks: list[DiffusionTasks],
|
||||
check_multioutput: bool,
|
||||
check_determinism: bool,
|
||||
tiny_model_paths: dict[str, str],
|
||||
run_level: str,
|
||||
subtests,
|
||||
):
|
||||
"""Smoke test: start a tiny model server and run each supported task via the API."""
|
||||
model_path = tiny_model_paths[model_name]
|
||||
server_args = build_server_args_from_diff_accelerations(accelerations)
|
||||
server_args.append("--enforce-eager")
|
||||
|
||||
with OmniServer(model_path, server_args) as server:
|
||||
# TODO: We may want to revisit run_level validation here,
|
||||
# because checks for things like image size etc should not
|
||||
# depend on whether or not the weights are real or random
|
||||
client = OpenAIClientHandler(
|
||||
host=server.host,
|
||||
port=server.port,
|
||||
api_key="EMPTY",
|
||||
run_level=run_level,
|
||||
log_stats=server.log_stats,
|
||||
)
|
||||
for task_type in supported_tasks:
|
||||
with subtests.test(msg=task_type):
|
||||
if task_type == DiffusionTasks.TEXT_TO_IMAGE:
|
||||
run_and_validate_online_text_to_image_request(server, client)
|
||||
elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
|
||||
run_and_validate_online_image_to_image_request(server, client)
|
||||
else:
|
||||
raise ValueError(f"Task type {task_type} is not yet supported")
|
||||
|
||||
# NOTE: For now, we only check determinism + multi output for the base case,
|
||||
# since checking it on every extra acceleration configuration is redundant
|
||||
# (see case_filtering).
|
||||
if check_determinism:
|
||||
with subtests.test(msg="determinism"):
|
||||
run_and_validate_online_text_to_image_determinism(server, client)
|
||||
if check_multioutput:
|
||||
with subtests.test(msg="multi_output"):
|
||||
run_and_validate_online_text_to_image_multi_output(server, client)
|
||||
@@ -0,0 +1,33 @@
|
||||
"""
|
||||
Utilities for resolving real models to their tiny model equivalents.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
|
||||
from vllm_omni.diffusion.data import resolve_model_class_name
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def resolve_tiny_model_path(model: str) -> str:
|
||||
"""Given a real model name/path, resolve it to a tiny model path.
|
||||
|
||||
Raises ValueError if the pipeline class cannot be determined (invalid
|
||||
model). Returns the original model path if no tiny builder exists yet."""
|
||||
pipeline_class = resolve_model_class_name(model)
|
||||
if pipeline_class is None:
|
||||
raise ValueError(
|
||||
f"Cannot resolve pipeline class for model: {model}. The model path may be invalid or its config unreadable."
|
||||
)
|
||||
|
||||
test_opts = DIFFUSION_TEST_SETTINGS.get(pipeline_class)
|
||||
if test_opts is None:
|
||||
logger.warning(
|
||||
"No tiny model builder for pipeline %s (model: %s). Using original model.",
|
||||
pipeline_class,
|
||||
model,
|
||||
)
|
||||
return model
|
||||
|
||||
return test_opts.builder()
|
||||
Reference in new issue
Block a user