Troubleshoot file_path and sub_dir_path debug

This commit is contained in:
Scorpinaus
2024-06-08 22:16:01 +08:00
parent 5dc37d2159
commit d8e4b5320c
5 changed files with 295 additions and 3 deletions
+11 -2
View File
@@ -7,6 +7,9 @@ from .sdxl_clip_loader import SDXLCLIPLoader
from .sd15_combined_loader import CombinedDiffusersSD15Loader
from .sdxl_combined_loader import CombinedDiffusersSDXLLoader
from .combined_diffusers_loader import CombinedDiffusersLoader
from .diffusers_clip_loader import DiffusersClipLoader
from .diffusers_unet_loader import DiffusersUNETLoader
from .diffusers_vae_loader import DiffusersVAELoader
NODE_CLASS_MAPPINGS = {
"SD15UNETLoader": SD15UNETLoader,
@@ -17,7 +20,10 @@ NODE_CLASS_MAPPINGS = {
"SDXLCLIPLoader": SDXLCLIPLoader,
"CombinedDiffusersSD15Loader": CombinedDiffusersSD15Loader,
"CombinedDiffusersSDXLLoader": CombinedDiffusersSDXLLoader,
"CombinedDiffusersLoader": CombinedDiffusersLoader
"CombinedDiffusersLoader": CombinedDiffusersLoader,
"DiffusersClipLoader": DiffusersClipLoader,
"DiffusersUNETLoader": DiffusersUNETLoader,
"DiffusersVAELoader": DiffusersVAELoader
}
@@ -30,5 +36,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"SDXLCLIPLoader": "SDXL CLIP Loader",
"CombinedDiffusersSD15Loader" : "Combined Diffusers SD15 Loader",
"CombinedDiffusersSDXLLoader" : "Combined Diffusers SDXL Loader",
"CombinedDiffusersLoader": "Combined Diffusers Loader"
"CombinedDiffusersLoader": "Combined Diffusers Loader",
"DiffusersClipLoader": "Diffusers CLIP Loader",
"DiffusersUNETLoader": "Diffusers UNET Loader",
"DiffusersVAELoader": "Diffusers VAE Loader"
}
+3 -1
View File
@@ -33,6 +33,7 @@ class CombinedDiffusersLoader:
def INPUT_TYPES(cls):
base_path = cls.get_base_path()
model_directories = cls.get_model_directories(base_path)
print(f"CombinedDiffusersLoader: model_directories: \n{model_directories}")
return {
"required": {
"sub_directory": (model_directories,),
@@ -46,10 +47,11 @@ class CombinedDiffusersLoader:
def load_models(self, sub_directory, clip_type="stable_diffusion"):
base_path = self.get_base_path()
print(f"CombinedDiffusersLoader: base_path: \n{base_path}")
sub_dir_path = os.path.join(base_path, sub_directory)
model_type = self.detect_model_type(sub_dir_path)
print(f"CombinedDiffusersLoader: Sub_dir_path: \n{sub_dir_path}")
vae_model = self.load_vae(sub_dir_path)
unet_model = self.load_unet(sub_dir_path)
clip_model = self.load_clip(sub_dir_path, clip_type, model_type)
+93
View File
@@ -0,0 +1,93 @@
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"]
+103
View File
@@ -0,0 +1,103 @@
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"]
+85
View File
@@ -0,0 +1,85 @@
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"]