3d039591fe
Previously before this commit, credits are already in entry and licenses are already in root. This commit will make info clearer.
39 lines
1.5 KiB
Python
39 lines
1.5 KiB
Python
# Taken from https://github.com/comfyanonymous/ComfyUI
|
|
|
|
import os
|
|
|
|
import ldm_patched.modules.sd
|
|
|
|
def first_file(path, filenames):
|
|
for f in filenames:
|
|
p = os.path.join(path, f)
|
|
if os.path.exists(p):
|
|
return p
|
|
return None
|
|
|
|
def load_diffusers(model_path, output_vae=True, output_clip=True, embedding_directory=None):
|
|
diffusion_model_names = ["diffusion_pytorch_model.fp16.safetensors", "diffusion_pytorch_model.safetensors", "diffusion_pytorch_model.fp16.bin", "diffusion_pytorch_model.bin"]
|
|
unet_path = first_file(os.path.join(model_path, "unet"), diffusion_model_names)
|
|
vae_path = first_file(os.path.join(model_path, "vae"), diffusion_model_names)
|
|
|
|
text_encoder_model_names = ["model.fp16.safetensors", "model.safetensors", "pytorch_model.fp16.bin", "pytorch_model.bin"]
|
|
text_encoder1_path = first_file(os.path.join(model_path, "text_encoder"), text_encoder_model_names)
|
|
text_encoder2_path = first_file(os.path.join(model_path, "text_encoder_2"), text_encoder_model_names)
|
|
|
|
text_encoder_paths = [text_encoder1_path]
|
|
if text_encoder2_path is not None:
|
|
text_encoder_paths.append(text_encoder2_path)
|
|
|
|
unet = ldm_patched.modules.sd.load_unet(unet_path)
|
|
|
|
clip = None
|
|
if output_clip:
|
|
clip = ldm_patched.modules.sd.load_clip(text_encoder_paths, embedding_directory=embedding_directory)
|
|
|
|
vae = None
|
|
if output_vae:
|
|
sd = ldm_patched.modules.utils.load_torch_file(vae_path)
|
|
vae = ldm_patched.modules.sd.VAE(sd=sd)
|
|
|
|
return (unet, clip, vae)
|