fix(model-manager): filter repo-id Diffusers files by index

invokeai/backend/model_manager/metadata/metadata_base.py

@@ -126,14 +126,18 @@ def download_urls(
 
         paths = filter_files([x.path for x in self.files], variant, subfolder, subfolders)  #  all files in the model
 
+        repo_root = Path(self.name) if paths and paths[0].parts[0] == self.name else Path()
+
         # Determine prefix for model_index.json check - only applies for single subfolder
         prefix = ""
         if subfolder and not subfolders:
             prefix = f"{subfolder}/"
 
         # the next step reads model_index.json to determine which subdirectories belong
         # to the model (only for single subfolder case)
-        if Path(f"{prefix}model_index.json") in paths:
+        model_root = repo_root / prefix
+        model_index_path = model_root / "model_index.json"
+        if model_index_path in paths:
             url = hf_hub_url(self.id, filename="model_index.json", subfolder=str(subfolder) if subfolder else None)
             try:
                 resp = session.get(url, timeout=HF_METADATA_REQUEST_TIMEOUT_SECONDS)
@@ -145,8 +149,8 @@ def download_urls(
             except RequestException as e:
                 raise ModelMetadataUnavailableError(f"Could not fetch model index for '{self.id}'") from e
             submodels = resp.json()
-            paths = [Path(subfolder or "", x) for x in paths if Path(x).parent.as_posix() in submodels]
-            paths.insert(0, Path(f"{prefix}model_index.json"))
+            paths = [x for x in paths if x.relative_to(model_root).parent.as_posix() in submodels]
+            paths.insert(0, model_index_path)
 
         return [x for x in self.files if x.path in paths]
 

tests/backend/model_manager/test_huggingface_metadata_download_urls.py

@@ -0,0 +1,87 @@
+from pathlib import Path
+from unittest.mock import Mock
+
+from invokeai.backend.model_manager.metadata.metadata_base import HuggingFaceMetadata, RemoteModelFile
+
+
+def test_repo_id_diffusers_download_uses_model_index_components() -> None:
+    session = Mock()
+    response = Mock(status_code=200)
+    response.json.return_value = {
+        "_class_name": "StableDiffusionPipeline",
+        "unet": ["diffusers", "UNet2DConditionModel"],
+        "vae": ["diffusers", "AutoencoderKL"],
+    }
+    session.get.return_value = response
+    repo_files = [
+        "model_index.json",
+        "model.safetensors",
+        "unet/config.json",
+        "unet/diffusion_pytorch_model.safetensors",
+        "vae/config.json",
+        "vae/diffusion_pytorch_model.safetensors",
+        "README.md",
+    ]
+    metadata = HuggingFaceMetadata(
+        id="org/diffusers-repo",
+        name="diffusers-repo",
+        files=[
+            RemoteModelFile(
+                url=f"https://huggingface.co/org/diffusers-repo/resolve/main/{path}",
+                path=Path("diffusers-repo", path),
+            )
+            for path in repo_files
+        ],
+    )
+
+    selected = {file.path.as_posix() for file in metadata.download_urls(session=session)}
+
+    assert selected == {
+        "diffusers-repo/model_index.json",
+        "diffusers-repo/unet/config.json",
+        "diffusers-repo/unet/diffusion_pytorch_model.safetensors",
+        "diffusers-repo/vae/config.json",
+        "diffusers-repo/vae/diffusion_pytorch_model.safetensors",
+    }
+    assert session.get.call_args.args[0].endswith("model_index.json")
+
+
+def test_repo_id_diffusers_subfolder_uses_its_model_index_components() -> None:
+    session = Mock()
+    response = Mock(status_code=200)
+    response.json.return_value = {
+        "_class_name": "StableDiffusionPipeline",
+        "unet": ["diffusers", "UNet2DConditionModel"],
+        "vae": ["diffusers", "AutoencoderKL"],
+    }
+    session.get.return_value = response
+    repo_files = [
+        "pipeline/model_index.json",
+        "pipeline/model.safetensors",
+        "pipeline/unet/config.json",
+        "pipeline/unet/diffusion_pytorch_model.safetensors",
+        "pipeline/vae/config.json",
+        "pipeline/vae/diffusion_pytorch_model.safetensors",
+    ]
+    metadata = HuggingFaceMetadata(
+        id="org/diffusers-repo",
+        name="diffusers-repo",
+        files=[
+            RemoteModelFile(
+                url=f"https://huggingface.co/org/diffusers-repo/resolve/main/{path}",
+                path=Path("diffusers-repo", path),
+            )
+            for path in repo_files
+        ],
+    )
+
+    selected = {file.path.as_posix() for file in metadata.download_urls(subfolder=Path("pipeline"), session=session)}
+
+    assert selected == {
+        "diffusers-repo/pipeline/model_index.json",
+        "diffusers-repo/pipeline/unet/config.json",
+        "diffusers-repo/pipeline/unet/diffusion_pytorch_model.safetensors",
+        "diffusers-repo/pipeline/vae/config.json",
+        "diffusers-repo/pipeline/vae/diffusion_pytorch_model.safetensors",
+    }
+    assert session.get.call_args.args[0].endswith("pipeline/model_index.json")