[Tests] Further extend the tiny model testing framework and enable multiple models (#5353)

Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/model_tests/diffusion/diff_model_builders.py

@@ -1,3 +1,4 @@
+from functools import partial
 from pathlib import Path
 
 from tests.helpers.tiny_model import build_tiny_from_configs
@@ -5,6 +6,25 @@ from tests.helpers.tiny_model import build_tiny_from_configs
 TINY_CONFIGS_DIR = Path(__file__).parent / "tiny_configs"
 
 
+def _shrink_dit_rope_config(
+    config: dict,
+    num_layers: int = 2,
+    num_single_layers: int | None = None,
+    target_head_dim: int = 32,
+    target_heads: int = 4,
+    default_axes_dims_rope: list[int] | None = None,
+) -> dict:
+    config["num_layers"] = num_layers
+    if num_single_layers is not None:
+        config["num_single_layers"] = num_single_layers
+    axes_dims_rope = config.get("axes_dims_rope", default_axes_dims_rope)
+    factor = config["attention_head_dim"] / target_head_dim
+    config["attention_head_dim"] = target_head_dim
+    config["num_attention_heads"] = target_heads
+    config["axes_dims_rope"] = [int(d / factor) for d in axes_dims_rope]
+    return config
+
+
 def tiny_flux2_klein_builder() -> str:
     """Build a tiny Flux2Klein model from vendored configs."""
     return build_tiny_from_configs(
@@ -15,3 +35,72 @@ def tiny_flux2_klein_builder() -> str:
 def tiny_ltx2_builder() -> str:
     """Build a tiny LTX2 model from vendored configs."""
     return build_tiny_from_configs("LTX2Pipeline", "Lightricks/LTX-2", TINY_CONFIGS_DIR / "LTX2Pipeline")
+
+
+def tiny_qwen_image_builder() -> str:
+    def shrink_text_encoder(config: dict) -> dict:
+        config["num_hidden_layers"] = 2
+        config["intermediate_size"] = 64
+        config["text_config"]["num_hidden_layers"] = 2
+        config["text_config"]["intermediate_size"] = 64
+        config["text_config"]["layer_types"] = config["text_config"]["layer_types"][:2]
+        config["vision_config"]["depth"] = 2
+        config["vision_config"]["intermediate_size"] = 64
+        config["vision_config"]["fullatt_block_indexes"] = [0, 1]
+        return config
+
+    return build_tiny_from_configs(
+        "QwenImagePipeline",
+        "Qwen/Qwen-Image",
+        transform={"text_encoder": shrink_text_encoder, "transformer": _shrink_dit_rope_config},
+    )
+
+
+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,
+            "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, default_axes_dims_rope=[16, 56, 56]),
+        },
+    )
+
+
+def tiny_flux2_builder() -> str:
+    def shrink_text_encoder(config: dict) -> dict:
+        # The real checkpoint ships language_model.lm_head.weight as its own
+        # tensor even though tie_word_embeddings defaults to True; keep it
+        # untied here too so transformers doesn't drop it on save.
+        config["tie_word_embeddings"] = False
+        config["text_config"]["num_hidden_layers"] = 31
+        config["text_config"]["intermediate_size"] = 64
+        config["text_config"]["num_attention_heads"] = 4
+        config["text_config"]["head_dim"] = 32
+        config["text_config"]["num_key_value_heads"] = 2
+        config["vision_config"]["num_hidden_layers"] = 2
+        config["vision_config"]["intermediate_size"] = 64
+        config["vision_config"]["num_attention_heads"] = 4
+        config["vision_config"]["head_dim"] = 32
+        return config
+
+    return build_tiny_from_configs(
+        "Flux2Pipeline",
+        "black-forest-labs/FLUX.2-dev",
+        transform={
+            "text_encoder": shrink_text_encoder,
+            "transformer": partial(_shrink_dit_rope_config, num_single_layers=2),
+        },
+    )