change to modeling level test.

Author: sayakpaul

tests/models/test_modeling_common.py

@@ -1714,6 +1714,35 @@ def test_push_to_hub_library_name(self):
         delete_repo(self.repo_id, token=TOKEN)
 
 
+class TorchCompileTesterMixin:
+    def setUp(self):
+        # clean up the VRAM before each test
+        super().setUp()
+        torch._dynamo.reset()
+        gc.collect()
+        backend_empty_cache(torch_device)
+
+    def tearDown(self):
+        # clean up the VRAM after each test in case of CUDA runtime errors
+        super().tearDown()
+        torch._dynamo.reset()
+        gc.collect()
+        backend_empty_cache(torch_device)
+
+    @require_torch_gpu
+    @require_torch_2
+    @slow
+    def test_torch_compile_recompilation_and_graph_break(self):
+        torch._dynamo.reset()
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+
+        model = self.model_class(**init_dict).to(torch_device)
+        model = torch.compile(model, fullgraph=True)
+
+        with torch._dynamo.config.patch(error_on_recompile=True), torch.no_grad():
+            _ = model(**inputs_dict)
+
+
 @slow
 @require_torch_2
 @require_torch_accelerator

tests/models/transformers/test_models_transformer_flux.py

@@ -22,7 +22,7 @@
 from diffusers.models.embeddings import ImageProjection
 from diffusers.utils.testing_utils import enable_full_determinism, torch_device
 
-from ..test_modeling_common import ModelTesterMixin
+from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin
 
 
 enable_full_determinism()
@@ -78,7 +78,7 @@ def create_flux_ip_adapter_state_dict(model):
     return ip_state_dict
 
 
-class FluxTransformerTests(ModelTesterMixin, unittest.TestCase):
+class FluxTransformerTests(ModelTesterMixin, TorchCompileTesterMixin, unittest.TestCase):
     model_class = FluxTransformer2DModel
     main_input_name = "hidden_states"
     # We override the items here because the transformer under consideration is small.

tests/pipelines/test_pipelines_common.py

@@ -56,7 +56,6 @@
     require_torch_gpu,
     require_transformers_version_greater,
     skip_mps,
-    slow,
     torch_device,
 )
 
@@ -2165,22 +2164,6 @@ def test_StableDiffusionMixin_component(self):
             )
         )
 
-    @require_torch_gpu
-    @slow
-    def test_torch_compile_recompilation_and_graph_break(self):
-        torch._dynamo.reset()
-        inputs = self.get_dummy_inputs(torch_device)
-        components = self.get_dummy_components()
-
-        pipe = self.pipeline_class(**components).to(torch_device)
-        if getattr(pipe, "unet", None) is not None:
-            pipe.unet = torch.compile(pipe.unet, fullgraph=True)
-        else:
-            pipe.transformer = torch.compile(pipe.transformer, fullgraph=True)
-
-        with torch._dynamo.config.patch(error_on_recompile=True):
-            _ = pipe(**inputs)
-
     @require_hf_hub_version_greater("0.26.5")
     @require_transformers_version_greater("4.47.1")
     def test_save_load_dduf(self, atol=1e-4, rtol=1e-4):