Updated VAE_Loader to allow user to select tae models from vae_approx folder

2024-08-21 18:10:45 +08:00
parent 99be89d2d3
commit e701d6a9d1
5 changed files with 201 additions and 69 deletions
@@ -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):
@@ -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
+        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:]
@@ -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

@@ -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
+                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
@@ -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
+            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)