[Tests] Enable tiny model testing for Qwen-Image Edit and EditPlus (#5656)

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

tests/model_tests/diffusion/diff_model_builders.py

@@ -13,10 +13,13 @@ def _shrink_dit_rope_config(
     target_head_dim: int = 32,
     target_heads: int = 4,
     default_axes_dims_rope: list[int] | None = None,
+    joint_attention_dim: int | None = None,
 ) -> dict:
     config["num_layers"] = num_layers
     if num_single_layers is not None:
         config["num_single_layers"] = num_single_layers
+    if joint_attention_dim is not None:
+        config["joint_attention_dim"] = joint_attention_dim
     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
@@ -50,22 +53,55 @@ def _shrink_flux_t5_text_encoder(config: dict) -> dict:
     return config
 
 
-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
+def _shrink_qwen_text_encoder_config(config: dict, hidden_size: int = 64) -> dict:
+    text_config = config["text_config"]
+    old_head_dim = text_config["hidden_size"] / text_config["num_attention_heads"]
+    text_config["num_hidden_layers"] = 2
+    text_config["intermediate_size"] = 64
+    text_config["hidden_size"] = hidden_size
+    text_config["num_attention_heads"] = 2
+    text_config["num_key_value_heads"] = 2
+    text_config["layer_types"] = text_config["layer_types"][:2]
+    factor = old_head_dim / (hidden_size / 2)
+    mrope_section = text_config["rope_scaling"]["mrope_section"]
+    text_config["rope_scaling"]["mrope_section"] = [round(d / factor) for d in mrope_section]
+    config["num_hidden_layers"] = 2
+    config["intermediate_size"] = 64
+    config["hidden_size"] = hidden_size
+    config["num_attention_heads"] = 2
+    config["num_key_value_heads"] = 2
+    config["vision_config"]["depth"] = 2
+    config["vision_config"]["intermediate_size"] = 64
+    config["vision_config"]["fullatt_block_indexes"] = [0, 1]
+    config["vision_config"]["out_hidden_size"] = hidden_size
+    return config
 
+
+def _shrink_qwen_transformer_config(config: dict, joint_attention_dim: int = 64) -> dict:
+    return _shrink_dit_rope_config(config, joint_attention_dim=joint_attention_dim)
+
+
+def tiny_qwen_image_builder() -> str:
     return build_tiny_from_configs(
         "QwenImagePipeline",
         "Qwen/Qwen-Image",
-        transform={"text_encoder": shrink_text_encoder, "transformer": _shrink_dit_rope_config},
+        transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
+    )
+
+
+def tiny_qwen_image_edit_builder() -> str:
+    return build_tiny_from_configs(
+        "QwenImageEditPipeline",
+        "Qwen/Qwen-Image-Edit",
+        transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
+    )
+
+
+def tiny_qwen_image_edit_plus_builder() -> str:
+    return build_tiny_from_configs(
+        "QwenImageEditPlusPipeline",
+        "Qwen/Qwen-Image-Edit-2511",
+        transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
     )
 
 

tests/model_tests/diffusion/model_settings.py

@@ -74,4 +74,14 @@ DIFFUSION_TEST_SETTINGS = {
         builder=diff_model_builders.tiny_flux2_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
     ),
+    "QwenImageEditPipeline": DiffusionModelTestOpts(
+        model="Qwen/Qwen-Image-Edit",
+        builder=diff_model_builders.tiny_qwen_image_edit_builder,
+        supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
+    ),
+    "QwenImageEditPlusPipeline": DiffusionModelTestOpts(
+        model="Qwen/Qwen-Image-Edit-2511",
+        builder=diff_model_builders.tiny_qwen_image_edit_plus_builder,
+        supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
+    ),
 }

tests/model_tests/diffusion/test_alignment.py

@@ -21,8 +21,6 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 # if you are adding a new model and see the tests below fail, please follow the pattern
 # for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
 EXCLUDED_MODELS = [
-    "QwenImageEditPipeline",
-    "QwenImageEditPlusPipeline",
     "QwenImageLayeredPipeline",
     "GlmImagePipeline",
     "ZImagePipeline",