[Bugfix] Apply function module marker per-model instead of hardcoding diffusion (#5875)
Signed-off-by: khairulkabir1661 <khairulkabir1661@users.noreply.github.com> Co-authored-by: khairulkabir1661 <khairulkabir1661@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
5 files changed
+50
-7
No files matched your search
@@ -7,6 +7,7 @@ from tests.model_tests.diffusion.config_types import (
|
||||
DiffusionAccs,
|
||||
DiffusionModelTestOpts,
|
||||
DiffusionTasks,
|
||||
ModelTypeMarker,
|
||||
)
|
||||
|
||||
# This object defines the (tiny) model configurations for common tests.
|
||||
@@ -37,12 +38,14 @@ DIFFUSION_TEST_SETTINGS = {
|
||||
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],
|
||||
@@ -53,6 +56,7 @@ DIFFUSION_TEST_SETTINGS = {
|
||||
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],
|
||||
@@ -63,6 +67,7 @@ DIFFUSION_TEST_SETTINGS = {
|
||||
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],
|
||||
@@ -82,6 +87,7 @@ DIFFUSION_TEST_SETTINGS = {
|
||||
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],
|
||||
@@ -95,11 +101,13 @@ DIFFUSION_TEST_SETTINGS = {
|
||||
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],
|
||||
@@ -110,6 +118,7 @@ DIFFUSION_TEST_SETTINGS = {
|
||||
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],
|
||||
@@ -120,11 +129,13 @@ DIFFUSION_TEST_SETTINGS = {
|
||||
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],
|
||||
@@ -135,6 +146,7 @@ DIFFUSION_TEST_SETTINGS = {
|
||||
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],
|
||||
@@ -145,16 +157,19 @@ DIFFUSION_TEST_SETTINGS = {
|
||||
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],
|
||||
|
||||
Reference in new issue
Block a user