Code refactoring for modularity and easier addition of new model types

2024-07-27 21:19:47 +08:00
parent ad1430a99b
commit dfeed59757
11 changed files with 279 additions and 456 deletions
@@ -0,0 +1,39 @@
+# vae_loader.py
+import os
+import comfy.sd
+import comfy.utils
+from .base_loader import DiffusersLoaderBase
+from .utils import DiffusersUtils
+
+class DiffusersVAELoader(DiffusersLoaderBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),),
+            }
+        }
+    
+    RETURN_TYPES = ("VAE",)
+    FUNCTION = "load_vae"
+    CATEGORY = "DiffusersLoader"
+
+    @classmethod
+    def load_vae(cls, sub_directory):
+        return (cls.load_model(sub_directory),)
+
+    @classmethod
+    def load_model(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)
+        print(f"DiffusersVAELoader: Attempting to load VAE model from: {vae_path}")
+
+        try:
+            vae_sd = comfy.utils.load_torch_file(vae_path)
+            return comfy.sd.VAE(sd=vae_sd)
+        except Exception as e:
+            print(f"DiffusersVAELoader: Error loading VAE model: {e}")
+            raise