update related unit test cases

Signed-off-by: Liu, Kaixuan <[email protected]>

Author: kaixuanliu

tests/pipelines/test_pipelines.py

@@ -167,9 +167,9 @@ def test_one_request_upon_cached(self):
             download_requests = [r.method for r in m.request_history]
             assert download_requests.count("HEAD") == 15, "15 calls to files"
             assert download_requests.count("GET") == 17, "15 calls to files + model_info + model_index.json"
-            assert len(download_requests) == 32, (
-                "2 calls per file (15 files) + send_telemetry, model_info and model_index.json"
-            )
+            assert (
+                len(download_requests) == 32
+            ), "2 calls per file (15 files) + send_telemetry, model_info and model_index.json"
 
             with requests_mock.mock(real_http=True) as m:
                 DiffusionPipeline.download(
@@ -179,9 +179,9 @@ def test_one_request_upon_cached(self):
             cache_requests = [r.method for r in m.request_history]
             assert cache_requests.count("HEAD") == 1, "model_index.json is only HEAD"
             assert cache_requests.count("GET") == 1, "model info is only GET"
-            assert len(cache_requests) == 2, (
-                "We should call only `model_info` to check for _commit hash and `send_telemetry`"
-            )
+            assert (
+                len(cache_requests) == 2
+            ), "We should call only `model_info` to check for _commit hash and `send_telemetry`"
 
     def test_less_downloads_passed_object(self):
         with tempfile.TemporaryDirectory() as tmpdirname:
@@ -217,9 +217,9 @@ def test_less_downloads_passed_object_calls(self):
             assert download_requests.count("HEAD") == 13, "13 calls to files"
             # 17 - 2 because no call to config or model file for `safety_checker`
             assert download_requests.count("GET") == 15, "13 calls to files + model_info + model_index.json"
-            assert len(download_requests) == 28, (
-                "2 calls per file (13 files) + send_telemetry, model_info and model_index.json"
-            )
+            assert (
+                len(download_requests) == 28
+            ), "2 calls per file (13 files) + send_telemetry, model_info and model_index.json"
 
             with requests_mock.mock(real_http=True) as m:
                 DiffusionPipeline.download(
@@ -229,9 +229,9 @@ def test_less_downloads_passed_object_calls(self):
             cache_requests = [r.method for r in m.request_history]
             assert cache_requests.count("HEAD") == 1, "model_index.json is only HEAD"
             assert cache_requests.count("GET") == 1, "model info is only GET"
-            assert len(cache_requests) == 2, (
-                "We should call only `model_info` to check for _commit hash and `send_telemetry`"
-            )
+            assert (
+                len(cache_requests) == 2
+            ), "We should call only `model_info` to check for _commit hash and `send_telemetry`"
 
     def test_download_only_pytorch(self):
         with tempfile.TemporaryDirectory() as tmpdirname:
@@ -538,26 +538,38 @@ def test_download_variant_partly(self):
             variant = "no_ema"
 
             with tempfile.TemporaryDirectory() as tmpdirname:
-                tmpdirname = StableDiffusionPipeline.download(
-                    "hf-internal-testing/stable-diffusion-all-variants",
-                    cache_dir=tmpdirname,
-                    variant=variant,
-                    use_safetensors=use_safetensors,
-                )
-                all_root_files = [t[-1] for t in os.walk(tmpdirname)]
-                files = [item for sublist in all_root_files for item in sublist]
-
-                unet_files = os.listdir(os.path.join(tmpdirname, "unet"))
+                if use_safetensors:
+                    with self.assertRaises(OSError) as error_context:
+                        tmpdirname = StableDiffusionPipeline.download(
+                            "hf-internal-testing/stable-diffusion-all-variants",
+                            cache_dir=tmpdirname,
+                            variant=variant,
+                            use_safetensors=use_safetensors,
+                        )
+                    assert "Could not find the necessary `safetensors` weights" in str(error_context.exception)
+                else:
+                    tmpdirname = StableDiffusionPipeline.download(
+                        "hf-internal-testing/stable-diffusion-all-variants",
+                        cache_dir=tmpdirname,
+                        variant=variant,
+                        use_safetensors=use_safetensors,
+                    )
+                    all_root_files = [t[-1] for t in os.walk(tmpdirname)]
+                    files = [item for sublist in all_root_files for item in sublist]
 
-                # Some of the downloaded files should be a non-variant file, check:
-                # https://huggingface.co/hf-internal-testing/stable-diffusion-all-variants/tree/main/unet
-                assert len(files) == 15, f"We should only download 15 files, not {len(files)}"
-                # only unet has "no_ema" variant
-                assert f"diffusion_pytorch_model.{variant}{this_format}" in unet_files
-                assert len([f for f in files if f.endswith(f"{variant}{this_format}")]) == 1
-                # vae, safety_checker and text_encoder should have no variant
-                assert sum(f.endswith(this_format) and not f.endswith(f"{variant}{this_format}") for f in files) == 3
-                assert not any(f.endswith(other_format) for f in files)
+                    unet_files = os.listdir(os.path.join(tmpdirname, "unet"))
+
+                    # Some of the downloaded files should be a non-variant file, check:
+                    # https://huggingface.co/hf-internal-testing/stable-diffusion-all-variants/tree/main/unet
+                    assert len(files) == 15, f"We should only download 15 files, not {len(files)}"
+                    # only unet has "no_ema" variant
+                    assert f"diffusion_pytorch_model.{variant}{this_format}" in unet_files
+                    assert len([f for f in files if f.endswith(f"{variant}{this_format}")]) == 1
+                    # vae, safety_checker and text_encoder should have no variant
+                    assert (
+                        sum(f.endswith(this_format) and not f.endswith(f"{variant}{this_format}") for f in files) == 3
+                    )
+                    assert not any(f.endswith(other_format) for f in files)
 
     def test_download_variants_with_sharded_checkpoints(self):
         # Here we test for downloading of "variant" files belonging to the `unet` and
@@ -588,19 +600,16 @@ def test_download_legacy_variants_with_sharded_ckpts_raises_warning(self):
         logger = logging.get_logger("diffusers.pipelines.pipeline_utils")
         deprecated_warning_msg = "Warning: The repository contains sharded checkpoints for variant"
 
-        for is_local in [True, False]:
-            with CaptureLogger(logger) as cap_logger:
-                with tempfile.TemporaryDirectory() as tmpdirname:
-                    local_repo_id = repo_id
-                    if is_local:
-                        local_repo_id = snapshot_download(repo_id, cache_dir=tmpdirname)
+        with CaptureLogger(logger) as cap_logger:
+            with tempfile.TemporaryDirectory() as tmpdirname:
+                local_repo_id = snapshot_download(repo_id, cache_dir=tmpdirname)
 
-                    _ = DiffusionPipeline.from_pretrained(
-                        local_repo_id,
-                        safety_checker=None,
-                        variant="fp16",
-                        use_safetensors=True,
-                    )
+                _ = DiffusionPipeline.from_pretrained(
+                    local_repo_id,
+                    safety_checker=None,
+                    variant="fp16",
+                    use_safetensors=True,
+                )
             assert deprecated_warning_msg in str(cap_logger), "Deprecation warning not found in logs"
 
     def test_download_safetensors_only_variant_exists_for_model(self):
@@ -616,7 +625,7 @@ def test_download_safetensors_only_variant_exists_for_model(self):
                     variant=variant,
                     use_safetensors=use_safetensors,
                 )
-            assert "Error no file name" in str(error_context.exception)
+            assert "Could not find the necessary `safetensors` weights" in str(error_context.exception)
 
         # text encoder has fp16 variants so we can load it
         with tempfile.TemporaryDirectory() as tmpdirname:
@@ -675,7 +684,7 @@ def test_download_safetensors_variant_does_not_exist_for_model(self):
                     use_safetensors=use_safetensors,
                 )
 
-            assert "Error no file name" in str(error_context.exception)
+            assert "Could not find the necessary `safetensors` weights" in str(error_context.exception)
 
     def test_download_bin_variant_does_not_exist_for_model(self):
         variant = "no_ema"