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

Improved on the model finding and naming process

Browse Source
This commit is contained in:
Scorpinaus
2024-08-26 22:48:52 +08:00
parent f9d8bb9592
commit 6263e91306
5 changed files with 112 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
clip_loader.py
+22 -14
View File
@@ -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
combined_diffusers_loader.py
+19 -13
View File
@@ -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
unet_loader.py
+21 -13
View File
@@ -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)
utils.py
+31 -2
View File
@@ -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]
@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}")
vae_loader.py
+19 -9
View File
@@ -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.")