diff --git a/clip_loader.py b/clip_loader.py index b5f05a4..32bbe49 100644 --- a/clip_loader.py +++ b/clip_loader.py @@ -10,10 +10,9 @@ from .utils import DiffusersUtils class DiffusersClipLoader(DiffusersLoaderBase): @classmethod def INPUT_TYPES(cls): - model_directories = DiffusersUtils.get_model_directories() return { "required": { - "sub_directory": (model_directories,), + "sub_directory": (DiffusersUtils.get_unique_display_names(DiffusersUtils.get_model_directories())[0],), "clip_type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "sdxl", "flux"],), } } @@ -28,18 +27,27 @@ class DiffusersClipLoader(DiffusersLoaderBase): @classmethod def load_model(cls, sub_directory, clip_type="stable_diffusion"): + if os.path.exists(sub_directory): + full_path = sub_directory + else: + model_directories = DiffusersUtils.get_model_directories() + _, unique_names = DiffusersUtils.get_unique_display_names(model_directories) - base_paths = DiffusersUtils.get_base_path() - sub_dir_path = next((os.path.join(base_path, sub_directory) for base_path in base_paths if os.path.exists(os.path.join(base_path, sub_directory))), None) + if "(" in sub_directory: + dir_name, index = sub_directory.rsplit(" (", 1) + index = int(index[:-1]) - 1 + full_path = unique_names[dir_name][index] + else: + full_path = unique_names[sub_directory][0] - if sub_dir_path is None: - raise ValueError(f"Subdirectory '{sub_directory}' not found in any of the diffusers paths.") + if not os.path.exists(full_path): + raise ValueError(f"Selected directory does not exist: {full_path}") - model_type = cls.detect_model_type(sub_dir_path) + model_type = cls.detect_model_type(full_path) clip_type_enum = cls.get_clip_type_enum(clip_type) - text_encoder_paths = cls.get_text_encoder_paths(sub_dir_path, model_type) + text_encoder_paths = cls.get_text_encoder_paths(full_path, model_type) clip_data = [] for path in text_encoder_paths: @@ -72,7 +80,7 @@ class DiffusersClipLoader(DiffusersLoaderBase): try: clip_model = comfy.sd.load_text_encoder_state_dicts(clip_data, - embedding_directory=os.path.join(sub_dir_path, "embeddings"), clip_type=clip_type_enum) + embedding_directory=os.path.join(full_path, "embeddings"), clip_type=clip_type_enum) except Exception as e: print(f"DiffusersClipLoader: Error loading clip model: {e}") raise @@ -99,12 +107,12 @@ class DiffusersClipLoader(DiffusersLoaderBase): return clip_type_map.get(clip_type, comfy.sd.CLIPType.STABLE_DIFFUSION) @classmethod - def get_text_encoder_paths(cls, sub_dir_path, model_type): + def get_text_encoder_paths(cls, full_path, model_type): text_encoder_paths = [] if model_type in ["SDXL", "SD3", "Flux"]: - 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_dir1 = os.path.join(full_path, "text_encoder") + text_encoder_dir2 = os.path.join(full_path, "text_encoder_2") text_encoder_paths.append(DiffusersUtils.find_model_file(text_encoder_dir1)) @@ -120,7 +128,7 @@ class DiffusersClipLoader(DiffusersLoaderBase): 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") + text_encoder_dir3 = os.path.join(full_path, "text_encoder_3") #For text_encoder_3, we will use the index file index_file = os.path.join(text_encoder_dir3, "text_encoder_3_model.safetensors.index.fp16.json") @@ -131,7 +139,7 @@ class DiffusersClipLoader(DiffusersLoaderBase): text_encoder_paths.append(DiffusersUtils.find_model_file(text_encoder_dir3)) else: #For SD15 or other single text encoder models - text_encoder_dir = os.path.join(sub_dir_path, "text_encoder") + text_encoder_dir = os.path.join(full_path, "text_encoder") text_encoder_paths = [DiffusersUtils.find_model_file(text_encoder_dir)] return text_encoder_paths diff --git a/combined_diffusers_loader.py b/combined_diffusers_loader.py index 538a19f..38a3248 100644 --- a/combined_diffusers_loader.py +++ b/combined_diffusers_loader.py @@ -9,10 +9,11 @@ import os class CombinedDiffusersLoader: @classmethod def INPUT_TYPES(cls): - model_directories = DiffusersUtils.get_model_directories() + display_names, _ = DiffusersUtils.get_unique_display_names(DiffusersUtils.get_model_directories()) + return { "required": { - "sub_directory": (model_directories,), + "sub_directory": (display_names,), "clip_type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "sdxl", "flux"],), "transformer_parts": (["all", "part_1", "part_2", "part_3"],), "vae_type": (["default", "taesd", "taesdxl", "taesd3", "taef1"],), @@ -27,19 +28,24 @@ class CombinedDiffusersLoader: @classmethod def load_models(cls, sub_directory, clip_type="stable_diffusion", transformer_parts="all", vae_type="default", weight_dtype="default"): - base_paths = DiffusersUtils.get_base_path() - sub_dir_path = next((os.path.join(base_path, sub_directory) for base_path in base_paths if os.path.exists(os.path.join(base_path, sub_directory))), None) - if sub_dir_path is None: - raise ValueError(f"Subdirectory '{sub_directory}' not found in any of the diffusers paths.") + model_directories = DiffusersUtils.get_model_directories() + _, unique_names = DiffusersUtils.get_unique_display_names(model_directories) - model_type = DiffusersLoaderBase.detect_model_type(sub_dir_path) - print(f"Detected model type: {model_type}") - - unet_model, = DiffusersUNETLoader.load_model(sub_directory, transformer_parts, weight_dtype) - - clip_model = DiffusersClipLoader.load_model(sub_directory, clip_type) + if "(" in sub_directory: + dir_name, index = sub_directory.rsplit(" (", 1) + index = int(index[:-1]) - 1 + full_path = unique_names[dir_name][index] + else: + full_path = unique_names[sub_directory][0] - vae_model = DiffusersVAELoader.load_model(sub_directory, vae_type) + if not os.path.exists(full_path): + raise ValueError(f"Model directory '{full_path}' not found.") + + unet_model, = DiffusersUNETLoader.load_model(full_path, transformer_parts, weight_dtype) + + clip_model = DiffusersClipLoader.load_model(full_path, clip_type) + + vae_model = DiffusersVAELoader.load_model(full_path, vae_type) return unet_model, clip_model, vae_model diff --git a/unet_loader.py b/unet_loader.py index 58613df..a2ffb4c 100644 --- a/unet_loader.py +++ b/unet_loader.py @@ -10,10 +10,9 @@ import torch class DiffusersUNETLoader(DiffusersLoaderBase): @classmethod def INPUT_TYPES(cls): - model_directories = DiffusersUtils.get_model_directories() return { "required": { - "sub_directory": (model_directories,), + "sub_directory": (DiffusersUtils.get_unique_display_names(DiffusersUtils.get_model_directories())[0],), "transformer_parts": (["all", "part_1", "part_2", "part_3"],), "weight_dtype": (["default", "fp8_e4m3fn", "fp8_e5m2"],), } @@ -29,17 +28,26 @@ class DiffusersUNETLoader(DiffusersLoaderBase): @classmethod def load_model(cls, sub_directory, transformer_parts="all", weight_dtype="default"): + if os.path.exists(sub_directory): + full_path = sub_directory + else: + model_directories = DiffusersUtils.get_model_directories() + _, unique_names = DiffusersUtils.get_unique_display_names(model_directories) + + if "(" in sub_directory: + dir_name, index = sub_directory.rsplit(" (", 1) + index = int(index[:-1]) - 1 + full_path = unique_names[dir_name][index] + else: + full_path = unique_names[sub_directory][0] - base_paths = DiffusersUtils.get_base_path() + if not os.path.exists(full_path): + raise ValueError(f"Selected directory does not exist: {full_path}") - sub_dir_path = next((os.path.join(base_path, sub_directory) for base_path in base_paths if os.path.exists(os.path.join(base_path, sub_directory))), None) - if sub_dir_path is None: - raise ValueError(f"Subdirectory '{sub_directory}' not found in any of the diffusers paths.") - - model_type = cls.detect_model_type(sub_dir_path) + model_type = cls.detect_model_type(full_path) print(f"DiffusersUNETLoader: Detected model type: {model_type}") - unet_path = cls.get_unet_path(sub_dir_path, model_type, transformer_parts) + unet_path = cls.get_unet_path(full_path, model_type, transformer_parts) DiffusersUtils.check_and_clear_cache('unet', unet_path) print(f"DiffusersUNETLoader: Attempting to load UNET model from: {unet_path}") @@ -123,9 +131,9 @@ class DiffusersUNETLoader(DiffusersLoaderBase): return model @classmethod - def get_unet_path(cls, sub_dir_path, model_type, transformer_parts): + def get_unet_path(cls, full_path, model_type, transformer_parts): if model_type in ["AuraFlow", "Flux", "SD3"]: - unet_folder = os.path.join(sub_dir_path, 'transformer') + unet_folder = os.path.join(full_path, 'transformer') print(f"Using transformer folder for {model_type}: {unet_folder}") if model_type == "SD3": return cls.handle_sd3_transformer(unet_folder) @@ -134,10 +142,10 @@ class DiffusersUNETLoader(DiffusersLoaderBase): elif model_type == "Flux": return cls.handle_flux_transformer(unet_folder, transformer_parts) else: # SD15, SDXL, etc. - unet_folder = os.path.join(sub_dir_path, 'unet') + unet_folder = os.path.join(full_path, 'unet') if not os.path.exists(unet_folder): #Fallback to transformer folder if unet folder does not exist - unet_folder = os.path.join(sub_dir_path, 'transformer') + unet_folder = os.path.join(full_path, 'transformer') if os.path.exists(unet_folder): print(f"Unet folder not found, using transformer folder: {unet_folder}") return cls.handle_flux_transformer(unet_folder, transformer_parts) diff --git a/utils.py b/utils.py index 604dca8..9f34f0f 100644 --- a/utils.py +++ b/utils.py @@ -60,7 +60,10 @@ class DiffusersUtils: if os.path.exists(base_path): for root, _, files in os.walk(base_path, followlinks=True): if "model_index.json" in files: - paths.append(os.path.relpath(root, start=base_path)) + relative_path = os.path.relpath(root, start=base_path) + full_path = os.path.join(base_path, relative_path) + dir_name = os.path.basename(relative_path) + paths.append((dir_name, full_path)) return paths @staticmethod @@ -145,4 +148,30 @@ class DiffusersUtils: print(f"Error in combine_safetensor_files: {e}") # If combining fails, return the path to the first file as a fallback return part_files[0] - \ No newline at end of file + + @staticmethod + def get_unique_display_names(model_directories): + unique_names = {} + for dir_name, full_path in model_directories: + if dir_name in unique_names: + unique_names[dir_name].append(full_path) + else: + unique_names[dir_name] = [full_path] + + display_names = [] + for dir_name, paths in unique_names.items(): + if len(paths) == 1: + display_names.append(dir_name) + else: + for i, path in enumerate(paths): + display_names.append(f"{dir_name} ({i+1})") + + return display_names, unique_names + + @staticmethod + def get_full_path(dir_name): + model_directories = DiffusersUtils.get_model_directories() + for name, path in model_directories: + if name == dir_name: + return path + raise ValueError(f"Directory not found: {dir_name}") \ No newline at end of file diff --git a/vae_loader.py b/vae_loader.py index ecdd825..a680d9e 100644 --- a/vae_loader.py +++ b/vae_loader.py @@ -16,10 +16,9 @@ class DiffusersVAELoader(DiffusersLoaderBase): } @classmethod def INPUT_TYPES(cls): - model_directories = DiffusersUtils.get_model_directories() return { "required": { - "sub_directory": (model_directories,), + "sub_directory": (DiffusersUtils.get_unique_display_names(DiffusersUtils.get_model_directories())[0],), "vae_type": (["default"] + list(cls.VAE_CONFIGS.keys()),), } } @@ -40,13 +39,24 @@ class DiffusersVAELoader(DiffusersLoaderBase): return cls.load_taesd(vae_type) @classmethod - def load_default_vae(cls, sub_directory): - base_paths = DiffusersUtils.get_base_path() - sub_dir_path = next((os.path.join(base_path, sub_directory) for base_path in base_paths if os.path.exists(os.path.join(base_path, sub_directory))), None) - if sub_dir_path is None: - raise ValueError(f"Subdirectory '{sub_directory}' not found in any of the diffusers paths.") + def load_default_vae(cls, sub_directory): + if os.path.exists(sub_directory): + full_path = sub_directory + else: + model_directories = DiffusersUtils.get_model_directories() + _, unique_names = DiffusersUtils.get_unique_display_names(model_directories) + + if "(" in sub_directory: + dir_name, index = sub_directory.rsplit(" (", 1) + index = int(index[:-1]) - 1 + full_path = unique_names[dir_name][index] + else: + full_path = unique_names[sub_directory][0] - vae_folder = os.path.join(sub_dir_path, "vae") + if not os.path.exists(full_path): + raise ValueError(f"Selected directory does not exist: {full_path}") + + vae_folder = os.path.join(full_path, "vae") vae_path = DiffusersUtils.find_model_file(vae_folder) DiffusersUtils.check_and_clear_cache('vae', vae_path) @@ -62,7 +72,7 @@ class DiffusersVAELoader(DiffusersLoaderBase): @staticmethod def load_taesd(vae_type): sd = {} - vae_approx_path = os.path.join(os.path.dirname(DiffusersUtils.get_base_path()), "vae_approx") + vae_approx_path = os.path.join(os.path.dirname(DiffusersUtils.get_base_path()[0]), "vae_approx") # Validation check for vae_approx_path if not os.path.exists(vae_approx_path): raise FileNotFoundError(f"The vae_approx path '{vae_approx_path}' does not exist.")