mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-04-13 12:23:30 +00:00
51 lines
2.1 KiB
Python
51 lines
2.1 KiB
Python
import os
|
|
import torch
|
|
import comfy.sd
|
|
import comfy.utils
|
|
|
|
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, weight_dtype=torch.float16):
|
|
"""
|
|
Load Stable Diffusion model components with custom precision.
|
|
|
|
:param model_path: Path to the model directory.
|
|
:param output_vae: Whether to load the VAE model.
|
|
:param output_clip: Whether to load the CLIP model (text encoder).
|
|
:param embedding_directory: Path to embedding directory.
|
|
:param weight_dtype: Data type for model weights (torch.float16, torch.float32, torch.bfloat16).
|
|
:return: (UNet, CLIP, VAE)
|
|
"""
|
|
|
|
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_encoder1_path else []
|
|
if text_encoder2_path:
|
|
text_encoder_paths.append(text_encoder2_path)
|
|
|
|
unet = comfy.sd.load_diffusion_model(unet_path, dtype=weight_dtype)
|
|
|
|
clip = None
|
|
if output_clip and text_encoder_paths:
|
|
clip = comfy.sd.load_clip(text_encoder_paths, embedding_directory=embedding_directory, dtype=weight_dtype)
|
|
|
|
vae = None
|
|
if output_vae and vae_path:
|
|
sd = comfy.utils.load_torch_file(vae_path, map_location="cpu")
|
|
vae = comfy.sd.VAE(sd=sd).to(dtype=weight_dtype)
|
|
|
|
return (unet, clip, vae)
|