108 lines
4.2 KiB
Python
108 lines
4.2 KiB
Python
import os
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import torch
|
|
|
|
class CombinedDiffusersSD15Loader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
script_dir = os.path.dirname(os.path.realpath(__file__))
|
|
base_path = os.path.join(script_dir, "..", "..", "models", "diffusers", "SD15")
|
|
try:
|
|
sub_dirs = [d for d in os.listdir(base_path) if os.path.isdir(os.path.join(base_path, d))]
|
|
except FileNotFoundError:
|
|
sub_dirs = []
|
|
|
|
return {"required": {"sub_directory": (sub_dirs,),
|
|
"clip_type": (["stable_diffusion", "stable_cascade"],)}}
|
|
|
|
RETURN_TYPES = ("CLIP", "MODEL", "VAE")
|
|
FUNCTION = "load_models"
|
|
CATEGORY = "testing_nodes"
|
|
|
|
def load_models(self, sub_directory, clip_type="stable_diffusion"):
|
|
clip = self.load_clip(sub_directory, clip_type)
|
|
unet = self.load_unet(sub_directory)
|
|
vae = self.load_vae(sub_directory)
|
|
return clip, unet, vae
|
|
|
|
def load_clip(self, sub_directory, clip_type="stable_diffusion"):
|
|
clip_type_enum = comfy.sd.CLIPType.STABLE_DIFFUSION
|
|
if clip_type == "stable_cascade":
|
|
clip_type_enum = comfy.sd.CLIPType.STABLE_CASCADE
|
|
|
|
script_dir = os.path.dirname(os.path.realpath(__file__))
|
|
text_encoder_dir = os.path.join(script_dir, "..", "..", "models", "diffusers", "SD15", sub_directory, "text_encoder")
|
|
text_encoder_file = self.find_safetensors_file(text_encoder_dir)
|
|
text_encoder_path = os.path.join(text_encoder_dir, text_encoder_file)
|
|
|
|
if not os.path.exists(text_encoder_path):
|
|
raise FileNotFoundError(f"File not found: {text_encoder_path}")
|
|
|
|
clip = comfy.sd.load_clip(ckpt_paths=[text_encoder_path], embedding_directory=os.path.join(script_dir, "..", "..", "models", "diffusers", "SD15", sub_directory, "embeddings"), clip_type=clip_type_enum)
|
|
return clip
|
|
|
|
def load_unet(self, sub_directory):
|
|
unet_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers', 'SD15', sub_directory, 'unet'))
|
|
diffusion_model_path = self.find_safetensors_file(unet_folder)
|
|
|
|
if not os.path.exists(diffusion_model_path):
|
|
raise FileNotFoundError(f"{os.path.basename(diffusion_model_path)} not found in the directory: {diffusion_model_path}")
|
|
|
|
try:
|
|
model = comfy.sd.load_unet(diffusion_model_path)
|
|
return model
|
|
except PermissionError as e:
|
|
self.handle_corrupted_file(unet_folder)
|
|
raise e
|
|
|
|
def load_vae(self, sub_directory):
|
|
base_path = self.get_base_path()
|
|
vae_folder = os.path.join(base_path, sub_directory, "vae")
|
|
vae_filename = self.find_safetensors_file(vae_folder)
|
|
vae_path = os.path.join(vae_folder, vae_filename)
|
|
full_path = os.path.abspath(vae_path)
|
|
|
|
print(f"Attempting to load VAE model from: {full_path}")
|
|
|
|
if not os.path.exists(full_path):
|
|
raise FileNotFoundError(f"The VAE file '{full_path}' does not exist.")
|
|
|
|
sd = comfy.utils.load_torch_file(full_path)
|
|
vae = comfy.sd.VAE(sd=sd)
|
|
return vae
|
|
|
|
@staticmethod
|
|
def get_base_path():
|
|
base_path = os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers', 'SD15')
|
|
base_path = os.path.abspath(base_path)
|
|
if not os.path.exists(base_path):
|
|
raise FileNotFoundError(f"The base path '{base_path}' does not exist.")
|
|
return base_path
|
|
|
|
@staticmethod
|
|
def find_safetensors_file(directory):
|
|
for file in os.listdir(directory):
|
|
if file.endswith(".safetensors"):
|
|
return os.path.join(directory, file)
|
|
raise FileNotFoundError(f"No .safetensors file found in {directory}")
|
|
|
|
@staticmethod
|
|
def handle_corrupted_file(unet_folder):
|
|
corrupted_path = unet_folder + ".corrupted"
|
|
try:
|
|
os.rename(unet_folder, corrupted_path)
|
|
except Exception as e:
|
|
raise e
|
|
|
|
# Register the node in the ComfyUI framework
|
|
NODE_CLASS_MAPPINGS = {
|
|
"CombinedDiffusersSD15Loader": CombinedDiffusersSD15Loader,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"CombinedDiffusersSD15Loader": "Combined Diffusers SD1.5 Loader",
|
|
}
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|