From 0b39e5d24f4d7b95a2fa169f6fedad0da60428bc Mon Sep 17 00:00:00 2001
From: zandm7 <zandmstudios@gmail.com>
Date: Sun, 19 May 2024 15:01:13 -0600
Subject: [PATCH 1/2] Implemented spandrel

---
 comfy/model_management.py           |  5 ++++-
 comfy_extras/nodes_upscale_model.py | 14 +++++++-------
 requirements.txt                    |  2 ++
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/comfy/model_management.py b/comfy/model_management.py
index 21ae8d291095..afad68d1b225 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -259,7 +259,10 @@ def get_torch_device_name(device):
 
 def module_size(module):
     module_mem = 0
-    sd = module.state_dict()
+    try:
+        sd = module.state_dict()
+    except:
+        sd = dict(module.__dict__.items())
     for k in sd:
         t = sd[k]
         module_mem += t.nelement() * t.element_size()
diff --git a/comfy_extras/nodes_upscale_model.py b/comfy_extras/nodes_upscale_model.py
index 52c95df23dc7..cacaf975561d 100644
--- a/comfy_extras/nodes_upscale_model.py
+++ b/comfy_extras/nodes_upscale_model.py
@@ -1,5 +1,5 @@
 import os
-from comfy_extras.chainner_models import model_loading
+from spandrel import ModelLoader, ImageModelDescriptor
 from comfy import model_management
 import torch
 import comfy.utils
@@ -20,7 +20,10 @@ def load_model(self, model_name):
         sd = comfy.utils.load_torch_file(model_path, safe_load=True)
         if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
             sd = comfy.utils.state_dict_prefix_replace(sd, {"module.":""})
-        out = model_loading.load_state_dict(sd).eval()
+        out = ModelLoader().load_from_state_dict(sd).eval()
+
+        if not isinstance(out, ImageModelDescriptor):
+            raise Exception("Upscale model must be a single-image model.")
         return (out, )
 
 
@@ -38,10 +41,7 @@ def INPUT_TYPES(s):
     def upscale(self, upscale_model, image):
         device = model_management.get_torch_device()
 
-        memory_required = model_management.module_size(upscale_model)
-        memory_required += (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate
-        memory_required += image.nelement() * image.element_size()
-        model_management.free_memory(memory_required, device)
+        model_management.get_free_memory(device)
 
         upscale_model.to(device)
         in_img = image.movedim(-1,-3).to(device)
@@ -61,7 +61,7 @@ def upscale(self, upscale_model, image):
                 if tile < 128:
                     raise e
 
-        upscale_model.cpu()
+        upscale_model.to("cpu")
         s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
         return (s,)
 
diff --git a/requirements.txt b/requirements.txt
index e7d8c0e9c7a7..acaa6b18d0ee 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -11,3 +11,5 @@ scipy
 tqdm
 psutil
 kornia>=0.7.1
+spandrel
+xformers

From 2e20a7616ad465785c9327a4d401a01cdecb78f6 Mon Sep 17 00:00:00 2001
From: zandm7 <zandmstudios@gmail.com>
Date: Mon, 20 May 2024 00:10:35 -0600
Subject: [PATCH 2/2] Changed default preview method back to TAESD

---
 latent_preview.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/latent_preview.py b/latent_preview.py
index 05d750e2c06b..b4e23c679ce5 100644
--- a/latent_preview.py
+++ b/latent_preview.py
@@ -64,6 +64,8 @@ def get_previewer(device, latent_format):
 
         if method == LatentPreviewMethod.Auto:
             method = LatentPreviewMethod.Latent2RGB
+            if taesd_decoder_path:
+                method = LatentPreviewMethod.TAESD
 
         if method == LatentPreviewMethod.TAESD:
             if taesd_decoder_path: