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

Added loading support for .bin format files

Browse Source
This commit is contained in:
Scorpinaus
2024-06-05 08:34:08 +08:00
parent c017edd120
commit 7739e26406
8 changed files with 38 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sd15_clip_loader.py
+5 -5
View File
@@ -36,8 +36,8 @@ class SD15CLIPLoader:
text_encoder_dir = os.path.join(script_dir, "..", "..", "models", "diffusers", "SD15", sub_directory, "text_encoder")
print(f"SD15CLIPLoader:Text encoder directory: {text_encoder_dir}") # Debug print
# Find the .safetensors file in the text_encoder directory
text_encoder_file = self.find_safetensors_file(text_encoder_dir)
# Find the .safetensors or .bin file in the text_encoder directory
text_encoder_file = self.find_model_file(text_encoder_dir)
text_encoder_path = os.path.join(text_encoder_dir, text_encoder_file)
print(f"SD15CLIPLoader:Text encoder path: {text_encoder_path}") # Debug print
@@ -50,11 +50,11 @@ class SD15CLIPLoader:
return (clip,)
@staticmethod
def find_safetensors_file(directory):
def find_model_file(directory):
for file in os.listdir(directory):
if file.endswith(".safetensors"):
if file.endswith(".safetensors") or file.endswith(".bin"):
return file
raise FileNotFoundError(f"SD15CLIPLoader:No .safetensors file found in {directory}")
raise FileNotFoundError(f"SD15CLIPLoader:No .safetensors file or .bin file found in {directory}")
# Register the node in the ComfyUI framework
NODE_CLASS_MAPPINGS = {
sd15_combined_loader.py
+6 -6
View File
@@ -33,7 +33,7 @@ class CombinedDiffusersSD15Loader:
script_dir = os.path.dirname(os.path.realpath(__file__))
text_encoder_dir = os.path.join(script_dir, "..", "..", "models", "diffusers", "SD15", sub_directory, "text_encoder")
text_encoder_file = self.find_safetensors_file(text_encoder_dir)
text_encoder_file = self.find_model_file(text_encoder_dir)
text_encoder_path = os.path.join(text_encoder_dir, text_encoder_file)
if not os.path.exists(text_encoder_path):
@@ -44,7 +44,7 @@ class CombinedDiffusersSD15Loader:
def load_unet(self, sub_directory):
unet_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers', 'SD15', sub_directory, 'unet'))
diffusion_model_path = self.find_safetensors_file(unet_folder)
diffusion_model_path = self.find_model_file(unet_folder)
if not os.path.exists(diffusion_model_path):
raise FileNotFoundError(f"{os.path.basename(diffusion_model_path)} not found in the directory: {diffusion_model_path}")
@@ -59,7 +59,7 @@ class CombinedDiffusersSD15Loader:
def load_vae(self, sub_directory):
base_path = self.get_base_path()
vae_folder = os.path.join(base_path, sub_directory, "vae")
vae_filename = self.find_safetensors_file(vae_folder)
vae_filename = self.find_model_file(vae_folder)
vae_path = os.path.join(vae_folder, vae_filename)
full_path = os.path.abspath(vae_path)
@@ -81,11 +81,11 @@ class CombinedDiffusersSD15Loader:
return base_path
@staticmethod
def find_safetensors_file(directory):
def find_model_file(directory):
for file in os.listdir(directory):
if file.endswith(".safetensors"):
if file.endswith(".safetensors") or file.endswith(".bin"):
return os.path.join(directory, file)
raise FileNotFoundError(f"No .safetensors file found in {directory}")
raise FileNotFoundError(f"No .safetensors file or . bin file found in {directory}")
@staticmethod
def handle_corrupted_file(unet_folder):
sd15_unet_loader.py
+4 -4
View File
@@ -42,7 +42,7 @@ class SD15UNETLoader:
def load_unet(self, sub_directory):
unet_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers', 'SD15', sub_directory, 'unet'))
diffusion_model_path = self.find_safetensors_file(unet_folder)
diffusion_model_path = self.find_model_file(unet_folder)
print(f"SD15UNETLoader:Attempting to load UNET model from: {diffusion_model_path}")
if not os.path.exists(diffusion_model_path):
@@ -60,11 +60,11 @@ class SD15UNETLoader:
raise e
@staticmethod
def find_safetensors_file(directory):
def find_model_file(directory):
for file in os.listdir(directory):
if file.endswith(".safetensors"):
if file.endswith(".safetensors") or file.endswith(".bin"):
return os.path.join(directory, file)
raise FileNotFoundError(f"SD15UNETLoader:No .safetensors file found in {directory}")
raise FileNotFoundError(f"SD15UNETLoader:No .safetensors file or .bin file found in {directory}")
# Ensure the node is registered properly
NODE_CLASS_MAPPINGS = {"SD15UNETLoader": SD15UNETLoader}
sd15_vae_loader.py
+4 -4
View File
@@ -18,11 +18,11 @@ class SD15VAELoader:
return [name for name in os.listdir(base_path) if os.path.isdir(os.path.join(base_path, name))]
@staticmethod
def find_safetensors_file(directory):
def find_model_file(directory):
for file in os.listdir(directory):
if file.endswith(".safetensors"):
if file.endswith(".safetensors") or file.endswith(".bin"):
return file
raise FileNotFoundError(f"No .safetensors file found in {directory}")
raise FileNotFoundError(f"No .safetensors file or .bin file found in {directory}")
@classmethod
def INPUT_TYPES(cls):
@@ -40,7 +40,7 @@ class SD15VAELoader:
def load_vae(self, sub_directory):
base_path = self.get_base_path()
vae_folder = os.path.join(base_path, sub_directory, "vae")
vae_filename = self.find_safetensors_file(vae_folder)
vae_filename = self.find_model_file(vae_folder)
vae_path = os.path.join(vae_folder, vae_filename)
full_path = os.path.abspath(vae_path)
sdxl_clip_loader.py
+4 -4
View File
@@ -30,8 +30,8 @@ class SDXLCLIPLoader:
text_encoder_dir1 = os.path.join(base_path, "text_encoder")
text_encoder_dir2 = os.path.join(base_path, "text_encoder_2")
text_encoder_file1 = self.find_safetensors_file(text_encoder_dir1)
text_encoder_file2 = self.find_safetensors_file(text_encoder_dir2)
text_encoder_file1 = self.find_model_file(text_encoder_dir1)
text_encoder_file2 = self.find_model_file(text_encoder_dir2)
text_encoder_path1 = os.path.join(text_encoder_dir1, text_encoder_file1)
text_encoder_path2 = os.path.join(text_encoder_dir2, text_encoder_file2)
@@ -51,9 +51,9 @@ class SDXLCLIPLoader:
return (clip,)
@staticmethod
def find_safetensors_file(directory):
def find_model_file(directory):
for file in os.listdir(directory):
if file.endswith(".safetensors"):
if file.endswith(".safetensors") or file.endswith(".bin"):
return file
raise FileNotFoundError(f"No .safetensors file found in {directory}")
sdxl_combined_loader.py
+7 -7
View File
@@ -17,11 +17,11 @@ class CombinedDiffusersSDXLLoader:
return [name for name in os.listdir(base_path) if os.path.isdir(os.path.join(base_path, name))]
@staticmethod
def find_safetensors_file(directory):
def find_model_file(directory):
for file in os.listdir(directory):
if file.endswith(".safetensors"):
if file.endswith(".safetensors") or file.endswith(".bin"):
return os.path.join(directory, file)
raise FileNotFoundError(f"CombinedDiffusersSDXLLoader: No .safetensors file found in {directory}")
raise FileNotFoundError(f"CombinedDiffusersSDXLLoader: No .safetensors file or .bin file found in {directory}")
@classmethod
def INPUT_TYPES(cls):
@@ -48,7 +48,7 @@ class CombinedDiffusersSDXLLoader:
def load_vae(self, base_path):
vae_folder = os.path.join(base_path, "vae")
vae_path = self.find_safetensors_file(vae_folder)
vae_path = self.find_model_file(vae_folder)
print(f"CombinedDiffusersSDXLLoader: Attempting to load VAE model from: {vae_path}")
vae_sd = comfy.utils.load_torch_file(vae_path)
@@ -56,7 +56,7 @@ class CombinedDiffusersSDXLLoader:
def load_unet(self, base_path):
unet_folder = os.path.join(base_path, 'unet')
unet_path = self.find_safetensors_file(unet_folder)
unet_path = self.find_model_file(unet_folder)
print(f"CombinedDiffusersSDXLLoader: Attempting to load UNET model from: {unet_path}")
try:
@@ -73,8 +73,8 @@ class CombinedDiffusersSDXLLoader:
text_encoder_dir1 = os.path.join(base_path, "text_encoder")
text_encoder_dir2 = os.path.join(base_path, "text_encoder_2")
text_encoder_path1 = self.find_safetensors_file(text_encoder_dir1)
text_encoder_path2 = self.find_safetensors_file(text_encoder_dir2)
text_encoder_path1 = self.find_model_file(text_encoder_dir1)
text_encoder_path2 = self.find_model_file(text_encoder_dir2)
print(f"CombinedDiffusersSDXLLoader: Checking paths:\n{text_encoder_path1}\n{text_encoder_path2}")
sdxl_unet_loader.py
+4 -4
View File
@@ -42,7 +42,7 @@ class SDXLUNETLoader:
def load_unet(self, sub_directory):
unet_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers', 'SDXL', sub_directory, 'unet'))
diffusion_model_path = self.find_safetensors_file(unet_folder)
diffusion_model_path = self.find_model_file(unet_folder)
print(f"SDXLUNETLoader:Attempting to load UNET model from: {diffusion_model_path}")
if not os.path.exists(diffusion_model_path):
@@ -60,11 +60,11 @@ class SDXLUNETLoader:
raise e
@staticmethod
def find_safetensors_file(directory):
def find_model_file(directory):
for file in os.listdir(directory):
if file.endswith(".safetensors"):
if file.endswith(".safetensors") or file.endswith(".bin"):
return os.path.join(directory, file)
raise FileNotFoundError(f"SDXLUNETLoader:No .safetensors file found in {directory}")
raise FileNotFoundError(f"SDXLUNETLoader:No .safetensors file or .bin file found in {directory}")
# Mapping the node to make it recognizable in the ComfyUI framework
NODE_CLASS_MAPPINGS = {
sdxl_vae_loader.py
+4 -4
View File
@@ -18,11 +18,11 @@ class SDXLVAELoader:
return [name for name in os.listdir(base_path) if os.path.isdir(os.path.join(base_path, name))]
@staticmethod
def find_safetensors_file(directory):
def find_model_file(directory):
for file in os.listdir(directory):
if file.endswith(".safetensors"):
if file.endswith(".safetensors") or file.endswith(".bin"):
return file
raise FileNotFoundError(f"SDXLVAELoader:No .safetensors file found in {directory}")
raise FileNotFoundError(f"SDXLVAELoader:No .safetensors file or .bin file found in {directory}")
@classmethod
def INPUT_TYPES(cls):
@@ -40,7 +40,7 @@ class SDXLVAELoader:
def load_vae(self, sub_directory):
base_path = self.get_base_path()
vae_folder = os.path.join(base_path, sub_directory, "vae")
vae_filename = self.find_safetensors_file(vae_folder)
vae_filename = self.find_model_file(vae_folder)
vae_path = os.path.join(vae_folder, vae_filename)
full_path = os.path.abspath(vae_path)