diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py index 9353d54f6..e47ea1895 100644 --- a/tests/model_tests/diffusion/diff_model_builders.py +++ b/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}, ) diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index 10500c688..cbb37729a 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/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], + ), } diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 17d1eb698..be064f603 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/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",