update

Author: DN6

tests/pipelines/test_pipeline_utils.py

@@ -224,6 +224,13 @@ def test_is_compatible_mixed_variants(self):
         ]
         self.assertTrue(is_safetensors_compatible(filenames, variant="fp16"))
 
+    def test_is_compatible_variant_and_non_safetensors(self):
+        filenames = [
+            "unet/diffusion_pytorch_model.fp16.safetensors",
+            "vae/diffusion_pytorch_model.bin",
+        ]
+        self.assertFalse(is_safetensors_compatible(filenames, variant="fp16"))
+
 
 class VariantCompatibleSiblingsTest(unittest.TestCase):
     def test_only_non_variants_downloaded(self):

tests/pipelines/test_pipelines.py

@@ -538,38 +538,26 @@ def test_download_variant_partly(self):
             variant = "no_ema"
 
             with tempfile.TemporaryDirectory() as tmpdirname:
-                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]
+                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"))
-
-                    # 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