Cannot save pytorch2.0-compiled models

[xingzhehe]

For inference, it works fine. However, if I want to save the model (for example, after dreambooth-fine-tune), it gives error.

import torch
from diffusers import StableDiffusionPipeline
import os


pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipe = pipe.to("cuda")
pipe.unet = torch.compile(pipe.unet)

prompt = "a photo of a face"
image = pipe(prompt=prompt).images[0]

import matplotlib.pyplot as plt
plt.imshow(image)
plt.show()

pipe.save_pretrained('./test_save')

Error: ModuleNotFoundError: No module named 'onnxruntime'

Full Error:
╭───────────────────── Traceback (most recent call last) ──────────────────────╮
│ xxxx/test_diffuser.py:17 in │
│ │
│ 14 plt.imshow(image) │
│ 15 plt.show() │
│ 16 │
│ ❱ 17 pipe.save_pretrained('./test_save') │
│ 18 │
│ │
│ xxxx/anaconda3/envs/xxxx/lib/python3.9/site-packages/diffusers │
│ /pipelines/pipeline_utils.py:320 in save_pretrained │
│ │
│ 317 │ │ │ save_method_name = None │
│ 318 │ │ │ # search for the model's base class in LOADABLE_CLASSES │
│ 319 │ │ │ for library_name, library_classes in LOADABLE_CLASSES.ite │
│ ❱ 320 │ │ │ │ library = importlib.import_module(library_name) │
│ 321 │ │ │ │ for base_class, save_load_methods in library_classes. │
│ 322 │ │ │ │ │ class_candidate = getattr(library, base_class, No │
│ 323 │ │ │ │ │ if class_candidate is not None and issubclass(mod │
│ │
│ xxxx/anaconda3/envs/xxxx/lib/python3.9/importlib/init.py:1 │
│ 27 in import_module │
│ │
│ 124 │ │ │ if character != '.': │
│ 125 │ │ │ │ break │
│ 126 │ │ │ level += 1 │
│ ❱ 127 │ return _bootstrap._gcd_import(name[level:], package, level) │
│ 128 │
│ 129 │
│ 130 _RELOADING = {} │
│ in _gcd_import:1030 │
│ in _find_and_load:1007 │
│ in _find_and_load_unlocked:972 │
│ in _call_with_frames_removed:228 │
│ in _gcd_import:1030 │
│ in _find_and_load:1007 │
│ in _find_and_load_unlocked:984 │
╰──────────────────────────────────────────────────────────────────────────────╯
ModuleNotFoundError: No module named 'onnxruntime'

Besides, I can neither create a new pipe line using the compiled unet. It says
is of type: <class 'torch._dynamo.eval_frame.OptimizedModule'>, but should be <class 'diffusers.models.modeling_utils.ModelMixin'>

[github-actions]

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.