[Test] Add tiny model builder for FluxKontextPipeline (#5823)

Signed-off-by: Nick Cao <ncao@redhat.com>

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,
         ),
     ]
 

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

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,

tests/model_tests/diffusion/test_alignment.py

@@ -48,7 +48,6 @@ EXCLUDED_MODELS = [
     "InternVLAA1Pipeline",
     "LongCatImageEditPipeline",
     "StableDiffusion3Pipeline",
-    "FluxKontextPipeline",
     "HunyuanImage3ForCausalMM",
     "ErnieImagePipeline",
     "NextStep11Pipeline",