From 437beaeafbe713e4bf1b42a3d339a38b75bd12fa Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 29 Feb 2024 15:26:11 +0200
Subject: [PATCH] Fix model changing

---
 SUPIR/models/SUPIR_model.py | 38 ++++++++++++++++++++++++-------------
 nodes.py                    | 11 +++++++----
 2 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/SUPIR/models/SUPIR_model.py b/SUPIR/models/SUPIR_model.py
index ce28b7e..c213509 100644
--- a/SUPIR/models/SUPIR_model.py
+++ b/SUPIR/models/SUPIR_model.py
@@ -152,19 +152,31 @@ class SUPIRModel(DiffusionEngine):
             samples = adaptive_instance_normalization(samples, x_stage1)
         return samples
 
-    def init_tile_vae(self, encoder_tile_size=512, decoder_tile_size=64):
-        self.first_stage_model.denoise_encoder.original_forward = self.first_stage_model.denoise_encoder.forward
-        self.first_stage_model.encoder.original_forward = self.first_stage_model.encoder.forward
-        self.first_stage_model.decoder.original_forward = self.first_stage_model.decoder.forward
-        self.first_stage_model.denoise_encoder.forward = VAEHook(
-            self.first_stage_model.denoise_encoder, encoder_tile_size, is_decoder=False, fast_decoder=False,
-            fast_encoder=False, color_fix=False, to_gpu=True)
-        self.first_stage_model.encoder.forward = VAEHook(
-            self.first_stage_model.encoder, encoder_tile_size, is_decoder=False, fast_decoder=False,
-            fast_encoder=False, color_fix=False, to_gpu=True)
-        self.first_stage_model.decoder.forward = VAEHook(
-            self.first_stage_model.decoder, decoder_tile_size, is_decoder=True, fast_decoder=False,
-            fast_encoder=False, color_fix=False, to_gpu=True)
+    def init_tile_vae(self, encoder_tile_size=512, decoder_tile_size=64, reset=False):
+        if reset:
+            # Reset the models to their original forward methods
+            if hasattr(self.first_stage_model.denoise_encoder, 'original_forward'):
+                self.first_stage_model.denoise_encoder.forward = self.first_stage_model.denoise_encoder.original_forward
+            if hasattr(self.first_stage_model.encoder, 'original_forward'):
+                self.first_stage_model.encoder.forward = self.first_stage_model.encoder.original_forward
+            if hasattr(self.first_stage_model.decoder, 'original_forward'):
+                self.first_stage_model.decoder.forward = self.first_stage_model.decoder.original_forward
+        else:
+            # Save the original forward methods
+            self.first_stage_model.denoise_encoder.original_forward = self.first_stage_model.denoise_encoder.forward
+            self.first_stage_model.encoder.original_forward = self.first_stage_model.encoder.forward
+            self.first_stage_model.decoder.original_forward = self.first_stage_model.decoder.forward
+
+            # Apply the VAEHook to the models
+            self.first_stage_model.denoise_encoder.forward = VAEHook(
+                self.first_stage_model.denoise_encoder, encoder_tile_size, is_decoder=False, fast_decoder=False,
+                fast_encoder=False, color_fix=False, to_gpu=True)
+            self.first_stage_model.encoder.forward = VAEHook(
+                self.first_stage_model.encoder, encoder_tile_size, is_decoder=False, fast_decoder=False,
+                fast_encoder=False, color_fix=False, to_gpu=True)
+            self.first_stage_model.decoder.forward = VAEHook(
+                self.first_stage_model.decoder, decoder_tile_size, is_decoder=True, fast_decoder=False,
+                fast_encoder=False, color_fix=False, to_gpu=True)
 
 
 if __name__ == '__main__':
diff --git a/nodes.py b/nodes.py
index b6af10d..b8abf56 100644
--- a/nodes.py
+++ b/nodes.py
@@ -65,7 +65,8 @@ class SUPIR_Upscale:
         
         device = comfy.model_management.get_torch_device()
         image = image.to(device)
-        batch_size = image.shape[0]
+        self.current_sdxl_model = None
+
         SUPIR_MODEL_PATH = folder_paths.get_full_path("checkpoints", supir_model)
         SDXL_MODEL_PATH = folder_paths.get_full_path("checkpoints", sdxl_model)
         
@@ -87,8 +88,8 @@ class SUPIR_Upscale:
             vae_dtype = 'fp32'
             model_dtype = 'fp32'
 
-        if not hasattr(self, "model") or self.model is None:
-            
+        if not hasattr(self, "model") or self.model is None or self.current_sdxl_model != sdxl_model:
+            self.current_sdxl_model = sdxl_model
             config = OmegaConf.load(config_path)
             config.model.params.ae_dtype = vae_dtype
             config.model.params.diffusion_dtype = model_dtype
@@ -101,7 +102,9 @@ class SUPIR_Upscale:
             self.model.load_state_dict(sdxl_state_dict, strict=False)
             self.model.to(device).to(dtype)
             if use_tiled_vae:
-                self.model.init_tile_vae(encoder_tile_size=encoder_tile_size_pixels, decoder_tile_size=decoder_tile_size_latent)
+                self.model.init_tile_vae(encoder_tile_size=encoder_tile_size_pixels, decoder_tile_size=decoder_tile_size_latent, reset=False)
+            else:
+                self.model.init_tile_vae(encoder_tile_size=encoder_tile_size_pixels, decoder_tile_size=decoder_tile_size_latent, reset=True)
    
         autocast_condition = dtype == torch.float16 or torch.bfloat16 and not comfy.model_management.is_device_mps(device)
         with torch.autocast(comfy.model_management.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext():