Custom VAE

2023-10-21 21:26:23 +02:00
parent bca8487272
commit 1c72bd4ae7
6 changed files with 680 additions and 1 deletions
@@ -0,0 +1,29 @@
+import folder_paths
+
+from .conf import vae_conf
+from .loader import EXVAE
+
+class ExtraVAELoader:
+	@classmethod
+	def INPUT_TYPES(s):
+		return {
+			"required": {
+				"vae_name": (folder_paths.get_filename_list("vae"),),
+				"vae_type": (list(vae_conf.keys()), {"default":"kl-f8"}),
+				"dtype"   : (["auto","fp32","fp16","bf16"],),
+			}
+		}
+	RETURN_TYPES = ("VAE",)
+	FUNCTION = "load_vae"
+	CATEGORY = "ExtraModels"
+	TITLE = "ExtraVAELoader"
+
+	def load_vae(self, vae_name, vae_type, dtype):
+		model_path = folder_paths.get_full_path("vae", vae_name)
+		model_conf = vae_conf[vae_type]
+		vae = EXVAE(model_path, model_conf, dtype)
+		return (vae,)
+
+NODE_CLASS_MAPPINGS = {
+	"ExtraVAELoader" : ExtraVAELoader,
+}