From d8e4b5320cb6d1de593547f1ea45ddfab4b797ae Mon Sep 17 00:00:00 2001 From: Scorpinaus Date: Sat, 8 Jun 2024 22:16:01 +0800 Subject: [PATCH] Troubleshoot file_path and sub_dir_path debug --- __init__.py | 13 ++++- combined_diffusers_loader.py | 4 +- diffusers_clip_loader.py | 93 +++++++++++++++++++++++++++++++ diffusers_unet_loader.py | 103 +++++++++++++++++++++++++++++++++++ diffusers_vae_loader.py | 85 +++++++++++++++++++++++++++++ 5 files changed, 295 insertions(+), 3 deletions(-) create mode 100644 diffusers_clip_loader.py create mode 100644 diffusers_unet_loader.py create mode 100644 diffusers_vae_loader.py diff --git a/__init__.py b/__init__.py index 088099b..38f9691 100644 --- a/__init__.py +++ b/__init__.py @@ -7,6 +7,9 @@ from .sdxl_clip_loader import SDXLCLIPLoader from .sd15_combined_loader import CombinedDiffusersSD15Loader from .sdxl_combined_loader import CombinedDiffusersSDXLLoader from .combined_diffusers_loader import CombinedDiffusersLoader +from .diffusers_clip_loader import DiffusersClipLoader +from .diffusers_unet_loader import DiffusersUNETLoader +from .diffusers_vae_loader import DiffusersVAELoader NODE_CLASS_MAPPINGS = { "SD15UNETLoader": SD15UNETLoader, @@ -17,7 +20,10 @@ NODE_CLASS_MAPPINGS = { "SDXLCLIPLoader": SDXLCLIPLoader, "CombinedDiffusersSD15Loader": CombinedDiffusersSD15Loader, "CombinedDiffusersSDXLLoader": CombinedDiffusersSDXLLoader, - "CombinedDiffusersLoader": CombinedDiffusersLoader + "CombinedDiffusersLoader": CombinedDiffusersLoader, + "DiffusersClipLoader": DiffusersClipLoader, + "DiffusersUNETLoader": DiffusersUNETLoader, + "DiffusersVAELoader": DiffusersVAELoader } @@ -30,5 +36,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "SDXLCLIPLoader": "SDXL CLIP Loader", "CombinedDiffusersSD15Loader" : "Combined Diffusers SD15 Loader", "CombinedDiffusersSDXLLoader" : "Combined Diffusers SDXL Loader", - "CombinedDiffusersLoader": "Combined Diffusers Loader" + "CombinedDiffusersLoader": "Combined Diffusers Loader", + "DiffusersClipLoader": "Diffusers CLIP Loader", + "DiffusersUNETLoader": "Diffusers UNET Loader", + "DiffusersVAELoader": "Diffusers VAE Loader" } diff --git a/combined_diffusers_loader.py b/combined_diffusers_loader.py index 15afe54..f5bdf22 100644 --- a/combined_diffusers_loader.py +++ b/combined_diffusers_loader.py @@ -33,6 +33,7 @@ class CombinedDiffusersLoader: def INPUT_TYPES(cls): base_path = cls.get_base_path() model_directories = cls.get_model_directories(base_path) + print(f"CombinedDiffusersLoader: model_directories: \n{model_directories}") return { "required": { "sub_directory": (model_directories,), @@ -46,10 +47,11 @@ class CombinedDiffusersLoader: def load_models(self, sub_directory, clip_type="stable_diffusion"): base_path = self.get_base_path() + print(f"CombinedDiffusersLoader: base_path: \n{base_path}") sub_dir_path = os.path.join(base_path, sub_directory) model_type = self.detect_model_type(sub_dir_path) - + print(f"CombinedDiffusersLoader: Sub_dir_path: \n{sub_dir_path}") vae_model = self.load_vae(sub_dir_path) unet_model = self.load_unet(sub_dir_path) clip_model = self.load_clip(sub_dir_path, clip_type, model_type) diff --git a/diffusers_clip_loader.py b/diffusers_clip_loader.py new file mode 100644 index 0000000..ebbae67 --- /dev/null +++ b/diffusers_clip_loader.py @@ -0,0 +1,93 @@ +import os +from comfy import sd +import comfy.utils +import torch + +class DiffusersClipLoader: + @staticmethod + def get_base_path(): + base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers')) + if not os.path.exists(base_path): + raise FileNotFoundError(f"DiffusersClipLoader: The base path '{base_path}' does not exist.") + return base_path + + @staticmethod + def get_model_directories(base_path): + model_dirs = [] + for root, dirs, files in os.walk(base_path): + if "model_index.json" in files: + model_dirs.append(os.path.relpath(root, base_path)) + # Remove unwanted subdirectories from dirs so that os.walk doesn't traverse them + dirs[:] = [d for d in dirs if d not in ["vae", "unet", "text_encoder", "text_encoder_2"]] + return model_dirs + + @staticmethod + def find_model_file(directory): + for file in os.listdir(directory): + if file.endswith(".safetensors") or file.endswith(".bin"): + return file + raise FileNotFoundError(f"DiffusersClipLoader: No .safetensors file or .bin file found in {directory}") + + @classmethod + def INPUT_TYPES(cls): + base_path = cls.get_base_path() + model_directories = cls.get_model_directories(base_path) + return { + "required": { + "sub_directory": (model_directories,), + "clip_type": (["stable_diffusion", "stable_cascade"],) + } + } + + RETURN_TYPES = ("CLIP",) + FUNCTION = "load_clip" + CATEGORY = "DiffusersLoader" + + def load_clip(self, sub_directory, clip_type="stable_diffusion"): + base_path = self.get_base_path() + sub_dir_path = os.path.join(base_path, sub_directory) + + model_type = self.detect_model_type(sub_dir_path) + + clip_model = self.load_clip_model(sub_dir_path, clip_type, model_type) + + return (clip_model,) + + def detect_model_type(self, sub_dir_path): + text_encoder_dir1 = os.path.join(sub_dir_path, "text_encoder") + text_encoder_dir2 = os.path.join(sub_dir_path, "text_encoder_2") + + if os.path.exists(text_encoder_dir1) and os.path.exists(text_encoder_dir2): + return "SDXL" + elif os.path.exists(text_encoder_dir1): + return "SD15" + else: + raise FileNotFoundError("No valid text_encoder directories found. This model is not SD15 or SDXL") + + def load_clip_model(self, sub_dir_path, clip_type="stable_diffusion", model_type="SD15"): + clip_type_enum = sd.CLIPType.STABLE_DIFFUSION + if clip_type == "stable_cascade": + clip_type_enum = sd.CLIPType.STABLE_CASCADE + + text_encoder_dir1 = os.path.join(sub_dir_path, "text_encoder") + text_encoder_paths = [os.path.join(text_encoder_dir1, self.find_model_file(text_encoder_dir1))] + + if model_type == "SDXL": + text_encoder_dir2 = os.path.join(sub_dir_path, "text_encoder_2") + text_encoder_paths.append(os.path.join(text_encoder_dir2, self.find_model_file(text_encoder_dir2))) + + print(f"DiffusersClipLoader: Checking paths: \n{text_encoder_paths}") + + clip_model = sd.load_clip(ckpt_paths=text_encoder_paths, embedding_directory=os.path.join(sub_dir_path, "embeddings"), clip_type=clip_type_enum) + + # Debugging statement to check if the loaded model has 'tokenize' method + if not hasattr(clip_model, 'tokenize'): + raise AttributeError("Loaded clip model does not have 'tokenize' method.") + + return clip_model + +# Ensure the node is registered properly +NODE_CLASS_MAPPINGS = {"DiffusersClipLoader": DiffusersClipLoader} +NODE_DISPLAY_NAME_MAPPINGS = {"DiffusersClipLoader": "Diffusers CLIP Loader"} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/diffusers_unet_loader.py b/diffusers_unet_loader.py new file mode 100644 index 0000000..ba3deda --- /dev/null +++ b/diffusers_unet_loader.py @@ -0,0 +1,103 @@ +import os +import comfy.sd + +def get_model_directories(base_path): + """Returns a list of sub-directories containing models within the given base path.""" + base_directory = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', base_path)) + print(f"DiffusersUNETLoader: Looking for model directories in: {base_directory}") + + model_dirs = [] + if os.path.exists(base_directory): + for root, dirs, files in os.walk(base_directory): + if "model_index.json" in files: + model_dirs.append(os.path.relpath(root, base_directory)) + print(f"DiffusersUNETLoader: Found model directory: {os.path.relpath(root, base_directory)}") + dirs[:] = [d for d in dirs if d not in ["vae", "unet", "text_encoder", "text_encoder_2"]] + else: + print(f"DiffusersUNETLoader: Path does not exist: {base_directory}") + return model_dirs + +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"DiffusersUNETLoader: Moved corrupted folder to: {corrupted_path}") + except PermissionError as e: + print(f"DiffusersUNETLoader: PermissionError while moving corrupted folder: {e}") + raise PermissionError(f"DiffusersUNETLoader: Permission denied while moving: {unet_folder}") + except Exception as e: + print(f"DiffusersUNETLoader: Error while moving corrupted folder: {e}") + raise e + +class DiffusersUNETLoader: + @classmethod + def INPUT_TYPES(cls): + base_path = "diffusers" + model_directories = get_model_directories(base_path) + return { + "required": { + "sub_directory": (model_directories,), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "load_unet" + CATEGORY = "DiffusersLoader" + + def load_unet(self, sub_directory): + base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers')) + sub_dir_path = os.path.join(base_path, sub_directory) + + model_type = self.detect_model_type(sub_dir_path) + + if model_type == "SD15": + unet_folder = os.path.join(sub_dir_path, 'unet') + elif model_type == "SDXL": + unet_folder = os.path.join(sub_dir_path, 'unet') + else: + raise ValueError("DiffusersUNETLoader: Unsupported model type") + + diffusion_model_path = self.find_model_file(unet_folder) + print(f"DiffusersUNETLoader: Attempting to load UNET model from: {diffusion_model_path}") + + if not os.path.exists(diffusion_model_path): + raise FileNotFoundError(f"DiffusersUNETLoader: {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"DiffusersUNETLoader: PermissionError: {e}") + handle_corrupted_file(unet_folder) + raise e + except Exception as e: + print(f"DiffusersUNETLoader: Error loading UNET model: {e}") + handle_corrupted_file(unet_folder) + raise e + + def detect_model_type(self, sub_dir_path): + """Detect the model type based on the directory structure.""" + text_encoder_dir1 = os.path.join(sub_dir_path, "text_encoder") + text_encoder_dir2 = os.path.join(sub_dir_path, "text_encoder_2") + + if os.path.exists(text_encoder_dir1) and os.path.exists(text_encoder_dir2): + return "SDXL" + elif os.path.exists(text_encoder_dir1): + return "SD15" + else: + raise FileNotFoundError("DiffusersUNETLoader: No valid text_encoder directories found. This model is not SD15 or SDXL") + + @staticmethod + def find_model_file(directory): + """Find the .safetensors or .bin file in the given directory.""" + for file in os.listdir(directory): + if file.endswith(".safetensors") or file.endswith(".bin"): + return os.path.join(directory, file) + raise FileNotFoundError(f"DiffusersUNETLoader: No .safetensors file or .bin file found in {directory}") + +# Ensure the node is registered properly +NODE_CLASS_MAPPINGS = {"DiffusersUNETLoader": DiffusersUNETLoader} +NODE_DISPLAY_NAME_MAPPINGS = {"DiffusersUNETLoader": "Diffusers UNET Loader"} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/diffusers_vae_loader.py b/diffusers_vae_loader.py new file mode 100644 index 0000000..7787a29 --- /dev/null +++ b/diffusers_vae_loader.py @@ -0,0 +1,85 @@ +import os +import comfy.utils +import comfy.sd +import torch + +class DiffusersVAELoader: + @staticmethod + def get_base_path(): + base_path = os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers') + base_path = os.path.abspath(base_path) + if not os.path.exists(base_path): + raise FileNotFoundError(f"DiffusersVAELoader: The base path '{base_path}' does not exist.") + return base_path + + @staticmethod + def get_model_directories(base_path): + model_dirs = [] + if os.path.exists(base_path): + for root, dirs, files in os.walk(base_path): + if "model_index.json" in files: + model_dirs.append(os.path.relpath(root, base_path)) + print(f"DiffusersVAELoader: Found model directory: {os.path.relpath(root, base_path)}") + dirs[:] = [d for d in dirs if d not in ["vae", "unet", "text_encoder", "text_encoder_2"]] + else: + print(f"DiffusersVAELoader: Path does not exist: {base_path}") + return model_dirs + + @staticmethod + def find_model_file(directory): + for file in os.listdir(directory): + if file.endswith(".safetensors") or file.endswith(".bin"): + return file + raise FileNotFoundError(f"DiffusersVAELoader: No .safetensors file or .bin file found in {directory}") + + @classmethod + def INPUT_TYPES(cls): + base_path = cls.get_base_path() + model_directories = cls.get_model_directories(base_path) + return { + "required": { + "sub_directory": (model_directories,), + } + } + + RETURN_TYPES = ("VAE",) + FUNCTION = "load_vae" + CATEGORY = "DiffusersLoader" + + def load_vae(self, sub_directory): + base_path = self.get_base_path() + sub_dir_path = os.path.join(base_path, sub_directory) + + model_type = self.detect_model_type(sub_dir_path) + + vae_folder = os.path.join(sub_dir_path, "vae") + vae_filename = self.find_model_file(vae_folder) + vae_path = os.path.join(vae_folder, vae_filename) + full_path = os.path.abspath(vae_path) + + print(f"DiffusersVAELoader: Attempting to load VAE model from: {full_path}") + + if not os.path.exists(full_path): + raise FileNotFoundError(f"DiffusersVAELoader: The VAE file '{full_path}' does not exist.") + + sd = comfy.utils.load_torch_file(full_path) + vae = comfy.sd.VAE(sd=sd) + return (vae,) + + def detect_model_type(self, sub_dir_path): + """Detect the model type based on the directory structure.""" + text_encoder_dir1 = os.path.join(sub_dir_path, "text_encoder") + text_encoder_dir2 = os.path.join(sub_dir_path, "text_encoder_2") + + if os.path.exists(text_encoder_dir1) and os.path.exists(text_encoder_dir2): + return "SDXL" + elif os.path.exists(text_encoder_dir1): + return "SD15" + else: + raise FileNotFoundError("DiffusersVAELoader: No valid text_encoder directories found. This model is not SD15 or SDXL") + +# Ensure the node is registered properly +NODE_CLASS_MAPPINGS = {"DiffusersVAELoader": DiffusersVAELoader} +NODE_DISPLAY_NAME_MAPPINGS = {"DiffusersVAELoader": "Diffusers VAE Loader"} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]