[tests] accept recompile_limit from the user in tests (#13150)

accept recompile_limit from the user in tests

Author: sayakpaul

tests/models/testing_utils/compile.py

@@ -81,7 +81,7 @@ def test_torch_compile_recompilation_and_graph_break(self):
             _ = model(**inputs_dict)
 
     @torch.no_grad()
-    def test_torch_compile_repeated_blocks(self):
+    def test_torch_compile_repeated_blocks(self, recompile_limit=1):
         if self.model_class._repeated_blocks is None:
             pytest.skip("Skipping test as the model class doesn't have `_repeated_blocks` set.")
 
@@ -92,7 +92,6 @@ def test_torch_compile_repeated_blocks(self):
         model.eval()
         model.compile_repeated_blocks(fullgraph=True)
 
-        recompile_limit = 1
         if self.model_class.__name__ == "UNet2DConditionModel":
             recompile_limit = 2
 

tests/models/transformers/test_models_transformer_wan_vace.py

@@ -147,22 +147,7 @@ class TestWanVACETransformer3DCompile(WanVACETransformer3DTesterConfig, TorchCom
     def test_torch_compile_repeated_blocks(self):
         # WanVACE has two block types (WanTransformerBlock and WanVACETransformerBlock),
         # so we need recompile_limit=2 instead of the default 1.
-        import torch._dynamo
-        import torch._inductor.utils
-
-        init_dict = self.get_init_dict()
-        inputs_dict = self.get_dummy_inputs()
-
-        model = self.model_class(**init_dict).to(torch_device)
-        model.eval()
-        model.compile_repeated_blocks(fullgraph=True)
-
-        with (
-            torch._inductor.utils.fresh_inductor_cache(),
-            torch._dynamo.config.patch(recompile_limit=2),
-        ):
-            _ = model(**inputs_dict)
-            _ = model(**inputs_dict)
+        super().test_torch_compile_repeated_blocks(recompile_limit=2)
 
 
 class TestWanVACETransformer3DBitsAndBytes(WanVACETransformer3DTesterConfig, BitsAndBytesTesterMixin):