[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)