[tests] refactor wan autoencoder tests (#13371)

* refactor wan autoencoder tests

* up

* address dhruv's feedback.

Author: sayakpaul

tests/models/autoencoders/test_models_autoencoder_wan.py

@@ -13,24 +13,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import unittest
+import pytest
+import torch
 
 from diffusers import AutoencoderKLWan
+from diffusers.utils.torch_utils import randn_tensor
 
-from ...testing_utils import enable_full_determinism, floats_tensor, torch_device
-from ..test_modeling_common import ModelTesterMixin
-from .testing_utils import AutoencoderTesterMixin
+from ...testing_utils import enable_full_determinism, torch_device
+from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin
+from .testing_utils import NewAutoencoderTesterMixin
 
 
 enable_full_determinism()
 
 
-class AutoencoderKLWanTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase):
-    model_class = AutoencoderKLWan
-    main_input_name = "sample"
-    base_precision = 1e-2
+class AutoencoderKLWanTesterConfig(BaseModelTesterConfig):
+    @property
+    def model_class(self):
+        return AutoencoderKLWan
 
-    def get_autoencoder_kl_wan_config(self):
+    @property
+    def output_shape(self):
+        return (3, 9, 16, 16)
+
+    @property
+    def generator(self):
+        return torch.Generator("cpu").manual_seed(0)
+
+    def get_init_dict(self):
         return {
             "base_dim": 3,
             "z_dim": 16,
@@ -39,54 +49,40 @@ def get_autoencoder_kl_wan_config(self):
             "temperal_downsample": [False, True, True],
         }
 
-    @property
-    def dummy_input(self):
+    def get_dummy_inputs(self):
         batch_size = 2
         num_frames = 9
         num_channels = 3
         sizes = (16, 16)
-        image = floats_tensor((batch_size, num_channels, num_frames) + sizes).to(torch_device)
+        image = randn_tensor(
+            (batch_size, num_channels, num_frames, *sizes), generator=self.generator, device=torch_device
+        )
         return {"sample": image}
 
-    @property
-    def dummy_input_tiling(self):
-        batch_size = 2
-        num_frames = 9
-        num_channels = 3
-        sizes = (128, 128)
-        image = floats_tensor((batch_size, num_channels, num_frames) + sizes).to(torch_device)
-        return {"sample": image}
-
-    @property
-    def input_shape(self):
-        return (3, 9, 16, 16)
 
-    @property
-    def output_shape(self):
-        return (3, 9, 16, 16)
+class TestAutoencoderKLWan(AutoencoderKLWanTesterConfig, ModelTesterMixin):
+    base_precision = 1e-2
 
-    def prepare_init_args_and_inputs_for_common(self):
-        init_dict = self.get_autoencoder_kl_wan_config()
-        inputs_dict = self.dummy_input
-        return init_dict, inputs_dict
 
-    def prepare_init_args_and_inputs_for_tiling(self):
-        init_dict = self.get_autoencoder_kl_wan_config()
-        inputs_dict = self.dummy_input_tiling
-        return init_dict, inputs_dict
+class TestAutoencoderKLWanTraining(AutoencoderKLWanTesterConfig, TrainingTesterMixin):
+    """Training tests for AutoencoderKLWan."""
 
-    @unittest.skip("Gradient checkpointing has not been implemented yet")
+    @pytest.mark.skip(reason="Gradient checkpointing has not been implemented yet")
     def test_gradient_checkpointing_is_applied(self):
         pass
 
-    @unittest.skip("Test not supported")
-    def test_forward_with_norm_groups(self):
-        pass
 
-    @unittest.skip("RuntimeError: fill_out not implemented for 'Float8_e4m3fn'")
-    def test_layerwise_casting_inference(self):
+class TestAutoencoderKLWanMemory(AutoencoderKLWanTesterConfig, MemoryTesterMixin):
+    """Memory optimization tests for AutoencoderKLWan."""
+
+    @pytest.mark.skip(reason="RuntimeError: fill_out not implemented for 'Float8_e4m3fn'")
+    def test_layerwise_casting_memory(self):
         pass
 
-    @unittest.skip("RuntimeError: fill_out not implemented for 'Float8_e4m3fn'")
+    @pytest.mark.skip(reason="RuntimeError: fill_out not implemented for 'Float8_e4m3fn'")
     def test_layerwise_casting_training(self):
         pass
+
+
+class TestAutoencoderKLWanSlicingTiling(AutoencoderKLWanTesterConfig, NewAutoencoderTesterMixin):
+    """Slicing and tiling tests for AutoencoderKLWan."""

tests/models/autoencoders/testing_utils.py

