# 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,
    ),
}