diff --git a/tests/others/test_check_copies.py b/tests/others/test_check_copies.py index ec3a5702ea9c..0ae32c1a66f1 100644 --- a/tests/others/test_check_copies.py +++ b/tests/others/test_check_copies.py @@ -16,8 +16,8 @@ import re import shutil import sys -import tempfile -import unittest + +import pytest git_repo_path = os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) @@ -45,42 +45,42 @@ """ -class CopyCheckTester(unittest.TestCase): - def setUp(self): - self.diffusers_dir = tempfile.mkdtemp() - os.makedirs(os.path.join(self.diffusers_dir, "schedulers/")) - check_copies.DIFFUSERS_PATH = self.diffusers_dir +class TestCopyCheck: + @pytest.fixture + def diffusers_dir(self, tmp_path, monkeypatch): + """A stand-in `src/diffusers` holding only `scheduling_ddpm.py`, pointed at by `check_copies`.""" + os.makedirs(tmp_path / "schedulers") shutil.copy( os.path.join(git_repo_path, "src/diffusers/schedulers/scheduling_ddpm.py"), - os.path.join(self.diffusers_dir, "schedulers/scheduling_ddpm.py"), + tmp_path / "schedulers" / "scheduling_ddpm.py", ) + monkeypatch.setattr(check_copies, "DIFFUSERS_PATH", str(tmp_path)) + return tmp_path - def tearDown(self): - check_copies.DIFFUSERS_PATH = "src/diffusers" - shutil.rmtree(self.diffusers_dir) - - def check_copy_consistency(self, comment, class_name, class_code, overwrite_result=None): + def check_copy_consistency(self, diffusers_dir, comment, class_name, class_code, overwrite_result=None): code = comment + f"\nclass {class_name}(nn.Module):\n" + class_code if overwrite_result is not None: expected = comment + f"\nclass {class_name}(nn.Module):\n" + overwrite_result code = check_copies.run_ruff(code) - fname = os.path.join(self.diffusers_dir, "new_code.py") + fname = diffusers_dir / "new_code.py" with open(fname, "w", newline="\n") as f: f.write(code) if overwrite_result is None: - self.assertTrue(len(check_copies.is_copy_consistent(fname)) == 0) + assert len(check_copies.is_copy_consistent(fname)) == 0 else: - check_copies.is_copy_consistent(f.name, overwrite=True) + check_copies.is_copy_consistent(fname, overwrite=True) with open(fname, "r") as f: - self.assertTrue(f.read(), expected) + assert f.read() == expected - def test_find_code_in_diffusers(self): + def test_find_code_in_diffusers(self, diffusers_dir): + # `diffusers_dir` is requested for its `DIFFUSERS_PATH` patch — the lookup below resolves against it. code = check_copies.find_code_in_diffusers("schedulers.scheduling_ddpm.DDPMSchedulerOutput") - self.assertEqual(code, REFERENCE_CODE) + assert code == REFERENCE_CODE - def test_is_copy_consistent(self): + def test_is_copy_consistent(self, diffusers_dir): # Base copy consistency self.check_copy_consistency( + diffusers_dir, "# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput", "DDPMSchedulerOutput", REFERENCE_CODE + "\n", @@ -88,6 +88,7 @@ def test_is_copy_consistent(self): # With no empty line at the end self.check_copy_consistency( + diffusers_dir, "# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput", "DDPMSchedulerOutput", REFERENCE_CODE, @@ -95,6 +96,7 @@ def test_is_copy_consistent(self): # Copy consistency with rename self.check_copy_consistency( + diffusers_dir, "# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->Test", "TestSchedulerOutput", re.sub("DDPM", "Test", REFERENCE_CODE), @@ -103,6 +105,7 @@ def test_is_copy_consistent(self): # Copy consistency with a really long name long_class_name = "TestClassWithAReallyLongNameBecauseSomePeopleLikeThatForSomeReason" self.check_copy_consistency( + diffusers_dir, f"# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->{long_class_name}", f"{long_class_name}SchedulerOutput", re.sub("Bert", long_class_name, REFERENCE_CODE), @@ -110,6 +113,7 @@ def test_is_copy_consistent(self): # Copy consistency with overwrite self.check_copy_consistency( + diffusers_dir, "# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->Test", "TestSchedulerOutput", REFERENCE_CODE, diff --git a/tests/others/test_check_dummies.py b/tests/others/test_check_dummies.py index b9ddc2764465..c9a3ba3d111d 100644 --- a/tests/others/test_check_dummies.py +++ b/tests/others/test_check_dummies.py @@ -14,7 +14,6 @@ import os import sys -import unittest git_repo_path = os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) @@ -28,48 +27,46 @@ check_dummies.PATH_TO_DIFFUSERS = os.path.join(git_repo_path, "src", "diffusers") -class CheckDummiesTester(unittest.TestCase): +class TestCheckDummies: def test_find_backend(self): simple_backend = find_backend(" if not is_torch_available():") - self.assertEqual(simple_backend, "torch") + assert simple_backend == "torch" # backend_with_underscore = find_backend(" if not is_tensorflow_text_available():") - # self.assertEqual(backend_with_underscore, "tensorflow_text") + # assert backend_with_underscore == "tensorflow_text" double_backend = find_backend(" if not (is_torch_available() and is_transformers_available()):") - self.assertEqual(double_backend, "torch_and_transformers") + assert double_backend == "torch_and_transformers" # double_backend_with_underscore = find_backend( # " if not (is_sentencepiece_available() and is_tensorflow_text_available()):" # ) - # self.assertEqual(double_backend_with_underscore, "sentencepiece_and_tensorflow_text") + # assert double_backend_with_underscore == "sentencepiece_and_tensorflow_text" triple_backend = find_backend( " if not (is_torch_available() and is_transformers_available() and is_onnx_available()):" ) - self.assertEqual(triple_backend, "torch_and_transformers_and_onnx") + assert triple_backend == "torch_and_transformers_and_onnx" def test_read_init(self): objects = read_init() # We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects - self.assertIn("torch", objects) - self.assertIn("torch_and_transformers", objects) - self.assertIn("torch_and_transformers_and_onnx", objects) + assert "torch" in objects + assert "torch_and_transformers" in objects + assert "torch_and_transformers_and_onnx" in objects # Likewise, we can't assert on the exact content of a key - self.assertIn("UNet2DModel", objects["torch"]) - self.assertIn("StableDiffusionPipeline", objects["torch_and_transformers"]) - self.assertIn("LMSDiscreteScheduler", objects["torch_and_scipy"]) - self.assertIn("OnnxStableDiffusionPipeline", objects["torch_and_transformers_and_onnx"]) + assert "UNet2DModel" in objects["torch"] + assert "StableDiffusionPipeline" in objects["torch_and_transformers"] + assert "LMSDiscreteScheduler" in objects["torch_and_scipy"] + assert "OnnxStableDiffusionPipeline" in objects["torch_and_transformers_and_onnx"] def test_create_dummy_object(self): dummy_constant = create_dummy_object("CONSTANT", "'torch'") - self.assertEqual(dummy_constant, "\nCONSTANT = None\n") + assert dummy_constant == "\nCONSTANT = None\n" dummy_function = create_dummy_object("function", "'torch'") - self.assertEqual( - dummy_function, "\ndef function(*args, **kwargs):\n requires_backends(function, 'torch')\n" - ) + assert dummy_function == "\ndef function(*args, **kwargs):\n requires_backends(function, 'torch')\n" expected_dummy_class = """ class FakeClass(metaclass=DummyObject): @@ -87,7 +84,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, 'torch') """ dummy_class = create_dummy_object("FakeClass", "'torch'") - self.assertEqual(dummy_class, expected_dummy_class) + assert dummy_class == expected_dummy_class def test_create_dummy_files(self): expected_dummy_pytorch_file = """# This file is autogenerated by the command `make fix-copies`, do not edit. @@ -116,4 +113,4 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) """ dummy_files = create_dummy_files({"torch": ["CONSTANT", "function", "FakeClass"]}) - self.assertEqual(dummy_files["torch"], expected_dummy_pytorch_file) + assert dummy_files["torch"] == expected_dummy_pytorch_file diff --git a/tests/others/test_check_support_list.py b/tests/others/test_check_support_list.py index 0f6b134aad49..a4b460933ef9 100644 --- a/tests/others/test_check_support_list.py +++ b/tests/others/test_check_support_list.py @@ -1,6 +1,5 @@ import os import sys -import unittest from unittest.mock import mock_open, patch @@ -10,10 +9,8 @@ from check_support_list import check_documentation # noqa: E402 -class TestCheckSupportList(unittest.TestCase): - def setUp(self): - # Mock doc and source contents that we can reuse - self.doc_content = """# Documentation +# Mock doc and source contents that we can reuse +DOC_CONTENT = """# Documentation ## FooProcessor [[autodoc]] module.FooProcessor @@ -22,7 +19,7 @@ def setUp(self): [[autodoc]] module.BarProcessor """ - self.source_content = """ +SOURCE_CONTENT = """ class FooProcessor(nn.Module): pass @@ -30,12 +27,14 @@ class BarProcessor(nn.Module): pass """ + +class TestCheckSupportList: def test_check_documentation_all_documented(self): # In this test, both FooProcessor and BarProcessor are documented - with patch("builtins.open", mock_open(read_data=self.doc_content)) as doc_file: + with patch("builtins.open", mock_open(read_data=DOC_CONTENT)) as doc_file: doc_file.side_effect = [ - mock_open(read_data=self.doc_content).return_value, - mock_open(read_data=self.source_content).return_value, + mock_open(read_data=DOC_CONTENT).return_value, + mock_open(read_data=SOURCE_CONTENT).return_value, ] undocumented = check_documentation( @@ -44,7 +43,7 @@ def test_check_documentation_all_documented(self): doc_regex=r"\[\[autodoc\]\]\s([^\n]+)", src_regex=r"class\s+(\w+Processor)\(.*?nn\.Module.*?\):", ) - self.assertEqual(len(undocumented), 0, f"Expected no undocumented classes, got {undocumented}") + assert len(undocumented) == 0, f"Expected no undocumented classes, got {undocumented}" def test_check_documentation_missing_class(self): # In this test, only FooProcessor is documented, but BarProcessor is missing from the docs @@ -56,7 +55,7 @@ def test_check_documentation_missing_class(self): with patch("builtins.open", mock_open(read_data=doc_content_missing)) as doc_file: doc_file.side_effect = [ mock_open(read_data=doc_content_missing).return_value, - mock_open(read_data=self.source_content).return_value, + mock_open(read_data=SOURCE_CONTENT).return_value, ] undocumented = check_documentation( @@ -65,4 +64,4 @@ def test_check_documentation_missing_class(self): doc_regex=r"\[\[autodoc\]\]\s([^\n]+)", src_regex=r"class\s+(\w+Processor)\(.*?nn\.Module.*?\):", ) - self.assertIn("BarProcessor", undocumented, f"BarProcessor should be undocumented, got {undocumented}") + assert "BarProcessor" in undocumented, f"BarProcessor should be undocumented, got {undocumented}" diff --git a/tests/others/test_config.py b/tests/others/test_config.py index 376633601f7e..58567f80f550 100644 --- a/tests/others/test_config.py +++ b/tests/others/test_config.py @@ -14,10 +14,10 @@ # limitations under the License. import json -import tempfile -import unittest from pathlib import Path +import pytest + from diffusers import ( DDIMScheduler, DDPMScheduler, @@ -102,9 +102,9 @@ def __init__(self, test_file_1=Path("foo/bar"), test_file_2=Path("foo bar\\bar") pass -class ConfigTester(unittest.TestCase): +class TestConfig: def test_load_not_from_mixin(self): - with self.assertRaises(ValueError): + with pytest.raises(ValueError): ConfigMixin.load_config("dummy_path") def test_register_to_config(self): @@ -143,7 +143,7 @@ def test_register_to_config(self): assert config["d"] == "for diffusion" assert config["e"] == [1, 3] - def test_save_load(self): + def test_save_load(self, tmp_path): obj = SampleObject() config = obj.config @@ -153,10 +153,9 @@ def test_save_load(self): assert config["d"] == "for diffusion" assert config["e"] == [1, 3] - with tempfile.TemporaryDirectory() as tmpdirname: - obj.save_config(tmpdirname) - new_obj = SampleObject.from_config(SampleObject.load_config(tmpdirname)) - new_config = new_obj.config + obj.save_config(tmp_path) + new_obj = SampleObject.from_config(SampleObject.load_config(tmp_path)) + new_config = new_obj.config # unfreeze configs config = dict(config) @@ -262,7 +261,7 @@ def test_load_dpmsolver(self): # no warning should be thrown assert cap_logger.out == "" - def test_use_default_values(self): + def test_use_default_values(self, tmp_path): # let's first save a config that should be in the form # a=2, # b=5, @@ -277,14 +276,13 @@ def test_use_default_values(self): # make sure that default config has all keys in `_use_default_values` assert set(config_dict.keys()) == set(config.config._use_default_values) - with tempfile.TemporaryDirectory() as tmpdirname: - config.save_config(tmpdirname) + config.save_config(tmp_path) - # now loading it with SampleObject2 should put f into `_use_default_values` - config = SampleObject2.from_config(SampleObject2.load_config(tmpdirname)) + # now loading it with SampleObject2 should put f into `_use_default_values` + config = SampleObject2.from_config(SampleObject2.load_config(tmp_path)) - assert "f" in config.config._use_default_values - assert config.config.f == [1, 3] + assert "f" in config.config._use_default_values + assert config.config.f == [1, 3] # now loading the config, should **NOT** use [1, 3] for `f`, but the default [1, 4] value # **BECAUSE** it is part of `config.config._use_default_values` diff --git a/tests/others/test_ema.py b/tests/others/test_ema.py index 87820ed6af84..05a082902b0d 100644 --- a/tests/others/test_ema.py +++ b/tests/others/test_ema.py @@ -13,9 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import tempfile -import unittest - import torch from diffusers import UNet2DConditionModel @@ -27,7 +24,7 @@ enable_full_determinism() -class EMAModelTests(unittest.TestCase): +class TestEMAModel: model_id = "hf-internal-testing/tiny-stable-diffusion-pipe" batch_size = 1 prompt_length = 77 @@ -60,15 +57,14 @@ def simulate_backprop(self, unet): unet.load_state_dict(updated_state_dict) return unet - def test_from_pretrained(self): + def test_from_pretrained(self, tmp_path): # Save the model parameters to a temporary directory unet, ema_unet = self.get_models() - with tempfile.TemporaryDirectory() as tmpdir: - ema_unet.save_pretrained(tmpdir) + ema_unet.save_pretrained(tmp_path) - # Load the EMA model from the saved directory - loaded_ema_unet = EMAModel.from_pretrained(tmpdir, model_cls=UNet2DConditionModel, foreach=False) - loaded_ema_unet.to(torch_device) + # Load the EMA model from the saved directory + loaded_ema_unet = EMAModel.from_pretrained(tmp_path, model_cls=UNet2DConditionModel, foreach=False) + loaded_ema_unet.to(torch_device) # Check that the shadow parameters of the loaded model match the original EMA model for original_param, loaded_param in zip(ema_unet.shadow_params, loaded_ema_unet.shadow_params): @@ -164,14 +160,13 @@ def test_zero_decay(self): assert torch.allclose(step_one, step_two) @skip_mps - def test_serialization(self): + def test_serialization(self, tmp_path): unet, ema_unet = self.get_models() noisy_latents, timesteps, encoder_hidden_states = self.get_dummy_inputs() - with tempfile.TemporaryDirectory() as tmpdir: - ema_unet.save_pretrained(tmpdir) - loaded_unet = UNet2DConditionModel.from_pretrained(tmpdir, model_cls=UNet2DConditionModel) - loaded_unet = loaded_unet.to(unet.device) + ema_unet.save_pretrained(tmp_path) + loaded_unet = UNet2DConditionModel.from_pretrained(tmp_path, model_cls=UNet2DConditionModel) + loaded_unet = loaded_unet.to(unet.device) # Since no EMA step has been performed the outputs should match. output = unet(noisy_latents, timesteps, encoder_hidden_states).sample @@ -180,7 +175,7 @@ def test_serialization(self): assert torch.allclose(output, output_loaded, atol=1e-4) -class EMAModelTestsForeach(unittest.TestCase): +class TestEMAModelForeach: model_id = "hf-internal-testing/tiny-stable-diffusion-pipe" batch_size = 1 prompt_length = 77 @@ -215,15 +210,14 @@ def simulate_backprop(self, unet): unet.load_state_dict(updated_state_dict) return unet - def test_from_pretrained(self): + def test_from_pretrained(self, tmp_path): # Save the model parameters to a temporary directory unet, ema_unet = self.get_models() - with tempfile.TemporaryDirectory() as tmpdir: - ema_unet.save_pretrained(tmpdir) + ema_unet.save_pretrained(tmp_path) - # Load the EMA model from the saved directory - loaded_ema_unet = EMAModel.from_pretrained(tmpdir, model_cls=UNet2DConditionModel, foreach=True) - loaded_ema_unet.to(torch_device) + # Load the EMA model from the saved directory + loaded_ema_unet = EMAModel.from_pretrained(tmp_path, model_cls=UNet2DConditionModel, foreach=True) + loaded_ema_unet.to(torch_device) # Check that the shadow parameters of the loaded model match the original EMA model for original_param, loaded_param in zip(ema_unet.shadow_params, loaded_ema_unet.shadow_params): @@ -319,14 +313,13 @@ def test_zero_decay(self): assert torch.allclose(step_one, step_two) @skip_mps - def test_serialization(self): + def test_serialization(self, tmp_path): unet, ema_unet = self.get_models() noisy_latents, timesteps, encoder_hidden_states = self.get_dummy_inputs() - with tempfile.TemporaryDirectory() as tmpdir: - ema_unet.save_pretrained(tmpdir) - loaded_unet = UNet2DConditionModel.from_pretrained(tmpdir, model_cls=UNet2DConditionModel) - loaded_unet = loaded_unet.to(unet.device) + ema_unet.save_pretrained(tmp_path) + loaded_unet = UNet2DConditionModel.from_pretrained(tmp_path, model_cls=UNet2DConditionModel) + loaded_unet = loaded_unet.to(unet.device) # Since no EMA step has been performed the outputs should match. output = unet(noisy_latents, timesteps, encoder_hidden_states).sample diff --git a/tests/others/test_flashpack.py b/tests/others/test_flashpack.py index c14410c0d8e0..c20836df1c98 100644 --- a/tests/others/test_flashpack.py +++ b/tests/others/test_flashpack.py @@ -13,9 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pathlib -import tempfile -import unittest +import pytest from diffusers import AutoPipelineForText2Image from diffusers.models.auto_model import AutoModel @@ -27,48 +25,46 @@ import torch -class FlashPackTests(unittest.TestCase): +class TestFlashPack: model_id: str = "hf-internal-testing/tiny-flux-pipe" + # `AutoModel.from_pretrained` builds `_diffusers_load_id` by string-joining the path it is given, + # so the `tmp_path` fixture has to be passed to it as a `str`. + @require_flashpack - def test_save_load_model(self): + def test_save_load_model(self, tmp_path): model = AutoModel.from_pretrained(self.model_id, subfolder="transformer") - with tempfile.TemporaryDirectory() as temp_dir: - model.save_pretrained(temp_dir, use_flashpack=True) - self.assertTrue((pathlib.Path(temp_dir) / "model.flashpack").exists()) - model = AutoModel.from_pretrained(temp_dir, use_flashpack=True) + model.save_pretrained(tmp_path, use_flashpack=True) + assert (tmp_path / "model.flashpack").exists() + model = AutoModel.from_pretrained(str(tmp_path), use_flashpack=True) @require_flashpack - def test_save_load_pipeline(self): + def test_save_load_pipeline(self, tmp_path): pipeline = AutoPipelineForText2Image.from_pretrained(self.model_id) - with tempfile.TemporaryDirectory() as temp_dir: - pipeline.save_pretrained(temp_dir, use_flashpack=True) - self.assertTrue((pathlib.Path(temp_dir) / "transformer" / "model.flashpack").exists()) - self.assertTrue((pathlib.Path(temp_dir) / "vae" / "model.flashpack").exists()) - pipeline = AutoPipelineForText2Image.from_pretrained(temp_dir, use_flashpack=True) + pipeline.save_pretrained(tmp_path, use_flashpack=True) + assert (tmp_path / "transformer" / "model.flashpack").exists() + assert (tmp_path / "vae" / "model.flashpack").exists() + pipeline = AutoPipelineForText2Image.from_pretrained(tmp_path, use_flashpack=True) @require_torch_gpu @require_flashpack - def test_load_model_device_str(self): + def test_load_model_device_str(self, tmp_path): model = AutoModel.from_pretrained(self.model_id, subfolder="transformer") - with tempfile.TemporaryDirectory() as temp_dir: - model.save_pretrained(temp_dir, use_flashpack=True) - model = AutoModel.from_pretrained(temp_dir, use_flashpack=True, device_map={"": "cuda"}) - self.assertTrue(model.device.type == "cuda") + model.save_pretrained(tmp_path, use_flashpack=True) + model = AutoModel.from_pretrained(str(tmp_path), use_flashpack=True, device_map={"": "cuda"}) + assert model.device.type == "cuda" @require_torch_gpu @require_flashpack - def test_load_model_device(self): + def test_load_model_device(self, tmp_path): model = AutoModel.from_pretrained(self.model_id, subfolder="transformer") - with tempfile.TemporaryDirectory() as temp_dir: - model.save_pretrained(temp_dir, use_flashpack=True) - model = AutoModel.from_pretrained(temp_dir, use_flashpack=True, device_map={"": torch.device("cuda")}) - self.assertTrue(model.device.type == "cuda") + model.save_pretrained(tmp_path, use_flashpack=True) + model = AutoModel.from_pretrained(str(tmp_path), use_flashpack=True, device_map={"": torch.device("cuda")}) + assert model.device.type == "cuda" @require_flashpack - def test_load_model_device_auto(self): + def test_load_model_device_auto(self, tmp_path): model = AutoModel.from_pretrained(self.model_id, subfolder="transformer") - with tempfile.TemporaryDirectory() as temp_dir: - model.save_pretrained(temp_dir, use_flashpack=True) - with self.assertRaises(ValueError): - model = AutoModel.from_pretrained(temp_dir, use_flashpack=True, device_map={"": "auto"}) + model.save_pretrained(tmp_path, use_flashpack=True) + with pytest.raises(ValueError): + model = AutoModel.from_pretrained(str(tmp_path), use_flashpack=True, device_map={"": "auto"}) diff --git a/tests/others/test_hub_utils.py b/tests/others/test_hub_utils.py index c897afd26db9..23c48b7f7f5e 100644 --- a/tests/others/test_hub_utils.py +++ b/tests/others/test_hub_utils.py @@ -14,9 +14,8 @@ # limitations under the License. import json import os -import unittest -from pathlib import Path -from tempfile import TemporaryDirectory + +import pytest from diffusers.utils.hub_utils import ( _get_checkpoint_shard_files, @@ -25,17 +24,16 @@ ) -class CreateModelCardTest(unittest.TestCase): - def test_generate_model_card_with_library_name(self): - with TemporaryDirectory() as tmpdir: - file_path = Path(tmpdir) / "README.md" - file_path.write_text("---\nlibrary_name: foo\n---\nContent\n") - model_card = load_or_create_model_card(file_path) - populate_model_card(model_card) - assert model_card.data.library_name == "foo" +class TestCreateModelCard: + def test_generate_model_card_with_library_name(self, tmp_path): + file_path = tmp_path / "README.md" + file_path.write_text("---\nlibrary_name: foo\n---\nContent\n") + model_card = load_or_create_model_card(file_path) + populate_model_card(model_card) + assert model_card.data.library_name == "foo" -class GetCheckpointShardFilesTest(unittest.TestCase): +class TestGetCheckpointShardFiles: def _write_index(self, model_dir, shard_filename): index = {"metadata": {"total_size": 1}, "weight_map": {"w": shard_filename}} index_filename = os.path.join(model_dir, "diffusion_pytorch_model.safetensors.index.json") @@ -43,26 +41,23 @@ def _write_index(self, model_dir, shard_filename): json.dump(index, f) return index_filename - def test_rejects_parent_directory_traversal(self): - with TemporaryDirectory() as tmpdir: - model_dir = os.path.join(tmpdir, "model") - os.makedirs(model_dir) - index_filename = self._write_index(model_dir, "../secret/SECRET.safetensors") - with self.assertRaises(ValueError): - _get_checkpoint_shard_files(model_dir, index_filename) - - def test_rejects_absolute_path(self): - with TemporaryDirectory() as tmpdir: - model_dir = os.path.join(tmpdir, "model") - os.makedirs(model_dir) - index_filename = self._write_index(model_dir, os.path.join(tmpdir, "secret", "SECRET.safetensors")) - with self.assertRaises(ValueError): - _get_checkpoint_shard_files(model_dir, index_filename) - - def test_rejects_subdirectory_component(self): - with TemporaryDirectory() as tmpdir: - model_dir = os.path.join(tmpdir, "model") - os.makedirs(model_dir) - index_filename = self._write_index(model_dir, "sub/shard.safetensors") - with self.assertRaises(ValueError): - _get_checkpoint_shard_files(model_dir, index_filename) + def test_rejects_parent_directory_traversal(self, tmp_path): + model_dir = os.path.join(tmp_path, "model") + os.makedirs(model_dir) + index_filename = self._write_index(model_dir, "../secret/SECRET.safetensors") + with pytest.raises(ValueError): + _get_checkpoint_shard_files(model_dir, index_filename) + + def test_rejects_absolute_path(self, tmp_path): + model_dir = os.path.join(tmp_path, "model") + os.makedirs(model_dir) + index_filename = self._write_index(model_dir, os.path.join(tmp_path, "secret", "SECRET.safetensors")) + with pytest.raises(ValueError): + _get_checkpoint_shard_files(model_dir, index_filename) + + def test_rejects_subdirectory_component(self, tmp_path): + model_dir = os.path.join(tmp_path, "model") + os.makedirs(model_dir) + index_filename = self._write_index(model_dir, "sub/shard.safetensors") + with pytest.raises(ValueError): + _get_checkpoint_shard_files(model_dir, index_filename) diff --git a/tests/others/test_image_processor.py b/tests/others/test_image_processor.py index 88e82ab54b82..0d358699f105 100644 --- a/tests/others/test_image_processor.py +++ b/tests/others/test_image_processor.py @@ -13,8 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - import numpy as np import PIL.Image import torch @@ -22,7 +20,7 @@ from diffusers.image_processor import VaeImageProcessor -class ImageProcessorTest(unittest.TestCase): +class TestImageProcessor: @property def dummy_sample(self): batch_size = 1 diff --git a/tests/others/test_outputs.py b/tests/others/test_outputs.py index 90b8bfe9464e..2c1a011523c2 100644 --- a/tests/others/test_outputs.py +++ b/tests/others/test_outputs.py @@ -1,5 +1,4 @@ import pickle as pkl -import unittest from dataclasses import dataclass import numpy as np @@ -15,7 +14,7 @@ class CustomOutput(BaseOutput): images: list[PIL.Image.Image] | np.ndarray -class ConfigTester(unittest.TestCase): +class TestOutputs: def test_outputs_single_attribute(self): outputs = CustomOutput(images=np.random.rand(1, 3, 4, 4)) @@ -80,14 +79,14 @@ def test_torch_pytree(self): data = np.random.rand(1, 3, 4, 4) x = CustomOutput(images=data) - self.assertFalse(torch.utils._pytree._is_leaf(x)) + assert not torch.utils._pytree._is_leaf(x) expected_flat_outs = [data] expected_tree_spec = torch.utils._pytree.TreeSpec(CustomOutput, ["images"], [torch.utils._pytree.LeafSpec()]) actual_flat_outs, actual_tree_spec = torch.utils._pytree.tree_flatten(x) - self.assertEqual(expected_flat_outs, actual_flat_outs) - self.assertEqual(expected_tree_spec, actual_tree_spec) + assert expected_flat_outs == actual_flat_outs + assert expected_tree_spec == actual_tree_spec unflattened_x = torch.utils._pytree.tree_unflatten(actual_flat_outs, actual_tree_spec) - self.assertEqual(x, unflattened_x) + assert x == unflattened_x diff --git a/tests/others/test_training.py b/tests/others/test_training.py index a339ee8a3c6b..1d7aa465863c 100644 --- a/tests/others/test_training.py +++ b/tests/others/test_training.py @@ -13,8 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - import torch from diffusers import DDIMScheduler, DDPMScheduler, UNet2DModel @@ -26,7 +24,7 @@ torch.backends.cuda.matmul.allow_tf32 = False -class TrainingTests(unittest.TestCase): +class TestTraining: def get_model_optimizer(self, resolution=32): set_seed(0) model = UNet2DModel(sample_size=resolution, in_channels=3, out_channels=3) @@ -83,8 +81,8 @@ def test_training_step_equality(self): optimizer.step() del model, optimizer - self.assertTrue(torch.allclose(ddpm_noisy_images, ddim_noisy_images, atol=1e-5)) - self.assertTrue(torch.allclose(ddpm_noise_pred, ddim_noise_pred, atol=1e-5)) + assert torch.allclose(ddpm_noisy_images, ddim_noisy_images, atol=1e-5) + assert torch.allclose(ddpm_noise_pred, ddim_noise_pred, atol=1e-5) def test_confidence_aware_loss(self): logits = torch.tensor([[[5.0, 0.0], [0.0, 5.0]]]) @@ -94,8 +92,8 @@ def test_confidence_aware_loss(self): loss, loss_sft, loss_conf = compute_confidence_aware_loss( logits, labels, lambda_conf=0.0, per_token_weights=weights ) - self.assertTrue(torch.allclose(loss, loss_sft)) - self.assertTrue(torch.allclose(loss_conf, torch.zeros_like(loss_conf))) + assert torch.allclose(loss, loss_sft) + assert torch.allclose(loss_conf, torch.zeros_like(loss_conf)) lambda_conf = 0.25 loss, loss_sft, loss_conf = compute_confidence_aware_loss( @@ -118,14 +116,14 @@ def test_confidence_aware_loss(self): ).sum().clamp_min(1) expected = expected_sft + lambda_conf * expected_conf - self.assertTrue(torch.allclose(loss_sft, expected_sft)) - self.assertTrue(torch.allclose(loss_conf, expected_conf)) - self.assertTrue(torch.allclose(loss, expected)) + assert torch.allclose(loss_sft, expected_sft) + assert torch.allclose(loss_conf, expected_conf) + assert torch.allclose(loss, expected) # Temperature affects only the confidence term. loss_t, loss_sft_t, loss_conf_t = compute_confidence_aware_loss( logits, labels, lambda_conf=lambda_conf, temperature=0.5, per_token_weights=weights ) - self.assertTrue(torch.allclose(loss_sft_t, expected_sft)) - self.assertFalse(torch.allclose(loss_conf_t, expected_conf)) - self.assertTrue(torch.allclose(loss_t, loss_sft_t + lambda_conf * loss_conf_t)) + assert torch.allclose(loss_sft_t, expected_sft) + assert not torch.allclose(loss_conf_t, expected_conf) + assert torch.allclose(loss_t, loss_sft_t + lambda_conf * loss_conf_t) diff --git a/tests/others/test_utils.py b/tests/others/test_utils.py index 412c7478c5a7..d1a59cec52f1 100755 --- a/tests/others/test_utils.py +++ b/tests/others/test_utils.py @@ -15,7 +15,6 @@ import importlib import os -import unittest import warnings import pytest @@ -34,19 +33,19 @@ TOKEN = "hf_94wBhPGp6KrrTH3KDchhKpRxZwd6dmHWLL" -class DeprecateTester(unittest.TestCase): +class TestDeprecate: higher_version = ".".join([str(int(__version__.split(".")[0]) + 1)] + __version__.split(".")[1:]) lower_version = "0.0.1" def test_deprecate_function_arg(self): kwargs = {"deprecated_arg": 4} - with self.assertWarns(FutureWarning) as warning: + with pytest.warns(FutureWarning) as warning: output = deprecate("deprecated_arg", self.higher_version, "message", take_from=kwargs) assert output == 4 assert ( - str(warning.warning) + str(warning[0].message) == f"The `deprecated_arg` argument is deprecated and will be removed in version {self.higher_version}." " message" ) @@ -54,19 +53,19 @@ def test_deprecate_function_arg(self): def test_deprecate_function_arg_tuple(self): kwargs = {"deprecated_arg": 4} - with self.assertWarns(FutureWarning) as warning: + with pytest.warns(FutureWarning) as warning: output = deprecate(("deprecated_arg", self.higher_version, "message"), take_from=kwargs) assert output == 4 assert ( - str(warning.warning) + str(warning[0].message) == f"The `deprecated_arg` argument is deprecated and will be removed in version {self.higher_version}." " message" ) def test_deprecate_function_args(self): kwargs = {"deprecated_arg_1": 4, "deprecated_arg_2": 8} - with self.assertWarns(FutureWarning) as warning: + with pytest.warns(FutureWarning) as warning: output_1, output_2 = deprecate( ("deprecated_arg_1", self.higher_version, "Hey"), ("deprecated_arg_2", self.higher_version, "Hey"), @@ -75,47 +74,45 @@ def test_deprecate_function_args(self): assert output_1 == 4 assert output_2 == 8 assert ( - str(warning.warnings[0].message) - == "The `deprecated_arg_1` argument is deprecated and will be removed in version" + str(warning[0].message) == "The `deprecated_arg_1` argument is deprecated and will be removed in version" f" {self.higher_version}. Hey" ) assert ( - str(warning.warnings[1].message) - == "The `deprecated_arg_2` argument is deprecated and will be removed in version" + str(warning[1].message) == "The `deprecated_arg_2` argument is deprecated and will be removed in version" f" {self.higher_version}. Hey" ) def test_deprecate_function_incorrect_arg(self): kwargs = {"deprecated_arg": 4} - with self.assertRaises(TypeError) as error: + with pytest.raises(TypeError) as error: deprecate(("wrong_arg", self.higher_version, "message"), take_from=kwargs) - assert "test_deprecate_function_incorrect_arg in" in str(error.exception) - assert "line" in str(error.exception) - assert "got an unexpected keyword argument `deprecated_arg`" in str(error.exception) + assert "test_deprecate_function_incorrect_arg in" in str(error.value) + assert "line" in str(error.value) + assert "got an unexpected keyword argument `deprecated_arg`" in str(error.value) def test_deprecate_arg_no_kwarg(self): - with self.assertWarns(FutureWarning) as warning: + with pytest.warns(FutureWarning) as warning: deprecate(("deprecated_arg", self.higher_version, "message")) assert ( - str(warning.warning) + str(warning[0].message) == f"`deprecated_arg` is deprecated and will be removed in version {self.higher_version}. message" ) def test_deprecate_args_no_kwarg(self): - with self.assertWarns(FutureWarning) as warning: + with pytest.warns(FutureWarning) as warning: deprecate( ("deprecated_arg_1", self.higher_version, "Hey"), ("deprecated_arg_2", self.higher_version, "Hey"), ) assert ( - str(warning.warnings[0].message) + str(warning[0].message) == f"`deprecated_arg_1` is deprecated and will be removed in version {self.higher_version}. Hey" ) assert ( - str(warning.warnings[1].message) + str(warning[1].message) == f"`deprecated_arg_2` is deprecated and will be removed in version {self.higher_version}. Hey" ) @@ -123,12 +120,12 @@ def test_deprecate_class_obj(self): class Args: arg = 5 - with self.assertWarns(FutureWarning) as warning: + with pytest.warns(FutureWarning) as warning: arg = deprecate(("arg", self.higher_version, "message"), take_from=Args()) assert arg == 5 assert ( - str(warning.warning) + str(warning[0].message) == f"The `arg` attribute is deprecated and will be removed in version {self.higher_version}. message" ) @@ -137,7 +134,7 @@ class Args: arg = 5 foo = 7 - with self.assertWarns(FutureWarning) as warning: + with pytest.warns(FutureWarning) as warning: arg_1, arg_2 = deprecate( ("arg", self.higher_version, "message"), ("foo", self.higher_version, "message"), @@ -148,41 +145,37 @@ class Args: assert arg_1 == 5 assert arg_2 == 7 assert ( - str(warning.warning) + str(warning[0].message) == f"The `arg` attribute is deprecated and will be removed in version {self.higher_version}. message" ) assert ( - str(warning.warnings[0].message) - == f"The `arg` attribute is deprecated and will be removed in version {self.higher_version}. message" - ) - assert ( - str(warning.warnings[1].message) + str(warning[1].message) == f"The `foo` attribute is deprecated and will be removed in version {self.higher_version}. message" ) def test_deprecate_incorrect_version(self): kwargs = {"deprecated_arg": 4} - with self.assertRaises(ValueError) as error: + with pytest.raises(ValueError) as error: deprecate(("wrong_arg", self.lower_version, "message"), take_from=kwargs) assert ( - str(error.exception) + str(error.value) == "The deprecation tuple ('wrong_arg', '0.0.1', 'message') should be removed since diffusers' version" f" {__version__} is >= {self.lower_version}" ) def test_deprecate_incorrect_no_standard_warn(self): - with self.assertWarns(FutureWarning) as warning: + with pytest.warns(FutureWarning) as warning: deprecate(("deprecated_arg", self.higher_version, "This message is better!!!"), standard_warn=False) - assert str(warning.warning) == "This message is better!!!" + assert str(warning[0].message) == "This message is better!!!" def test_deprecate_stacklevel(self): - with self.assertWarns(FutureWarning) as warning: + with pytest.warns(FutureWarning) as warning: deprecate(("deprecated_arg", self.higher_version, "This message is better!!!"), standard_warn=False) - assert str(warning.warning) == "This message is better!!!" - assert "diffusers/tests/others/test_utils.py" in warning.filename + assert str(warning[0].message) == "This message is better!!!" + assert "diffusers/tests/others/test_utils.py" in warning[0].filename def test_deprecate_testing_utils_module(self): import diffusers.utils.testing_utils @@ -204,7 +197,7 @@ def test_deprecate_testing_utils_module(self): ), f"Expected deprecation message substring not found, got: {messages}" -class FourierFilterTester(unittest.TestCase): +class TestFourierFilter: """Tests for :func:`diffusers.utils.torch_utils.fourier_filter` (FreeU helper).""" def _run_without_complexhalf_warning(self, dtype): @@ -247,7 +240,7 @@ def test_fourier_filter_preserves_dtype_and_shape(self): assert out.shape == x.shape -class RandnTensorTester(unittest.TestCase): +class TestRandnTensor: """Tests for :func:`diffusers.utils.torch_utils.randn_tensor`.""" def test_mps_suppresses_cpu_generator_info_log(self): @@ -270,22 +263,14 @@ def _capture(target_device): return cl.out mps_out = _capture("mps") - self.assertNotIn( - "moved to", - mps_out, - f"MPS target should not emit the CPU-fallback info log, got: {mps_out}", - ) + assert "moved to" not in mps_out, f"MPS target should not emit the CPU-fallback info log, got: {mps_out}" cuda_out = _capture("cuda") - self.assertIn( - "moved to", - cuda_out, - f"Non-MPS target should still emit the CPU-fallback info log, got: {cuda_out}", - ) + assert "moved to" in cuda_out, f"Non-MPS target should still emit the CPU-fallback info log, got: {cuda_out}" # Copied from https://github.com/huggingface/transformers/blob/main/tests/utils/test_expectations.py -class ExpectationsTester(unittest.TestCase): +class TestExpectations: def test_expectations(self): expectations = Expectations( { @@ -312,7 +297,7 @@ def check(value, key): check(2, ("cuda", 2)) expectations = Expectations({("cuda", 8): 1}) - with self.assertRaises(ValueError): + with pytest.raises(ValueError): expectations.find_expectation(("xpu", None)) diff --git a/tests/others/test_video_processor.py b/tests/others/test_video_processor.py index 7fd98f26652f..882b3891de89 100644 --- a/tests/others/test_video_processor.py +++ b/tests/others/test_video_processor.py @@ -13,12 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - import numpy as np import PIL.Image +import pytest import torch -from parameterized import parameterized from diffusers.video_processor import VideoProcessor @@ -27,7 +25,7 @@ torch.manual_seed(0) -class VideoProcessorTest(unittest.TestCase): +class TestVideoProcessor: def get_dummy_sample(self, input_type): batch_size = 1 num_frames = 5 @@ -128,7 +126,7 @@ def to_np(self, video): return video - @parameterized.expand(["list_images", "list_list_images"]) + @pytest.mark.parametrize("input_type", ["list_images", "list_list_images"]) def test_video_processor_pil(self, input_type): video_processor = VideoProcessor(do_resize=False, do_normalize=True) @@ -140,7 +138,7 @@ def test_video_processor_pil(self, input_type): input_np = self.to_np(input).astype("float32") / 255.0 if output_type != "pil" else self.to_np(input) assert np.abs(input_np - out_np).max() < 1e-6, f"Decoded output does not match input for {output_type=}" - @parameterized.expand(["list_4d_np", "list_5d_np", "5d_np"]) + @pytest.mark.parametrize("input_type", ["list_4d_np", "list_5d_np", "5d_np"]) def test_video_processor_np(self, input_type): video_processor = VideoProcessor(do_resize=False, do_normalize=True) @@ -154,7 +152,7 @@ def test_video_processor_np(self, input_type): ) assert np.abs(input_np - out_np).max() < 1e-6, f"Decoded output does not match input for {output_type=}" - @parameterized.expand(["list_4d_pt", "list_5d_pt", "5d_pt"]) + @pytest.mark.parametrize("input_type", ["list_4d_pt", "list_5d_pt", "5d_pt"]) def test_video_processor_pt(self, input_type): video_processor = VideoProcessor(do_resize=False, do_normalize=True)