From bcb4ae288f77dbee431c0649b72ef43a4722bb2b Mon Sep 17 00:00:00 2001 From: Khairul Kabir <250204754+khairulkabir1661@users.noreply.github.com> Date: Fri, 25 Sep 2026 14:54:38 -0700 Subject: [PATCH] [CI][Diffusion] Add tiny model builder for Krea2Pipeline (#6421) Signed-off-by: khairulkabir1661 Co-authored-by: khairulkabir1661 --- tests/helpers/tiny_model.py | 24 ++++--- .../diffusion/diff_model_builders.py | 66 +++++++++++++++++++ tests/model_tests/diffusion/model_settings.py | 5 ++ tests/model_tests/diffusion/test_alignment.py | 1 - 4 files changed, 85 insertions(+), 11 deletions(-) diff --git a/tests/helpers/tiny_model.py b/tests/helpers/tiny_model.py index 7ada9f987..aa273824a 100644 --- a/tests/helpers/tiny_model.py +++ b/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) diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py index cf8b4c92c..14a0a0cff 100644 --- a/tests/model_tests/diffusion/diff_model_builders.py +++ b/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, + }, + ) diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index e2a202c89..0a7571b7b 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/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], + ), } diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index c2489777a..0bf19a1a3 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -60,7 +60,6 @@ EXCLUDED_MODELS = [ "ErnieImagePipeline", "NextStep11Pipeline", "FluxDMD2Pipeline", - "Krea2Pipeline", "QwenImageDMD2Pipeline", "OmniGen2Pipeline", "HeliosPipeline",