[Feature] AutoModel can load components using model_index.json

src/diffusers/models/auto_model.py

@@ -16,9 +16,11 @@
 import os
 from typing import Optional, Union
 
-from huggingface_hub.utils import validate_hf_hub_args
+from huggingface_hub import hf_hub_download
+from huggingface_hub.utils import EntryNotFoundError, validate_hf_hub_args
 
 from ..configuration_utils import ConfigMixin
+from ..pipelines.pipeline_loading_utils import ALL_IMPORTABLE_CLASSES, get_class_obj_and_candidates
 
 
 class AutoModel(ConfigMixin):
@@ -156,12 +158,38 @@ def from_pretrained(cls, pretrained_model_or_path: Optional[Union[str, os.PathLi
             "subfolder": subfolder,
         }
 
-        config = cls.load_config(pretrained_model_or_path, **load_config_kwargs)
-        orig_class_name = config["_class_name"]
+        if subfolder is not None:
+            try:
+                # To avoid circular import problem.
+                from diffusers import pipelines
+
+                mindex_kwargs = {k: v for k, v in load_config_kwargs.items() if k != "subfolder"}
+                mindex_kwargs["filename"] = "model_index.json"
+                config_path = hf_hub_download(pretrained_model_or_path, **mindex_kwargs)
+                config = cls.load_config(config_path, **load_config_kwargs)
+                library, orig_class_name = config[subfolder]
+                model_cls, _ = get_class_obj_and_candidates(
+                    library_name=library,
+                    class_name=orig_class_name,
+                    importable_classes=ALL_IMPORTABLE_CLASSES,
+                    pipelines=pipelines,
+                    is_pipeline_module=hasattr(pipelines, library),
+                )
+            except EntryNotFoundError:
+                # If `model_index.json` is not found, we load the model from the
+                # `config.json` file and `diffusers` library.
+                config = cls.load_config(pretrained_model_or_path, **load_config_kwargs)
+                library = importlib.import_module("diffusers")
+                orig_class_name = config["_class_name"]
+                model_cls = getattr(library, orig_class_name, None)
+        else:
+            # If `subfolder` is not provided, we load the model from the
+            # `config.json` file and `diffusers` library.
+            config = cls.load_config(pretrained_model_or_path, **load_config_kwargs)
+            library = importlib.import_module("diffusers")
+            orig_class_name = config["_class_name"]
+            model_cls = getattr(library, orig_class_name, None)
 
-        library = importlib.import_module("diffusers")
-
-        model_cls = getattr(library, orig_class_name, None)
         if model_cls is None:
             raise ValueError(f"AutoModel can't find a model linked to {orig_class_name}.")
 

src/diffusers/pipelines/pipeline_loading_utils.py

@@ -335,14 +335,14 @@ def get_class_obj_and_candidates(
     library_name, class_name, importable_classes, pipelines, is_pipeline_module, component_name=None, cache_dir=None
 ):
     """Simple helper method to retrieve class object of module as well as potential parent class objects"""
-    component_folder = os.path.join(cache_dir, component_name)
+    component_folder = os.path.join(cache_dir, component_name) if component_name and cache_dir else None
 
     if is_pipeline_module:
         pipeline_module = getattr(pipelines, library_name)
 
         class_obj = getattr(pipeline_module, class_name)
         class_candidates = dict.fromkeys(importable_classes.keys(), class_obj)
-    elif os.path.isfile(os.path.join(component_folder, library_name + ".py")):
+    elif component_folder and os.path.isfile(os.path.join(component_folder, library_name + ".py")):
         # load custom component
         class_obj = get_class_from_dynamic_module(
             component_folder, module_file=library_name + ".py", class_name=class_name

tests/models/test_models_auto.py

@@ -0,0 +1,20 @@
+import unittest
+from unittest.mock import patch
+
+from huggingface_hub.utils import EntryNotFoundError
+from transformers import CLIPTextModel
+
+from diffusers.models import AutoModel, UNet2DConditionModel
+
+
+class TestAutoModel(unittest.TestCase):
+    @patch("diffusers.models.auto_model.hf_hub_download", side_effect=EntryNotFoundError("File not found"))
+    def test_from_pretrained_falls_back_on_entry_error(self, mock_hf_hub_download):
+        model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-torch", subfolder="unet")
+        assert isinstance(model, UNet2DConditionModel)
+
+    def test_from_pretrained_loads_successfully(
+        self
+    ):
+        model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-torch", subfolder="text_encoder")
+        assert isinstance(model, CLIPTextModel)