diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index 9a6b4b9b6fb4..f4fee82ecba5 100644 --- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py +++ b/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) diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index e97b543ff85d..41cd5f344779 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -1,5 +1,4 @@ import gc -import tempfile from typing import Callable import pytest @@ -328,16 +327,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) diff --git a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py index 9c5fd5be326d..422fa3323e21 100644 --- a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py +++ b/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"