From 51c724b6d0298d3c9d4da19ecdbf36c7ece58a08 Mon Sep 17 00:00:00 2001 From: Scorpinaus Date: Sun, 28 Jul 2024 12:10:09 +0800 Subject: [PATCH] Added loading support for AuraFlow and memory cache clearing between renders for optimization --- base_loader.py | 2 + clip_loader.py | 23 +++++--- combined_diffusers_loader.py | 14 +++-- unet_loader.py | 68 +++++++++++++++------ utils.py | 111 +++++++++++++++++++++++++++++------ vae_loader.py | 4 ++ 6 files changed, 173 insertions(+), 49 deletions(-) diff --git a/base_loader.py b/base_loader.py index eb2cdb4..c545f57 100644 --- a/base_loader.py +++ b/base_loader.py @@ -14,6 +14,8 @@ class DiffusersLoaderBase: 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) and os.path.exists(transformer): + return "AuraFlow" elif os.path.exists(text_encoder_dir1): return "SD15" else: diff --git a/clip_loader.py b/clip_loader.py index 6ff5c6b..a9b4e91 100644 --- a/clip_loader.py +++ b/clip_loader.py @@ -11,7 +11,7 @@ class DiffusersClipLoader(DiffusersLoaderBase): "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"],) + "clip_parts": (["all", "part_1", "part_2"],) } } @@ -20,11 +20,12 @@ class DiffusersClipLoader(DiffusersLoaderBase): 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),) + def load_clip(cls, sub_directory, clip_type="stable_diffusion", clip_parts="all"): + return (cls.load_model(sub_directory, clip_type, clip_parts),) @classmethod - def load_model(cls, sub_directory, clip_type="stable_diffusion", file_parts="all"): + def load_model(cls, sub_directory, clip_type="stable_diffusion", clip_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) @@ -42,16 +43,20 @@ class DiffusersClipLoader(DiffusersLoaderBase): if model_type == "SD3": text_encoder_dir3 = os.path.join(sub_dir_path, "text_encoder_3") - if file_parts == "all": + print("clip_parts Selected:", clip_parts) + if clip_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"]) + elif clip_parts == "part_1": + text_encoder_paths += DiffusersUtils.find_model_files(text_encoder_dir3, num_parts=1) + elif clip_parts == "part_2": + text_encoder_paths += DiffusersUtils.find_model_files(text_encoder_dir3, num_parts=2)[1:] + for path in text_encoder_paths: + DiffusersUtils.check_and_clear_cache('clip', path) + print(f"DiffusersClipLoader: Checking paths: \n{text_encoder_paths}") try: diff --git a/combined_diffusers_loader.py b/combined_diffusers_loader.py index 0ce38d9..a804400 100644 --- a/combined_diffusers_loader.py +++ b/combined_diffusers_loader.py @@ -13,7 +13,8 @@ class CombinedDiffusersLoader: "required": { "sub_directory": (model_directories,), "clip_type": (["stable_diffusion", "stable_cascade"],), - "file_parts": (["none", "all", "part_1", "part_2"],) + "clip_parts": (["all", "part_1", "part_2"],), + "transformer_parts": (["all", "part_1", "part_2", "part_3"],), } } @@ -22,9 +23,14 @@ class CombinedDiffusersLoader: CATEGORY = "DiffusersLoader/Combined" @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) + def load_models(cls, sub_directory, clip_type="stable_diffusion", clip_parts="all", transformer_parts="all"): + + unet_model = DiffusersUNETLoader.load_model(sub_directory, transformer_parts) + + + clip_model = DiffusersClipLoader.load_model(sub_directory, clip_type, clip_parts) + + vae_model = DiffusersVAELoader.load_model(sub_directory) return unet_model, clip_model, vae_model diff --git a/unet_loader.py b/unet_loader.py index ac3c71c..46b96a3 100644 --- a/unet_loader.py +++ b/unet_loader.py @@ -3,6 +3,7 @@ import os import comfy.sd from .base_loader import DiffusersLoaderBase from .utils import DiffusersUtils +import torch class DiffusersUNETLoader(DiffusersLoaderBase): @classmethod @@ -10,6 +11,7 @@ class DiffusersUNETLoader(DiffusersLoaderBase): return { "required": { "sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),), + "transformer_parts": (["all", "part_1", "part_2", "part_3"],), } } @@ -18,34 +20,62 @@ class DiffusersUNETLoader(DiffusersLoaderBase): CATEGORY = "DiffusersLoader" @classmethod - def load_unet(cls, sub_directory): - return (cls.load_model(sub_directory),) + def load_unet(cls, sub_directory, transformer_parts="all"): + return (cls.load_model(sub_directory, transformer_parts),) @classmethod - def load_model(cls, sub_directory): + def load_model(cls, sub_directory, transformer_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) + print(f"DiffusersUNETLoader: Detected model type: {model_type}") - if model_type == "SD3": + if model_type in ["AuraFlow"]: unet_folder = os.path.join(sub_dir_path, 'transformer') + print(f"Using transformer folder: {unet_folder}") + if transformer_parts == "all": + combined_file_path = os.path.join(unet_folder, "combined_transformer.safetensors") + if not os.path.exists(combined_file_path): + print(f"Combined file not found, creating at: {combined_file_path}") + try: + combined_file_path = DiffusersUtils.combine_safetensor_files(unet_folder, base_path, num_parts=3) + except Exception as e: + print(f"Error combining safetensor files: {e}") + # If combining fails, use the first part file + combined_file_path = DiffusersUtils.find_model_files(unet_folder, num_parts=1)[0] + unet_path = combined_file_path + else: + part_num = int(transformer_parts.split('_')[1]) + unet_path = DiffusersUtils.find_model_files(unet_folder, num_parts=part_num)[-1] + print(f"Selected UNET/Transformer path: {unet_path}") + elif model_type in ["SD3"]: + unet_folder = os.path.join(sub_dir_path, 'transformer') + print(f"Using transformer folder: {unet_folder}") + unet_path = DiffusersUtils.find_model_file(unet_folder) else: unet_folder = os.path.join(sub_dir_path, 'unet') + unet_path = DiffusersUtils.find_model_file(unet_folder) + + DiffusersUtils.check_and_clear_cache('unet', unet_path) - 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}") + print(f"DiffusersUNETLoader: Attempting to load UNET/Transformer 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 + model = comfy.sd.load_unet(unet_path) + print("UNET/Transformer model loaded successfully") + return model + except RuntimeError as e: + if "out of memory" in str(e): + print("Out of memory error. Attempting to clear memory and retry...") + cls.clear_memory() + try: + model = comfy.sd.load_unet(unet_path) + print("UNET/Transformer model loaded successfully after memory clear") + return model + except Exception as retry_e: + print(f"Error loading model after memory clear: {retry_e}") + raise + else: + print(f"DiffusersUNETLoader: Error loading UNET/Transformer model: {e}") + raise \ No newline at end of file diff --git a/utils.py b/utils.py index 44acc0d..8d09148 100644 --- a/utils.py +++ b/utils.py @@ -1,8 +1,46 @@ import os import safetensors.torch import torch +import gc class DiffusersUtils: + _model_cache = {'clip': None, 'unet': None, 'vae': None} + _current_model_hashes = {'clip': None, 'unet': None, 'vae': None} + + @staticmethod + def clear_memory(): + gc.collect() + torch.cuda.empty_cache() + print("Cleared CUDA Memory and ran garbage collection") + + @classmethod + def clear_model_cache(cls): + cls._model_cache = {'clip': None, 'unet': None, 'vae': None} + cls._current_model_hashes = {'clip': None, 'unet': None, 'vae': None} + cls.clear_memory() + print("Cleared model cache") + + @classmethod + def get_model_hash(cls, model_path): + import hashlib + with open(model_path, "rb") as f: + file_hash = hashlib.md5() + chunk = f.read(8192) + while chunk: + file_hash.update(chunk) + chunk = f.read(8192) + return file_hash.hexdigest() + + @classmethod + def check_and_clear_cache(cls, model_type, model_path): + new_hash = cls.get_model_hash(model_path) + if cls._current_model_hashes[model_type] != new_hash: + print(f"Detected change in {model_type} model. Clearing cache.") + cls.clear_model_cache() + cls._current_model_hashes[model_type] = new_hash + else: + print(f"No change detected in {model_type} model.") + @staticmethod def get_base_path(): base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers')) @@ -20,7 +58,7 @@ class DiffusersUtils: return model_dirs @staticmethod - def find_model_files(directory, file_parts=None): + def find_model_files(directory, file_parts=None, num_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) @@ -31,6 +69,8 @@ class DiffusersUtils: if file_parts: return [file for file in files if any(part in file for part in file_parts)] + elif num_parts: + return [file for file in files if any(f"0000{i+1}-of-" in file for i in range(num_parts))] return files @staticmethod @@ -42,21 +82,58 @@ class DiffusersUtils: @staticmethod def load_safetensor_paths(file_paths): - tensors = [safetensors.torch.load_file(file) for file in file_paths] + print("Running Loading Safetensors") 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) + for file_path in file_paths: + print(f"Loading file: {file_path}") + try: + tensor_dict = safetensors.torch.load_file(file_path) + for key, value in tensor_dict.items(): + print(f"Processing key: {key}") + if key in combined_tensors: + try: + combined_tensors[key] = torch.cat((combined_tensors[key], value), dim=0) + except RuntimeError as e: + print(f"Error concatenating tensors for key {key}: {e}") + print(f"Shapes: combined {combined_tensors[key].shape}, new {value.shape}") + # Instead of raising, we'll skip this tensor and continue + print(f"Skipping problematic tensor: {key}") + else: + combined_tensors[key] = value + print(f"Finished processing file: {file_path}") + # Free up memory + del tensor_dict + torch.cuda.empty_cache() + except Exception as e: + print(f"Error processing file {file_path}: {e}") + # Continue to the next file instead of stopping + continue + print("Combining Safetensors completed") + return combined_tensors - return combined_file_path \ No newline at end of file + @staticmethod + def combine_safetensor_files(directory, base_path, num_parts): + print(f"Running combine_safetensor_files for directory: {directory}") + part_files = DiffusersUtils.find_model_files(directory, num_parts=num_parts) + print("Path to part_files:", part_files) + try: + combined_tensors = DiffusersUtils.load_safetensor_paths(part_files) + print("Combined tensors keys:", list(combined_tensors.keys())) + + if "text_encoder" in directory: + combined_file_path = os.path.join(directory, "combined_text_encoder.safetensors") + elif "transformer" in directory: + combined_file_path = os.path.join(directory, "combined_transformer.safetensors") + else: + raise ValueError(f"Unsupported directory for combining files: {directory}") + + print(f"Saving combined file to: {combined_file_path}") + safetensors.torch.save_file(combined_tensors, combined_file_path) + print("Combined file saved successfully") + + return combined_file_path + except Exception as e: + 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 diff --git a/vae_loader.py b/vae_loader.py index a4cd92f..21afedc 100644 --- a/vae_loader.py +++ b/vae_loader.py @@ -24,11 +24,15 @@ class DiffusersVAELoader(DiffusersLoaderBase): @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) + + DiffusersUtils.check_and_clear_cache('vae', vae_path) + print(f"DiffusersVAELoader: Attempting to load VAE model from: {vae_path}") try: