propagate

Author: sayakpaul

tests/models/transformers/test_models_transformer_hunyuan_video.py

@@ -28,7 +28,7 @@
 enable_full_determinism()
 
 
-class HunyuanVideoTransformer3DTests(ModelTesterMixin, TorchCompileTesterMixin, unittest.TestCase):
+class HunyuanVideoTransformer3DTests(ModelTesterMixin, unittest.TestCase):
     model_class = HunyuanVideoTransformer3DModel
     main_input_name = "hidden_states"
     uses_custom_attn_processor = True
@@ -93,7 +93,14 @@ def test_gradient_checkpointing_is_applied(self):
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
 
-class HunyuanSkyreelsImageToVideoTransformer3DTests(ModelTesterMixin, TorchCompileTesterMixin, unittest.TestCase):
+class HunyuanTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase):
+    model_class = HunyuanVideoTransformer3DModel
+
+    def prepare_init_args_and_inputs_for_common(self):
+        return HunyuanVideoTransformer3DTests().prepare_init_args_and_inputs_for_common()
+
+
+class HunyuanSkyreelsImageToVideoTransformer3DTests(ModelTesterMixin, unittest.TestCase):
     model_class = HunyuanVideoTransformer3DModel
     main_input_name = "hidden_states"
     uses_custom_attn_processor = True
@@ -161,7 +168,14 @@ def test_gradient_checkpointing_is_applied(self):
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
 
-class HunyuanVideoImageToVideoTransformer3DTests(ModelTesterMixin, TorchCompileTesterMixin, unittest.TestCase):
+class HunyuanSkyreelsImageToVideoCompileTests(TorchCompileTesterMixin, unittest.TestCase):
+    model_class = HunyuanVideoTransformer3DModel
+
+    def prepare_init_args_and_inputs_for_common(self):
+        return HunyuanSkyreelsImageToVideoTransformer3DTests().prepare_init_args_and_inputs_for_common()
+
+
+class HunyuanVideoImageToVideoTransformer3DTests(ModelTesterMixin, unittest.TestCase):
     model_class = HunyuanVideoTransformer3DModel
     main_input_name = "hidden_states"
     uses_custom_attn_processor = True
@@ -227,9 +241,14 @@ def test_gradient_checkpointing_is_applied(self):
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
 
-class HunyuanVideoTokenReplaceImageToVideoTransformer3DTests(
-    ModelTesterMixin, TorchCompileTesterMixin, unittest.TestCase
-):
+class HunyuanImageToVideoCompileTests(TorchCompileTesterMixin, unittest.TestCase):
+    model_class = HunyuanVideoTransformer3DModel
+
+    def prepare_init_args_and_inputs_for_common(self):
+        return HunyuanVideoImageToVideoTransformer3DTests().prepare_init_args_and_inputs_for_common()
+
+
+class HunyuanVideoTokenReplaceImageToVideoTransformer3DTests(ModelTesterMixin, unittest.TestCase):
     model_class = HunyuanVideoTransformer3DModel
     main_input_name = "hidden_states"
     uses_custom_attn_processor = True
@@ -295,3 +314,10 @@ def test_output(self):
     def test_gradient_checkpointing_is_applied(self):
         expected_set = {"HunyuanVideoTransformer3DModel"}
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
+
+class HunyuanVideoTokenReplaceCompileTests(TorchCompileTesterMixin, unittest.TestCase):
+    model_class = HunyuanVideoTransformer3DModel
+
+    def prepare_init_args_and_inputs_for_common(self):
+        return HunyuanVideoTokenReplaceImageToVideoTransformer3DTests().prepare_init_args_and_inputs_for_common()

tests/models/transformers/test_models_transformer_ltx.py

@@ -26,7 +26,7 @@
 enable_full_determinism()
 
 
-class LTXTransformerTests(ModelTesterMixin, TorchCompileTesterMixin, unittest.TestCase):
+class LTXTransformerTests(ModelTesterMixin, unittest.TestCase):
     model_class = LTXVideoTransformer3DModel
     main_input_name = "hidden_states"
     uses_custom_attn_processor = True
@@ -81,3 +81,10 @@ def prepare_init_args_and_inputs_for_common(self):
     def test_gradient_checkpointing_is_applied(self):
         expected_set = {"LTXVideoTransformer3DModel"}
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
+
+class LTXTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase):
+    model_class = LTXVideoTransformer3DModel
+
+    def prepare_init_args_and_inputs_for_common(self):
+        return LTXTransformerTests().prepare_init_args_and_inputs_for_common()

tests/models/transformers/test_models_transformer_wan.py

@@ -28,7 +28,7 @@
 enable_full_determinism()
 
 
-class WanTransformer3DTests(ModelTesterMixin, TorchCompileTesterMixin, unittest.TestCase):
+class WanTransformer3DTests(ModelTesterMixin, unittest.TestCase):
     model_class = WanTransformer3DModel
     main_input_name = "hidden_states"
     uses_custom_attn_processor = True
@@ -82,3 +82,10 @@ def prepare_init_args_and_inputs_for_common(self):
     def test_gradient_checkpointing_is_applied(self):
         expected_set = {"WanTransformer3DModel"}
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
+
+class WanTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase):
+    model_class = WanTransformer3DModel
+
+    def prepare_init_args_and_inputs_for_common(self):
+        return WanTransformer3DTests().prepare_init_args_and_inputs_for_common()