1 files changed, 7 insertions, 1 deletions
diff --git a/modules/ui_extra_networks_checkpoints_user_metadata.py b/modules/ui_extra_networks_checkpoints_user_metadata.py
index 2c69aab8..25df0a80 100644
--- a/modules/ui_extra_networks_checkpoints_user_metadata.py
+++ b/modules/ui_extra_networks_checkpoints_user_metadata.py
@@ -1,6 +1,6 @@
 import gradio as gr
 
-from modules import ui_extra_networks_user_metadata, sd_vae
+from modules import ui_extra_networks_user_metadata, sd_vae, shared
 from modules.ui_common import create_refresh_button
 
 
@@ -18,6 +18,10 @@ class CheckpointUserMetadataEditor(ui_extra_networks_user_metadata.UserMetadataE
 
         self.write_user_metadata(name, user_metadata)
 
+    def update_vae(self, name):
+        if name == shared.sd_model.sd_checkpoint_info.name_for_extra:
+            sd_vae.reload_vae_weights()
+
     def put_values_into_components(self, name):
         user_metadata = self.get_user_metadata(name)
         values = super().put_values_into_components(name)
@@ -58,3 +62,5 @@ class CheckpointUserMetadataEditor(ui_extra_networks_user_metadata.UserMetadataE
         ]
 
         self.setup_save_handler(self.button_save, self.save_user_metadata, edited_components)
+        self.button_save.click(fn=self.update_vae, inputs=[self.edit_name_input])
+