[Feature] AutoModel can load components using model_index.json

[ishan-modi]

What does this PR do?

Allows AutoModel to load components using model_index.json. If model_index.json is not found config.json(existing) is used as a fallback

Fixes #11388

Who can review?

@yiyixuxu

[yiyixuxu]

thanks @ishan-modi !
can you run on the testing script here #11388 (comment)

[ishan-modi]

On it

[yiyixuxu]

probably this?

Suggested change

	library, orig_class_name = config["subfolder"]
	library, orig_class_name = config[subfolder]

[ishan-modi]

@yiyixuxu, should be working now !

[yiyixuxu]

currently test1 and 3 still fail

[yiyixuxu]

Suggested change

config_name = "model_index.json"

this should not changed, the model config should still be saved in a config.json
maybe, you can try to use this to download model_index https://github.com/huggingface/diffusers/blob/main/src/diffusers/loaders/single_file.py#L242

[ishan-modi]

My bad, completely missed it. Let me know if the testcases work now. I think they were failing earlier because I didn't have latest changes in the branch

[yiyixuxu]

does this work for custom modules? like this

# test3: load a custom diffusers model 
# https://huggingface.co/Kwai-Kolors/Kolors-diffusers/blob/main/model_index.json#L9
try:
    model = AutoModel.from_pretrained("Kwai-Kolors/Kolors-diffusers", subfolder="text_encoder", torch_dtype=torch.bfloat16)
    print(f"test3 passed!")
except Exception as e:
    print(f"test3 failed: {e}")

you can probably use this

diffusers/src/diffusers/pipelines/pipeline_loading_utils.py

Line 334 in 4b60f4b

def get_class_obj_and_candidates(

[ishan-modi]

Thanks for the patience, all test cases work now

Script

from diffusers import AutoModel 
import torch


# test1: load a diffusers model
try:
    model = AutoModel.from_pretrained(
        "Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers", subfolder="transformer", torch_dtype=torch.bfloat16, 
    )
    print(f"test1 passed!")
except Exception as e:
    print(f"test1 failed: {e}")


# test2: load a non-diffusers model
try:
    model = AutoModel.from_pretrained(
        "HiDream-ai/HiDream-I1-Full", subfolder="text_encoder", torch_dtype=torch.bfloat16, 
    )
    print(f"test2 passed!")
except Exception as e:
    print(f"test2 failed: {e}")


# test3: load a custom diffusers model 
# https://huggingface.co/Kwai-Kolors/Kolors-diffusers/blob/main/model_index.json#L9
try:
    model = AutoModel.from_pretrained("Kwai-Kolors/Kolors-diffusers", subfolder="text_encoder", variant="fp16")
    print(f"test3 passed!")
except Exception as e:
    print(f"test3 failed: {e}")


# test4: load a model directly (not subfolder)
controlnet_repo = "InstantX/SD3-Controlnet-Canny"
try:
    controlnet_model = AutoModel.from_pretrained(
        controlnet_repo, revision="refs/pr/12"
    )
    print(f"test4 passed!")
except Exception as e:
    print(f"test4 failed: {e}")

Please let me know if anything seems incorrect !