From 4ddc28100521c06f500e85eca814f58c81094580 Mon Sep 17 00:00:00 2001 From: Scorpinaus Date: Thu, 20 Jun 2024 23:37:17 +0800 Subject: [PATCH] Cleaned up repo & Added SD3 support for CombinedDiffusersLoader --- .gitignore.txt | 175 ++++++++++++++++++ .../sd15_clip_loader.py | 0 .../sd15_combined_loader.py | 0 .../sd15_unet_loader.py | 0 sd15_vae_loader.py => Old/sd15_vae_loader.py | 0 .../sdxl_clip_loader.py | 0 .../sdxl_combined_loader.py | 0 .../sdxl_unet_loader.py | 0 sdxl_vae_loader.py => Old/sdxl_vae_loader.py | 0 combined_diffusers_loader.py | 121 +++++++++--- 10 files changed, 271 insertions(+), 25 deletions(-) create mode 100644 .gitignore.txt rename sd15_clip_loader.py => Old/sd15_clip_loader.py (100%) rename sd15_combined_loader.py => Old/sd15_combined_loader.py (100%) rename sd15_unet_loader.py => Old/sd15_unet_loader.py (100%) rename sd15_vae_loader.py => Old/sd15_vae_loader.py (100%) rename sdxl_clip_loader.py => Old/sdxl_clip_loader.py (100%) rename sdxl_combined_loader.py => Old/sdxl_combined_loader.py (100%) rename sdxl_unet_loader.py => Old/sdxl_unet_loader.py (100%) rename sdxl_vae_loader.py => Old/sdxl_vae_loader.py (100%) diff --git a/.gitignore.txt b/.gitignore.txt new file mode 100644 index 0000000..071aaa0 --- /dev/null +++ b/.gitignore.txt @@ -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/ \ No newline at end of file diff --git a/sd15_clip_loader.py b/Old/sd15_clip_loader.py similarity index 100% rename from sd15_clip_loader.py rename to Old/sd15_clip_loader.py diff --git a/sd15_combined_loader.py b/Old/sd15_combined_loader.py similarity index 100% rename from sd15_combined_loader.py rename to Old/sd15_combined_loader.py diff --git a/sd15_unet_loader.py b/Old/sd15_unet_loader.py similarity index 100% rename from sd15_unet_loader.py rename to Old/sd15_unet_loader.py diff --git a/sd15_vae_loader.py b/Old/sd15_vae_loader.py similarity index 100% rename from sd15_vae_loader.py rename to Old/sd15_vae_loader.py diff --git a/sdxl_clip_loader.py b/Old/sdxl_clip_loader.py similarity index 100% rename from sdxl_clip_loader.py rename to Old/sdxl_clip_loader.py diff --git a/sdxl_combined_loader.py b/Old/sdxl_combined_loader.py similarity index 100% rename from sdxl_combined_loader.py rename to Old/sdxl_combined_loader.py diff --git a/sdxl_unet_loader.py b/Old/sdxl_unet_loader.py similarity index 100% rename from sdxl_unet_loader.py rename to Old/sdxl_unet_loader.py diff --git a/sdxl_vae_loader.py b/Old/sdxl_vae_loader.py similarity index 100% rename from sdxl_vae_loader.py rename to Old/sdxl_vae_loader.py diff --git a/combined_diffusers_loader.py b/combined_diffusers_loader.py index 5b434db..5b20b8c 100644 --- a/combined_diffusers_loader.py +++ b/combined_diffusers_loader.py @@ -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)