diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index b1bdaabbad7a..84d3679eb6a9 100644 --- a/tests/models/test_modeling_common.py +++ b/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() diff --git a/tests/models/unets/test_models_unet_2d_condition.py b/tests/models/unets/test_models_unet_2d_condition.py index f31d148caf2a..ab367b749191 100644 --- a/tests/models/unets/test_models_unet_2d_condition.py +++ b/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() diff --git a/utils/generate_model_tests.py b/utils/generate_model_tests.py index c006a71580db..75762d01e17c 100644 --- a/utils/generate_model_tests.py +++ b/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 (",