From dfeed597573c7ceaec53c6372d6460bd8492d6d9 Mon Sep 17 00:00:00 2001 From: Scorpinaus Date: Sat, 27 Jul 2024 21:19:47 +0800 Subject: [PATCH] Code refactoring for modularity and easier addition of new model types --- __init__.py | 14 +-- base_loader.py | 24 +++++ clip_loader.py | 66 +++++++++++++ combined_diffusers_loader.py | 182 +++-------------------------------- diffusers_clip_loader.py | 93 ------------------ diffusers_unet_loader.py | 103 -------------------- diffusers_vae_loader.py | 85 ---------------- loader_factory.py | 16 +++ unet_loader.py | 51 ++++++++++ utils.py | 62 ++++++++++++ vae_loader.py | 39 ++++++++ 11 files changed, 279 insertions(+), 456 deletions(-) create mode 100644 base_loader.py create mode 100644 clip_loader.py delete mode 100644 diffusers_clip_loader.py delete mode 100644 diffusers_unet_loader.py delete mode 100644 diffusers_vae_loader.py create mode 100644 loader_factory.py create mode 100644 unet_loader.py create mode 100644 utils.py create mode 100644 vae_loader.py diff --git a/__init__.py b/__init__.py index 95bb074..f2ef0f9 100644 --- a/__init__.py +++ b/__init__.py @@ -1,19 +1,19 @@ +# __init__.py from .combined_diffusers_loader import CombinedDiffusersLoader -from .diffusers_clip_loader import DiffusersClipLoader -from .diffusers_unet_loader import DiffusersUNETLoader -from .diffusers_vae_loader import DiffusersVAELoader +from .unet_loader import DiffusersUNETLoader +from .clip_loader import DiffusersClipLoader +from .vae_loader import DiffusersVAELoader NODE_CLASS_MAPPINGS = { "CombinedDiffusersLoader": CombinedDiffusersLoader, - "DiffusersClipLoader": DiffusersClipLoader, "DiffusersUNETLoader": DiffusersUNETLoader, + "DiffusersClipLoader": DiffusersClipLoader, "DiffusersVAELoader": DiffusersVAELoader - } NODE_DISPLAY_NAME_MAPPINGS = { "CombinedDiffusersLoader": "Combined Diffusers Loader", - "DiffusersClipLoader": "Diffusers CLIP Loader", "DiffusersUNETLoader": "Diffusers UNET Loader", + "DiffusersClipLoader": "Diffusers CLIP Loader", "DiffusersVAELoader": "Diffusers VAE Loader" -} +} \ No newline at end of file diff --git a/base_loader.py b/base_loader.py new file mode 100644 index 0000000..eb2cdb4 --- /dev/null +++ b/base_loader.py @@ -0,0 +1,24 @@ +# base_loader.py +import os +from .utils import DiffusersUtils + +class DiffusersLoaderBase: + @classmethod + def detect_model_type(cls, 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") + text_encoder_dir3 = os.path.join(sub_dir_path, "text_encoder_3") + transformer = os.path.join(sub_dir_path, "transformer") + + if os.path.exists(text_encoder_dir3) and os.path.exists(transformer): + return "SD3" + elif 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, SDXL, or SD3") + + @classmethod + def load_model(cls, sub_directory): + raise NotImplementedError \ No newline at end of file diff --git a/clip_loader.py b/clip_loader.py new file mode 100644 index 0000000..6ff5c6b --- /dev/null +++ b/clip_loader.py @@ -0,0 +1,66 @@ +# clip_loader.py +import os +import comfy.sd +from .base_loader import DiffusersLoaderBase +from .utils import DiffusersUtils + +class DiffusersClipLoader(DiffusersLoaderBase): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),), + "clip_type": (["stable_diffusion", "stable_cascade"],), + "file_parts": (["none", "all", "part_1", "part_2"],) + } + } + + RETURN_TYPES = ("CLIP",) + FUNCTION = "load_clip" + CATEGORY = "DiffusersLoader" + + @classmethod + def load_clip(cls, sub_directory, clip_type="stable_diffusion", file_parts="all"): + return (cls.load_model(sub_directory, clip_type, file_parts),) + + @classmethod + def load_model(cls, sub_directory, clip_type="stable_diffusion", file_parts="all"): + base_path = DiffusersUtils.get_base_path() + sub_dir_path = os.path.join(base_path, sub_directory) + model_type = cls.detect_model_type(sub_dir_path) + + clip_type_enum = comfy.sd.CLIPType.STABLE_DIFFUSION + if clip_type == "stable_cascade": + clip_type_enum = comfy.sd.CLIPType.STABLE_CASCADE + + text_encoder_dir1 = os.path.join(sub_dir_path, "text_encoder") + text_encoder_paths = [DiffusersUtils.find_model_file(text_encoder_dir1)] + + if model_type in ["SDXL", "SD3"]: + text_encoder_dir2 = os.path.join(sub_dir_path, "text_encoder_2") + text_encoder_paths.append(DiffusersUtils.find_model_file(text_encoder_dir2)) + + if model_type == "SD3": + text_encoder_dir3 = os.path.join(sub_dir_path, "text_encoder_3") + if file_parts == "all": + combined_file_path = os.path.join(text_encoder_dir3, "combined_text_encoder.safetensors") + if not os.path.exists(combined_file_path): + combined_file_path = DiffusersUtils.combine_safetensor_files(text_encoder_dir3, base_path) + text_encoder_paths.append(combined_file_path) + elif file_parts == "part_1": + text_encoder_paths += DiffusersUtils.find_model_files(text_encoder_dir3, ["00001-of-00002"]) + elif file_parts == "part_2": + text_encoder_paths += DiffusersUtils.find_model_files(text_encoder_dir3, ["00002-of-00002"]) + + print(f"DiffusersClipLoader: Checking paths: \n{text_encoder_paths}") + + try: + clip_model = comfy.sd.load_clip(ckpt_paths=text_encoder_paths, embedding_directory=os.path.join(sub_dir_path, "embeddings"), clip_type=clip_type_enum) + except Exception as e: + print(f"DiffusersClipLoader: Error loading clip model: {e}") + raise + + if not hasattr(clip_model, 'tokenize'): + raise AttributeError("DiffusersClipLoader: Loaded clip model does not have 'tokenize' method.") + + return clip_model \ No newline at end of file diff --git a/combined_diffusers_loader.py b/combined_diffusers_loader.py index 5b20b8c..0ce38d9 100644 --- a/combined_diffusers_loader.py +++ b/combined_diffusers_loader.py @@ -1,80 +1,19 @@ -import os -import comfy.sd -import comfy.utils -import torch -import safetensors.torch +# combined_diffusers_loader.py +from .utils import DiffusersUtils +from .unet_loader import DiffusersUNETLoader +from .clip_loader import DiffusersClipLoader +from .vae_loader import DiffusersVAELoader class CombinedDiffusersLoader: - @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"CombinedDiffusersLoader: 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 sub-directories from dirs - dirs[:] = [d for d in dirs if d not in ["vae", "unet", "text_encoder", "text_encoder_2", "text_encoder_3", "transformer"]] - return model_dirs - - @staticmethod - def find_model_files(directory, file_parts=None): - if not os.path.exists(directory): - raise FileNotFoundError(f"The directory '{directory}' does not exist.") - files = [] - for file in os.listdir(directory): - if file.endswith(".safetensors") or file.endswith(".bin"): - files.append(os.path.join(directory, file)) - - if not files: - raise FileNotFoundError(f"CombinedDiffusersLoader: No .safetensors file or .bin file found in {directory}") - - if file_parts: - return [file for file in files if any(part in file for part in file_parts)] - return files - - @staticmethod - def find_model_file(directory): - files = CombinedDiffusersLoader.find_model_files(directory) - if files: - return files[0] - raise FileNotFoundError(f"CombinedDiffusersLoader: No .safetensors or .bin file found in {directory}") - - @staticmethod - def load_safetensor_paths(file_paths): - tensors = [safetensors.torch.load_file(file) for file in file_paths] - combined_tensors = {} - for tensor_dict in tensors: - for key, value in tensor_dict.items(): - if key in combined_tensors: - combined_tensors[key] = torch.cat((combined_tensors[key], value), dim=0) - else: - combined_tensors[key] = value - return combined_tensors - - @staticmethod - def combine_safetensor_files(text_encoder_dir3, base_path): - part_files = CombinedDiffusersLoader.find_model_files(text_encoder_dir3, ["00001-of-00002", "00002-of-00002"]) - combined_tensors = CombinedDiffusersLoader.load_safetensor_paths(part_files) - combined_file_path = os.path.join(text_encoder_dir3, "combined_text_encoder.safetensors") - safetensors.torch.save_file(combined_tensors, combined_file_path) - - return combined_file_path - @classmethod def INPUT_TYPES(cls): - base_path = cls.get_base_path() - model_directories = cls.get_model_directories(base_path) + base_path = DiffusersUtils.get_base_path() + model_directories = DiffusersUtils.get_model_directories(base_path) return { "required": { "sub_directory": (model_directories,), "clip_type": (["stable_diffusion", "stable_cascade"],), - "file_parts": (["none", "all", "part_1", "part_2",],) + "file_parts": (["none", "all", "part_1", "part_2"],) } } @@ -82,107 +21,14 @@ class CombinedDiffusersLoader: FUNCTION = "load_models" CATEGORY = "DiffusersLoader/Combined" - def load_models(self, sub_directory, clip_type="stable_diffusion", file_parts="all"): - 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_model = self.load_vae(sub_dir_path, model_type) - unet_model = self.load_unet(sub_dir_path, model_type) - clip_model = self.load_clip(sub_dir_path, clip_type, model_type, file_parts) + @classmethod + def load_models(cls, sub_directory, clip_type="stable_diffusion", file_parts="all"): + unet_model = DiffusersUNETLoader.load_model(sub_directory) + clip_model = DiffusersClipLoader.load_model(sub_directory, clip_type, file_parts) + vae_model = DiffusersVAELoader.load_model(sub_directory) return unet_model, clip_model, vae_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") - text_encoder_dir3 = os.path.join(sub_dir_path, "text_encoder_3") - transformer = os.path.join(sub_dir_path, "transformer") - - if os.path.exists(text_encoder_dir3) and os.path.exists(transformer): - print("This model is SD3") - return "SD3" - elif os.path.exists(text_encoder_dir1) and os.path.exists(text_encoder_dir2): - print("This model is SDXL") - return "SDXL" - elif os.path.exists(text_encoder_dir1): - print("This model is SD1.5") - return "SD15" - else: - print(f"CombinedDiffusersLoader: Sub_dir_path: \n{sub_dir_path}") - raise FileNotFoundError("No valid text_encoder directories found. This model is not SD15 or SDXL or SD3") - - def load_clip(self, base_path, clip_type="stable_diffusion", model_type="SD15", file_parts="all"): - clip_type_enum = comfy.sd.CLIPType.STABLE_DIFFUSION - if clip_type == "stable_cascade": - clip_type_enum = comfy.sd.CLIPType.STABLE_CASCADE - - text_encoder_dir1 = os.path.join(base_path, "text_encoder") - text_encoder_paths = [CombinedDiffusersLoader.find_model_file(text_encoder_dir1)] - - if model_type in ["SDXL", "SD3"]: - text_encoder_dir2 = os.path.join(base_path, "text_encoder_2") - text_encoder_paths.append(CombinedDiffusersLoader.find_model_file(text_encoder_dir2)) - - if model_type == "SD3": - text_encoder_dir3 = os.path.join(base_path, "text_encoder_3") - if file_parts == "all": - combined_file_path = os.path.join(text_encoder_dir3, "combined_text_encoder.safetensors") - if not os.path.exists(combined_file_path): - combined_file_path = CombinedDiffusersLoader.combine_safetensor_files(text_encoder_dir3, base_path) - text_encoder_paths.append(combined_file_path) - elif file_parts == "part_1": - text_encoder_paths += CombinedDiffusersLoader.find_model_files(text_encoder_dir3, ["00001-of-00002"]) - elif file_parts == "part_2": - text_encoder_paths += CombinedDiffusersLoader.find_model_files(text_encoder_dir3, ["00002-of-00002"]) - - print(f"CombinedDiffusersLoader: Checking paths: \n{text_encoder_paths}") - - try: - clip_model = comfy.sd.load_clip(ckpt_paths=text_encoder_paths, embedding_directory=os.path.join(base_path, "embeddings"), clip_type=clip_type_enum) - except Exception as e: - print(f"Error loading clip model: {e}") - raise - - # 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 - - def load_unet(self, base_path, model_type): - if model_type == "SD3": - unet_folder = os.path.join(base_path, 'transformer') - else: - unet_folder = os.path.join(base_path, 'unet') - - if not os.path.exists(unet_folder): - raise FileNotFoundError(f"CombinedDiffusersLoader: The directory '{unet_folder}' does not exist for model type '{model_type}'") - - unet_path = CombinedDiffusersLoader.find_model_file(unet_folder) - print(f"CombinedDiffusersLoader: 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"Permission denied: {unet_path}") - except FileNotFoundError as e: - print(f"FileNotFoundError: {e}") - raise FileNotFoundError(f"UNET file not found: {unet_path}") - except Exception as e: - print(f"Error loading UNET model: {e}") - raise e - - def load_vae(self, base_path, model_type): - vae_folder = os.path.join(base_path, "vae") - vae_path = CombinedDiffusersLoader.find_model_file(vae_folder) - print(f"Attempting to load VAE model from: {vae_path}") - - vae_sd = comfy.utils.load_torch_file(vae_path) - return comfy.sd.VAE(sd=vae_sd) - # Ensure the node is registered properly NODE_CLASS_MAPPINGS = {"CombinedDiffusersLoader": CombinedDiffusersLoader} -NODE_DISPLAY_NAME_MAPPINGS = {"CombinedDiffusersLoader": "Combined Diffusers Loader"} +NODE_DISPLAY_NAME_MAPPINGS = {"CombinedDiffusersLoader": "Combined Diffusers Loader"} \ No newline at end of file diff --git a/diffusers_clip_loader.py b/diffusers_clip_loader.py deleted file mode 100644 index ebbae67..0000000 --- a/diffusers_clip_loader.py +++ /dev/null @@ -1,93 +0,0 @@ -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 deleted file mode 100644 index ba3deda..0000000 --- a/diffusers_unet_loader.py +++ /dev/null @@ -1,103 +0,0 @@ -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 deleted file mode 100644 index 7787a29..0000000 --- a/diffusers_vae_loader.py +++ /dev/null @@ -1,85 +0,0 @@ -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"] diff --git a/loader_factory.py b/loader_factory.py new file mode 100644 index 0000000..e889700 --- /dev/null +++ b/loader_factory.py @@ -0,0 +1,16 @@ +# loader_factory.py +from .unet_loader import DiffusersUNETLoader +from .clip_loader import DiffusersClipLoader +from .vae_loader import DiffusersVAELoader + +class LoaderFactory: + @staticmethod + def create_loader(loader_type, sub_directory): + if loader_type == "UNET": + return DiffusersUNETLoader(sub_directory) + elif loader_type == "CLIP": + return DiffusersClipLoader(sub_directory) + elif loader_type == "VAE": + return DiffusersVAELoader(sub_directory) + else: + raise ValueError(f"Unknown loader type: {loader_type}") \ No newline at end of file diff --git a/unet_loader.py b/unet_loader.py new file mode 100644 index 0000000..ac3c71c --- /dev/null +++ b/unet_loader.py @@ -0,0 +1,51 @@ +# unet_loader.py +import os +import comfy.sd +from .base_loader import DiffusersLoaderBase +from .utils import DiffusersUtils + +class DiffusersUNETLoader(DiffusersLoaderBase): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "load_unet" + CATEGORY = "DiffusersLoader" + + @classmethod + def load_unet(cls, sub_directory): + return (cls.load_model(sub_directory),) + + @classmethod + def load_model(cls, sub_directory): + base_path = DiffusersUtils.get_base_path() + sub_dir_path = os.path.join(base_path, sub_directory) + model_type = cls.detect_model_type(sub_dir_path) + + if model_type == "SD3": + unet_folder = os.path.join(sub_dir_path, 'transformer') + else: + unet_folder = os.path.join(sub_dir_path, 'unet') + + if not os.path.exists(unet_folder): + raise FileNotFoundError(f"The directory '{unet_folder}' does not exist for model type '{model_type}'") + + unet_path = DiffusersUtils.find_model_file(unet_folder) + print(f"DiffusersUNETLoader: Attempting to load UNET model from {unet_path}") + + try: + return comfy.sd.load_unet(unet_path) + except PermissionError as e: + print(f"DiffusersUNETLoader: PermissionError: {e}") + raise PermissionError(f"Permission denied: {unet_path}") + except FileNotFoundError as e: + print(f"DiffusersUNETLoader: FileNotFoundError: {e}") + raise FileNotFoundError(f"UNET file not found: {unet_path}") + except Exception as e: + print(f"DiffusersUNETLoader: Error loading UNET model: {e}") + raise e \ No newline at end of file diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..44acc0d --- /dev/null +++ b/utils.py @@ -0,0 +1,62 @@ +import os +import safetensors.torch +import torch + +class DiffusersUtils: + @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"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)) + dirs[:] = [d for d in dirs if d not in ["vae", "unet", "text_encoder", "text_encoder_2", "text_encoder_3", "transformer"]] + return model_dirs + + @staticmethod + def find_model_files(directory, file_parts=None): + if not os.path.exists(directory): + raise FileNotFoundError(f"The directory '{directory}' does not exist.") + files = [os.path.join(directory, f) for f in os.listdir(directory) + if f.endswith((".safetensors", ".bin"))] + + if not files: + raise FileNotFoundError(f"No .safetensors or .bin file found in {directory}") + + if file_parts: + return [file for file in files if any(part in file for part in file_parts)] + return files + + @staticmethod + def find_model_file(directory): + files = DiffusersUtils.find_model_files(directory) + if files: + return files[0] + raise FileNotFoundError(f"No .safetensors or .bin file found in {directory}") + + @staticmethod + def load_safetensor_paths(file_paths): + tensors = [safetensors.torch.load_file(file) for file in file_paths] + combined_tensors = {} + for tensor_dict in tensors: + for key, value in tensor_dict.items(): + if key in combined_tensors: + combined_tensors[key] = torch.cat((combined_tensors[key], value), dim=0) + else: + combined_tensors[key] = value + return combined_tensors + + @staticmethod + def combine_safetensor_files(text_encoder_dir3, base_path): + part_files = DiffusersUtils.find_model_files(text_encoder_dir3, ["00001-of-00002", "00002-of-00002"]) + combined_tensors = DiffusersUtils.load_safetensor_paths(part_files) + combined_file_path = os.path.join(text_encoder_dir3, "combined_text_encoder.safetensors") + safetensors.torch.save_file(combined_tensors, combined_file_path) + + return combined_file_path \ No newline at end of file diff --git a/vae_loader.py b/vae_loader.py new file mode 100644 index 0000000..a4cd92f --- /dev/null +++ b/vae_loader.py @@ -0,0 +1,39 @@ +# vae_loader.py +import os +import comfy.sd +import comfy.utils +from .base_loader import DiffusersLoaderBase +from .utils import DiffusersUtils + +class DiffusersVAELoader(DiffusersLoaderBase): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),), + } + } + + RETURN_TYPES = ("VAE",) + FUNCTION = "load_vae" + CATEGORY = "DiffusersLoader" + + @classmethod + def load_vae(cls, sub_directory): + return (cls.load_model(sub_directory),) + + @classmethod + def load_model(cls, sub_directory): + base_path = DiffusersUtils.get_base_path() + sub_dir_path = os.path.join(base_path, sub_directory) + + vae_folder = os.path.join(sub_dir_path, "vae") + vae_path = DiffusersUtils.find_model_file(vae_folder) + print(f"DiffusersVAELoader: Attempting to load VAE model from: {vae_path}") + + try: + vae_sd = comfy.utils.load_torch_file(vae_path) + return comfy.sd.VAE(sd=vae_sd) + except Exception as e: + print(f"DiffusersVAELoader: Error loading VAE model: {e}") + raise \ No newline at end of file