From f9d8bb959286f3f07079802d635e083c41cbfa8f Mon Sep 17 00:00:00 2001 From: Scorpinaus Date: Mon, 26 Aug 2024 21:43:54 +0800 Subject: [PATCH] Improved on model detection to include extra_model_paths.yaml --- README.md | 20 +++++++--- clip_loader.py | 40 +++++++------------- combined_diffusers_loader.py | 12 +++--- unet_loader.py | 73 +++++++++++++++++++----------------- utils.py | 24 +++++++----- vae_loader.py | 10 +++-- 6 files changed, 95 insertions(+), 84 deletions(-) diff --git a/README.md b/README.md index 8a00a26..f9a4879 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,7 @@ Follow these steps to install and set up the project: - Add Nodes -> DiffusersLoader -> Combined -> CombinedDiffusersLoader - ![image](https://github.com/user-attachments/assets/79a576b2-dc27-49e4-a7c7-ec0e01bc5bad) -3. Connect the DiffusersLoader +3. Connect the DiffusersLoader to your workflow @@ -48,12 +48,16 @@ The combined loader work in the same manner as existing checkpoints loader as se Select the name of the diffusers checkpoint folder from the dropdown list and connect the nodes. -Take note for SD3, AuraFlow and Flux: -- The first run through will be slightly slower due to creation of a combined_safetensors.safetensors ## Limitations & Future Improvements - Add support for other compatible diffusers format checkpoints -- For AuraFlow and SD3, the loading and combining for the combined_model.safetensors will have a delay for low-vram or low-system ram machines + - Future model_types: + - HunyuanDiT + - PixArt + - Kolors + Note: This will need to raise a PR on main comfyUI repo +- For Flux, use of default weight_dtype would OOM (Use fp8-based d-types) + ## Contributing @@ -67,5 +71,11 @@ Hafiz Saffie - hafiz.safs@gmail.com Project Link: https://github.com/Scorpinaus/ComfyUI-DiffusersLoader -## Version History +## Update History +26/8 - Improved on model detection to include extra_model_paths.yaml. +25/8 - Implemented multi-file loading support where the combined_safetensor.safetensor file method is used as a back-up 21/8 - Compatibility Update With ComfyUI, Addition of Flux Support +30/7 - Added repo to comfyUI registry +28/7 - Added AuraFlow support +20/6 - Added SD3 support +2/6 - Initial commit diff --git a/clip_loader.py b/clip_loader.py index bc2a5cd..b5f05a4 100644 --- a/clip_loader.py +++ b/clip_loader.py @@ -10,9 +10,10 @@ from .utils import DiffusersUtils class DiffusersClipLoader(DiffusersLoaderBase): @classmethod def INPUT_TYPES(cls): + model_directories = DiffusersUtils.get_model_directories() return { "required": { - "sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),), + "sub_directory": (model_directories,), "clip_type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "sdxl", "flux"],), } } @@ -28,8 +29,12 @@ class DiffusersClipLoader(DiffusersLoaderBase): @classmethod def load_model(cls, sub_directory, clip_type="stable_diffusion"): - base_path = DiffusersUtils.get_base_path() - sub_dir_path = os.path.join(base_path, 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.") + model_type = cls.detect_model_type(sub_dir_path) clip_type_enum = cls.get_clip_type_enum(clip_type) @@ -105,10 +110,9 @@ class DiffusersClipLoader(DiffusersLoaderBase): if model_type == "Flux": # use the index file - index_files = [f for f in os.listdir(text_encoder_dir2) if f.endswith('index.json') or 'index.' in f and f.endswith('.json')] + index_files = [f for f in os.listdir(text_encoder_dir2) if f.endswith('index.json')] if index_files: - index_file = os.path.join(text_encoder_dir2, index_files[0]) - text_encoder_paths.append(cls.load_flux_text_encoder_2(index_file)) + text_encoder_paths.append(os.path.join(text_encoder_dir2, index_files[0])) else: print(f"No index file found in {text_encoder_dir2}. Checking for combined text encoder step") text_encoder_paths.append(DiffusersUtils.find_model_file(text_encoder_dir2)) @@ -118,10 +122,9 @@ class DiffusersClipLoader(DiffusersLoaderBase): if model_type == "SD3": text_encoder_dir3 = os.path.join(sub_dir_path, "text_encoder_3") - #For text_encoder_3, we will use the index file. Look up for any file ending with 'index.json' or 'index.*.json' - index_files = [f for f in os.listdir(text_encoder_dir3) if f.endswith('index.json') or 'index.' in f and f.endswith('.json')] - if index_files: - index_file = os.path.join(text_encoder_dir3, index_files[0]) + #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") + if os.path.exists(index_file): text_encoder_paths.append(cls.load_sd3_text_encoder_3(index_file)) else: #If index file not found @@ -141,23 +144,6 @@ class DiffusersClipLoader(DiffusersLoaderBase): weight_map = index_data['weight_map'] base_path = os.path.dirname(index_file) - sd = {} - for key, file_name in weight_map.items(): - file_path = os.path.join(base_path, file_name) - if os.path.exists(file_path): - part_sd = comfy.utils.load_torch_file(file_path, safe_load=True) - sd.update({key: part_sd[key] for key in part_sd if key in weight_map}) - - return sd - - @classmethod - def load_flux_text_encoder_2(cls, index_file): - with open(index_file, 'r') as f: - index_data = json.load(f) - - weight_map = index_data['weight_map'] - base_path = os.path.dirname(index_file) - sd = {} for key, file_name in weight_map.items(): file_path = os.path.join(base_path, file_name) diff --git a/combined_diffusers_loader.py b/combined_diffusers_loader.py index 1cd76b1..538a19f 100644 --- a/combined_diffusers_loader.py +++ b/combined_diffusers_loader.py @@ -9,8 +9,7 @@ import os class CombinedDiffusersLoader: @classmethod def INPUT_TYPES(cls): - base_path = DiffusersUtils.get_base_path() - model_directories = DiffusersUtils.get_model_directories(base_path) + model_directories = DiffusersUtils.get_model_directories() return { "required": { "sub_directory": (model_directories,), @@ -28,12 +27,15 @@ class CombinedDiffusersLoader: @classmethod def load_models(cls, sub_directory, clip_type="stable_diffusion", transformer_parts="all", vae_type="default", weight_dtype="default"): - base_path = DiffusersUtils.get_base_path() - sub_dir_path = os.path.join(base_path, 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.") + 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) + unet_model, = DiffusersUNETLoader.load_model(sub_directory, transformer_parts, weight_dtype) clip_model = DiffusersClipLoader.load_model(sub_directory, clip_type) diff --git a/unet_loader.py b/unet_loader.py index f19987f..58613df 100644 --- a/unet_loader.py +++ b/unet_loader.py @@ -10,16 +10,17 @@ import torch class DiffusersUNETLoader(DiffusersLoaderBase): @classmethod def INPUT_TYPES(cls): + model_directories = DiffusersUtils.get_model_directories() return { "required": { - "sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),), + "sub_directory": (model_directories,), "transformer_parts": (["all", "part_1", "part_2", "part_3"],), "weight_dtype": (["default", "fp8_e4m3fn", "fp8_e5m2"],), } } RETURN_TYPES = ("MODEL",) - FUNCTION = "load_unet/transformer" + FUNCTION = "load_model" CATEGORY = "DiffusersLoader" @classmethod @@ -28,8 +29,13 @@ class DiffusersUNETLoader(DiffusersLoaderBase): @classmethod def load_model(cls, sub_directory, transformer_parts="all", weight_dtype="default"): - base_path = DiffusersUtils.get_base_path() - sub_dir_path = os.path.join(base_path, 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.") + model_type = cls.detect_model_type(sub_dir_path) print(f"DiffusersUNETLoader: Detected model type: {model_type}") @@ -46,23 +52,42 @@ class DiffusersUNETLoader(DiffusersLoaderBase): print(f"model_options: {model_options}") + def load_from_index(index_path): + print(f"Load_from_index: index_path: {index_path}") + + with open(index_path, 'r') as f: + index_data = json.load(f) + weight_map = index_data['weight_map'] + state_dict = {} + + for key, file in weight_map.items(): + file_path = os.path.join(os.path.dirname(index_path), file) + part_dict = comfy.utils.load_torch_file(file_path, safe_load=True) + state_dict.update(part_dict) + del part_dict + #torch.cuda.empty_cache() + + print(f"Load from index function completed. Returning state_dict") + return state_dict + try: if model_type == "AuraFlow" and transformer_parts == "all" and unet_path.endswith('index.json'): - state_dict = cls.load_from_index(unet_path) + state_dict = load_from_index(unet_path) print(f"State_dict obtained") model = cls.load_diffusion_model_from_state_dict(state_dict, model_options=model_options) elif model_type == "Flux" and transformer_parts == "all" and unet_path.endswith('index.json'): - state_dict = cls.load_from_index(unet_path) + state_dict = load_from_index(unet_path) print(f"State_dict obtained") model = cls.load_diffusion_model_from_state_dict(state_dict, model_options=model_options) else: model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options) print("UNET/Transformer model loaded successfully") - return model + #Return model as a tuple + return (model,) except RuntimeError as e: if "out of memory" in str(e): @@ -71,18 +96,19 @@ class DiffusersUNETLoader(DiffusersLoaderBase): try: if model_type == "AuraFlow" and transformer_parts == "all" and unet_path.endswith('index.json'): # Load model using the index file (retry) - state_dict = cls.load_from_index(unet_path) + state_dict = load_from_index(unet_path) model = cls.load_diffusion_model_from_state_dict(state_dict, model_options=model_options) elif model_type == "Flux" and transformer_parts == "all" and unet_path.endswith('index.json'): - state_dict = cls.load_from_index(unet_path) + state_dict = load_from_index(unet_path) model = cls.load_diffusion_model_from_state_dict(state_dict, model_options=model_options) else: model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options) print("UNET/Transformer model loaded successfully after memory clear") - return model + #Return model as a tuple + return (model,) except Exception as retry_e: print(f"Error loading model after memory clear: {retry_e}") raise @@ -90,28 +116,6 @@ class DiffusersUNETLoader(DiffusersLoaderBase): print(f"DiffusersUNETLoader: Error loading UNET/Transformer model: {e}") raise - @classmethod - def load_from_index(cls, index_path): - print(f"Load_from_index: index_path: {index_path}") - - with open(index_path, 'r') as f: - index_data = json.load(f) - - weight_map = index_data['weight_map'] - base_path = os.path.dirname(index_path) - - state_dict = {} - for key, file_name in weight_map.items(): - file_path = os.path.join(base_path, file_name) - if os.path.exists(file_path): - part_dict = comfy.utils.load_torch_file(file_path, safe_load=True) - state_dict.update({k: part_dict[k] for k in part_dict if k in weight_map}) - else: - print(f"WARNING: File not found: {file_path}") - - print(f"Load from index function completed. Returning state_dict") - return state_dict - @classmethod def load_diffusion_model_from_state_dict(cls, state_dict, model_options): print("Running load_diffusion_model_from_state_dict function...") @@ -158,7 +162,7 @@ class DiffusersUNETLoader(DiffusersLoaderBase): @classmethod def handle_auraflow_transformer(cls, unet_folder, transformer_parts): #Find file ending with 'index.json' - index_files = [f for f in os.listdir(unet_folder) if f.endswith('index.json') or 'index.' in f and f.endswith('.json')] + index_files = [f for f in os.listdir(unet_folder) if f.endswith('index.json')] print(f"Index files: {index_files}") if index_files: index_file = os.path.join(unet_folder, index_files[0]) @@ -189,7 +193,7 @@ class DiffusersUNETLoader(DiffusersLoaderBase): @classmethod def handle_flux_transformer(cls, unet_folder, transformer_parts): - index_files = [f for f in os.listdir(unet_folder) if f.endswith('index.json') or 'index.' in f and f.endswith('.json')] + index_files = [f for f in os.listdir(unet_folder) if f.endswith('index.json')] if index_files: index_file = os.path.join(unet_folder, index_files[0]) print(f"Found index file: {index_file}") @@ -202,7 +206,6 @@ class DiffusersUNETLoader(DiffusersLoaderBase): part_num = int(transformer_parts.split('_')[1]) part_files = [file for file in index_data['weight_map'].values() if f"-0000{part_num}-of-" in file] - print(f"Part files: {part_files}") if part_files: return os.path.join(unet_folder, part_files[0]) diff --git a/utils.py b/utils.py index 170e0cf..604dca8 100644 --- a/utils.py +++ b/utils.py @@ -2,6 +2,9 @@ import os import safetensors.torch import torch import gc +import sys +sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +import folder_paths class DiffusersUtils: _model_cache = {'clip': None, 'unet': None, 'vae': None} @@ -43,19 +46,22 @@ 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): + base_path = folder_paths.get_folder_paths("diffusers") + #print(f"Base path: {base_path}") + if not 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 + def get_model_directories(): + paths = [] + base_paths = DiffusersUtils.get_base_path() + for base_path in base_paths: + 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)) + return paths @staticmethod def find_model_files(directory, file_parts=None): diff --git a/vae_loader.py b/vae_loader.py index e1998b0..ecdd825 100644 --- a/vae_loader.py +++ b/vae_loader.py @@ -16,9 +16,10 @@ class DiffusersVAELoader(DiffusersLoaderBase): } @classmethod def INPUT_TYPES(cls): + model_directories = DiffusersUtils.get_model_directories() return { "required": { - "sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),), + "sub_directory": (model_directories,), "vae_type": (["default"] + list(cls.VAE_CONFIGS.keys()),), } } @@ -40,8 +41,11 @@ class DiffusersVAELoader(DiffusersLoaderBase): @classmethod def load_default_vae(cls, sub_directory): - base_path = DiffusersUtils.get_base_path() - sub_dir_path = os.path.join(base_path, 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.") + vae_folder = os.path.join(sub_dir_path, "vae") vae_path = DiffusersUtils.find_model_file(vae_folder) DiffusersUtils.check_and_clear_cache('vae', vae_path)