Updated and split loaders into combined, clip, unet and vae for both SDXL and SD1.5
This commit is contained in:
Scorpinaus
+25
-7
@@ -1,13 +1,31 @@
|
||||
from .loader import DiffusersLoader, SDXLDiffusersLoader
|
||||
from .sd15_unet_loader import SD15UNETLoader
|
||||
from .sd15_vae_loader import SD15VAELoader
|
||||
from .sd15_clip_loader import SD15CLIPLoader
|
||||
from .sdxl_unet_loader import SDXLUNETLoader
|
||||
from .sdxl_vae_loader import SDXLVAELoader
|
||||
from .sdxl_clip_loader import SDXLCLIPLoader
|
||||
from .sd15_combined_loader import CombinedDiffusersSD15Loader
|
||||
from .sdxl_combined_loader import CombinedDiffusersSDXLLoader
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Diffusers_Loader": DiffusersLoader,
|
||||
"SDXL_Diffusers_Loader": SDXLDiffusersLoader,
|
||||
"SD15UNETLoader": SD15UNETLoader,
|
||||
"SD15VAELoader": SD15VAELoader,
|
||||
"SD15CLIPLoader": SD15CLIPLoader,
|
||||
"SDXLUNETLoader": SDXLUNETLoader,
|
||||
"SDXLVAELoader": SDXLVAELoader,
|
||||
"SDXLCLIPLoader": SDXLCLIPLoader,
|
||||
"CombinedDiffusersSD15Loader": CombinedDiffusersSD15Loader,
|
||||
"CombinedDiffusersSDXLLoader": CombinedDiffusersSDXLLoader,
|
||||
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Diffusers_Loader": "Diffusers Loader",
|
||||
"SDXL_Diffusers_Loader": "SDXL Diffusers Loader",
|
||||
"SD15UNETLoader": "SD1.5 UNET Loader",
|
||||
"SD15VAELoader": "SD1.5 VAE Loader",
|
||||
"SD15CLIPLoader": "SD1.5 CLIP Loader",
|
||||
"SDXLUNETLoader": "SDXL UNET Loader",
|
||||
"SDXLVAELoader": "SDXL VAE Loader",
|
||||
"SDXLCLIPLoader": "SDXL CLIP Loader",
|
||||
"CombinedDiffusersSD15Loader" : "Combined Diffusers SD15 Loader",
|
||||
"CombinedDiffusersSDXLLoader" : "Combined Diffusers SDXL Loader"
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -1,20 +0,0 @@
|
||||
import os
|
||||
|
||||
def get_comfyui_base_path():
|
||||
# Assuming the script is in the standard directory structure
|
||||
return os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
|
||||
def get_folder_paths(folder_type):
|
||||
base_path = get_comfyui_base_path()
|
||||
|
||||
if folder_type == "diffusers":
|
||||
return [os.path.join(base_path, "models", "diffusers")]
|
||||
if folder_type == "embeddings":
|
||||
return [os.path.join(base_path, "models", "embeddings")]
|
||||
|
||||
return []
|
||||
|
||||
# Example usage:
|
||||
if __name__ == "__main__":
|
||||
print(get_folder_paths("diffusers")) # Should print the path to the diffusers models
|
||||
print(get_folder_paths("embeddings")) # Should print the path to the embeddings
|
||||
@@ -1,56 +0,0 @@
|
||||
import os
|
||||
from .utils import load_diffusers_model, load_diffusers_sdxl_model
|
||||
from .folder_paths import get_folder_paths
|
||||
|
||||
class DiffusersLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
paths = []
|
||||
for search_path in get_folder_paths("diffusers"):
|
||||
if os.path.exists(search_path):
|
||||
for root, subdir, files in os.walk(search_path, followlinks=True):
|
||||
if "model_index.json" in files:
|
||||
paths.append(os.path.relpath(root, start=search_path))
|
||||
|
||||
return {"required": {"model_path": (paths,)}}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
||||
FUNCTION = "load_checkpoint"
|
||||
CATEGORY = "advanced/loaders/diffusers"
|
||||
|
||||
def load_checkpoint(self, model_path, output_vae=True, output_clip=True):
|
||||
for search_path in get_folder_paths("diffusers"):
|
||||
if os.path.exists(search_path):
|
||||
path = os.path.join(search_path, model_path)
|
||||
if os.path.exists(path):
|
||||
model_path = path
|
||||
break
|
||||
|
||||
return load_diffusers_model(model_path, output_vae=output_vae, output_clip=output_clip)
|
||||
|
||||
|
||||
class SDXLDiffusersLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
paths = []
|
||||
for search_path in get_folder_paths("diffusers"):
|
||||
if os.path.exists(search_path):
|
||||
for root, subdir, files in os.walk(search_path, followlinks=True):
|
||||
if "model_index.json" in files:
|
||||
paths.append(os.path.relpath(root, start=search_path))
|
||||
|
||||
return {"required": {"model_path": (paths,)}}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE", "TEXT_ENCODER_2")
|
||||
FUNCTION = "load_checkpoint"
|
||||
CATEGORY = "advanced/loaders/sdxl"
|
||||
|
||||
def load_checkpoint(self, model_path, output_vae=True, output_clip=True):
|
||||
for search_path in get_folder_paths("diffusers"):
|
||||
if os.path.exists(search_path):
|
||||
path = os.path.join(search_path, model_path)
|
||||
if os.path.exists(path):
|
||||
model_path = path
|
||||
break
|
||||
|
||||
return load_diffusers_sdxl_model(model_path, output_vae=output_vae, output_clip=output_clip)
|
||||
@@ -0,0 +1,68 @@
|
||||
import os
|
||||
from comfy import sd
|
||||
|
||||
class SD15CLIPLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
# Path to the 'diffusers/SD15' folder, relative to the script location
|
||||
script_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
base_path = os.path.join(script_dir, "..", "..", "models", "diffusers", "SD15")
|
||||
print(f"SD15CLIPLoader:Base path: {base_path}") # Debug print
|
||||
|
||||
# Get the list of sub-directories in the 'diffusers/SD15' folder
|
||||
try:
|
||||
sub_dirs = [d for d in os.listdir(base_path) if os.path.isdir(os.path.join(base_path, d))]
|
||||
print(f"Sub-directories: {sub_dirs}") # Debug print
|
||||
except FileNotFoundError:
|
||||
sub_dirs = []
|
||||
print("SD15CLIPLoader:Base path not found.") # Debug print
|
||||
|
||||
return {"required": {"sub_directory": (sub_dirs,),
|
||||
"clip_type": (["stable_diffusion", "stable_cascade"],),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("CLIP",)
|
||||
FUNCTION = "load_clip"
|
||||
CATEGORY = "testing_nodes"
|
||||
|
||||
def load_clip(self, sub_directory, clip_type="stable_diffusion"):
|
||||
# Determine the clip type
|
||||
clip_type_enum = sd.CLIPType.STABLE_DIFFUSION
|
||||
if clip_type == "stable_cascade":
|
||||
clip_type_enum = sd.CLIPType.STABLE_CASCADE
|
||||
|
||||
# Construct the path to the text_encoder directory
|
||||
script_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
text_encoder_dir = os.path.join(script_dir, "..", "..", "models", "diffusers", "SD15", sub_directory, "text_encoder")
|
||||
print(f"SD15CLIPLoader:Text encoder directory: {text_encoder_dir}") # Debug print
|
||||
|
||||
# Find the .safetensors file in the text_encoder directory
|
||||
text_encoder_file = self.find_safetensors_file(text_encoder_dir)
|
||||
text_encoder_path = os.path.join(text_encoder_dir, text_encoder_file)
|
||||
print(f"SD15CLIPLoader:Text encoder path: {text_encoder_path}") # Debug print
|
||||
|
||||
# Ensure the file exists
|
||||
if not os.path.exists(text_encoder_path):
|
||||
raise FileNotFoundError(f"SD15CLIPLoader:File not found: {text_encoder_path}")
|
||||
|
||||
# Load the CLIP model using the constructed path
|
||||
clip = 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,)
|
||||
|
||||
@staticmethod
|
||||
def find_safetensors_file(directory):
|
||||
for file in os.listdir(directory):
|
||||
if file.endswith(".safetensors"):
|
||||
return file
|
||||
raise FileNotFoundError(f"SD15CLIPLoader:No .safetensors file found in {directory}")
|
||||
|
||||
# Register the node in the ComfyUI framework
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SD15CLIPLoader": SD15CLIPLoader,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SD15CLIPLoader": "SD1.5 CLIP Loader",
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,107 @@
|
||||
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"]
|
||||
@@ -0,0 +1,73 @@
|
||||
import os
|
||||
import comfy.sd
|
||||
|
||||
def get_sub_directory_list(base_path):
|
||||
"""Returns a list of sub-directories within the given base path."""
|
||||
base_directory = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', base_path))
|
||||
print(f"SD15UNETLoader:Looking for sub-directories in: {base_directory}")
|
||||
|
||||
if os.path.exists(base_directory):
|
||||
sub_dirs = [d for d in os.listdir(base_directory) if os.path.isdir(os.path.join(base_directory, d))]
|
||||
print(f"SD15UNETLoader:Found sub-directories: {sub_dirs}")
|
||||
return sub_dirs
|
||||
else:
|
||||
print(f"SD15UNETLoader:Path does not exist: {base_directory}")
|
||||
return []
|
||||
|
||||
def handle_corrupted_file(unet_folder):
|
||||
"""Handle corrupted UNET folder by renaming it and logging the action."""
|
||||
corrupted_path = unet_folder + ".corrupted"
|
||||
try:
|
||||
os.rename(unet_folder, corrupted_path)
|
||||
print(f"SD15UNETLoader:Moved corrupted folder to: {corrupted_path}")
|
||||
except PermissionError as e:
|
||||
print(f"SD15UNETLoader:PermissionError while moving corrupted folder: {e}")
|
||||
raise PermissionError(f"SD15UNETLoader:Permission denied while moving: {unet_folder}")
|
||||
except Exception as e:
|
||||
print(f"SD15UNETLoader:Error while moving corrupted folder: {e}")
|
||||
raise e
|
||||
|
||||
class SD15UNETLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"sub_directory": (get_sub_directory_list("diffusers/SD15"),),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "load_unet"
|
||||
CATEGORY = "testing_nodes"
|
||||
|
||||
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)
|
||||
print(f"SD15UNETLoader:Attempting to load UNET model from: {diffusion_model_path}")
|
||||
|
||||
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:
|
||||
print(f"SD15UNETLoader:PermissionError: {e}")
|
||||
raise PermissionError(f"SD15UNETLoader:Permission denied: {diffusion_model_path}")
|
||||
except Exception as e:
|
||||
print(f"SD15UNETLoader:Error loading UNET model: {e}")
|
||||
handle_corrupted_file(unet_folder)
|
||||
raise e
|
||||
|
||||
@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"SD15UNETLoader:No .safetensors file found in {directory}")
|
||||
|
||||
# Ensure the node is registered properly
|
||||
NODE_CLASS_MAPPINGS = {"SD15UNETLoader": SD15UNETLoader}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"SD15UNETLoader": "SD1.5 UNET Loader"}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,60 @@
|
||||
import os
|
||||
import comfy.utils
|
||||
import torch
|
||||
|
||||
class SD15VAELoader:
|
||||
@staticmethod
|
||||
def get_base_path():
|
||||
# Adjust the base path to the location where the SD15 subdirectory is located
|
||||
base_path = os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers', 'SD15')
|
||||
base_path = os.path.abspath(base_path) # Convert to absolute path
|
||||
if not os.path.exists(base_path):
|
||||
raise FileNotFoundError(f"The base path '{base_path}' does not exist.")
|
||||
return base_path
|
||||
|
||||
@staticmethod
|
||||
def get_subdirectories():
|
||||
base_path = SD15VAELoader.get_base_path()
|
||||
return [name for name in os.listdir(base_path) if os.path.isdir(os.path.join(base_path, name))]
|
||||
|
||||
@staticmethod
|
||||
def find_safetensors_file(directory):
|
||||
for file in os.listdir(directory):
|
||||
if file.endswith(".safetensors"):
|
||||
return file
|
||||
raise FileNotFoundError(f"No .safetensors file found in {directory}")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
subdirectories = cls.get_subdirectories()
|
||||
return {
|
||||
"required": {
|
||||
"sub_directory": (subdirectories, ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VAE",)
|
||||
FUNCTION = "load_vae"
|
||||
CATEGORY = "testing_nodes"
|
||||
|
||||
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,)
|
||||
|
||||
# Ensure the node is registered properly
|
||||
NODE_CLASS_MAPPINGS = {"SD15VAELoader": SD15VAELoader}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"SD15VAELoader": "SD1.5 VAE Loader"}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,69 @@
|
||||
import os
|
||||
from comfy import sd
|
||||
|
||||
class SDXLCLIPLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
script_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
base_path = os.path.abspath(os.path.join(script_dir, "../../models/diffusers/SDXL"))
|
||||
print(f"SDXLCLIPLoader: Base path for SDXL models: {base_path}")
|
||||
|
||||
if not os.path.exists(base_path):
|
||||
raise FileNotFoundError(f"SDXLCLIPLoader: The base path {base_path} does not exist.")
|
||||
|
||||
sub_dirs = [d for d in os.listdir(base_path) if os.path.isdir(os.path.join(base_path, d))]
|
||||
print(f"SDXLCLIPLoader: Available subdirectories: {sub_dirs}")
|
||||
return {
|
||||
"required": {
|
||||
"sub_directory": (sub_dirs, ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CLIP",)
|
||||
FUNCTION = "load_sdxl_clip"
|
||||
CATEGORY = "testing_nodes"
|
||||
|
||||
def load_sdxl_clip(self, sub_directory):
|
||||
script_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
base_path = os.path.abspath(os.path.join(script_dir, "../../models/diffusers/SDXL", sub_directory))
|
||||
|
||||
text_encoder_dir1 = os.path.join(base_path, "text_encoder")
|
||||
text_encoder_dir2 = os.path.join(base_path, "text_encoder_2")
|
||||
|
||||
text_encoder_file1 = self.find_safetensors_file(text_encoder_dir1)
|
||||
text_encoder_file2 = self.find_safetensors_file(text_encoder_dir2)
|
||||
|
||||
text_encoder_path1 = os.path.join(text_encoder_dir1, text_encoder_file1)
|
||||
text_encoder_path2 = os.path.join(text_encoder_dir2, text_encoder_file2)
|
||||
|
||||
print(f"SDXLCLIPLoader: Checking paths:\n{text_encoder_path1}\n{text_encoder_path2}")
|
||||
|
||||
path1_exists = os.path.exists(text_encoder_path1)
|
||||
path2_exists = os.path.exists(text_encoder_path2)
|
||||
|
||||
print(f"SDXLCLIPLoader: Path 1 exists: {path1_exists}")
|
||||
print(f"SDXLCLIPLoader: Path 2 exists: {path2_exists}")
|
||||
|
||||
if not path1_exists or not path2_exists:
|
||||
raise FileNotFoundError(f"SDXLCLIPLoader: One or both text encoder files not found in {base_path}")
|
||||
|
||||
clip = sd.load_clip(ckpt_paths=[text_encoder_path1, text_encoder_path2], embedding_directory=os.path.join(base_path, "embeddings"))
|
||||
return (clip,)
|
||||
|
||||
@staticmethod
|
||||
def find_safetensors_file(directory):
|
||||
for file in os.listdir(directory):
|
||||
if file.endswith(".safetensors"):
|
||||
return file
|
||||
raise FileNotFoundError(f"No .safetensors file found in {directory}")
|
||||
|
||||
# Mapping the node to make it recognizable in the ComfyUI framework
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SDXLCLIPLoader": SDXLCLIPLoader,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SDXLCLIPLoader": "SDXL CLIP Loader",
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,101 @@
|
||||
import os
|
||||
import comfy.utils
|
||||
import comfy.sd
|
||||
import torch
|
||||
|
||||
class CombinedDiffusersSDXLLoader:
|
||||
@staticmethod
|
||||
def get_base_path():
|
||||
base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers', 'SDXL'))
|
||||
if not os.path.exists(base_path):
|
||||
raise FileNotFoundError(f"CombinedDiffusersSDXLLoader: The base path '{base_path}' does not exist.")
|
||||
return base_path
|
||||
|
||||
@staticmethod
|
||||
def get_subdirectories():
|
||||
base_path = CombinedDiffusersSDXLLoader.get_base_path()
|
||||
return [name for name in os.listdir(base_path) if os.path.isdir(os.path.join(base_path, name))]
|
||||
|
||||
@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"CombinedDiffusersSDXLLoader: No .safetensors file found in {directory}")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
subdirectories = cls.get_subdirectories()
|
||||
return {
|
||||
"required": {
|
||||
"sub_directory": (subdirectories, ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
||||
FUNCTION = "load_models"
|
||||
CATEGORY = "testing_nodes"
|
||||
|
||||
def load_models(self, sub_directory):
|
||||
base_path = self.get_base_path()
|
||||
sub_dir_path = os.path.join(base_path, sub_directory)
|
||||
|
||||
vae_model = self.load_vae(sub_dir_path)
|
||||
unet_model = self.load_unet(sub_dir_path)
|
||||
clip_model = self.load_clip(sub_dir_path)
|
||||
|
||||
return unet_model, clip_model, vae_model
|
||||
|
||||
def load_vae(self, base_path):
|
||||
vae_folder = os.path.join(base_path, "vae")
|
||||
vae_path = self.find_safetensors_file(vae_folder)
|
||||
print(f"CombinedDiffusersSDXLLoader: Attempting to load VAE model from: {vae_path}")
|
||||
|
||||
vae_sd = comfy.utils.load_torch_file(vae_path)
|
||||
return comfy.sd.VAE(sd=vae_sd)
|
||||
|
||||
def load_unet(self, base_path):
|
||||
unet_folder = os.path.join(base_path, 'unet')
|
||||
unet_path = self.find_safetensors_file(unet_folder)
|
||||
print(f"CombinedDiffusersSDXLLoader: Attempting to load UNET model from: {unet_path}")
|
||||
|
||||
try:
|
||||
return comfy.sd.load_unet(unet_path)
|
||||
except PermissionError as e:
|
||||
print(f"PermissionError: {e}")
|
||||
raise PermissionError(f"CombinedDiffusersSDXLLoader: Permission denied: {unet_path}")
|
||||
except Exception as e:
|
||||
print(f"CombinedDiffusersSDXLLoader: Error loading UNET model: {e}")
|
||||
self.handle_corrupted_file(unet_folder)
|
||||
raise e
|
||||
|
||||
def load_clip(self, base_path):
|
||||
text_encoder_dir1 = os.path.join(base_path, "text_encoder")
|
||||
text_encoder_dir2 = os.path.join(base_path, "text_encoder_2")
|
||||
|
||||
text_encoder_path1 = self.find_safetensors_file(text_encoder_dir1)
|
||||
text_encoder_path2 = self.find_safetensors_file(text_encoder_dir2)
|
||||
|
||||
print(f"CombinedDiffusersSDXLLoader: Checking paths:\n{text_encoder_path1}\n{text_encoder_path2}")
|
||||
|
||||
return comfy.sd.load_clip(ckpt_paths=[text_encoder_path1, text_encoder_path2], embedding_directory=os.path.join(base_path, "embeddings"))
|
||||
|
||||
@staticmethod
|
||||
def handle_corrupted_file(unet_folder):
|
||||
"""Handle corrupted UNET folder by renaming it and logging the action."""
|
||||
corrupted_path = unet_folder + ".corrupted"
|
||||
try:
|
||||
os.rename(unet_folder, corrupted_path)
|
||||
print(f"CombinedDiffusersSDXLLoader: Moved corrupted folder to: {corrupted_path}")
|
||||
except PermissionError as e:
|
||||
print(f"CombinedDiffusersSDXLLoader: PermissionError while moving corrupted folder: {e}")
|
||||
raise PermissionError(f"Permission denied while moving: {unet_folder}")
|
||||
except Exception as e:
|
||||
print(f"CombinedDiffusersSDXLLoader: Error while moving corrupted folder: {e}")
|
||||
raise e
|
||||
|
||||
# Ensure the node is registered properly
|
||||
NODE_CLASS_MAPPINGS = {"CombinedDiffusersSDXLLoader": CombinedDiffusersSDXLLoader}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"CombinedDiffusersSDXLLoader": "Combined Diffusers SDXL Loader"}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,78 @@
|
||||
import os
|
||||
import comfy.sd
|
||||
|
||||
def get_sub_directory_list(base_path):
|
||||
"""Returns a list of sub-directories within the given base path."""
|
||||
base_directory = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', base_path))
|
||||
print(f"SDXLUNETLoader:Looking for sub-directories in: {base_directory}")
|
||||
|
||||
if os.path.exists(base_directory):
|
||||
sub_dirs = [d for d in os.listdir(base_directory) if os.path.isdir(os.path.join(base_directory, d))]
|
||||
print(f"SDXLUNETLoader:Found sub-directories: {sub_dirs}")
|
||||
return sub_dirs
|
||||
else:
|
||||
print(f"SDXLUNETLoader:Path does not exist: {base_directory}")
|
||||
return []
|
||||
|
||||
def handle_corrupted_file(unet_folder):
|
||||
"""Handle corrupted UNET folder by renaming it and logging the action."""
|
||||
corrupted_path = unet_folder + ".corrupted"
|
||||
try:
|
||||
os.rename(unet_folder, corrupted_path)
|
||||
print(f"SDXLUNETLoader:Moved corrupted folder to: {corrupted_path}")
|
||||
except PermissionError as e:
|
||||
print(f"SDXLUNETLoader:PermissionError while moving corrupted folder: {e}")
|
||||
raise PermissionError(f"Permission denied while moving: {unet_folder}")
|
||||
except Exception as e:
|
||||
print(f"SDXLUNETLoader:Error while moving corrupted folder: {e}")
|
||||
raise e
|
||||
|
||||
class SDXLUNETLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"sub_directory": (get_sub_directory_list("diffusers/SDXL"),),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "load_unet"
|
||||
CATEGORY = "testing_nodes"
|
||||
|
||||
def load_unet(self, sub_directory):
|
||||
unet_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers', 'SDXL', sub_directory, 'unet'))
|
||||
diffusion_model_path = self.find_safetensors_file(unet_folder)
|
||||
print(f"SDXLUNETLoader:Attempting to load UNET model from: {diffusion_model_path}")
|
||||
|
||||
if not os.path.exists(diffusion_model_path):
|
||||
raise FileNotFoundError(f"SDXLUNETLoader:{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:
|
||||
print(f"PermissionError: {e}")
|
||||
raise PermissionError(f"SDXLUNETLoader:Permission denied: {diffusion_model_path}")
|
||||
except Exception as e:
|
||||
print(f"SDXLUNETLoader:Error loading UNET model: {e}")
|
||||
handle_corrupted_file(unet_folder)
|
||||
raise e
|
||||
|
||||
@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"SDXLUNETLoader:No .safetensors file found in {directory}")
|
||||
|
||||
# Mapping the node to make it recognizable in the ComfyUI framework
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SDXLUNETLoader": SDXLUNETLoader,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SDXLUNETLoader": "SDXL UNET Loader",
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,60 @@
|
||||
import os
|
||||
import comfy.utils
|
||||
import torch
|
||||
|
||||
class SDXLVAELoader:
|
||||
@staticmethod
|
||||
def get_base_path():
|
||||
# Adjust the base path to the location where the SDXL subdirectory is located
|
||||
base_path = os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers', 'SDXL')
|
||||
base_path = os.path.abspath(base_path) # Convert to absolute path
|
||||
if not os.path.exists(base_path):
|
||||
raise FileNotFoundError(f"SDXLVAELoader:The base path '{base_path}' does not exist.")
|
||||
return base_path
|
||||
|
||||
@staticmethod
|
||||
def get_subdirectories():
|
||||
base_path = SDXLVAELoader.get_base_path()
|
||||
return [name for name in os.listdir(base_path) if os.path.isdir(os.path.join(base_path, name))]
|
||||
|
||||
@staticmethod
|
||||
def find_safetensors_file(directory):
|
||||
for file in os.listdir(directory):
|
||||
if file.endswith(".safetensors"):
|
||||
return file
|
||||
raise FileNotFoundError(f"SDXLVAELoader:No .safetensors file found in {directory}")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
subdirectories = cls.get_subdirectories()
|
||||
return {
|
||||
"required": {
|
||||
"sub_directory": (subdirectories, ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VAE",)
|
||||
FUNCTION = "load_vae"
|
||||
CATEGORY = "testing_nodes"
|
||||
|
||||
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"SDXLVAELoader:Attempting to load VAE model from: {full_path}")
|
||||
|
||||
if not os.path.exists(full_path):
|
||||
raise FileNotFoundError(f"SDXLVAELoader:The VAE file '{full_path}' does not exist.")
|
||||
|
||||
sd = comfy.utils.load_torch_file(full_path)
|
||||
vae = comfy.sd.VAE(sd=sd)
|
||||
return (vae,)
|
||||
|
||||
# Ensure the node is registered properly
|
||||
NODE_CLASS_MAPPINGS = {"SDXLVAELoader": SDXLVAELoader}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"SDXLVAELoader": "SDXL VAE Loader"}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -1,34 +0,0 @@
|
||||
import os
|
||||
import torch
|
||||
|
||||
def check_file_permissions(file_path):
|
||||
try:
|
||||
if os.access(file_path, os.R_OK | os.W_OK):
|
||||
print(f"Read and write permissions are available for {file_path}")
|
||||
else:
|
||||
raise PermissionError(f"Permission issue with {file_path}")
|
||||
except Exception as e:
|
||||
print(f"Error checking permissions for {file_path}: {e}")
|
||||
raise
|
||||
|
||||
def test_loading_file(file_path):
|
||||
try:
|
||||
check_file_permissions(file_path)
|
||||
print(f"Attempting to load file: {file_path}")
|
||||
data = torch.load(file_path)
|
||||
print(f"Successfully loaded file: {file_path}")
|
||||
except PermissionError as e:
|
||||
print(f"Permission denied when accessing {file_path}. Please check the file permissions.")
|
||||
except Exception as e:
|
||||
print(f"Error loading file {file_path}: {e}")
|
||||
|
||||
# Paths to test
|
||||
test_paths = [
|
||||
"C:\\Users\\Admin\\Desktop\\ComfyUI_windows_portable\\ComfyUI\\models\\diffusers\\raemumix_v90\\model_index.json",
|
||||
"C:\\Users\\Admin\\Desktop\\ComfyUI_windows_portable\\ComfyUI\\models\\diffusers\\raemumix_v90\\vae\\diffusion_pytorch_model.safetensors",
|
||||
"C:\\Users\\Admin\\Desktop\\ComfyUI_windows_portable\\ComfyUI\\models\\diffusers\\raemumix_v90\\unet\\diffusion_pytorch_model.safetensors",
|
||||
"C:\\Users\\Admin\\Desktop\\ComfyUI_windows_portable\\ComfyUI\\models\\diffusers\\raemumix_v90\\text_encoder\\model.safetensors",
|
||||
]
|
||||
|
||||
for path in test_paths:
|
||||
test_loading_file(path)
|
||||
@@ -1,39 +0,0 @@
|
||||
import os
|
||||
import utils
|
||||
|
||||
# Define the base model path
|
||||
base_model_path = "C:\\Users\\Admin\\Desktop\\ComfyUI_windows_portable\\ComfyUI\\models\\diffusers\\raemumix_v90"
|
||||
|
||||
# Test loading the main model
|
||||
try:
|
||||
model = utils.load_model_from_path(os.path.join(base_model_path, "model_index.json"))
|
||||
print("Main model loaded successfully.")
|
||||
except Exception as e:
|
||||
print(f"Error loading main model: {e}")
|
||||
|
||||
# Test loading the VAE model
|
||||
try:
|
||||
vae = utils.load_vae_from_path(base_model_path)
|
||||
print("VAE model loaded successfully.")
|
||||
except Exception as e:
|
||||
print(f"Error loading VAE model: {e}")
|
||||
|
||||
# Test loading the CLIP model
|
||||
try:
|
||||
clip = utils.load_text_encoder_from_path(base_model_path, encoder_id=None)
|
||||
print("CLIP model loaded successfully.")
|
||||
except Exception as e:
|
||||
print(f"Error loading CLIP model: {e}")
|
||||
|
||||
# Test loading the text encoders for SDXL
|
||||
try:
|
||||
text_encoder_1 = utils.load_text_encoder_from_path(base_model_path, 1)
|
||||
print("Text encoder 1 loaded successfully.")
|
||||
except Exception as e:
|
||||
print(f"Error loading text encoder 1: {e}")
|
||||
|
||||
try:
|
||||
text_encoder_2 = utils.load_text_encoder_from_path(base_model_path, 2)
|
||||
print("Text encoder 2 loaded successfully.")
|
||||
except Exception as e:
|
||||
print(f"Error loading text encoder 2: {e}")
|
||||
@@ -1,121 +0,0 @@
|
||||
import os
|
||||
import json
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
|
||||
def load_diffusers_model(model_path, output_vae=True, output_clip=True):
|
||||
try:
|
||||
model = load_model_from_path(model_path)
|
||||
vae = load_vae_from_path(model_path) if output_vae else None
|
||||
clip = load_text_encoder_from_path(model_path, encoder_id=None) if output_clip else None
|
||||
return model, vae, clip
|
||||
except Exception as e:
|
||||
print(f"Error loading diffusers model from {model_path}: {e}")
|
||||
raise
|
||||
|
||||
def load_diffusers_sdxl_model(model_path, output_vae=True, output_clip=True, embedding_directory=None):
|
||||
try:
|
||||
model = load_model_from_path(model_path)
|
||||
vae = load_vae_from_path(model_path) if output_vae else None
|
||||
clip = load_text_encoder_from_path(model_path, encoder_id=1) if output_clip else None
|
||||
text_encoder_2 = load_text_encoder_from_path(model_path, encoder_id=2)
|
||||
|
||||
if embedding_directory:
|
||||
# Load or handle embeddings if applicable
|
||||
pass
|
||||
|
||||
return model, vae, clip, text_encoder_2
|
||||
except Exception as e:
|
||||
print(f"Error loading SDXL diffusers model from {model_path}: {e}")
|
||||
raise
|
||||
|
||||
def load_model_from_path(model_path):
|
||||
try:
|
||||
print(f"Loading model from {model_path}")
|
||||
check_file_permissions(model_path)
|
||||
if model_path.endswith('.safetensors'):
|
||||
model = load_safetensors(model_path)
|
||||
else:
|
||||
model = load_json(model_path)
|
||||
return model
|
||||
except PermissionError as e:
|
||||
print(f"Permission denied when accessing {model_path}. Please check the file permissions.")
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"Error loading model from {model_path}: {e}")
|
||||
raise
|
||||
|
||||
def load_vae_from_path(model_path):
|
||||
try:
|
||||
print(f"Loading VAE from {model_path}")
|
||||
vae_path = os.path.join(model_path, "vae", "diffusion_pytorch_model.safetensors")
|
||||
check_file_permissions(vae_path)
|
||||
vae = load_safetensors(vae_path)
|
||||
return vae
|
||||
except PermissionError as e:
|
||||
print(f"Permission denied when accessing {vae_path}. Please check the file permissions.")
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"Error loading VAE from {vae_path}: {e}")
|
||||
raise
|
||||
|
||||
def load_clip_from_path(model_path):
|
||||
try:
|
||||
print(f"Loading CLIP model from {model_path}")
|
||||
clip_path = os.path.join(model_path, "clip", "diffusion_pytorch_model.safetensors")
|
||||
check_file_permissions(clip_path)
|
||||
clip = load_safetensors(clip_path)
|
||||
return clip
|
||||
except PermissionError as e:
|
||||
print(f"Permission denied when accessing {clip_path}. Please check the file permissions.")
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"Error loading CLIP model from {clip_path}: {e}")
|
||||
raise
|
||||
|
||||
def load_text_encoder_from_path(model_path, encoder_id=None):
|
||||
try:
|
||||
if encoder_id is not None:
|
||||
encoder_path = os.path.join(model_path, f"text_encoder_{encoder_id}", "model.safetensors")
|
||||
print(f"Loading text encoder {encoder_id} from {encoder_path}")
|
||||
else:
|
||||
encoder_path = os.path.join(model_path, "text_encoder", "model.safetensors")
|
||||
print(f"Loading text encoder from {encoder_path}")
|
||||
|
||||
check_file_permissions(encoder_path)
|
||||
text_encoder = load_safetensors(encoder_path)
|
||||
return text_encoder
|
||||
except PermissionError as e:
|
||||
print(f"Permission denied when accessing {encoder_path}. Please check the file permissions.")
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"Error loading text encoder from {encoder_path}: {e}")
|
||||
raise
|
||||
|
||||
def load_json(json_path):
|
||||
try:
|
||||
print(f"Loading JSON from {json_path}")
|
||||
with open(json_path, 'r') as f:
|
||||
return json.load(f)
|
||||
except Exception as e:
|
||||
print(f"Error loading JSON from {json_path}: {e}")
|
||||
raise
|
||||
|
||||
def load_safetensors(file_path):
|
||||
try:
|
||||
print(f"Loading safetensors from {file_path}")
|
||||
with safe_open(file_path, framework="pt", device="cpu") as f:
|
||||
return {k: f.get_tensor(k) for k in f.keys()}
|
||||
except Exception as e:
|
||||
print(f"Error loading safetensors from {file_path}: {e}")
|
||||
raise
|
||||
|
||||
def check_file_permissions(file_path):
|
||||
try:
|
||||
if os.access(file_path, os.R_OK | os.W_OK):
|
||||
print(f"Read and write permissions are available for {file_path}")
|
||||
else:
|
||||
raise PermissionError(f"Permission issue with {file_path}")
|
||||
except Exception as e:
|
||||
print(f"Error checking permissions for {file_path}: {e}")
|
||||
raise
|
||||
Reference in New Issue
Block a user