From 7db09264c4e120992b3d47793da238f9e75875c7 Mon Sep 17 00:00:00 2001 From: lorenzozanee Date: Sun, 18 Oct 2026 14:20:53 +0800 Subject: [PATCH] fix(model-manager): filter repo-id Diffusers files by index --- .../model_manager/metadata/metadata_base.py | 10 ++- ...test_huggingface_metadata_download_urls.py | 87 +++++++++++++++++++ 2 files changed, 94 insertions(+), 3 deletions(-) create mode 100644 tests/backend/model_manager/test_huggingface_metadata_download_urls.py diff --git a/invokeai/backend/model_manager/metadata/metadata_base.py b/invokeai/backend/model_manager/metadata/metadata_base.py index 3a011850f5d..e100f1c7b06 100644 --- a/invokeai/backend/model_manager/metadata/metadata_base.py +++ b/invokeai/backend/model_manager/metadata/metadata_base.py @@ -119,6 +119,8 @@ 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: @@ -126,7 +128,9 @@ def download_urls( # 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) @@ -134,8 +138,8 @@ def download_urls( raise ModelMetadataUnavailableError(f"Could not fetch model index for '{self.id}'") from e resp.raise_for_status() 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] diff --git a/tests/backend/model_manager/test_huggingface_metadata_download_urls.py b/tests/backend/model_manager/test_huggingface_metadata_download_urls.py new file mode 100644 index 00000000000..da14d4977b3 --- /dev/null +++ b/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")