Code refactoring for modularity and easier addition of new model types
This commit is contained in:
Scorpinaus
+7
-7
@@ -1,19 +1,19 @@
|
||||
# __init__.py
|
||||
from .combined_diffusers_loader import CombinedDiffusersLoader
|
||||
from .diffusers_clip_loader import DiffusersClipLoader
|
||||
from .diffusers_unet_loader import DiffusersUNETLoader
|
||||
from .diffusers_vae_loader import DiffusersVAELoader
|
||||
from .unet_loader import DiffusersUNETLoader
|
||||
from .clip_loader import DiffusersClipLoader
|
||||
from .vae_loader import DiffusersVAELoader
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"CombinedDiffusersLoader": CombinedDiffusersLoader,
|
||||
"DiffusersClipLoader": DiffusersClipLoader,
|
||||
"DiffusersUNETLoader": DiffusersUNETLoader,
|
||||
"DiffusersClipLoader": DiffusersClipLoader,
|
||||
"DiffusersVAELoader": DiffusersVAELoader
|
||||
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"CombinedDiffusersLoader": "Combined Diffusers Loader",
|
||||
"DiffusersClipLoader": "Diffusers CLIP Loader",
|
||||
"DiffusersUNETLoader": "Diffusers UNET Loader",
|
||||
"DiffusersClipLoader": "Diffusers CLIP Loader",
|
||||
"DiffusersVAELoader": "Diffusers VAE Loader"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
# base_loader.py
|
||||
import os
|
||||
from .utils import DiffusersUtils
|
||||
|
||||
class DiffusersLoaderBase:
|
||||
@classmethod
|
||||
def detect_model_type(cls, sub_dir_path):
|
||||
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_dir3 = os.path.join(sub_dir_path, "text_encoder_3")
|
||||
transformer = os.path.join(sub_dir_path, "transformer")
|
||||
|
||||
if os.path.exists(text_encoder_dir3) and os.path.exists(transformer):
|
||||
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):
|
||||
return "SD15"
|
||||
else:
|
||||
raise FileNotFoundError("No valid text_encoder directories found. This model is not SD15, SDXL, or SD3")
|
||||
|
||||
@classmethod
|
||||
def load_model(cls, sub_directory):
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,66 @@
|
||||
# clip_loader.py
|
||||
import os
|
||||
import comfy.sd
|
||||
from .base_loader import DiffusersLoaderBase
|
||||
from .utils import DiffusersUtils
|
||||
|
||||
class DiffusersClipLoader(DiffusersLoaderBase):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"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"],)
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CLIP",)
|
||||
FUNCTION = "load_clip"
|
||||
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),)
|
||||
|
||||
@classmethod
|
||||
def load_model(cls, sub_directory, clip_type="stable_diffusion", file_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)
|
||||
|
||||
clip_type_enum = comfy.sd.CLIPType.STABLE_DIFFUSION
|
||||
if clip_type == "stable_cascade":
|
||||
clip_type_enum = comfy.sd.CLIPType.STABLE_CASCADE
|
||||
|
||||
text_encoder_dir1 = os.path.join(sub_dir_path, "text_encoder")
|
||||
text_encoder_paths = [DiffusersUtils.find_model_file(text_encoder_dir1)]
|
||||
|
||||
if model_type in ["SDXL", "SD3"]:
|
||||
text_encoder_dir2 = os.path.join(sub_dir_path, "text_encoder_2")
|
||||
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")
|
||||
if file_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"])
|
||||
|
||||
print(f"DiffusersClipLoader: Checking paths: \n{text_encoder_paths}")
|
||||
|
||||
try:
|
||||
clip_model = comfy.sd.load_clip(ckpt_paths=text_encoder_paths, embedding_directory=os.path.join(sub_dir_path, "embeddings"), clip_type=clip_type_enum)
|
||||
except Exception as e:
|
||||
print(f"DiffusersClipLoader: Error loading clip model: {e}")
|
||||
raise
|
||||
|
||||
if not hasattr(clip_model, 'tokenize'):
|
||||
raise AttributeError("DiffusersClipLoader: Loaded clip model does not have 'tokenize' method.")
|
||||
|
||||
return clip_model
|
||||
+14
-168
@@ -1,80 +1,19 @@
|
||||
import os
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import torch
|
||||
import safetensors.torch
|
||||
# combined_diffusers_loader.py
|
||||
from .utils import DiffusersUtils
|
||||
from .unet_loader import DiffusersUNETLoader
|
||||
from .clip_loader import DiffusersClipLoader
|
||||
from .vae_loader import DiffusersVAELoader
|
||||
|
||||
class CombinedDiffusersLoader:
|
||||
@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):
|
||||
raise FileNotFoundError(f"CombinedDiffusersLoader: 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))
|
||||
#Remove unwanted sub-directories from dirs
|
||||
dirs[:] = [d for d in dirs if d not in ["vae", "unet", "text_encoder", "text_encoder_2", "text_encoder_3", "transformer"]]
|
||||
return model_dirs
|
||||
|
||||
@staticmethod
|
||||
def find_model_files(directory, file_parts=None):
|
||||
if not os.path.exists(directory):
|
||||
raise FileNotFoundError(f"The directory '{directory}' does not exist.")
|
||||
files = []
|
||||
for file in os.listdir(directory):
|
||||
if file.endswith(".safetensors") or file.endswith(".bin"):
|
||||
files.append(os.path.join(directory, file))
|
||||
|
||||
if not files:
|
||||
raise FileNotFoundError(f"CombinedDiffusersLoader: No .safetensors file or .bin file found in {directory}")
|
||||
|
||||
if file_parts:
|
||||
return [file for file in files if any(part in file for part in file_parts)]
|
||||
return files
|
||||
|
||||
@staticmethod
|
||||
def find_model_file(directory):
|
||||
files = CombinedDiffusersLoader.find_model_files(directory)
|
||||
if files:
|
||||
return files[0]
|
||||
raise FileNotFoundError(f"CombinedDiffusersLoader: No .safetensors or .bin file found in {directory}")
|
||||
|
||||
@staticmethod
|
||||
def load_safetensor_paths(file_paths):
|
||||
tensors = [safetensors.torch.load_file(file) for file in file_paths]
|
||||
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 = CombinedDiffusersLoader.find_model_files(text_encoder_dir3, ["00001-of-00002", "00002-of-00002"])
|
||||
combined_tensors = CombinedDiffusersLoader.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)
|
||||
|
||||
return combined_file_path
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
base_path = cls.get_base_path()
|
||||
model_directories = cls.get_model_directories(base_path)
|
||||
base_path = DiffusersUtils.get_base_path()
|
||||
model_directories = DiffusersUtils.get_model_directories(base_path)
|
||||
return {
|
||||
"required": {
|
||||
"sub_directory": (model_directories,),
|
||||
"clip_type": (["stable_diffusion", "stable_cascade"],),
|
||||
"file_parts": (["none", "all", "part_1", "part_2",],)
|
||||
"file_parts": (["none", "all", "part_1", "part_2"],)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -82,107 +21,14 @@ class CombinedDiffusersLoader:
|
||||
FUNCTION = "load_models"
|
||||
CATEGORY = "DiffusersLoader/Combined"
|
||||
|
||||
def load_models(self, sub_directory, clip_type="stable_diffusion", file_parts="all"):
|
||||
base_path = self.get_base_path()
|
||||
sub_dir_path = os.path.join(base_path, sub_directory)
|
||||
|
||||
model_type = self.detect_model_type(sub_dir_path)
|
||||
|
||||
vae_model = self.load_vae(sub_dir_path, model_type)
|
||||
unet_model = self.load_unet(sub_dir_path, model_type)
|
||||
clip_model = self.load_clip(sub_dir_path, clip_type, model_type, file_parts)
|
||||
@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)
|
||||
vae_model = DiffusersVAELoader.load_model(sub_directory)
|
||||
|
||||
return unet_model, clip_model, vae_model
|
||||
|
||||
def detect_model_type(self, sub_dir_path):
|
||||
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_dir3 = os.path.join(sub_dir_path, "text_encoder_3")
|
||||
transformer = os.path.join(sub_dir_path, "transformer")
|
||||
|
||||
if os.path.exists(text_encoder_dir3) and os.path.exists(transformer):
|
||||
print("This model is SD3")
|
||||
return "SD3"
|
||||
elif os.path.exists(text_encoder_dir1) and os.path.exists(text_encoder_dir2):
|
||||
print("This model is SDXL")
|
||||
return "SDXL"
|
||||
elif os.path.exists(text_encoder_dir1):
|
||||
print("This model is SD1.5")
|
||||
return "SD15"
|
||||
else:
|
||||
print(f"CombinedDiffusersLoader: Sub_dir_path: \n{sub_dir_path}")
|
||||
raise FileNotFoundError("No valid text_encoder directories found. This model is not SD15 or SDXL or SD3")
|
||||
|
||||
def load_clip(self, base_path, clip_type="stable_diffusion", model_type="SD15", file_parts="all"):
|
||||
clip_type_enum = comfy.sd.CLIPType.STABLE_DIFFUSION
|
||||
if clip_type == "stable_cascade":
|
||||
clip_type_enum = comfy.sd.CLIPType.STABLE_CASCADE
|
||||
|
||||
text_encoder_dir1 = os.path.join(base_path, "text_encoder")
|
||||
text_encoder_paths = [CombinedDiffusersLoader.find_model_file(text_encoder_dir1)]
|
||||
|
||||
if model_type in ["SDXL", "SD3"]:
|
||||
text_encoder_dir2 = os.path.join(base_path, "text_encoder_2")
|
||||
text_encoder_paths.append(CombinedDiffusersLoader.find_model_file(text_encoder_dir2))
|
||||
|
||||
if model_type == "SD3":
|
||||
text_encoder_dir3 = os.path.join(base_path, "text_encoder_3")
|
||||
if file_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 = CombinedDiffusersLoader.combine_safetensor_files(text_encoder_dir3, base_path)
|
||||
text_encoder_paths.append(combined_file_path)
|
||||
elif file_parts == "part_1":
|
||||
text_encoder_paths += CombinedDiffusersLoader.find_model_files(text_encoder_dir3, ["00001-of-00002"])
|
||||
elif file_parts == "part_2":
|
||||
text_encoder_paths += CombinedDiffusersLoader.find_model_files(text_encoder_dir3, ["00002-of-00002"])
|
||||
|
||||
print(f"CombinedDiffusersLoader: Checking paths: \n{text_encoder_paths}")
|
||||
|
||||
try:
|
||||
clip_model = comfy.sd.load_clip(ckpt_paths=text_encoder_paths, embedding_directory=os.path.join(base_path, "embeddings"), clip_type=clip_type_enum)
|
||||
except Exception as e:
|
||||
print(f"Error loading clip model: {e}")
|
||||
raise
|
||||
|
||||
# Debugging statement to check if the loaded model has 'tokenize' method
|
||||
if not hasattr(clip_model, 'tokenize'):
|
||||
raise AttributeError("Loaded clip model does not have 'tokenize' method.")
|
||||
|
||||
return clip_model
|
||||
|
||||
def load_unet(self, base_path, model_type):
|
||||
if model_type == "SD3":
|
||||
unet_folder = os.path.join(base_path, 'transformer')
|
||||
else:
|
||||
unet_folder = os.path.join(base_path, 'unet')
|
||||
|
||||
if not os.path.exists(unet_folder):
|
||||
raise FileNotFoundError(f"CombinedDiffusersLoader: The directory '{unet_folder}' does not exist for model type '{model_type}'")
|
||||
|
||||
unet_path = CombinedDiffusersLoader.find_model_file(unet_folder)
|
||||
print(f"CombinedDiffusersLoader: Attempting to load UNET model from {unet_path}")
|
||||
|
||||
try:
|
||||
return comfy.sd.load_unet(unet_path)
|
||||
except PermissionError as e:
|
||||
print(f"Permissionerror:{e}")
|
||||
raise PermissionError(f"Permission denied: {unet_path}")
|
||||
except FileNotFoundError as e:
|
||||
print(f"FileNotFoundError: {e}")
|
||||
raise FileNotFoundError(f"UNET file not found: {unet_path}")
|
||||
except Exception as e:
|
||||
print(f"Error loading UNET model: {e}")
|
||||
raise e
|
||||
|
||||
def load_vae(self, base_path, model_type):
|
||||
vae_folder = os.path.join(base_path, "vae")
|
||||
vae_path = CombinedDiffusersLoader.find_model_file(vae_folder)
|
||||
print(f"Attempting to load VAE model from: {vae_path}")
|
||||
|
||||
vae_sd = comfy.utils.load_torch_file(vae_path)
|
||||
return comfy.sd.VAE(sd=vae_sd)
|
||||
|
||||
# Ensure the node is registered properly
|
||||
NODE_CLASS_MAPPINGS = {"CombinedDiffusersLoader": CombinedDiffusersLoader}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"CombinedDiffusersLoader": "Combined Diffusers Loader"}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"CombinedDiffusersLoader": "Combined Diffusers Loader"}
|
||||
@@ -1,93 +0,0 @@
|
||||
import os
|
||||
from comfy import sd
|
||||
import comfy.utils
|
||||
import torch
|
||||
|
||||
class DiffusersClipLoader:
|
||||
@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):
|
||||
raise FileNotFoundError(f"DiffusersClipLoader: 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))
|
||||
# Remove unwanted subdirectories from dirs so that os.walk doesn't traverse them
|
||||
dirs[:] = [d for d in dirs if d not in ["vae", "unet", "text_encoder", "text_encoder_2"]]
|
||||
return model_dirs
|
||||
|
||||
@staticmethod
|
||||
def find_model_file(directory):
|
||||
for file in os.listdir(directory):
|
||||
if file.endswith(".safetensors") or file.endswith(".bin"):
|
||||
return file
|
||||
raise FileNotFoundError(f"DiffusersClipLoader: No .safetensors file or .bin file found in {directory}")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
base_path = cls.get_base_path()
|
||||
model_directories = cls.get_model_directories(base_path)
|
||||
return {
|
||||
"required": {
|
||||
"sub_directory": (model_directories,),
|
||||
"clip_type": (["stable_diffusion", "stable_cascade"],)
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CLIP",)
|
||||
FUNCTION = "load_clip"
|
||||
CATEGORY = "DiffusersLoader"
|
||||
|
||||
def load_clip(self, sub_directory, clip_type="stable_diffusion"):
|
||||
base_path = self.get_base_path()
|
||||
sub_dir_path = os.path.join(base_path, sub_directory)
|
||||
|
||||
model_type = self.detect_model_type(sub_dir_path)
|
||||
|
||||
clip_model = self.load_clip_model(sub_dir_path, clip_type, model_type)
|
||||
|
||||
return (clip_model,)
|
||||
|
||||
def detect_model_type(self, sub_dir_path):
|
||||
text_encoder_dir1 = os.path.join(sub_dir_path, "text_encoder")
|
||||
text_encoder_dir2 = os.path.join(sub_dir_path, "text_encoder_2")
|
||||
|
||||
if os.path.exists(text_encoder_dir1) and os.path.exists(text_encoder_dir2):
|
||||
return "SDXL"
|
||||
elif os.path.exists(text_encoder_dir1):
|
||||
return "SD15"
|
||||
else:
|
||||
raise FileNotFoundError("No valid text_encoder directories found. This model is not SD15 or SDXL")
|
||||
|
||||
def load_clip_model(self, sub_dir_path, clip_type="stable_diffusion", model_type="SD15"):
|
||||
clip_type_enum = sd.CLIPType.STABLE_DIFFUSION
|
||||
if clip_type == "stable_cascade":
|
||||
clip_type_enum = sd.CLIPType.STABLE_CASCADE
|
||||
|
||||
text_encoder_dir1 = os.path.join(sub_dir_path, "text_encoder")
|
||||
text_encoder_paths = [os.path.join(text_encoder_dir1, self.find_model_file(text_encoder_dir1))]
|
||||
|
||||
if model_type == "SDXL":
|
||||
text_encoder_dir2 = os.path.join(sub_dir_path, "text_encoder_2")
|
||||
text_encoder_paths.append(os.path.join(text_encoder_dir2, self.find_model_file(text_encoder_dir2)))
|
||||
|
||||
print(f"DiffusersClipLoader: Checking paths: \n{text_encoder_paths}")
|
||||
|
||||
clip_model = sd.load_clip(ckpt_paths=text_encoder_paths, embedding_directory=os.path.join(sub_dir_path, "embeddings"), clip_type=clip_type_enum)
|
||||
|
||||
# Debugging statement to check if the loaded model has 'tokenize' method
|
||||
if not hasattr(clip_model, 'tokenize'):
|
||||
raise AttributeError("Loaded clip model does not have 'tokenize' method.")
|
||||
|
||||
return clip_model
|
||||
|
||||
# Ensure the node is registered properly
|
||||
NODE_CLASS_MAPPINGS = {"DiffusersClipLoader": DiffusersClipLoader}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"DiffusersClipLoader": "Diffusers CLIP Loader"}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -1,103 +0,0 @@
|
||||
import os
|
||||
import comfy.sd
|
||||
|
||||
def get_model_directories(base_path):
|
||||
"""Returns a list of sub-directories containing models within the given base path."""
|
||||
base_directory = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', base_path))
|
||||
print(f"DiffusersUNETLoader: Looking for model directories in: {base_directory}")
|
||||
|
||||
model_dirs = []
|
||||
if os.path.exists(base_directory):
|
||||
for root, dirs, files in os.walk(base_directory):
|
||||
if "model_index.json" in files:
|
||||
model_dirs.append(os.path.relpath(root, base_directory))
|
||||
print(f"DiffusersUNETLoader: Found model directory: {os.path.relpath(root, base_directory)}")
|
||||
dirs[:] = [d for d in dirs if d not in ["vae", "unet", "text_encoder", "text_encoder_2"]]
|
||||
else:
|
||||
print(f"DiffusersUNETLoader: Path does not exist: {base_directory}")
|
||||
return model_dirs
|
||||
|
||||
def handle_corrupted_file(unet_folder):
|
||||
"""Handle corrupted UNET folder by renaming it and logging the action."""
|
||||
corrupted_path = unet_folder + ".corrupted"
|
||||
try:
|
||||
os.rename(unet_folder, corrupted_path)
|
||||
print(f"DiffusersUNETLoader: Moved corrupted folder to: {corrupted_path}")
|
||||
except PermissionError as e:
|
||||
print(f"DiffusersUNETLoader: PermissionError while moving corrupted folder: {e}")
|
||||
raise PermissionError(f"DiffusersUNETLoader: Permission denied while moving: {unet_folder}")
|
||||
except Exception as e:
|
||||
print(f"DiffusersUNETLoader: Error while moving corrupted folder: {e}")
|
||||
raise e
|
||||
|
||||
class DiffusersUNETLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
base_path = "diffusers"
|
||||
model_directories = get_model_directories(base_path)
|
||||
return {
|
||||
"required": {
|
||||
"sub_directory": (model_directories,),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "load_unet"
|
||||
CATEGORY = "DiffusersLoader"
|
||||
|
||||
def load_unet(self, sub_directory):
|
||||
base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers'))
|
||||
sub_dir_path = os.path.join(base_path, sub_directory)
|
||||
|
||||
model_type = self.detect_model_type(sub_dir_path)
|
||||
|
||||
if model_type == "SD15":
|
||||
unet_folder = os.path.join(sub_dir_path, 'unet')
|
||||
elif model_type == "SDXL":
|
||||
unet_folder = os.path.join(sub_dir_path, 'unet')
|
||||
else:
|
||||
raise ValueError("DiffusersUNETLoader: Unsupported model type")
|
||||
|
||||
diffusion_model_path = self.find_model_file(unet_folder)
|
||||
print(f"DiffusersUNETLoader: Attempting to load UNET model from: {diffusion_model_path}")
|
||||
|
||||
if not os.path.exists(diffusion_model_path):
|
||||
raise FileNotFoundError(f"DiffusersUNETLoader: {os.path.basename(diffusion_model_path)} not found in the directory: {diffusion_model_path}")
|
||||
|
||||
try:
|
||||
model = comfy.sd.load_unet(diffusion_model_path)
|
||||
return (model,)
|
||||
except PermissionError as e:
|
||||
print(f"DiffusersUNETLoader: PermissionError: {e}")
|
||||
handle_corrupted_file(unet_folder)
|
||||
raise e
|
||||
except Exception as e:
|
||||
print(f"DiffusersUNETLoader: Error loading UNET model: {e}")
|
||||
handle_corrupted_file(unet_folder)
|
||||
raise e
|
||||
|
||||
def detect_model_type(self, sub_dir_path):
|
||||
"""Detect the model type based on the directory structure."""
|
||||
text_encoder_dir1 = os.path.join(sub_dir_path, "text_encoder")
|
||||
text_encoder_dir2 = os.path.join(sub_dir_path, "text_encoder_2")
|
||||
|
||||
if os.path.exists(text_encoder_dir1) and os.path.exists(text_encoder_dir2):
|
||||
return "SDXL"
|
||||
elif os.path.exists(text_encoder_dir1):
|
||||
return "SD15"
|
||||
else:
|
||||
raise FileNotFoundError("DiffusersUNETLoader: No valid text_encoder directories found. This model is not SD15 or SDXL")
|
||||
|
||||
@staticmethod
|
||||
def find_model_file(directory):
|
||||
"""Find the .safetensors or .bin file in the given directory."""
|
||||
for file in os.listdir(directory):
|
||||
if file.endswith(".safetensors") or file.endswith(".bin"):
|
||||
return os.path.join(directory, file)
|
||||
raise FileNotFoundError(f"DiffusersUNETLoader: No .safetensors file or .bin file found in {directory}")
|
||||
|
||||
# Ensure the node is registered properly
|
||||
NODE_CLASS_MAPPINGS = {"DiffusersUNETLoader": DiffusersUNETLoader}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"DiffusersUNETLoader": "Diffusers UNET Loader"}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -1,85 +0,0 @@
|
||||
import os
|
||||
import comfy.utils
|
||||
import comfy.sd
|
||||
import torch
|
||||
|
||||
class DiffusersVAELoader:
|
||||
@staticmethod
|
||||
def get_base_path():
|
||||
base_path = os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers')
|
||||
base_path = os.path.abspath(base_path)
|
||||
if not os.path.exists(base_path):
|
||||
raise FileNotFoundError(f"DiffusersVAELoader: The base path '{base_path}' does not exist.")
|
||||
return base_path
|
||||
|
||||
@staticmethod
|
||||
def get_model_directories(base_path):
|
||||
model_dirs = []
|
||||
if os.path.exists(base_path):
|
||||
for root, dirs, files in os.walk(base_path):
|
||||
if "model_index.json" in files:
|
||||
model_dirs.append(os.path.relpath(root, base_path))
|
||||
print(f"DiffusersVAELoader: Found model directory: {os.path.relpath(root, base_path)}")
|
||||
dirs[:] = [d for d in dirs if d not in ["vae", "unet", "text_encoder", "text_encoder_2"]]
|
||||
else:
|
||||
print(f"DiffusersVAELoader: Path does not exist: {base_path}")
|
||||
return model_dirs
|
||||
|
||||
@staticmethod
|
||||
def find_model_file(directory):
|
||||
for file in os.listdir(directory):
|
||||
if file.endswith(".safetensors") or file.endswith(".bin"):
|
||||
return file
|
||||
raise FileNotFoundError(f"DiffusersVAELoader: No .safetensors file or .bin file found in {directory}")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
base_path = cls.get_base_path()
|
||||
model_directories = cls.get_model_directories(base_path)
|
||||
return {
|
||||
"required": {
|
||||
"sub_directory": (model_directories,),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VAE",)
|
||||
FUNCTION = "load_vae"
|
||||
CATEGORY = "DiffusersLoader"
|
||||
|
||||
def load_vae(self, sub_directory):
|
||||
base_path = self.get_base_path()
|
||||
sub_dir_path = os.path.join(base_path, sub_directory)
|
||||
|
||||
model_type = self.detect_model_type(sub_dir_path)
|
||||
|
||||
vae_folder = os.path.join(sub_dir_path, "vae")
|
||||
vae_filename = self.find_model_file(vae_folder)
|
||||
vae_path = os.path.join(vae_folder, vae_filename)
|
||||
full_path = os.path.abspath(vae_path)
|
||||
|
||||
print(f"DiffusersVAELoader: Attempting to load VAE model from: {full_path}")
|
||||
|
||||
if not os.path.exists(full_path):
|
||||
raise FileNotFoundError(f"DiffusersVAELoader: The VAE file '{full_path}' does not exist.")
|
||||
|
||||
sd = comfy.utils.load_torch_file(full_path)
|
||||
vae = comfy.sd.VAE(sd=sd)
|
||||
return (vae,)
|
||||
|
||||
def detect_model_type(self, sub_dir_path):
|
||||
"""Detect the model type based on the directory structure."""
|
||||
text_encoder_dir1 = os.path.join(sub_dir_path, "text_encoder")
|
||||
text_encoder_dir2 = os.path.join(sub_dir_path, "text_encoder_2")
|
||||
|
||||
if os.path.exists(text_encoder_dir1) and os.path.exists(text_encoder_dir2):
|
||||
return "SDXL"
|
||||
elif os.path.exists(text_encoder_dir1):
|
||||
return "SD15"
|
||||
else:
|
||||
raise FileNotFoundError("DiffusersVAELoader: No valid text_encoder directories found. This model is not SD15 or SDXL")
|
||||
|
||||
# Ensure the node is registered properly
|
||||
NODE_CLASS_MAPPINGS = {"DiffusersVAELoader": DiffusersVAELoader}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"DiffusersVAELoader": "Diffusers VAE Loader"}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,16 @@
|
||||
# loader_factory.py
|
||||
from .unet_loader import DiffusersUNETLoader
|
||||
from .clip_loader import DiffusersClipLoader
|
||||
from .vae_loader import DiffusersVAELoader
|
||||
|
||||
class LoaderFactory:
|
||||
@staticmethod
|
||||
def create_loader(loader_type, sub_directory):
|
||||
if loader_type == "UNET":
|
||||
return DiffusersUNETLoader(sub_directory)
|
||||
elif loader_type == "CLIP":
|
||||
return DiffusersClipLoader(sub_directory)
|
||||
elif loader_type == "VAE":
|
||||
return DiffusersVAELoader(sub_directory)
|
||||
else:
|
||||
raise ValueError(f"Unknown loader type: {loader_type}")
|
||||
@@ -0,0 +1,51 @@
|
||||
# unet_loader.py
|
||||
import os
|
||||
import comfy.sd
|
||||
from .base_loader import DiffusersLoaderBase
|
||||
from .utils import DiffusersUtils
|
||||
|
||||
class DiffusersUNETLoader(DiffusersLoaderBase):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "load_unet"
|
||||
CATEGORY = "DiffusersLoader"
|
||||
|
||||
@classmethod
|
||||
def load_unet(cls, sub_directory):
|
||||
return (cls.load_model(sub_directory),)
|
||||
|
||||
@classmethod
|
||||
def load_model(cls, sub_directory):
|
||||
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)
|
||||
|
||||
if model_type == "SD3":
|
||||
unet_folder = os.path.join(sub_dir_path, 'transformer')
|
||||
else:
|
||||
unet_folder = os.path.join(sub_dir_path, 'unet')
|
||||
|
||||
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}")
|
||||
|
||||
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
|
||||
@@ -0,0 +1,62 @@
|
||||
import os
|
||||
import safetensors.torch
|
||||
import torch
|
||||
|
||||
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):
|
||||
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
|
||||
|
||||
@staticmethod
|
||||
def find_model_files(directory, file_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)
|
||||
if f.endswith((".safetensors", ".bin"))]
|
||||
|
||||
if not files:
|
||||
raise FileNotFoundError(f"No .safetensors or .bin file found in {directory}")
|
||||
|
||||
if file_parts:
|
||||
return [file for file in files if any(part in file for part in file_parts)]
|
||||
return files
|
||||
|
||||
@staticmethod
|
||||
def find_model_file(directory):
|
||||
files = DiffusersUtils.find_model_files(directory)
|
||||
if files:
|
||||
return files[0]
|
||||
raise FileNotFoundError(f"No .safetensors or .bin file found in {directory}")
|
||||
|
||||
@staticmethod
|
||||
def load_safetensor_paths(file_paths):
|
||||
tensors = [safetensors.torch.load_file(file) for file in file_paths]
|
||||
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)
|
||||
|
||||
return combined_file_path
|
||||
@@ -0,0 +1,39 @@
|
||||
# vae_loader.py
|
||||
import os
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
from .base_loader import DiffusersLoaderBase
|
||||
from .utils import DiffusersUtils
|
||||
|
||||
class DiffusersVAELoader(DiffusersLoaderBase):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VAE",)
|
||||
FUNCTION = "load_vae"
|
||||
CATEGORY = "DiffusersLoader"
|
||||
|
||||
@classmethod
|
||||
def load_vae(cls, sub_directory):
|
||||
return (cls.load_model(sub_directory),)
|
||||
|
||||
@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)
|
||||
print(f"DiffusersVAELoader: Attempting to load VAE model from: {vae_path}")
|
||||
|
||||
try:
|
||||
vae_sd = comfy.utils.load_torch_file(vae_path)
|
||||
return comfy.sd.VAE(sd=vae_sd)
|
||||
except Exception as e:
|
||||
print(f"DiffusersVAELoader: Error loading VAE model: {e}")
|
||||
raise
|
||||
Reference in New Issue
Block a user