From f2b60a42c3b91397837af70f92e90b2d530fd463 Mon Sep 17 00:00:00 2001
From: Kijai <40791699+kijai@users.noreply.github.com>
Date: Tue, 28 May 2024 17:08:31 +0300
Subject: [PATCH] Diffusers 0.28 compatibility

---
 configs/vae_config.yaml | 30 ++++++++++++++++++++++++++++++
 nodes.py                |  8 ++++----
 2 files changed, 34 insertions(+), 4 deletions(-)
 create mode 100644 configs/vae_config.yaml

diff --git a/configs/vae_config.yaml b/configs/vae_config.yaml
new file mode 100644
index 0000000..6681ef0
--- /dev/null
+++ b/configs/vae_config.yaml
@@ -0,0 +1,30 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.8.0",
+  "_name_or_path": "hf-models/stable-diffusion-v2-768x768/vae",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "in_channels": 3,
+  "latent_channels": 4,
+  "layers_per_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 512,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ]
+}
diff --git a/nodes.py b/nodes.py
index a7be3e9..da154ab 100644
--- a/nodes.py
+++ b/nodes.py
@@ -74,15 +74,15 @@ class geowizard_model_loader:
             self.current_config = custom_config
             # setup pretrained models
             original_config = OmegaConf.load(os.path.join(script_directory, f"configs/v1-inference.yaml"))
+            vae_config = OmegaConf.load(os.path.join(script_directory, f"configs/vae_config.yaml"))
 
             dtype = convert_dtype(dtype)
             
-            from diffusers.loaders.single_file_utils import (convert_ldm_vae_checkpoint, create_vae_diffusers_config)
+            from diffusers.loaders.single_file_utils import (convert_ldm_vae_checkpoint)
             
             sd = vae.get_sd()
-            converted_vae_config = create_vae_diffusers_config(original_config, image_size=512)
-            converted_vae = convert_ldm_vae_checkpoint(sd, converted_vae_config)
-            self.vae = AutoencoderKL(**converted_vae_config)
+            converted_vae = convert_ldm_vae_checkpoint(sd, vae_config)
+            self.vae = AutoencoderKL(**vae_config)
             self.vae.load_state_dict(converted_vae, strict=False)