update

Author: DN6

tests/models/testing_utils/common.py

@@ -260,7 +260,7 @@ def test_output(self, expected_output_shape=None):
 
         assert output is not None, "Model output is None"
         assert (
-            output.shape == expected_output_shape
+            output[0].shape == expected_output_shape or self.output_shape
         ), f"Output shape does not match expected. Expected {expected_output_shape}, got {output.shape}"
 
     def test_outputs_equivalence(self):
@@ -302,15 +302,11 @@ def recursive_check(tuple_object, dict_object):
 
         recursive_check(outputs_tuple, outputs_dict)
 
-    def test_model_config_to_json_string(self):
-        model = self.model_class(**self.get_init_dict())
-
-        json_string = model.config.to_json_string()
-        assert isinstance(json_string, str), "Config to_json_string should return a string"
-        assert len(json_string) > 0, "JSON string should not be empty"
-
     @require_accelerator
-    @pytest.mark.skipif(torch_device not in ["cuda", "xpu"])
+    @pytest.mark.skipif(
+        torch_device not in ["cuda", "xpu"],
+        reason="float16 and bfloat16 can only be use for inference with an accelerator",
+    )
     def test_from_save_pretrained_float16_bfloat16(self):
         model = self.model_class(**self.get_init_dict())
         model.to(torch_device)

tests/models/testing_utils/ip_adapter.py

@@ -100,6 +100,7 @@ def test_load_ip_adapter(self):
         init_dict = self.get_init_dict()
         inputs_dict = self.get_dummy_inputs()
         model = self.model_class(**init_dict).to(torch_device)
+        self.prepare_model(model)
 
         torch.manual_seed(0)
         output_no_adapter = model(**inputs_dict, return_dict=False)[0]
@@ -128,9 +129,10 @@ def test_ip_adapter_scale(self):
         init_dict = self.get_init_dict()
         inputs_dict = self.get_dummy_inputs()
         model = self.model_class(**init_dict).to(torch_device)
+        # self.prepare_model(model)
 
         # Create and load dummy IP adapter state dict
-        ip_adapter_state_dict = create_ip_adapter_state_dict(model)
+        ip_adapter_state_dict = self.create_ip_adapter_state_dict(model)
         model._load_ip_adapter_weights([ip_adapter_state_dict])
 
         # Test scale = 0.0 (no effect)
@@ -151,12 +153,13 @@ def test_ip_adapter_scale(self):
     def test_unload_ip_adapter(self):
         init_dict = self.get_init_dict()
         model = self.model_class(**init_dict).to(torch_device)
+        self.prepare_model(model)
 
         # Save original processors
         original_processors = {k: type(v).__name__ for k, v in model.attn_processors.items()}
 
         # Create and load IP adapter
-        ip_adapter_state_dict = create_ip_adapter_state_dict(model)
+        ip_adapter_state_dict = self.create_ip_adapter_state_dict(model)
         model._load_ip_adapter_weights([ip_adapter_state_dict])
         assert check_if_ip_adapter_correctly_set(model), "IP Adapter should be set"
 
@@ -172,9 +175,10 @@ def test_ip_adapter_save_load(self):
         init_dict = self.get_init_dict()
         inputs_dict = self.get_dummy_inputs()
         model = self.model_class(**init_dict).to(torch_device)
+        self.prepare_model(model)
 
         # Create and load IP adapter
-        ip_adapter_state_dict = self.create_ip_adapter_state_dict()
+        ip_adapter_state_dict = self.create_ip_adapter_state_dict(model)
         model._load_ip_adapter_weights([ip_adapter_state_dict])
 
         torch.manual_seed(0)

tests/models/testing_utils/quantization.py

@@ -40,9 +40,9 @@
     require_accelerator,
     require_bitsandbytes_version_greater,
     require_gguf_version_greater_or_equal,
+    require_modelopt_version_greater_or_equal,
     require_quanto,
     require_torchao_version_greater_or_equal,
-    require_modelopt_version_greater_or_equal,
     torch_device,
 )
 

tests/models/transformers/test_models_transformer_flux_.py

@@ -19,6 +19,7 @@
 
 from diffusers import FluxTransformer2DModel
 from diffusers.models.embeddings import ImageProjection
+from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor
 from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import enable_full_determinism, torch_device
@@ -87,11 +88,11 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor]:
 
     @property
     def input_shape(self) -> tuple[int, int]:
-        return (1, 16, 4)
+        return (16, 4)
 
     @property
     def output_shape(self) -> tuple[int, int]:
-        return (1, 16, 4)
+        return (16, 4)
 
 
 class TestFluxTransformer(FluxTransformerTesterConfig, ModelTesterMixin):
@@ -148,6 +149,11 @@ class TestFluxTransformerAttention(FluxTransformerTesterConfig, AttentionTesterM
 class TestFluxTransformerIPAdapter(FluxTransformerTesterConfig, IPAdapterTesterMixin):
     """IP Adapter tests for Flux Transformer."""
 
+    def prepare_model(self, model):
+        joint_attention_dim = model.config["joint_attention_dim"]
+        hidden_size = model.config["num_attention_heads"] * model.config["attention_head_dim"]
+        model.set_attn_processor(FluxIPAdapterAttnProcessor(hidden_size, joint_attention_dim, scale=1.0))
+
     def create_ip_adapter_state_dict(self, model: Any) -> dict[str, dict[str, Any]]:
         from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor
 

tests/testing_utils.py

@@ -37,8 +37,8 @@
     is_flax_available,
     is_gguf_available,
     is_kernels_available,
-    is_nvidia_modelopt_available,
     is_note_seq_available,
+    is_nvidia_modelopt_available,
     is_onnx_available,
     is_opencv_available,
     is_optimum_quanto_available,