fred-dev/comfy_ui_ali

1import os2 3import comfy.sd4 5def first_file(path, filenames):6    for f in filenames:7        p = os.path.join(path, f)8        if os.path.exists(p):9            return p10    return None11 12def load_diffusers(model_path, output_vae=True, output_clip=True, embedding_directory=None):13    diffusion_model_names = ["diffusion_pytorch_model.fp16.safetensors", "diffusion_pytorch_model.safetensors", "diffusion_pytorch_model.fp16.bin", "diffusion_pytorch_model.bin"]14    unet_path = first_file(os.path.join(model_path, "unet"), diffusion_model_names)15    vae_path = first_file(os.path.join(model_path, "vae"), diffusion_model_names)16 17    text_encoder_model_names = ["model.fp16.safetensors", "model.safetensors", "pytorch_model.fp16.bin", "pytorch_model.bin"]18    text_encoder1_path = first_file(os.path.join(model_path, "text_encoder"), text_encoder_model_names)19    text_encoder2_path = first_file(os.path.join(model_path, "text_encoder_2"), text_encoder_model_names)20 21    text_encoder_paths = [text_encoder1_path]22    if text_encoder2_path is not None:23        text_encoder_paths.append(text_encoder2_path)24 25    unet = comfy.sd.load_diffusion_model(unet_path)26 27    clip = None28    if output_clip:29        clip = comfy.sd.load_clip(text_encoder_paths, embedding_directory=embedding_directory)30 31    vae = None32    if output_vae:33        sd = comfy.utils.load_torch_file(vae_path)34        vae = comfy.sd.VAE(sd=sd)35 36    return (unet, clip, vae)37