nodes/Scorpinaus-ComfyUI-DiffusersLoader
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added loading support for AuraFlow and memory cache clearing between renders for optimization

Browse Source
This commit is contained in:
Scorpinaus
2024-07-28 12:10:09 +08:00
parent 9bf5db2533
commit 51c724b6d0
6 changed files with 173 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
base_loader.py
+2
View File
@@ -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:
clip_loader.py
+14 -9
View File
@@ -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:
combined_diffusers_loader.py
+10 -4
View File
@@ -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
unet_loader.py
+49 -19
View File
@@ -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
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
utils.py
+94 -17
View File
@@ -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
@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]
vae_loader.py
+4
View File
@@ -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: