[CI][Diffusion] Add tiny model builder for Krea2Pipeline (#6421)

Signed-off-by: khairulkabir1661 <khairulkabir1661@users.noreply.github.com>
Co-authored-by: khairulkabir1661 <khairulkabir1661@users.noreply.github.com>

tests/helpers/tiny_model.py

@@ -44,14 +44,16 @@ def build_tiny_from_configs(
             config or transform) that don't have vendored configs.
         configs_dir: Path to a directory containing vendored config files
             (model_index.json and per-component config.json).
-        transform: Maps a component (subfolder) name to a function that
-            takes that component's raw config dict (fetched from model_id)
-            and returns a shrunk one, applied only when that component has
-            no vendored config in configs_dir. Lets each pipeline encode
-            exactly the shrink it needs (layer counts, and any component-
-            internal width that isn't read by another component), including
-            keeping sibling config fields in sync (e.g. layer_types matching
-            num_hidden_layers).
+        transform: Maps a component (subfolder) name or None to a function that
+            transforms a config dict. For subfolder keys, the function takes
+            that component's raw config dict (fetched from model_id) and returns
+            a shrunk one, applied only when that component has no vendored config
+            in configs_dir. For None key, the function transforms the pipeline's
+            model_index.json config dict. Lets each pipeline encode exactly the
+            transforms it needs (layer counts, config-level fields, and any
+            component-internal width that isn't read by another component),
+            including keeping sibling config fields in sync (e.g. layer_types
+            matching num_hidden_layers).
 
     Returns:
         Path to the saved tiny model directory with safetensors weights.
@@ -72,7 +74,7 @@ def build_tiny_from_configs(
         subfolder = kwargs.get("subfolder")
         load_kwargs = {"subfolder": subfolder} if subfolder else {}
         config = cls.load_config(pretrained_model_name_or_path, **load_kwargs)
-        if subfolder in transform:
+        if subfolder is not None and subfolder in transform:
             config = transform[subfolder](dict(config))
         return cls.from_config(config)
 
@@ -86,7 +88,7 @@ def build_tiny_from_configs(
         subfolder = kwargs.get("subfolder")
         load_kwargs = {"subfolder": subfolder} if subfolder else {}
         config_dict, _ = PretrainedConfig.get_config_dict(pretrained_model_name_or_path, **load_kwargs)
-        if subfolder in transform:
+        if subfolder is not None and subfolder in transform:
             config_dict = transform[subfolder](config_dict)
         config = cls.config_class.from_dict(config_dict)
         return cls(config)
@@ -96,6 +98,8 @@ def build_tiny_from_configs(
     # can vendor just one component without also needing its own model_index.json.
     vendored_index = config_dir is not None and (config_dir / "model_index.json").exists()
     config_dict = DiffusionPipeline.load_config(config_dir if vendored_index else model_id)
+    if None in transform:
+        config_dict = transform[None](dict(config_dict))
     pipeline_cls = _get_pipeline_class(DiffusionPipeline, config=config_dict)
 
     init_dict, _, _ = pipeline_cls.extract_init_dict(config_dict)

tests/model_tests/diffusion/diff_model_builders.py

@@ -302,3 +302,69 @@ def tiny_flux2_builder() -> str:
             "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, joint_attention_dim=96),
         },
     )
+
+
+def _shrink_krea2_text_encoder(config: dict) -> dict:
+    tc = config["text_config"]
+    old_head_dim = tc["head_dim"]
+    tc["num_hidden_layers"] = 3
+    tc["hidden_size"] = 32
+    tc["intermediate_size"] = 64
+    tc["num_attention_heads"] = 2
+    tc["num_key_value_heads"] = 2
+    tc["head_dim"] = 16
+    if "rope_parameters" in tc:
+        rp = tc["rope_parameters"]
+        if "mrope_section" in rp:
+            factor = old_head_dim / 16
+            rp["mrope_section"] = [max(1, round(d / factor)) for d in rp["mrope_section"]]
+    vc = config["vision_config"]
+    vc["depth"] = 2
+    vc["intermediate_size"] = 64
+    vc["num_heads"] = 2
+    vc["out_hidden_size"] = 32
+    if "deepstack_visual_indexes" in vc:
+        vc["deepstack_visual_indexes"] = [0, 1]
+    return config
+
+
+def _shrink_krea2_transformer(config: dict) -> dict:
+    config["num_layers"] = 2
+    config["attention_head_dim"] = 32
+    config["num_attention_heads"] = 4
+    config["num_key_value_heads"] = 2
+    config["intermediate_size"] = 128
+    config["timestep_embed_dim"] = 32
+    config["text_hidden_dim"] = 32
+    config["num_text_layers"] = 3
+    config["text_num_attention_heads"] = 2
+    config["text_num_key_value_heads"] = 2
+    config["text_intermediate_size"] = 64
+    config["num_layerwise_text_blocks"] = 1
+    config["num_refiner_text_blocks"] = 1
+    config["axes_dims_rope"] = [8, 12, 12]
+    return config
+
+
+def _shrink_krea2_vae(config: dict) -> dict:
+    config["base_dim"] = 32
+    return config
+
+
+def _patch_krea2_model_index(config: dict) -> dict:
+    config["text_encoder_select_layers"] = [0, 1, 2]
+    return config
+
+
+def tiny_krea2_builder() -> str:
+    """Build a tiny Krea2 model with random weights using diffusers building blocks."""
+    return build_tiny_from_configs(
+        "Krea2Pipeline",
+        "krea/Krea-2-Turbo",
+        transform={
+            None: _patch_krea2_model_index,
+            "text_encoder": _shrink_krea2_text_encoder,
+            "transformer": _shrink_krea2_transformer,
+            "vae": _shrink_krea2_vae,
+        },
+    )

tests/model_tests/diffusion/model_settings.py

@@ -175,4 +175,9 @@ DIFFUSION_TEST_SETTINGS = {
             [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],
+    ),
 }

tests/model_tests/diffusion/test_alignment.py

@@ -60,7 +60,6 @@ EXCLUDED_MODELS = [
     "ErnieImagePipeline",
     "NextStep11Pipeline",
     "FluxDMD2Pipeline",
-    "Krea2Pipeline",
     "QwenImageDMD2Pipeline",
     "OmniGen2Pipeline",
     "HeliosPipeline",