From e701d6a9d169d27d3e9a58dd5f28c5b436b8e29f Mon Sep 17 00:00:00 2001 From: Scorpinaus Date: Wed, 21 Aug 2024 18:10:45 +0800 Subject: [PATCH] Updated VAE_Loader to allow user to select tae models from vae_approx folder --- base_loader.py | 36 +++++++------ clip_loader.py | 67 ++++++++++++++++++------- combined_diffusers_loader.py | 9 ++-- unet_loader.py | 97 ++++++++++++++++++++++++++++-------- vae_loader.py | 61 ++++++++++++++++++++--- 5 files changed, 201 insertions(+), 69 deletions(-) diff --git a/base_loader.py b/base_loader.py index c545f57..d4570e5 100644 --- a/base_loader.py +++ b/base_loader.py @@ -1,25 +1,31 @@ -# base_loader.py import os +import json from .utils import DiffusersUtils class DiffusersLoaderBase: @classmethod def detect_model_type(cls, 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") - text_encoder_dir3 = os.path.join(sub_dir_path, "text_encoder_3") - transformer = os.path.join(sub_dir_path, "transformer") + model_index_path = os.path.join(sub_dir_path, "model_index.json") - if os.path.exists(text_encoder_dir3) and os.path.exists(transformer): - return "SD3" - elif os.path.exists(text_encoder_dir1) and os.path.exists(text_encoder_dir2): - return "SDXL" - elif os.path.exists(text_encoder_dir1) and os.path.exists(transformer): - return "AuraFlow" - elif os.path.exists(text_encoder_dir1): - return "SD15" - else: - raise FileNotFoundError("No valid text_encoder directories found. This model is not SD15, SDXL, or SD3") + if os.path.exists(model_index_path): + with open(model_index_path, 'r') as f: + model_info = json.load(f) + + class_name = model_info.get("_class_name", "") + + if class_name == "StableDiffusionPipeline": + if os.path.exists(os.path.join(sub_dir_path, "text_encoder_2")): + return "SDXL" + else: + return "SD15" + elif class_name == "StableDiffusion3Pipeline": + return "SD3" + elif class_name == "AuraFlowPipeline": + return "AuraFlow" + elif class_name == "FluxPipeline": + return "Flux" + else: + return "Unknown" @classmethod def load_model(cls, sub_directory): diff --git a/clip_loader.py b/clip_loader.py index a9b4e91..a7da94e 100644 --- a/clip_loader.py +++ b/clip_loader.py @@ -25,7 +25,6 @@ class DiffusersClipLoader(DiffusersLoaderBase): @classmethod def load_model(cls, sub_directory, clip_type="stable_diffusion", clip_parts="all"): - base_path = DiffusersUtils.get_base_path() sub_dir_path = os.path.join(base_path, sub_directory) model_type = cls.detect_model_type(sub_dir_path) @@ -37,23 +36,18 @@ class DiffusersClipLoader(DiffusersLoaderBase): text_encoder_dir1 = os.path.join(sub_dir_path, "text_encoder") text_encoder_paths = [DiffusersUtils.find_model_file(text_encoder_dir1)] - if model_type in ["SDXL", "SD3"]: - text_encoder_dir2 = os.path.join(sub_dir_path, "text_encoder_2") - text_encoder_paths.append(DiffusersUtils.find_model_file(text_encoder_dir2)) - - if model_type == "SD3": - text_encoder_dir3 = os.path.join(sub_dir_path, "text_encoder_3") - print("clip_parts Selected:", clip_parts) - if clip_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 = DiffusersUtils.combine_safetensor_files(text_encoder_dir3, base_path) - text_encoder_paths.append(combined_file_path) - elif clip_parts == "part_1": - text_encoder_paths += DiffusersUtils.find_model_files(text_encoder_dir3, num_parts=1) - elif clip_parts == "part_2": - text_encoder_paths += DiffusersUtils.find_model_files(text_encoder_dir3, num_parts=2)[1:] - + if model_type == "SDXL": + cls.handle_sdxl_clip(sub_dir_path, text_encoder_paths) + elif model_type == "SD3": + cls.handle_sd3_clip(sub_dir_path, text_encoder_paths, clip_parts) + elif model_type == "Flux": + cls.handle_flux_clip(sub_dir_path, text_encoder_paths, clip_parts) + elif model_type == "AuraFlow": + cls.handle_auraflow_clip(sub_dir_path, text_encoder_paths) + else: # SD15 and other single text encoder models + text_encoder_dir = os.path.join(sub_dir_path, "text_encoder") + text_encoder_paths.append(DiffusersUtils.find_model_file(text_encoder_dir)) + for path in text_encoder_paths: DiffusersUtils.check_and_clear_cache('clip', path) @@ -68,4 +62,39 @@ class DiffusersClipLoader(DiffusersLoaderBase): if not hasattr(clip_model, 'tokenize'): raise AttributeError("DiffusersClipLoader: Loaded clip model does not have 'tokenize' method.") - return clip_model \ No newline at end of file + return clip_model + + @classmethod + def handle_sdxl_clip(cls, sub_dir_path, text_encoder_paths): + text_encoder_dir2 = os.path.join(sub_dir_path, "text_encoder_2") + text_encoder_paths.append(DiffusersUtils.find_model_file(text_encoder_dir2)) + + @classmethod + def handle_sd3_clip(cls, sub_dir_path, text_encoder_paths, clip_parts): + cls.handle_sdxl_clip(sub_dir_path, text_encoder_paths) + text_encoder_dir3 = os.path.join(sub_dir_path, "text_encoder_3") + cls.handle_multi_part_clip(text_encoder_dir3, text_encoder_paths, clip_parts) + + @classmethod + def handle_flux_clip(cls, sub_dir_path, text_encoder_paths, clip_parts): + cls.handle_sdxl_clip(sub_dir_path, text_encoder_paths) + text_encoder_dir2 = os.path.join(sub_dir_path, "text_encoder_2") + cls.handle_multi_part_clip(text_encoder_dir2, text_encoder_paths, clip_parts) + + @classmethod + def handle_auraflow_clip(cls, sub_dir_path, text_encoder_paths): + text_encoder_dir = os.path.join(sub_dir_path, "text_encoder") + text_encoder_paths.append(DiffusersUtils.find_model_file(text_encoder_dir)) + print(f"AuraFlow CLIP model path: {text_encoder_paths[-1]}") + + @classmethod + def handle_multi_part_clip(cls, encoder_dir, text_encoder_paths, clip_parts): + if clip_parts == "all": + combined_file_path = os.path.join(encoder_dir, "combined_text_encoder.safetensors") + if not os.path.exists(combined_file_path): + combined_file_path = DiffusersUtils.combine_safetensor_files(encoder_dir, os.path.dirname(encoder_dir)) + text_encoder_paths.append(combined_file_path) + elif clip_parts == "part_1": + text_encoder_paths += DiffusersUtils.find_model_files(encoder_dir, num_parts=1) + elif clip_parts == "part_2": + text_encoder_paths += DiffusersUtils.find_model_files(encoder_dir, num_parts=2)[1:] \ No newline at end of file diff --git a/combined_diffusers_loader.py b/combined_diffusers_loader.py index a804400..d21c877 100644 --- a/combined_diffusers_loader.py +++ b/combined_diffusers_loader.py @@ -15,6 +15,7 @@ class CombinedDiffusersLoader: "clip_type": (["stable_diffusion", "stable_cascade"],), "clip_parts": (["all", "part_1", "part_2"],), "transformer_parts": (["all", "part_1", "part_2", "part_3"],), + "vae_type": (["default", "taesd", "taesdxl", "taesd3", "taef1"],) } } @@ -23,15 +24,11 @@ class CombinedDiffusersLoader: CATEGORY = "DiffusersLoader/Combined" @classmethod - def load_models(cls, sub_directory, clip_type="stable_diffusion", clip_parts="all", transformer_parts="all"): + def load_models(cls, sub_directory, clip_type="stable_diffusion", clip_parts="all", transformer_parts="all", vae_type="default"): unet_model = DiffusersUNETLoader.load_model(sub_directory, transformer_parts) - - clip_model = DiffusersClipLoader.load_model(sub_directory, clip_type, clip_parts) - - - vae_model = DiffusersVAELoader.load_model(sub_directory) + vae_model = DiffusersVAELoader.load_model(sub_directory, vae_type) return unet_model, clip_model, vae_model diff --git a/unet_loader.py b/unet_loader.py index 46b96a3..e8d020c 100644 --- a/unet_loader.py +++ b/unet_loader.py @@ -31,29 +31,19 @@ class DiffusersUNETLoader(DiffusersLoaderBase): model_type = cls.detect_model_type(sub_dir_path) print(f"DiffusersUNETLoader: Detected model type: {model_type}") - if model_type in ["AuraFlow"]: + if model_type == "AuraFlow": unet_folder = os.path.join(sub_dir_path, 'transformer') - print(f"Using transformer folder: {unet_folder}") - if transformer_parts == "all": - combined_file_path = os.path.join(unet_folder, "combined_transformer.safetensors") - if not os.path.exists(combined_file_path): - print(f"Combined file not found, creating at: {combined_file_path}") - try: - combined_file_path = DiffusersUtils.combine_safetensor_files(unet_folder, base_path, num_parts=3) - except Exception as e: - print(f"Error combining safetensor files: {e}") - # If combining fails, use the first part file - combined_file_path = DiffusersUtils.find_model_files(unet_folder, num_parts=1)[0] - unet_path = combined_file_path - else: - part_num = int(transformer_parts.split('_')[1]) - unet_path = DiffusersUtils.find_model_files(unet_folder, num_parts=part_num)[-1] - print(f"Selected UNET/Transformer path: {unet_path}") - elif model_type in ["SD3"]: + print(f"Using transformer folder for AuraFlow: {unet_folder}") + unet_path = cls.handle_auraflow_transformer(unet_folder, transformer_parts) + elif model_type == "Flux": unet_folder = os.path.join(sub_dir_path, 'transformer') - print(f"Using transformer folder: {unet_folder}") - unet_path = DiffusersUtils.find_model_file(unet_folder) - else: + print(f"Using transformer folder for Flux: {unet_folder}") + unet_path = cls.handle_flux_transformer(unet_folder, transformer_parts) + elif model_type == "SD3": + unet_folder = os.path.join(sub_dir_path, 'transformer') + print(f"Using transformer folder for SD3: {unet_folder}") + unet_path = cls.handle_sd3_transformer(unet_folder, transformer_parts) + else: # SD15, SDXL, etc. unet_folder = os.path.join(sub_dir_path, 'unet') unet_path = DiffusersUtils.find_model_file(unet_folder) @@ -78,4 +68,67 @@ class DiffusersUNETLoader(DiffusersLoaderBase): raise else: print(f"DiffusersUNETLoader: Error loading UNET/Transformer model: {e}") - raise \ No newline at end of file + raise + + @classmethod + def handle_auraflow_transformer(cls, unet_folder, transformer_parts): + print(f"Handling AuraFlow transformer in folder: {unet_folder}") + if transformer_parts == "all": + combined_file_path = os.path.join(unet_folder, "combined_transformer.safetensors") + if not os.path.exists(combined_file_path): + print(f"Combined file not found, creating at: {combined_file_path}") + try: + combined_file_path = DiffusersUtils.combine_safetensor_files(unet_folder, os.path.dirname(unet_folder), num_parts=3) + except Exception as e: + print(f"Error combining safetensor files: {e}") + # If combining fails, use the first part file + combined_file_path = DiffusersUtils.find_model_files(unet_folder, num_parts=1)[0] + unet_path = combined_file_path + else: + part_num = int(transformer_parts.split('_')[1]) + unet_path = DiffusersUtils.find_model_files(unet_folder, num_parts=part_num)[-1] + + print(f"Selected AuraFlow transformer path: {unet_path}") + return unet_path + + @classmethod + def handle_sd3_transformer(cls, unet_folder, transformer_parts): + print(f"Handling SD3 transformer in folder: {unet_folder}") + if transformer_parts == "all": + combined_file_path = os.path.join(unet_folder, "combined_transformer.safetensors") + if not os.path.exists(combined_file_path): + print(f"Combined file not found, creating at: {combined_file_path}") + try: + combined_file_path = DiffusersUtils.combine_safetensor_files(unet_folder, os.path.dirname(unet_folder), num_parts=3) + except Exception as e: + print(f"Error combining safetensor files: {e}") + # If combining fails, use the first part file + combined_file_path = DiffusersUtils.find_model_files(unet_folder, num_parts=1)[0] + unet_path = combined_file_path + else: + part_num = int(transformer_parts.split('_')[1]) + unet_path = DiffusersUtils.find_model_files(unet_folder, num_parts=part_num)[-1] + + print(f"Selected SD3 transformer path: {unet_path}") + return unet_path + + @classmethod + def handle_flux_transformer(cls, unet_folder, transformer_parts): + print(f"Handling Flux transformer in folder: {unet_folder}") + if transformer_parts == "all": + combined_file_path = os.path.join(unet_folder, "combined_transformer.safetensors") + if not os.path.exists(combined_file_path): + print(f"Combined file not found, creating at: {combined_file_path}") + try: + combined_file_path = DiffusersUtils.combine_safetensor_files(unet_folder, os.path.dirname(unet_folder), num_parts=3) + except Exception as e: + print(f"Error combining safetensor files: {e}") + # If combining fails, use the first part file + combined_file_path = DiffusersUtils.find_model_files(unet_folder, num_parts=1)[0] + unet_path = combined_file_path + else: + part_num = int(transformer_parts.split('_')[1]) + unet_path = DiffusersUtils.find_model_files(unet_folder, num_parts=part_num)[-1] + + print(f"Selected Flux transformer path: {unet_path}") + return unet_path \ No newline at end of file diff --git a/vae_loader.py b/vae_loader.py index 21afedc..e1998b0 100644 --- a/vae_loader.py +++ b/vae_loader.py @@ -1,16 +1,25 @@ # vae_loader.py import os +import torch import comfy.sd import comfy.utils from .base_loader import DiffusersLoaderBase from .utils import DiffusersUtils class DiffusersVAELoader(DiffusersLoaderBase): + # Moving the VAE_Configs to a Json file might be for a future refactor! + VAE_CONFIGS = { + "taesd": {"scale": 0.18215, "shift": 0.0}, + "taesdxl": {"scale": 0.13025, "shift": 0.0}, + "taesd3": {"scale": 1.5305, "shift": 0.0609}, + "taef1": {"scale": 0.3611, "shift": 0.1159}, + } @classmethod def INPUT_TYPES(cls): return { "required": { "sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),), + "vae_type": (["default"] + list(cls.VAE_CONFIGS.keys()),), } } @@ -19,18 +28,22 @@ class DiffusersVAELoader(DiffusersLoaderBase): CATEGORY = "DiffusersLoader" @classmethod - def load_vae(cls, sub_directory): - return (cls.load_model(sub_directory),) + def load_vae(cls, sub_directory, vae_type="default"): + return (cls.load_model(sub_directory, vae_type),) @classmethod - def load_model(cls, sub_directory): - + def load_model(cls, sub_directory, vae_type="default"): + if vae_type == "default": + return cls.load_default_vae(sub_directory) + else: + return cls.load_taesd(vae_type) + + @classmethod + def load_default_vae(cls, sub_directory): base_path = DiffusersUtils.get_base_path() sub_dir_path = os.path.join(base_path, sub_directory) - vae_folder = os.path.join(sub_dir_path, "vae") vae_path = DiffusersUtils.find_model_file(vae_folder) - DiffusersUtils.check_and_clear_cache('vae', vae_path) print(f"DiffusersVAELoader: Attempting to load VAE model from: {vae_path}") @@ -40,4 +53,38 @@ class DiffusersVAELoader(DiffusersLoaderBase): return comfy.sd.VAE(sd=vae_sd) except Exception as e: print(f"DiffusersVAELoader: Error loading VAE model: {e}") - raise \ No newline at end of file + raise + + @staticmethod + def load_taesd(vae_type): + sd = {} + vae_approx_path = os.path.join(os.path.dirname(DiffusersUtils.get_base_path()), "vae_approx") + # Validation check for vae_approx_path + if not os.path.exists(vae_approx_path): + raise FileNotFoundError(f"The vae_approx path '{vae_approx_path}' does not exist.") + + encoder_file = next(f for f in os.listdir(vae_approx_path) if f.startswith(f"{vae_type}_encoder.")) + decoder_file = next(f for f in os.listdir(vae_approx_path) if f.startswith(f"{vae_type}_decoder.")) + + #Validation check for presence of encoder and decoder_file + if encoder_file is None or decoder_file is None: + raise FileNotFoundError(f"No encoder or decoder file found in {vae_approx_path}.") + + enc = comfy.utils.load_torch_file(os.path.join(vae_approx_path, encoder_file)) + for k in enc: + sd[f"taesd_encoder.{k}"] = enc[k] + + dec = comfy.utils.load_torch_file(os.path.join(vae_approx_path, decoder_file)) + for k in dec: + sd[f"taesd_decoder.{k}"] = dec[k] + + if vae_type in DiffusersVAELoader.VAE_CONFIGS: + config = DiffusersVAELoader.VAE_CONFIGS[vae_type] + sd["vae_scale"] = torch.tensor(config["scale"]) + sd["vae_shift"] = torch.tensor(config["shift"]) + else: + print(f"Warning: No configuration found for {vae_type}. Using default values.") + sd["vae_scale"] = torch.tensor(1.0) + sd["vae_shift"] = torch.tensor(0.0) + + return comfy.sd.VAE(sd=sd) \ No newline at end of file