diff --git a/__init__.py b/__init__.py index 8ee1a42..c5d97b1 100644 --- a/__init__.py +++ b/__init__.py @@ -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"] diff --git a/__pycache__/__init__.cpython-311.pyc b/__pycache__/__init__.cpython-311.pyc deleted file mode 100644 index 35da619..0000000 Binary files a/__pycache__/__init__.cpython-311.pyc and /dev/null differ diff --git a/__pycache__/folder_paths.cpython-311.pyc b/__pycache__/folder_paths.cpython-311.pyc deleted file mode 100644 index 3d73fe4..0000000 Binary files a/__pycache__/folder_paths.cpython-311.pyc and /dev/null differ diff --git a/__pycache__/loader.cpython-311.pyc b/__pycache__/loader.cpython-311.pyc deleted file mode 100644 index 03250e1..0000000 Binary files a/__pycache__/loader.cpython-311.pyc and /dev/null differ diff --git a/__pycache__/utils.cpython-310.pyc b/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index d8be3a3..0000000 Binary files a/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/__pycache__/utils.cpython-311.pyc b/__pycache__/utils.cpython-311.pyc deleted file mode 100644 index 35cd8d8..0000000 Binary files a/__pycache__/utils.cpython-311.pyc and /dev/null differ diff --git a/folder_paths.py b/folder_paths.py deleted file mode 100644 index 6a465d7..0000000 --- a/folder_paths.py +++ /dev/null @@ -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 diff --git a/loader.py b/loader.py deleted file mode 100644 index cbd69f5..0000000 --- a/loader.py +++ /dev/null @@ -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) diff --git a/sd15_clip_loader.py b/sd15_clip_loader.py new file mode 100644 index 0000000..39ca341 --- /dev/null +++ b/sd15_clip_loader.py @@ -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"] diff --git a/sd15_combined_loader.py b/sd15_combined_loader.py new file mode 100644 index 0000000..f82a58b --- /dev/null +++ b/sd15_combined_loader.py @@ -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"] diff --git a/sd15_unet_loader.py b/sd15_unet_loader.py new file mode 100644 index 0000000..63017f9 --- /dev/null +++ b/sd15_unet_loader.py @@ -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"] diff --git a/sd15_vae_loader.py b/sd15_vae_loader.py new file mode 100644 index 0000000..0c00d2a --- /dev/null +++ b/sd15_vae_loader.py @@ -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"] diff --git a/sdxl_clip_loader.py b/sdxl_clip_loader.py new file mode 100644 index 0000000..17d426f --- /dev/null +++ b/sdxl_clip_loader.py @@ -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"] diff --git a/sdxl_combined_loader.py b/sdxl_combined_loader.py new file mode 100644 index 0000000..1f41cf7 --- /dev/null +++ b/sdxl_combined_loader.py @@ -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"] diff --git a/sdxl_unet_loader.py b/sdxl_unet_loader.py new file mode 100644 index 0000000..1bd8aea --- /dev/null +++ b/sdxl_unet_loader.py @@ -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"] diff --git a/sdxl_vae_loader.py b/sdxl_vae_loader.py new file mode 100644 index 0000000..9059d53 --- /dev/null +++ b/sdxl_vae_loader.py @@ -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"] diff --git a/test.py b/test.py deleted file mode 100644 index f2a1c2b..0000000 --- a/test.py +++ /dev/null @@ -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) diff --git a/test2.py b/test2.py deleted file mode 100644 index 2986b47..0000000 --- a/test2.py +++ /dev/null @@ -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}") diff --git a/utils.py b/utils.py deleted file mode 100644 index d3879bd..0000000 --- a/utils.py +++ /dev/null @@ -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