[CI] Clean up test_modeling_common.py

tests/models/test_modeling_common.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import inspect
 import logging
 import os
 import tempfile
@@ -277,37 +276,6 @@ def get_dummy_inputs():
         SD3Transformer2DModel._keep_in_fp32_modules = fp32_modules
 
 
-class UNetTesterMixin:
-    @staticmethod
-    def _accepts_norm_num_groups(model_class):
-        model_sig = inspect.signature(model_class.__init__)
-        accepts_norm_groups = "norm_num_groups" in model_sig.parameters
-        return accepts_norm_groups
-
-    def test_forward_with_norm_groups(self):
-        if not self._accepts_norm_num_groups(self.model_class):
-            pytest.skip(f"Test not supported for {self.model_class.__name__}")
-        init_dict = self.get_init_dict()
-        inputs_dict = self.get_dummy_inputs()
-
-        init_dict["norm_num_groups"] = 16
-        init_dict["block_out_channels"] = (16, 32)
-
-        model = self.model_class(**init_dict)
-        model.to(torch_device)
-        model.eval()
-
-        with torch.no_grad():
-            output = model(**inputs_dict)
-
-            if isinstance(output, dict):
-                output = output.to_tuple()[0]
-
-        assert output is not None
-        expected_shape = inputs_dict["sample"].shape
-        assert output.shape == expected_shape, "Input and output shapes do not match"
-
-
 @is_staging_test
 class TestModelPushToHub:
     identifier = uuid.uuid4()

tests/models/unets/test_models_unet_2d_condition.py

@@ -51,7 +51,6 @@
     torch_all_close,
     torch_device,
 )
-from ..test_modeling_common import UNetTesterMixin
 from ..testing_utils import (
     AttentionTesterMixin,
     BaseModelTesterConfig,
@@ -382,7 +381,28 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor]:
         }
 
 
-class TestUNet2DCondition(UNet2DConditionTesterConfig, ModelTesterMixin, UNetTesterMixin):
+class TestUNet2DCondition(UNet2DConditionTesterConfig, ModelTesterMixin):
+    def test_forward_with_norm_groups(self):
+        init_dict = self.get_init_dict()
+        inputs_dict = self.get_dummy_inputs()
+
+        init_dict["norm_num_groups"] = 16
+        init_dict["block_out_channels"] = (16, 32)
+
+        model = self.model_class(**init_dict)
+        model.to(torch_device)
+        model.eval()
+
+        with torch.no_grad():
+            output = model(**inputs_dict)
+
+            if isinstance(output, dict):
+                output = output.to_tuple()[0]
+
+        assert output is not None
+        expected_shape = inputs_dict["sample"].shape
+        assert output.shape == expected_shape, "Input and output shapes do not match"
+
     def test_model_with_attention_head_dim_tuple(self):
         init_dict = self.get_init_dict()
         inputs_dict = self.get_dummy_inputs()

utils/generate_model_tests.py

@@ -430,7 +430,9 @@ def generate_test_class(model_name: str, config_class: str, tester: str) -> str:
 def generate_test_file(model_info: dict, model_filepath: str, include_optional: list[str], imports: set[str]) -> str:
     model_name = model_info["name"].replace("2DModel", "").replace("3DModel", "").replace("Model", "")
     testers = determine_testers(model_info, include_optional, imports)
-    tester_imports = sorted(set(testers) - {"LoraHotSwappingForModelTesterMixin"})
+    tester_imports = set(testers)
+    if "LoraTesterMixin" in testers:
+        tester_imports.add("LoraHotSwappingForModelTesterMixin")
 
     lines = [
         "# coding=utf-8",
@@ -456,9 +458,6 @@ def generate_test_file(model_info: dict, model_filepath: str, include_optional:
         "from ...testing_utils import enable_full_determinism, torch_device",
     ]
 
-    if "LoraTesterMixin" in testers:
-        lines.append("from ..test_modeling_common import LoraHotSwappingForModelTesterMixin")
-
     lines.extend(
         [
             "from ..testing_utils import (",