From 7739e26406467bf345636443d9df470d38dd947e Mon Sep 17 00:00:00 2001 From: Scorpinaus Date: Wed, 5 Jun 2024 08:34:08 +0800 Subject: [PATCH] Added loading support for .bin format files --- sd15_clip_loader.py | 10 +++++----- sd15_combined_loader.py | 12 ++++++------ sd15_unet_loader.py | 8 ++++---- sd15_vae_loader.py | 8 ++++---- sdxl_clip_loader.py | 8 ++++---- sdxl_combined_loader.py | 14 +++++++------- sdxl_unet_loader.py | 8 ++++---- sdxl_vae_loader.py | 8 ++++---- 8 files changed, 38 insertions(+), 38 deletions(-) diff --git a/sd15_clip_loader.py b/sd15_clip_loader.py index 3323d3b..ff6b231 100644 --- a/sd15_clip_loader.py +++ b/sd15_clip_loader.py @@ -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 = { diff --git a/sd15_combined_loader.py b/sd15_combined_loader.py index 2ae6fc4..5f160fc 100644 --- a/sd15_combined_loader.py +++ b/sd15_combined_loader.py @@ -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): diff --git a/sd15_unet_loader.py b/sd15_unet_loader.py index e203776..688eb66 100644 --- a/sd15_unet_loader.py +++ b/sd15_unet_loader.py @@ -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} diff --git a/sd15_vae_loader.py b/sd15_vae_loader.py index f1ea65d..c9433ec 100644 --- a/sd15_vae_loader.py +++ b/sd15_vae_loader.py @@ -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) diff --git a/sdxl_clip_loader.py b/sdxl_clip_loader.py index 4c5ff08..07ae938 100644 --- a/sdxl_clip_loader.py +++ b/sdxl_clip_loader.py @@ -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}") diff --git a/sdxl_combined_loader.py b/sdxl_combined_loader.py index acfb50a..a0307a1 100644 --- a/sdxl_combined_loader.py +++ b/sdxl_combined_loader.py @@ -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}") diff --git a/sdxl_unet_loader.py b/sdxl_unet_loader.py index f96a97e..4b9f0fb 100644 --- a/sdxl_unet_loader.py +++ b/sdxl_unet_loader.py @@ -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 = { diff --git a/sdxl_vae_loader.py b/sdxl_vae_loader.py index e25e283..283b868 100644 --- a/sdxl_vae_loader.py +++ b/sdxl_vae_loader.py @@ -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)