use fixture for tmp_path in modular tests.

Author: sayakpaul

tests/modular_pipelines/flux/test_modular_pipeline_flux.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 import random
-import tempfile
 
 import numpy as np
 import PIL
@@ -129,18 +128,16 @@ def get_dummy_inputs(self, seed=0):
 
         return inputs
 
-    def test_save_from_pretrained(self):
+    def test_save_from_pretrained(self, tmp_path):
         pipes = []
         base_pipe = self.get_pipeline().to(torch_device)
         pipes.append(base_pipe)
 
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            base_pipe.save_pretrained(tmpdirname)
-
-            pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device)
-            pipe.load_components(torch_dtype=torch.float32)
-            pipe.to(torch_device)
-            pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
+        base_pipe.save_pretrained(tmp_path)
+        pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
+        pipe.load_components(torch_dtype=torch.float32)
+        pipe.to(torch_device)
+        pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
 
         pipes.append(pipe)
 
@@ -212,18 +209,16 @@ def get_dummy_inputs(self, seed=0):
 
         return inputs
 
-    def test_save_from_pretrained(self):
+    def test_save_from_pretrained(self, tmp_path):
         pipes = []
         base_pipe = self.get_pipeline().to(torch_device)
         pipes.append(base_pipe)
 
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            base_pipe.save_pretrained(tmpdirname)
-
-            pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device)
-            pipe.load_components(torch_dtype=torch.float32)
-            pipe.to(torch_device)
-            pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
+        base_pipe.save_pretrained(tmp_path)
+        pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
+        pipe.load_components(torch_dtype=torch.float32)
+        pipe.to(torch_device)
+        pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
 
         pipes.append(pipe)
 

tests/modular_pipelines/test_modular_pipelines_common.py

@@ -1,7 +1,6 @@
 import gc
 import json
 import os
-import tempfile
 from typing import Callable
 
 import pytest
@@ -330,16 +329,15 @@ def test_components_auto_cpu_offload_inference_consistent(self):
 
         assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3
 
-    def test_save_from_pretrained(self):
+    def test_save_from_pretrained(self, tmp_path):
         pipes = []
         base_pipe = self.get_pipeline().to(torch_device)
         pipes.append(base_pipe)
 
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            base_pipe.save_pretrained(tmpdirname)
-            pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device)
-            pipe.load_components(torch_dtype=torch.float32)
-            pipe.to(torch_device)
+        base_pipe.save_pretrained(tmp_path)
+        pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
+        pipe.load_components(torch_dtype=torch.float32)
+        pipe.to(torch_device)
 
         pipes.append(pipe)
 
@@ -351,32 +349,31 @@ def test_save_from_pretrained(self):
 
         assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3
 
-    def test_modular_index_consistency(self):
+    def test_modular_index_consistency(self, tmp_path):
         pipe = self.get_pipeline()
         components_spec = pipe._component_specs
         components = sorted(components_spec.keys())
 
-        with tempfile.TemporaryDirectory() as tmpdir:
-            pipe.save_pretrained(tmpdir)
-            index_file = os.path.join(tmpdir, "modular_model_index.json")
-            assert os.path.exists(index_file)
-
-            with open(index_file) as f:
-                index_contents = json.load(f)
-
-            compulsory_keys = {"_blocks_class_name", "_class_name", "_diffusers_version"}
-            for k in compulsory_keys:
-                assert k in index_contents
-
-            to_check_attrs = {"pretrained_model_name_or_path", "revision", "subfolder"}
-            for component in components:
-                spec = components_spec[component]
-                for attr in to_check_attrs:
-                    if getattr(spec, "pretrained_model_name_or_path", None) is not None:
-                        for attr in to_check_attrs:
-                            assert component in index_contents, f"{component} should be present in index but isn't."
-                            attr_value_from_index = index_contents[component][2][attr]
-                            assert getattr(spec, attr) == attr_value_from_index
+        pipe.save_pretrained(tmp_path)
+        index_file = os.path.join(tmp_path, "modular_model_index.json")
+        assert os.path.exists(index_file)
+
+        with open(index_file) as f:
+            index_contents = json.load(f)
+
+        compulsory_keys = {"_blocks_class_name", "_class_name", "_diffusers_version"}
+        for k in compulsory_keys:
+            assert k in index_contents
+
+        to_check_attrs = {"pretrained_model_name_or_path", "revision", "subfolder"}
+        for component in components:
+            spec = components_spec[component]
+            for attr in to_check_attrs:
+                if getattr(spec, "pretrained_model_name_or_path", None) is not None:
+                    for attr in to_check_attrs:
+                        assert component in index_contents, f"{component} should be present in index but isn't."
+                        attr_value_from_index = index_contents[component][2][attr]
+                        assert getattr(spec, attr) == attr_value_from_index
 
     def test_workflow_map(self):
         blocks = self.pipeline_blocks_class()

tests/modular_pipelines/test_modular_pipelines_custom_blocks.py

@@ -14,7 +14,6 @@
 
 import json
 import os
-import tempfile
 from collections import deque
 from typing import List
 
@@ -153,25 +152,24 @@ def test_custom_block_output(self):
         output_prompt = output.values["output_prompt"]
         assert output_prompt.startswith("Modular diffusers + ")
 
-    def test_custom_block_saving_loading(self):
+    def test_custom_block_saving_loading(self, tmp_path):
         custom_block = DummyCustomBlockSimple()
 
-        with tempfile.TemporaryDirectory() as tmpdir:
-            custom_block.save_pretrained(tmpdir)
-            assert any("modular_config.json" in k for k in os.listdir(tmpdir))
+        custom_block.save_pretrained(tmp_path)
+        assert any("modular_config.json" in k for k in os.listdir(tmp_path))
 
-            with open(os.path.join(tmpdir, "modular_config.json"), "r") as f:
-                config = json.load(f)
-            auto_map = config["auto_map"]
-            assert auto_map == {"ModularPipelineBlocks": "test_modular_pipelines_custom_blocks.DummyCustomBlockSimple"}
+        with open(os.path.join(tmp_path, "modular_config.json"), "r") as f:
+            config = json.load(f)
+        auto_map = config["auto_map"]
+        assert auto_map == {"ModularPipelineBlocks": "test_modular_pipelines_custom_blocks.DummyCustomBlockSimple"}
 
-            # For now, the Python script that implements the custom block has to be manually pushed to the Hub.
-            # This is why, we have to separately save the Python script here.
-            code_path = os.path.join(tmpdir, "test_modular_pipelines_custom_blocks.py")
-            with open(code_path, "w") as f:
-                f.write(CODE_STR)
+        # For now, the Python script that implements the custom block has to be manually pushed to the Hub.
+        # This is why, we have to separately save the Python script here.
+        code_path = os.path.join(tmp_path, "test_modular_pipelines_custom_blocks.py")
+        with open(code_path, "w") as f:
+            f.write(CODE_STR)
 
-            loaded_custom_block = ModularPipelineBlocks.from_pretrained(tmpdir, trust_remote_code=True)
+        loaded_custom_block = ModularPipelineBlocks.from_pretrained(tmp_path, trust_remote_code=True)
 
         pipe = loaded_custom_block.init_pipeline()
         prompt = "Diffusers is nice"