@@ -145,3 +145,138 @@ def test_enable_disable_slicing(self):
             output_without_slicing.detach().cpu().numpy().all(),
             output_without_slicing_2.detach().cpu().numpy().all(),
         ), "Without slicing outputs should match with the outputs when slicing is manually disabled."
+
+
+class NewAutoencoderTesterMixin:
+    @staticmethod
+    def _accepts_generator(model):
+        model_sig = inspect.signature(model.forward)
+        accepts_generator = "generator" in model_sig.parameters
+        return accepts_generator
+
+    @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"
+
+    def test_enable_disable_tiling(self):
+        if not hasattr(self.model_class, "enable_tiling"):
+            pytest.skip(f"Skipping test as {self.model_class.__name__} doesn't support tiling.")
+
+        init_dict = self.get_init_dict()
+        inputs_dict = self.get_dummy_inputs()
+
+        torch.manual_seed(0)
+        model = self.model_class(**init_dict).to(torch_device)
+
+        if not hasattr(model, "use_tiling"):
+            pytest.skip(f"Skipping test as {self.model_class.__name__} doesn't support tiling.")
+
+        inputs_dict.update({"return_dict": False})
+        _ = inputs_dict.pop("generator", None)
+        accepts_generator = self._accepts_generator(model)
+
+        with torch.no_grad():
+            torch.manual_seed(0)
+            if accepts_generator:
+                inputs_dict["generator"] = torch.manual_seed(0)
+            output_without_tiling = model(**inputs_dict)[0]
+            if isinstance(output_without_tiling, DecoderOutput):
+                output_without_tiling = output_without_tiling.sample
+
+            torch.manual_seed(0)
+            model.enable_tiling()
+            if accepts_generator:
+                inputs_dict["generator"] = torch.manual_seed(0)
+            output_with_tiling = model(**inputs_dict)[0]
+            if isinstance(output_with_tiling, DecoderOutput):
+                output_with_tiling = output_with_tiling.sample
+
+            assert (output_without_tiling.cpu() - output_with_tiling.cpu()).max() < 0.5, (
+                "VAE tiling should not affect the inference results"
+            )
+
+            torch.manual_seed(0)
+            model.disable_tiling()
+            if accepts_generator:
+                inputs_dict["generator"] = torch.manual_seed(0)
+            output_without_tiling_2 = model(**inputs_dict)[0]
+            if isinstance(output_without_tiling_2, DecoderOutput):
+                output_without_tiling_2 = output_without_tiling_2.sample
+
+            assert torch.allclose(output_without_tiling.cpu(), output_without_tiling_2.cpu()), (
+                "Without tiling outputs should match with the outputs when tiling is manually disabled."
+            )
+
+    def test_enable_disable_slicing(self):
+        if not hasattr(self.model_class, "enable_slicing"):
+            pytest.skip(f"Skipping test as {self.model_class.__name__} doesn't support slicing.")
+
+        init_dict = self.get_init_dict()
+        inputs_dict = self.get_dummy_inputs()
+
+        torch.manual_seed(0)
+        model = self.model_class(**init_dict).to(torch_device)
+        if not hasattr(model, "use_slicing"):
+            pytest.skip(f"Skipping test as {self.model_class.__name__} doesn't support tiling.")
+
+        inputs_dict.update({"return_dict": False})
+        _ = inputs_dict.pop("generator", None)
+        accepts_generator = self._accepts_generator(model)
+
+        with torch.no_grad():
+            if accepts_generator:
+                inputs_dict["generator"] = torch.manual_seed(0)
+
+            torch.manual_seed(0)
+            output_without_slicing = model(**inputs_dict)[0]
+            if isinstance(output_without_slicing, DecoderOutput):
+                output_without_slicing = output_without_slicing.sample
+
+            torch.manual_seed(0)
+            model.enable_slicing()
+            if accepts_generator:
+                inputs_dict["generator"] = torch.manual_seed(0)
+            output_with_slicing = model(**inputs_dict)[0]
+            if isinstance(output_with_slicing, DecoderOutput):
+                output_with_slicing = output_with_slicing.sample
+
+            assert (output_without_slicing.cpu() - output_with_slicing.cpu()).max() < 0.5, (
+                "VAE slicing should not affect the inference results"
+            )
+
+            torch.manual_seed(0)
+            model.disable_slicing()
+            if accepts_generator:
+                inputs_dict["generator"] = torch.manual_seed(0)
+            output_without_slicing_2 = model(**inputs_dict)[0]
+            if isinstance(output_without_slicing_2, DecoderOutput):
+                output_without_slicing_2 = output_without_slicing_2.sample
+
+            assert torch.allclose(output_without_slicing.cpu(), output_without_slicing_2.cpu()), (
+                "Without slicing outputs should match with the outputs when slicing is manually disabled."
+            )