diff --git a/__init__.py b/__init__.py index c5d97b1..dd977da 100644 --- a/__init__.py +++ b/__init__.py @@ -6,6 +6,7 @@ from .sdxl_vae_loader import SDXLVAELoader from .sdxl_clip_loader import SDXLCLIPLoader from .sd15_combined_loader import CombinedDiffusersSD15Loader from .sdxl_combined_loader import CombinedDiffusersSDXLLoader +from .combined_diffusers_loader import CombinedDiffusersLoader NODE_CLASS_MAPPINGS = { "SD15UNETLoader": SD15UNETLoader, @@ -16,6 +17,7 @@ NODE_CLASS_MAPPINGS = { "SDXLCLIPLoader": SDXLCLIPLoader, "CombinedDiffusersSD15Loader": CombinedDiffusersSD15Loader, "CombinedDiffusersSDXLLoader": CombinedDiffusersSDXLLoader, + "CombinedDiffusersLoader": CombinedDiffusersLoader } @@ -28,4 +30,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "SDXLCLIPLoader": "SDXL CLIP Loader", "CombinedDiffusersSD15Loader" : "Combined Diffusers SD15 Loader", "CombinedDiffusersSDXLLoader" : "Combined Diffusers SDXL Loader" + "CombinedDiffusersLoader": "Combined Diffusers Loader" } diff --git a/combined_diffusers_loader.py b/combined_diffusers_loader.py new file mode 100644 index 0000000..fcdc96e --- /dev/null +++ b/combined_diffusers_loader.py @@ -0,0 +1,109 @@ +import os +import comfy.sd +import comfy.utils +import torch + +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_subdirectories(base_path): + return [name for name in os.listdir(base_path) if os.path.isdir(os.path.join(base_path, name))] + + @staticmethod + def find_model_file(directory): + for file in os.listdir(directory): + if file.endswith(".safetensors") or file.endswith(".bin"): + return os.path.join(directory, file) + raise FileNotFoundError(f"CombinedDiffusersLoader: No .safetensors file or .bin file found in {directory}") + + @classmethod + def INPUT_TYPES(cls): + base_path = cls.get_base_path() + subdirectories = cls.get_subdirectories(base_path) + return { + "required": { + "sub_directory": (subdirectories,), + "clip_type": (["stable_diffusion", "stable_cascade"],) + } + } + + RETURN_TYPES = ("MODEL", "CLIP", "VAE") + FUNCTION = "load_models" + CATEGORY = "DiffusersLoader/Combined" + + def load_models(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) + + 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) + + 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") + + 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(self, base_path, clip_type="stable_diffusion", model_type="SD15"): + 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 = [self.find_model_file(text_encoder_dir1)] + + if model_type == "SDXL": + text_encoder_dir2 = os.path.join(base_path, "text_encoder_2") + text_encoder_paths.append(self.find_model_file(text_encoder_dir2)) + + print(f"CombinedDiffusersLoader: Checking paths: \n{text_encoder_paths}") + + clip_model = comfy.sd.load_clip(ckpt_paths=text_encoder_paths, embedding_directory=os.path.join(base_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 + + def load_unet(self, base_path): + unet_folder = os.path.join(base_path, 'unet') + unet_path = self.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 Exception as e: + print(f"Error loading UNET model:{e}") + raise e + + def load_vae(self, base_path): + vae_folder = os.path.join(base_path, "vae") + vae_path = self.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"}