298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-04-13 15:03:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update diffusers_load.py

Browse Source
This commit is contained in:
tongyu 2025-03-19 21:38:15 +08:00 committed by GitHub
parent 3b19fc76e3
commit e4d69eb8df
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 26 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
comfy/diffusers_load.py
View File

@ -1,6 +1,7 @@
import os import os
import torch
import comfy.sd import comfy.sd
import comfy.utils
def first_file(path, filenames): def first_file(path, filenames):
for f in filenames: for f in filenames:
@ -9,28 +10,41 @@ def first_file(path, filenames):
return p return p
return None return None
def load_diffusers(model_path, output_vae=True, output_clip=True, embedding_directory=None): def load_diffusers(model_path, output_vae=True, output_clip=True, embedding_directory=None, weight_dtype=torch.float16):
diffusion_model_names = ["diffusion_pytorch_model.fp16.safetensors", "diffusion_pytorch_model.safetensors", "diffusion_pytorch_model.fp16.bin", "diffusion_pytorch_model.bin"] """
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) 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) 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_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_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_encoder2_path = first_file(os.path.join(model_path, "text_encoder_2"), text_encoder_model_names)
text_encoder_paths = [text_encoder1_path] text_encoder_paths = [text_encoder1_path] if text_encoder1_path else []
if text_encoder2_path is not None: if text_encoder2_path:
text_encoder_paths.append(text_encoder2_path) text_encoder_paths.append(text_encoder2_path)
unet = comfy.sd.load_diffusion_model(unet_path) unet = comfy.sd.load_diffusion_model(unet_path, dtype=weight_dtype)
clip = None clip = None
if output_clip: if output_clip and text_encoder_paths:
clip = comfy.sd.load_clip(text_encoder_paths, embedding_directory=embedding_directory) clip = comfy.sd.load_clip(text_encoder_paths, embedding_directory=embedding_directory, dtype=weight_dtype)
vae = None vae = None
if output_vae: if output_vae and vae_path:
sd = comfy.utils.load_torch_file(vae_path) sd = comfy.utils.load_torch_file(vae_path, map_location="cpu")
vae = comfy.sd.VAE(sd=sd) vae = comfy.sd.VAE(sd=sd).to(dtype=weight_dtype)
return (unet, clip, vae) return (unet, clip, vae)