diff --git a/tests/e2e/online_serving/test_flux_kontext_expansion.py b/tests/e2e/online_serving/test_flux_kontext_expansion.py index 5714ef711..ede5ff56b 100644 --- a/tests/e2e/online_serving/test_flux_kontext_expansion.py +++ b/tests/e2e/online_serving/test_flux_kontext_expansion.py @@ -14,47 +14,15 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.slow] EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting." NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist" MODEL = "black-forest-labs/FLUX.1-Kontext-dev" -PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=2) +BASE_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}) def _get_diffusion_feature_cases(model: str): return [ pytest.param( - OmniServerParams( - model=model, - server_args=[ - "--tensor-parallel-size", - "2", - "--enable-cpu-offload", - ], - ), - id="parallel_tp_2", - marks=PARALLEL_FEATURE_MARKS, - ), - pytest.param( - OmniServerParams( - model=model, - server_args=[ - "--enable-cpu-offload", - "--cfg-parallel-size", - "2", - ], - ), - id="parallel_cfg_2", - marks=PARALLEL_FEATURE_MARKS, - ), - pytest.param( - OmniServerParams( - model=model, - server_args=[ - "--tensor-parallel-size", - "2", - "--enable-cpu-offload", - "--cache-backend", - "cache_dit", - ], - ), - id="parallel_002", + OmniServerParams(model=model), + id="base", + marks=BASE_FEATURE_MARKS, ), ] diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py index ea69d95e5..9353d54f6 100644 --- a/tests/model_tests/diffusion/diff_model_builders.py +++ b/tests/model_tests/diffusion/diff_model_builders.py @@ -37,6 +37,19 @@ def tiny_ltx2_builder() -> str: return build_tiny_from_configs("LTX2Pipeline", "Lightricks/LTX-2", TINY_CONFIGS_DIR / "LTX2Pipeline") +def _shrink_flux_clip_text_encoder(config: dict) -> dict: + config["num_hidden_layers"] = 2 + return config + + +def _shrink_flux_t5_text_encoder(config: dict) -> dict: + config["num_layers"] = 2 + config["num_decoder_layers"] = 2 + config["d_ff"] = 64 + config["num_heads"] = 4 + return config + + def tiny_qwen_image_builder() -> str: def shrink_text_encoder(config: dict) -> dict: config["num_hidden_layers"] = 2 @@ -57,23 +70,24 @@ def tiny_qwen_image_builder() -> str: def tiny_flux_builder() -> str: - def shrink_clip_text_encoder(config: dict) -> dict: - config["num_hidden_layers"] = 2 - return config - - def shrink_t5_text_encoder(config: dict) -> dict: - config["num_layers"] = 2 - config["num_decoder_layers"] = 2 - config["d_ff"] = 64 - config["num_heads"] = 4 - return config - return build_tiny_from_configs( "FluxPipeline", "black-forest-labs/FLUX.1-schnell", transform={ - "text_encoder": shrink_clip_text_encoder, - "text_encoder_2": shrink_t5_text_encoder, + "text_encoder": _shrink_flux_clip_text_encoder, + "text_encoder_2": _shrink_flux_t5_text_encoder, + "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, default_axes_dims_rope=[16, 56, 56]), + }, + ) + + +def tiny_flux_kontext_builder() -> str: + return build_tiny_from_configs( + "FluxKontextPipeline", + "black-forest-labs/FLUX.1-Kontext-dev", + transform={ + "text_encoder": _shrink_flux_clip_text_encoder, + "text_encoder_2": _shrink_flux_t5_text_encoder, "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, default_axes_dims_rope=[16, 56, 56]), }, ) diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index 037faf67d..10500c688 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/tests/model_tests/diffusion/model_settings.py @@ -59,6 +59,16 @@ DIFFUSION_TEST_SETTINGS = { builder=diff_model_builders.tiny_flux_builder, supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE], ), + "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], + 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, diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 539f25c72..17d1eb698 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -48,7 +48,6 @@ EXCLUDED_MODELS = [ "InternVLAA1Pipeline", "LongCatImageEditPipeline", "StableDiffusion3Pipeline", - "FluxKontextPipeline", "HunyuanImage3ForCausalMM", "ErnieImagePipeline", "NextStep11Pipeline",