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

Cleaned up repo & Added SD3 support for CombinedDiffusersLoader

Browse Source
This commit is contained in:
Scorpinaus
2024-06-20 23:37:17 +08:00
parent 651cb8dbf5
commit 4ddc281005
10 changed files with 271 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore.txt
+175
View File
@@ -0,0 +1,175 @@
#Files to ignore
sd15_clip_loader.py
sd15_combined_loader.py
sd15_unet_loader.py
sd15_vae_loader.py
sdxl_clip_loader.py
sdxl_combined_loader.py
sdxl_unet_loader.py
sdxl_vae_loader.py
#Ignore Old directory
Old/
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# poetry
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
#poetry.lock
# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control.
# https://pdm.fming.dev/latest/usage/project/#working-with-version-control
.pdm.toml
.pdm-python
.pdm-build/
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/
# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
#.idea/
sd15_clip_loader.py → Old/sd15_clip_loader.py
View File
sd15_combined_loader.py → Old/sd15_combined_loader.py
View File
sd15_unet_loader.py → Old/sd15_unet_loader.py
View File
sd15_vae_loader.py → Old/sd15_vae_loader.py
View File
sdxl_clip_loader.py → Old/sdxl_clip_loader.py
View File
sdxl_combined_loader.py → Old/sdxl_combined_loader.py
View File
sdxl_unet_loader.py → Old/sdxl_unet_loader.py
View File
sdxl_vae_loader.py → Old/sdxl_vae_loader.py
View File
combined_diffusers_loader.py
+96 -25
View File
@@ -2,6 +2,7 @@ import os
import comfy.sd
import comfy.utils
import torch
import safetensors.torch
class CombinedDiffusersLoader:
@staticmethod
@@ -18,16 +19,52 @@ class CombinedDiffusersLoader:
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"]]
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_file(directory):
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"):
return os.path.join(directory, file)
raise FileNotFoundError(f"CombinedDiffusersLoader: No .safetensors file or .bin file found in {directory}")
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):
@@ -36,7 +73,8 @@ class CombinedDiffusersLoader:
return {
"required": {
"sub_directory": (model_directories,),
"clip_type": (["stable_diffusion", "stable_cascade"],)
"clip_type": (["stable_diffusion", "stable_cascade"],),
"file_parts": (["none", "all", "part_1", "part_2",],)
}
}
@@ -44,45 +82,68 @@ class CombinedDiffusersLoader:
FUNCTION = "load_models"
CATEGORY = "DiffusersLoader/Combined"
def load_models(self, sub_directory, clip_type="stable_diffusion"):
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)
unet_model = self.load_unet(sub_dir_path)
clip_model = self.load_clip(sub_dir_path, clip_type, model_type)
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)
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")
if os.path.exists(text_encoder_dir1) and os.path.exists(text_encoder_dir2):
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")
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"):
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 = [self.find_model_file(text_encoder_dir1)]
text_encoder_paths = [CombinedDiffusersLoader.find_model_file(text_encoder_dir1)]
if model_type == "SDXL":
if model_type in ["SDXL", "SD3"]:
text_encoder_dir2 = os.path.join(base_path, "text_encoder_2")
text_encoder_paths.append(self.find_model_file(text_encoder_dir2))
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}")
clip_model = comfy.sd.load_clip(ckpt_paths=text_encoder_paths, embedding_directory=os.path.join(base_path, "embeddings"), clip_type=clip_type_enum)
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'):
@@ -90,9 +151,16 @@ class CombinedDiffusersLoader:
return clip_model
def load_unet(self, base_path):
unet_folder = os.path.join(base_path, 'unet')
unet_path = self.find_model_file(unet_folder)
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:
@@ -100,13 +168,16 @@ class CombinedDiffusersLoader:
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}")
print(f"Error loading UNET model: {e}")
raise e
def load_vae(self, base_path):
def load_vae(self, base_path, model_type):
vae_folder = os.path.join(base_path, "vae")
vae_path = self.find_model_file(vae_folder)
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)