use a custom model in pipeline. unable to load directly by from_pretrained()

[Weifeng-Chen]

I define a new unet class (eg. CustomUnet)and a new pipeline(CustomPipeline) that is slightly different from StableDiffusionPipeline. After training, I use save_pretrained() to save it and I find the model_index is like(models is where I put the unet I define, and but not in the system path? ):

CustomPipeline {
  "_class_name": "CustomPipeline",
  "_diffusers_version": "0.11.1",
  "scheduler": [
    "diffusers",
    "PNDMScheduler"
  ],
  "text_encoder": [
    "transformers",
    "CLIPTextModel"
  ],
  "tokenizer": [
    "transformers",
    "CLIPTokenizer"
  ],
  "unet": [
    "models",
    "CustomUnet"
  ],
  "vae": [
    "diffusers",
    "AutoencoderKL"
  ]
}

Then I use from_pretrained() to load the pipeline, but got:

~/anaconda3/lib/python3.9/site-packages/diffusers/pipeline_utils.py in from_pretrained(cls, pretrained_model_name_or_path, **kwargs)
    657             else:
    658                 # else we just import it from the library.
--> 659                 library = importlib.import_module(library_name)
    660 
    661                 class_obj = getattr(library, class_name)

~/anaconda3/lib/python3.9/importlib/__init__.py in import_module(name, package)
    125                 break
    126             level += 1
--> 127     return _bootstrap._gcd_import(name[level:], package, level)
    128 
    129 

~/anaconda3/lib/python3.9/importlib/_bootstrap.py in _gcd_import(name, package, level)

~/anaconda3/lib/python3.9/importlib/_bootstrap.py in _find_and_load(name, import_)

~/anaconda3/lib/python3.9/importlib/_bootstrap.py in _find_and_load_unlocked(name, import_)

ModuleNotFoundError: No module named 'models'

I think the library doesn't have the model I define.

I can load the model by the following code but it's not elegant.

unet = CustomUnet.from_pretrained(PIPELINE_SAVING_PATH, subfolder='unet')
pipeline = CustomPipeline.from_pretrained('runwayml/stable-diffusion-v1-5', unet=unet).to("cuda")

How can I use pipeline to load my custom model?

[patrickvonplaten]

That's a very interesting use case! @Weifeng-Chen, you're spot on. Currently it is not possible to load custom models via:

pipeline = CustomPipeline.from_pretrained('runwayml/stable-diffusion-v1-5').to("cuda")

but:

unet = CustomUnet.from_pretrained(PIPELINE_SAVING_PATH, subfolder='unet')
pipeline = CustomPipeline.from_pretrained('runwayml/stable-diffusion-v1-5', unet=unet).to("cuda")

should indeed work. We could add some functionality to allow loading custom unets similar to how it's done for Transformers: https://huggingface.co/docs/transformers/custom_models

Wdyt @patil-suraj @pcuenca ?

[pcuenca]

I think it's an interesting use case indeed! What would the solution entail, uploading the model files to the Hub, and then have from_pretrained use them? Sounds good to me!

[patil-suraj]

Sounds good to me!

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.