diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index b073bc542..51c4223f2 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -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: diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 85f078533..e3a5aa353 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -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 diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index d12267fe0..7e0942bf6 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -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: diff --git a/pyproject.toml b/pyproject.toml index 2a8fc9c8f..0da2d7b5f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/tests/e2e/offline_inference/test_flux2_klein.py b/tests/e2e/offline_inference/test_flux2_klein.py index 3bd815229..9656772c9 100644 --- a/tests/e2e/offline_inference/test_flux2_klein.py +++ b/tests/e2e/offline_inference/test_flux2_klein.py @@ -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) diff --git a/tests/e2e/online_serving/test_flux2_klein.py b/tests/e2e/online_serving/test_flux2_klein.py deleted file mode 100644 index f8c705df2..000000000 --- a/tests/e2e/online_serving/test_flux2_klein.py +++ /dev/null @@ -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) diff --git a/tests/e2e/online_serving/test_flux2_klein_expansion.py b/tests/e2e/online_serving/test_flux2_klein_expansion.py deleted file mode 100644 index 98ccf0014..000000000 --- a/tests/e2e/online_serving/test_flux2_klein_expansion.py +++ /dev/null @@ -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 diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index f33489578..6ca8d2608 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -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 diff --git a/tests/model_tests/__init__.py b/tests/model_tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_tests/conftest.py b/tests/model_tests/conftest.py new file mode 100644 index 000000000..502f3e685 --- /dev/null +++ b/tests/model_tests/conftest.py @@ -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) diff --git a/tests/model_tests/diffusion/__init__.py b/tests/model_tests/diffusion/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_tests/diffusion/case_filtering.py b/tests/model_tests/diffusion/case_filtering.py new file mode 100644 index 000000000..a4fa5dc13 --- /dev/null +++ b/tests/model_tests/diffusion/case_filtering.py @@ -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)) diff --git a/tests/model_tests/diffusion/config_types.py b/tests/model_tests/diffusion/config_types.py new file mode 100644 index 000000000..4120900fe --- /dev/null +++ b/tests/model_tests/diffusion/config_types.py @@ -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 diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py new file mode 100644 index 000000000..9bee71f9b --- /dev/null +++ b/tests/model_tests/diffusion/diff_model_builders.py @@ -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 diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py new file mode 100644 index 000000000..3f44afcae --- /dev/null +++ b/tests/model_tests/diffusion/model_settings.py @@ -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], + ], + ) +} diff --git a/tests/model_tests/diffusion/task_runners.py b/tests/model_tests/diffusion/task_runners.py new file mode 100644 index 000000000..6b9c7ff99 --- /dev/null +++ b/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, + ) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py new file mode 100644 index 000000000..beb553994 --- /dev/null +++ b/tests/model_tests/diffusion/test_alignment.py @@ -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 diff --git a/tests/model_tests/diffusion/test_common_offline.py b/tests/model_tests/diffusion/test_common_offline.py new file mode 100644 index 000000000..f24ed3780 --- /dev/null +++ b/tests/model_tests/diffusion/test_common_offline.py @@ -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() diff --git a/tests/model_tests/diffusion/test_common_online.py b/tests/model_tests/diffusion/test_common_online.py new file mode 100644 index 000000000..e4c6be0d4 --- /dev/null +++ b/tests/model_tests/diffusion/test_common_online.py @@ -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) diff --git a/tests/model_tests/diffusion/utils.py b/tests/model_tests/diffusion/utils.py new file mode 100644 index 000000000..35577b68f --- /dev/null +++ b/tests/model_tests/diffusion/utils.py @@ -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()