Skip to content

[Feature] AutoModel can load components using model_index.json #11401

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 22 commits into from
May 26, 2025
Merged

[Feature] AutoModel can load components using model_index.json #11401

Changes from 3 commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
e506314
update
ishan-modi Apr 23, 2025
85024b0
update
ishan-modi Apr 23, 2025
6a0d0be
update
ishan-modi Apr 23, 2025
d86b0f2
update
ishan-modi Apr 23, 2025
314b6cc
Merge branch 'main' into fixes-issue-11388
ishan-modi Apr 23, 2025
528e002
addressed PR comments
ishan-modi Apr 29, 2025
6e92f40
update
ishan-modi Apr 29, 2025
76ea98d
Merge branch 'main' into fixes-issue-11388
ishan-modi Apr 29, 2025
0e53ad0
Merge branch 'main' into fixes-issue-11388
sayakpaul May 2, 2025
f697631
addressed PR comments
ishan-modi May 2, 2025
5614a15
added tests
ishan-modi May 2, 2025
f6b6b42
addressed PR comments
ishan-modi May 2, 2025
4e5cac1
updates
ishan-modi May 2, 2025
24f16f6
update
ishan-modi May 4, 2025
684384c
Merge branch 'main' into fixes-issue-11388
ishan-modi May 4, 2025
0fe68cd
addressed PR comments
ishan-modi May 13, 2025
2950372
Merge branch 'fixes-issue-11388' of https://github.com/ishan-modi/dif…
ishan-modi May 13, 2025
67e3404
Merge branch 'main' into fixes-issue-11388
ishan-modi May 13, 2025
694b81c
update
ishan-modi May 20, 2025
3bf51cd
Merge branch 'fixes-issue-11388' of https://github.com/ishan-modi/dif…
ishan-modi May 20, 2025
13420fb
fix style
ishan-modi May 26, 2025
af007ab
Merge branch 'main' into fixes-issue-11388
DN6 May 26, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 12 additions & 5 deletions src/diffusers/models/auto_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@


class AutoModel(ConfigMixin):
config_name = "config.json"
config_name = "model_index.json"

def __init__(self, *args, **kwargs):
raise EnvironmentError(
Expand Down Expand Up @@ -156,10 +156,17 @@ 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)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would avoid using exceptions for control flow and simplify this a bit

        load_config_kwargs = {
            "cache_dir": cache_dir,
            "force_download": force_download,
            "proxies": proxies,
            "token": token,
            "local_files_only": local_files_only,
            "revision": revision,
        }

        library = None
        orig_class_name = None
        from diffusers import pipelines

        # Always attempt to fetch model_index.json first
        try:
            cls.config_name = "model_index.json"
            config = cls.load_config(pretrained_model_or_path, **load_config_kwargs)

            if subfolder is not None and subfolder in config:
                library, orig_class_name = config[subfolder]

        except EntryNotFoundError as e:
            logger.debug(e)

        # Unable to load from model_index.json so fallback to loading from config
        if library is None and orig_class_name is None:
            cls.config_name = "config.json"
            load_config_kwargs.update({"subfolder": subfolder})

            config = cls.load_config(pretrained_model_or_path, **load_config_kwargs)
            orig_class_name = config["_class_name"]
            library = "diffusers"

        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), 
        )

orig_class_name = config["_class_name"]

library = importlib.import_module("diffusers")
try:
mindex_kwargs = {k: v for k, v in load_config_kwargs.items() if k != "subfolder"}
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

move this section under a if subfolder is not None:

if subfolder is not None:
     try: 
          ...

I think we are not supporting local path for pretrained_model_or_path here, are we? it would be a bit more complex if we do

Copy link
Contributor Author

@ishan-modi ishan-modi May 2, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What would be the complexity related to supporting local path here ?

Shouldn't we also be able to load models from structure where there is no folder heirarchy using this code, Following is a working example

try:
    control_net = AutoModel.from_pretrained(
        "ishan24/Sana_600M_1024px_ControlNet_diffusers",
        torch_dtype=torch.float16
    )
    print(f"test passed!")
except Exception as e:
    print(f"test failed: {e}")

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think for local path, you can append subfolder to the path as well ( we don't have to consider that for now)

like unet = AutoModel.from_pretrained("ishan24/SDXL_diffusers/unet")

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ohk, so we dont want to support flat repo's like this to load models

Because doing the following would mean that we dont support above case

if subfolder is not None:
     try: 
          ...

Copy link
Collaborator

@yiyixuxu yiyixuxu May 2, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ohh no we need to support flat repos, we need to all repos (we don't need to support some edge cases with local path when useer include subfolder directly in the file path, not as subfolder argument)

so basically here
if subfolder is not None -> we try model_index.json approach first, if that fails, we try the config.json approach
if subfolder is None -> we try the config.json approach directly(I think it is not meaningful to use model_index.json when subfolder is None because you need use subfolder as name to locate the info there, no?)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

got it, made the change

config = cls.load_config(os.path.join(pretrained_model_or_path), **mindex_kwargs)
library, orig_class_name = config[subfolder]
library = importlib.import_module(library)
except Exception:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should catch the specific Exception here instead of making it generic. This will help eliminate other side-effects.

# Fallback to loading the config from the config.json file
cls.config_name = "config.json"
config = cls.load_config(os.path.join(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)
if model_cls is None:
Expand Down