From 1cbc1971b0ec9d4d2e5c0909d1597284168f7427 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 21 Aug 2026 08:07:52 +0000 Subject: [PATCH] remove autoencodertestermixin and rename new one. --- .../test_models_asymmetric_autoencoder_kl.py | 4 +- .../test_models_autoencoder_cosmos.py | 4 +- .../test_models_autoencoder_dc.py | 4 +- .../test_models_autoencoder_hunyuan_video.py | 4 +- .../test_models_autoencoder_kl.py | 4 +- .../test_models_autoencoder_kl_cogvideox.py | 4 +- .../test_models_autoencoder_kl_kvae.py | 4 +- .../test_models_autoencoder_kl_kvae_video.py | 4 +- .../test_models_autoencoder_kl_ltx2_audio.py | 4 +- .../test_models_autoencoder_kl_minimax_h3.py | 4 +- ..._models_autoencoder_kl_temporal_decoder.py | 6 +- .../test_models_autoencoder_ltx2_video.py | 4 +- .../test_models_autoencoder_ltx_video.py | 4 +- .../test_models_autoencoder_magvit.py | 4 +- .../test_models_autoencoder_mochi.py | 4 +- .../test_models_autoencoder_oobleck.py | 4 +- .../test_models_autoencoder_rae.py | 4 - .../test_models_autoencoder_tiny.py | 4 +- .../test_models_autoencoder_vidtok.py | 4 +- .../test_models_autoencoder_wan.py | 4 +- .../test_models_consistency_decoder_vae.py | 4 +- tests/models/autoencoders/test_models_vq.py | 4 +- tests/models/autoencoders/testing_utils.py | 135 ------------------ 23 files changed, 42 insertions(+), 183 deletions(-) diff --git a/tests/models/autoencoders/test_models_asymmetric_autoencoder_kl.py b/tests/models/autoencoders/test_models_asymmetric_autoencoder_kl.py index 425848efd700..2aa1114a5b1a 100644 --- a/tests/models/autoencoders/test_models_asymmetric_autoencoder_kl.py +++ b/tests/models/autoencoders/test_models_asymmetric_autoencoder_kl.py @@ -37,7 +37,7 @@ torch_device, ) from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -99,7 +99,7 @@ class TestAsymmetricAutoencoderKLMemory(AsymmetricAutoencoderKLTesterConfig, Mem """Memory optimization tests for AsymmetricAutoencoderKL.""" -class TestAsymmetricAutoencoderKLSlicingTiling(AsymmetricAutoencoderKLTesterConfig, NewAutoencoderTesterMixin): +class TestAsymmetricAutoencoderKLSlicingTiling(AsymmetricAutoencoderKLTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AsymmetricAutoencoderKL.""" @pytest.mark.skip("Unsupported test.") diff --git a/tests/models/autoencoders/test_models_autoencoder_cosmos.py b/tests/models/autoencoders/test_models_autoencoder_cosmos.py index b70f05dc47eb..a7a37c2dbf47 100644 --- a/tests/models/autoencoders/test_models_autoencoder_cosmos.py +++ b/tests/models/autoencoders/test_models_autoencoder_cosmos.py @@ -20,7 +20,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -92,5 +92,5 @@ class TestAutoencoderKLCosmosMemory(AutoencoderKLCosmosTesterConfig, MemoryTeste """Memory optimization tests for AutoencoderKLCosmos.""" -class TestAutoencoderKLCosmosSlicingTiling(AutoencoderKLCosmosTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLCosmosSlicingTiling(AutoencoderKLCosmosTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLCosmos.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_dc.py b/tests/models/autoencoders/test_models_autoencoder_dc.py index c1743623021d..c1c304d81155 100644 --- a/tests/models/autoencoders/test_models_autoencoder_dc.py +++ b/tests/models/autoencoders/test_models_autoencoder_dc.py @@ -27,7 +27,7 @@ SingleFileTesterMixin, TrainingTesterMixin, ) -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -106,7 +106,7 @@ def test_layerwise_casting_memory(self): super().test_layerwise_casting_memory() -class TestAutoencoderDCSlicingTiling(AutoencoderDCTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderDCSlicingTiling(AutoencoderDCTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderDC.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py b/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py index cdaebcf2eea9..c575162ba91a 100644 --- a/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py @@ -21,7 +21,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -124,7 +124,7 @@ class TestAutoencoderKLHunyuanVideoMemory(AutoencoderKLHunyuanVideoTesterConfig, """Memory optimization tests for AutoencoderKLHunyuanVideo.""" -class TestAutoencoderKLHunyuanVideoSlicingTiling(AutoencoderKLHunyuanVideoTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLHunyuanVideoSlicingTiling(AutoencoderKLHunyuanVideoTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLHunyuanVideo.""" # Overwritten because the base test's block_out_channels doesn't account for the length of down_block_types. diff --git a/tests/models/autoencoders/test_models_autoencoder_kl.py b/tests/models/autoencoders/test_models_autoencoder_kl.py index 2ab286ffb8c2..8be893cc8213 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl.py @@ -42,7 +42,7 @@ SingleFileTesterMixin, TrainingTesterMixin, ) -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -203,7 +203,7 @@ class TestAutoencoderKLMemory(AutoencoderKLTesterConfig, MemoryTesterMixin): """Memory optimization tests for AutoencoderKL.""" -class TestAutoencoderKLSlicingTiling(AutoencoderKLTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLSlicingTiling(AutoencoderKLTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKL.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_cogvideox.py b/tests/models/autoencoders/test_models_autoencoder_kl_cogvideox.py index d4c141685115..e574dd895aca 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_cogvideox.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_cogvideox.py @@ -21,7 +21,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -107,7 +107,7 @@ class TestAutoencoderKLCogVideoXMemory(AutoencoderKLCogVideoXTesterConfig, Memor """Memory optimization tests for AutoencoderKLCogVideoX.""" -class TestAutoencoderKLCogVideoXSlicingTiling(AutoencoderKLCogVideoXTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLCogVideoXSlicingTiling(AutoencoderKLCogVideoXTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLCogVideoX.""" # Overwritten because the base test's block_out_channels doesn't account for the length of down_block_types. diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_kvae.py b/tests/models/autoencoders/test_models_autoencoder_kl_kvae.py index 066864abb4ff..a28fcffe06bf 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_kvae.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_kvae.py @@ -20,7 +20,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -79,5 +79,5 @@ class TestAutoencoderKLKVAEMemory(AutoencoderKLKVAETesterConfig, MemoryTesterMix """Memory optimization tests for AutoencoderKLKVAE.""" -class TestAutoencoderKLKVAESlicingTiling(AutoencoderKLKVAETesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLKVAESlicingTiling(AutoencoderKLKVAETesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLKVAE.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py index 6923def0d94d..d5263649d722 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py @@ -27,7 +27,7 @@ TrainingTesterMixin, run_nondeterministic, ) -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -116,5 +116,5 @@ def test_layerwise_casting_training(self): run_nondeterministic(super().test_layerwise_casting_training) -class TestAutoencoderKLKVAEVideoSlicingTiling(AutoencoderKLKVAEVideoTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLKVAEVideoSlicingTiling(AutoencoderKLKVAEVideoTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLKVAEVideo.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_ltx2_audio.py b/tests/models/autoencoders/test_models_autoencoder_kl_ltx2_audio.py index c10de84f04d1..562c6b8c47b9 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_ltx2_audio.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_ltx2_audio.py @@ -21,7 +21,7 @@ from ...testing_utils import is_flaky, torch_device from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin class AutoencoderKLLTX2AudioTesterConfig(BaseModelTesterConfig): @@ -104,5 +104,5 @@ def test_group_offloading_with_disk(self, tmp_path, record_stream, offload_type, super().test_group_offloading_with_disk(tmp_path, record_stream, offload_type, atol=atol, rtol=rtol) -class TestAutoencoderKLLTX2AudioSlicingTiling(AutoencoderKLLTX2AudioTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLLTX2AudioSlicingTiling(AutoencoderKLLTX2AudioTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLLTX2Audio.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py index 676299e812d5..2ffcd43af216 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py @@ -28,7 +28,7 @@ TrainingTesterMixin, run_nondeterministic, ) -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -183,5 +183,5 @@ class TestAutoencoderKLMiniMaxH3Attention(AutoencoderKLMiniMaxH3TesterConfig, At """Attention processor tests for the MiniMax-H3 video autoencoder.""" -class TestAutoencoderKLMiniMaxH3SlicingTiling(AutoencoderKLMiniMaxH3TesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLMiniMaxH3SlicingTiling(AutoencoderKLMiniMaxH3TesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for the MiniMax-H3 video autoencoder.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py b/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py index 56297b7540bb..61984b2e0fb6 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py @@ -21,7 +21,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -88,7 +88,5 @@ class TestAutoencoderKLTemporalDecoderMemory(AutoencoderKLTemporalDecoderTesterC """Memory optimization tests for AutoencoderKLTemporalDecoder.""" -class TestAutoencoderKLTemporalDecoderSlicingTiling( - AutoencoderKLTemporalDecoderTesterConfig, NewAutoencoderTesterMixin -): +class TestAutoencoderKLTemporalDecoderSlicingTiling(AutoencoderKLTemporalDecoderTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLTemporalDecoder.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_ltx2_video.py b/tests/models/autoencoders/test_models_autoencoder_ltx2_video.py index dac3f8fdcc6b..c9cd0ce8c0af 100644 --- a/tests/models/autoencoders/test_models_autoencoder_ltx2_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_ltx2_video.py @@ -21,7 +21,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -105,5 +105,5 @@ class TestAutoencoderKLLTX2VideoMemory(AutoencoderKLLTX2VideoTesterConfig, Memor """Memory optimization tests for AutoencoderKLLTX2Video.""" -class TestAutoencoderKLLTX2VideoSlicingTiling(AutoencoderKLLTX2VideoTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLLTX2VideoSlicingTiling(AutoencoderKLLTX2VideoTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLLTX2Video.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_ltx_video.py b/tests/models/autoencoders/test_models_autoencoder_ltx_video.py index b8d992c63d33..a29a8ba922f4 100644 --- a/tests/models/autoencoders/test_models_autoencoder_ltx_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_ltx_video.py @@ -27,7 +27,7 @@ SingleFileTesterMixin, TrainingTesterMixin, ) -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -110,7 +110,7 @@ class TestAutoencoderKLLTXVideo090Memory(AutoencoderKLLTXVideo090TesterConfig, M """Memory optimization tests for AutoencoderKLLTXVideo (0.9.0 config).""" -class TestAutoencoderKLLTXVideo090SlicingTiling(AutoencoderKLLTXVideo090TesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLLTXVideo090SlicingTiling(AutoencoderKLLTXVideo090TesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLLTXVideo (0.9.0 config).""" @pytest.mark.skip("AutoencoderKLLTXVideo does not support `norm_num_groups` because it does not use GroupNorm.") diff --git a/tests/models/autoencoders/test_models_autoencoder_magvit.py b/tests/models/autoencoders/test_models_autoencoder_magvit.py index 2f3ad90d40e1..bd16c0e03dad 100644 --- a/tests/models/autoencoders/test_models_autoencoder_magvit.py +++ b/tests/models/autoencoders/test_models_autoencoder_magvit.py @@ -21,7 +21,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -99,7 +99,7 @@ class TestAutoencoderKLMagvitMemory(AutoencoderKLMagvitTesterConfig, MemoryTeste """Memory optimization tests for AutoencoderKLMagvit.""" -class TestAutoencoderKLMagvitSlicingTiling(AutoencoderKLMagvitTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLMagvitSlicingTiling(AutoencoderKLMagvitTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLMagvit.""" @pytest.mark.skip("Unsupported test.") diff --git a/tests/models/autoencoders/test_models_autoencoder_mochi.py b/tests/models/autoencoders/test_models_autoencoder_mochi.py index 3beb0c94066b..9134d298d4c7 100755 --- a/tests/models/autoencoders/test_models_autoencoder_mochi.py +++ b/tests/models/autoencoders/test_models_autoencoder_mochi.py @@ -21,7 +21,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -91,5 +91,5 @@ class TestAutoencoderKLMochiMemory(AutoencoderKLMochiTesterConfig, MemoryTesterM """Memory optimization tests for AutoencoderKLMochi.""" -class TestAutoencoderKLMochiSlicingTiling(AutoencoderKLMochiTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLMochiSlicingTiling(AutoencoderKLMochiTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLMochi.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_oobleck.py b/tests/models/autoencoders/test_models_autoencoder_oobleck.py index bf603da829ce..3371ecfb0d99 100644 --- a/tests/models/autoencoders/test_models_autoencoder_oobleck.py +++ b/tests/models/autoencoders/test_models_autoencoder_oobleck.py @@ -32,7 +32,7 @@ torch_device, ) from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -98,7 +98,7 @@ def test_layerwise_casting_memory(self): super().test_layerwise_casting_memory() -class TestAutoencoderOobleckSlicingTiling(AutoencoderOobleckTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderOobleckSlicingTiling(AutoencoderOobleckTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderOobleck.""" def test_enable_disable_slicing(self): diff --git a/tests/models/autoencoders/test_models_autoencoder_rae.py b/tests/models/autoencoders/test_models_autoencoder_rae.py index c2b801ec348c..ce6e58f841ed 100644 --- a/tests/models/autoencoders/test_models_autoencoder_rae.py +++ b/tests/models/autoencoders/test_models_autoencoder_rae.py @@ -124,10 +124,6 @@ def generator(self): def get_dummy_inputs(self): return {"sample": torch.randn(2, 3, 32, 32, generator=self.generator, device="cpu").to(torch_device)} - # Bridge for AutoencoderTesterMixin which still uses the old interface - def prepare_init_args_and_inputs_for_common(self): - return self.get_init_dict(), self.get_dummy_inputs() - def _make_model(self, **overrides) -> AutoencoderRAE: config = self.get_init_dict() config.update(overrides) diff --git a/tests/models/autoencoders/test_models_autoencoder_tiny.py b/tests/models/autoencoders/test_models_autoencoder_tiny.py index 8f0740a8a70e..77f8f644d903 100644 --- a/tests/models/autoencoders/test_models_autoencoder_tiny.py +++ b/tests/models/autoencoders/test_models_autoencoder_tiny.py @@ -33,7 +33,7 @@ torch_device, ) from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -143,7 +143,7 @@ def test_layerwise_casting_memory(self): super().test_layerwise_casting_memory() -class TestAutoencoderTinySlicingTiling(AutoencoderTinyTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderTinySlicingTiling(AutoencoderTinyTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderTiny.""" @pytest.mark.skip("Model does not yet support smaller resolution.") diff --git a/tests/models/autoencoders/test_models_autoencoder_vidtok.py b/tests/models/autoencoders/test_models_autoencoder_vidtok.py index 323d92b40572..5d810141ee8b 100644 --- a/tests/models/autoencoders/test_models_autoencoder_vidtok.py +++ b/tests/models/autoencoders/test_models_autoencoder_vidtok.py @@ -26,7 +26,7 @@ TrainingTesterMixin, run_nondeterministic, ) -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -110,7 +110,7 @@ def test_layerwise_casting_training(self): run_nondeterministic(super().test_layerwise_casting_training) -class TestAutoencoderVidTokSlicingTiling(AutoencoderVidTokTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderVidTokSlicingTiling(AutoencoderVidTokTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderVidTok.""" def test_enable_disable_tiling(self): diff --git a/tests/models/autoencoders/test_models_autoencoder_wan.py b/tests/models/autoencoders/test_models_autoencoder_wan.py index ee67cc87a9c3..b8026f575e49 100644 --- a/tests/models/autoencoders/test_models_autoencoder_wan.py +++ b/tests/models/autoencoders/test_models_autoencoder_wan.py @@ -27,7 +27,7 @@ SingleFileTesterMixin, TrainingTesterMixin, ) -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -90,7 +90,7 @@ def test_layerwise_casting_training(self): pass -class TestAutoencoderKLWanSlicingTiling(AutoencoderKLWanTesterConfig, NewAutoencoderTesterMixin): +class TestAutoencoderKLWanSlicingTiling(AutoencoderKLWanTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLWan.""" diff --git a/tests/models/autoencoders/test_models_consistency_decoder_vae.py b/tests/models/autoencoders/test_models_consistency_decoder_vae.py index b5831d6ba09d..d28543031176 100644 --- a/tests/models/autoencoders/test_models_consistency_decoder_vae.py +++ b/tests/models/autoencoders/test_models_consistency_decoder_vae.py @@ -31,7 +31,7 @@ torch_device, ) from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -103,7 +103,7 @@ class TestConsistencyDecoderVAEMemory(ConsistencyDecoderVAETesterConfig, MemoryT """Memory optimization tests for ConsistencyDecoderVAE.""" -class TestConsistencyDecoderVAESlicingTiling(ConsistencyDecoderVAETesterConfig, NewAutoencoderTesterMixin): +class TestConsistencyDecoderVAESlicingTiling(ConsistencyDecoderVAETesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for ConsistencyDecoderVAE.""" diff --git a/tests/models/autoencoders/test_models_vq.py b/tests/models/autoencoders/test_models_vq.py index b6dbcfa64b19..5b1b97d19a6c 100644 --- a/tests/models/autoencoders/test_models_vq.py +++ b/tests/models/autoencoders/test_models_vq.py @@ -21,7 +21,7 @@ from ...testing_utils import backend_manual_seed, enable_full_determinism, torch_device from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -from .testing_utils import NewAutoencoderTesterMixin +from .testing_utils import AutoencoderTesterMixin enable_full_determinism() @@ -136,5 +136,5 @@ class TestVQModelMemory(VQModelTesterConfig, MemoryTesterMixin): """Memory optimization tests for VQModel.""" -class TestVQModelSlicingTiling(VQModelTesterConfig, NewAutoencoderTesterMixin): +class TestVQModelSlicingTiling(VQModelTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for VQModel.""" diff --git a/tests/models/autoencoders/testing_utils.py b/tests/models/autoencoders/testing_utils.py index c2a5eb5b906f..dcc086fdcf10 100644 --- a/tests/models/autoencoders/testing_utils.py +++ b/tests/models/autoencoders/testing_utils.py @@ -1,6 +1,5 @@ import inspect -import numpy as np import pytest import torch @@ -15,140 +14,6 @@ class AutoencoderTesterMixin: usually don't do slicing and tiling. """ - @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, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - 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, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - 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) - - torch.manual_seed(0) - if accepts_generator: - inputs_dict["generator"] = torch.manual_seed(0) - output_without_tiling = model(**inputs_dict)[0] - # Mochi-1 - 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.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy() - ).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 np.allclose( - output_without_tiling.detach().cpu().numpy().all(), - output_without_tiling_2.detach().cpu().numpy().all(), - ), "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, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - 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) - - if accepts_generator: - inputs_dict["generator"] = torch.manual_seed(0) - - torch.manual_seed(0) - output_without_slicing = model(**inputs_dict)[0] - # Mochi-1 - 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.detach().cpu().numpy() - output_with_slicing.detach().cpu().numpy() - ).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 np.allclose( - 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)