185 lines
8.7 KiB
Python
185 lines
8.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from tests.model_tests.diffusion import diff_model_builders
|
|
from tests.model_tests.diffusion.anima_builder import CHECKPOINT_FILENAME, real_anima_model, tiny_anima_builder
|
|
from tests.model_tests.diffusion.config_types import (
|
|
DiffusionAccs,
|
|
DiffusionModelTestOpts,
|
|
DiffusionTasks,
|
|
ModelTypeMarker,
|
|
)
|
|
|
|
# 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 = {
|
|
"AnimaPipeline": DiffusionModelTestOpts(
|
|
model=real_anima_model,
|
|
builder=tiny_anima_builder,
|
|
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
checkpoint_filename=CHECKPOINT_FILENAME,
|
|
),
|
|
"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],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
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],
|
|
],
|
|
),
|
|
"LTX2Pipeline": DiffusionModelTestOpts(
|
|
model="Lightricks/LTX-2",
|
|
builder=diff_model_builders.tiny_ltx2_builder,
|
|
supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO, DiffusionTasks.IMAGE_TO_VIDEO],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
extra_test_groups=[
|
|
[DiffusionAccs.HSDP, DiffusionAccs.CACHE_DIT],
|
|
[DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
|
|
[DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
|
|
],
|
|
),
|
|
"SanaVideoPipeline": DiffusionModelTestOpts(
|
|
model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
|
|
builder=diff_model_builders.tiny_sana_video_builder,
|
|
supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
extra_test_groups=[
|
|
[DiffusionAccs.TENSOR_PARALLEL],
|
|
[DiffusionAccs.CFG_PARALLEL],
|
|
[DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL],
|
|
[DiffusionAccs.SEQUENCE_PARALLEL],
|
|
[DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CFG_PARALLEL],
|
|
[DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.SEQUENCE_PARALLEL],
|
|
[DiffusionAccs.CACHE_DIT],
|
|
[DiffusionAccs.CPU_OFFLOAD],
|
|
[DiffusionAccs.LAYERWISE_OFFLOAD],
|
|
[DiffusionAccs.CACHE_DIT, DiffusionAccs.CPU_OFFLOAD],
|
|
[DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
|
|
],
|
|
check_multi_output=False,
|
|
),
|
|
"SanaImageToVideoPipeline": DiffusionModelTestOpts(
|
|
model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
|
|
builder=diff_model_builders.tiny_sana_video_i2v_builder,
|
|
supported_tasks=[DiffusionTasks.IMAGE_TO_VIDEO],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
extra_test_groups=[
|
|
[DiffusionAccs.CFG_PARALLEL],
|
|
[DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL],
|
|
[DiffusionAccs.SEQUENCE_PARALLEL],
|
|
[DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CFG_PARALLEL],
|
|
],
|
|
check_i2v_t2v_divergence=False,
|
|
check_multi_output=False,
|
|
),
|
|
"QwenImagePipeline": DiffusionModelTestOpts(
|
|
model="Qwen/Qwen-Image",
|
|
builder=diff_model_builders.tiny_qwen_image_builder,
|
|
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
),
|
|
"LongCatImagePipeline": DiffusionModelTestOpts(
|
|
model="meituan-longcat/LongCat-Image",
|
|
builder=diff_model_builders.tiny_longcat_image_builder,
|
|
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
extra_test_groups=[
|
|
[DiffusionAccs.TEA_CACHE],
|
|
[DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
|
|
[DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
|
|
],
|
|
),
|
|
"LongCatImageEditPipeline": DiffusionModelTestOpts(
|
|
model="meituan-longcat/LongCat-Image-Edit",
|
|
builder=diff_model_builders.tiny_longcat_image_edit_builder,
|
|
supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
extra_test_groups=[
|
|
[DiffusionAccs.TEA_CACHE],
|
|
[DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
|
|
[DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
|
|
],
|
|
),
|
|
"FluxPipeline": DiffusionModelTestOpts(
|
|
model="black-forest-labs/FLUX.1-schnell",
|
|
builder=diff_model_builders.tiny_flux_builder,
|
|
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
),
|
|
"FluxKontextPipeline": DiffusionModelTestOpts(
|
|
model="black-forest-labs/FLUX.1-Kontext-dev",
|
|
builder=diff_model_builders.tiny_flux_kontext_builder,
|
|
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
extra_test_groups=[
|
|
[DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
|
|
[DiffusionAccs.CFG_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
|
|
[DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD, DiffusionAccs.CACHE_DIT],
|
|
],
|
|
),
|
|
"Flux2Pipeline": DiffusionModelTestOpts(
|
|
model="black-forest-labs/FLUX.2-dev",
|
|
builder=diff_model_builders.tiny_flux2_builder,
|
|
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
extra_test_groups=[
|
|
[DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
|
|
[DiffusionAccs.CFG_PARALLEL, DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
|
|
[DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.VAE_PATCH_PARALLEL],
|
|
],
|
|
),
|
|
"QwenImageEditPipeline": DiffusionModelTestOpts(
|
|
model="Qwen/Qwen-Image-Edit",
|
|
builder=diff_model_builders.tiny_qwen_image_edit_builder,
|
|
supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
),
|
|
"QwenImageEditPlusPipeline": DiffusionModelTestOpts(
|
|
model="Qwen/Qwen-Image-Edit-2511",
|
|
builder=diff_model_builders.tiny_qwen_image_edit_plus_builder,
|
|
supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
),
|
|
"StableDiffusion3Pipeline": DiffusionModelTestOpts(
|
|
model="stabilityai/stable-diffusion-3.5-medium",
|
|
builder=diff_model_builders.tiny_sd3_builder,
|
|
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
extra_test_groups=[
|
|
[DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
|
|
[DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
|
|
],
|
|
),
|
|
"Krea2Pipeline": DiffusionModelTestOpts(
|
|
model="krea/Krea-2-Turbo",
|
|
builder=diff_model_builders.tiny_krea2_builder,
|
|
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
|
|
model_type_marker=ModelTypeMarker.DIFFUSION,
|
|
),
|
|
}
|