| from diffusers import AutoencoderKL, UNet2DConditionModel, StableDiffusionPipeline | |
| from transformers import CLIPTextModel, CLIPTokenizer | |
| def load_models(pretrained_model_name_or_path): | |
| text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder="text_encoder") | |
| vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae") | |
| unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder="unet") | |
| tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer") | |
| return text_encoder, vae, unet, tokenizer | |