update

Author: DN6

src/diffusers/pipelines/pipeline_loading_utils.py

@@ -197,29 +197,16 @@ def find_component(filename):
         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)
 
-    def has_variant(component, variant_filenames):
-        component = component + "/" if component else ""
         # Check for any variant file in this component
         return any(f.startswith(component) for f in variant_filenames)
 
     for filename in non_variant_filenames:
-        component = find_component(filename)
-
-        # Determine if sharded variant exists based on index file
-        if has_sharded_variant(component, variant, variant_filenames):
-            continue
-
         # If a variant exists skip adding to allowed patterns
-        if has_variant(component, variant_filenames):
+        if has_variant(filename, variant_filenames):
             continue
 
         usable_filenames.add(filename)

tests/pipelines/test_pipeline_utils.py

@@ -313,6 +313,20 @@ def test_mixed_sharded_and_variant_in_main_dir_downloaded(self):
         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 = [