Fix for fetching variants only

src/diffusers/pipelines/pipeline_loading_utils.py

@@ -191,37 +191,21 @@ def variant_compatible_siblings(filenames, variant=None) -> Union[List[os.PathLi
     # all variant filenames will be used by default
     usable_filenames = set(variant_filenames)
 
-    def convert_to_variant(filename):
-        if "index" in filename:
-            variant_filename = filename.replace("index", f"index.{variant}")
-        elif re.compile(f"^(.*?){transformers_index_format}").match(filename) is not None:
-            variant_filename = f"{filename.split('-')[0]}.{variant}-{'-'.join(filename.split('-')[1:])}"
-        else:
-            variant_filename = f"{filename.split('.')[0]}.{variant}.{filename.split('.')[1]}"
-        return variant_filename
-
     def find_component(filename):
         if not len(filename.split("/")) == 2:
             return
         component = filename.split("/")[0]
         return component
 
-    def has_sharded_variant(component, variant, variant_filenames):
-        # If component exists check for sharded variant index filename
-        # If component doesn't exist check main dir for sharded variant index filename
+    def has_variant(filename, variant_filenames):
+        component = find_component(filename)
         component = component + "/" if component else ""
-        variant_index_re = re.compile(
-            rf"{component}({'|'.join(weight_prefixes)})\.({'|'.join(weight_suffixs)})\.index\.{variant}\.json$"
-        )
-        return any(f for f in variant_filenames if variant_index_re.match(f) is not None)
 
-    for filename in non_variant_filenames:
-        if convert_to_variant(filename) in variant_filenames:
-            continue
+        return any(f.startswith(component) for f in variant_filenames)
 
-        component = find_component(filename)
-        # If a sharded variant exists skip adding to allowed patterns
-        if has_sharded_variant(component, variant, variant_filenames):
+    for filename in non_variant_filenames:
+        # If a variant exists skip adding to allowed patterns
+        if has_variant(filename, variant_filenames):
             continue
 
         usable_filenames.add(filename)

tests/pipelines/test_pipeline_utils.py

@@ -259,8 +259,6 @@ def test_non_variants_in_main_dir_downloaded(self):
             "diffusion_pytorch_model.safetensors",
             "model.safetensors",
             f"model.{variant}.safetensors",
-            f"diffusion_pytorch_model.{variant}.safetensors",
-            "diffusion_pytorch_model.safetensors",
         ]
         model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=None)
         assert all(variant not in f for f in model_filenames)
@@ -285,12 +283,50 @@ def test_mixed_variants_in_main_dir_downloaded(self):
             f"diffusion_pytorch_model.{variant}.safetensors",
             "diffusion_pytorch_model.safetensors",
             "model.safetensors",
-            f"diffusion_pytorch_model.{variant}.safetensors",
-            "diffusion_pytorch_model.safetensors",
         ]
         model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=variant)
         assert all(variant in f if f != non_variant_file else variant not in f for f in model_filenames)
 
+    def test_sharded_variants_in_main_dir_downloaded(self):
+        variant = "fp16"
+        filenames = [
+            "diffusion_pytorch_model.safetensors.index.json",
+            "diffusion_pytorch_model-00001-of-00003.safetensors",
+            "diffusion_pytorch_model-00002-of-00003.safetensors",
+            "diffusion_pytorch_model-00003-of-00003.safetensors",
+            f"diffusion_pytorch_model.{variant}-00001-of-00002.safetensors",
+            f"diffusion_pytorch_model.{variant}-00002-of-00002.safetensors",
+            f"diffusion_pytorch_model.safetensors.index.{variant}.json",
+        ]
+        model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=variant)
+        assert all(variant in f for f in model_filenames)
+
+    def test_mixed_sharded_and_variant_in_main_dir_downloaded(self):
+        variant = "fp16"
+        filenames = [
+            "diffusion_pytorch_model.safetensors.index.json",
+            "diffusion_pytorch_model-00001-of-00003.safetensors",
+            "diffusion_pytorch_model-00002-of-00003.safetensors",
+            "diffusion_pytorch_model-00003-of-00003.safetensors",
+            f"diffusion_pytorch_model.{variant}.safetensors",
+        ]
+        model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=variant)
+        assert all(variant in f for f in model_filenames)
+
+    def test_mixed_sharded_non_variants_in_main_dir_downloaded(self):
+        variant = "fp16"
+        filenames = [
+            f"diffusion_pytorch_model.safetensors.index.{variant}.json",
+            "diffusion_pytorch_model.safetensors.index.json",
+            "diffusion_pytorch_model-00001-of-00003.safetensors",
+            "diffusion_pytorch_model-00002-of-00003.safetensors",
+            "diffusion_pytorch_model-00003-of-00003.safetensors",
+            f"diffusion_pytorch_model.{variant}-00001-of-00002.safetensors",
+            f"diffusion_pytorch_model.{variant}-00002-of-00002.safetensors",
+        ]
+        model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=None)
+        assert all(variant not in f for f in model_filenames)
+
     def test_sharded_non_variants_downloaded(self):
         variant = "fp16"
         filenames = [
@@ -319,6 +355,35 @@ def test_sharded_variants_downloaded(self):
         model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=variant)
         assert all(variant in f for f in model_filenames)
 
+    def test_single_variant_with_sharded_non_variant_downloaded(self):
+        variant = "fp16"
+        filenames = [
+            "unet/diffusion_pytorch_model.safetensors.index.json",
+            "unet/diffusion_pytorch_model-00001-of-00003.safetensors",
+            "unet/diffusion_pytorch_model-00002-of-00003.safetensors",
+            "unet/diffusion_pytorch_model-00003-of-00003.safetensors",
+            f"unet/diffusion_pytorch_model.{variant}.safetensors",
+        ]
+        model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=variant)
+        assert all(variant in f for f in model_filenames)
+
+    def test_mixed_single_variant_with_sharded_non_variant_downloaded(self):
+        variant = "fp16"
+        allowed_non_variant = "unet"
+        filenames = [
+            "vae/diffusion_pytorch_model.safetensors.index.json",
+            "vae/diffusion_pytorch_model-00001-of-00003.safetensors",
+            "vae/diffusion_pytorch_model-00002-of-00003.safetensors",
+            "vae/diffusion_pytorch_model-00003-of-00003.safetensors",
+            f"vae/diffusion_pytorch_model.{variant}.safetensors",
+            "unet/diffusion_pytorch_model.safetensors.index.json",
+            "unet/diffusion_pytorch_model-00001-of-00003.safetensors",
+            "unet/diffusion_pytorch_model-00002-of-00003.safetensors",
+            "unet/diffusion_pytorch_model-00003-of-00003.safetensors",
+        ]
+        model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=variant)
+        assert all(variant in f if allowed_non_variant not in f else variant not in f for f in model_filenames)
+
     def test_sharded_mixed_variants_downloaded(self):
         variant = "fp16"
         allowed_non_variant = "unet"
@@ -338,6 +403,12 @@ def test_sharded_mixed_variants_downloaded(self):
         model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=variant)
         assert all(variant in f if allowed_non_variant not in f else variant not in f for f in model_filenames)
 
+    def test_downloading_when_no_variant_exists(self):
+        variant = "fp16"
+        filenames = ["model.safetensors", "diffusion_pytorch_model.safetensors"]
+        model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=variant)
+        assert len(model_filenames) != 0
+
 
 class ProgressBarTests(unittest.TestCase):
     def get_dummy_components_image_generation(self):