From d09b5ef4ef150adab31195761725eaba409f6343 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Mon, 22 Apr 2024 18:42:41 -0400
Subject: [PATCH] Free some memory before loading upscale models.

---
 comfy_extras/nodes_upscale_model.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/comfy_extras/nodes_upscale_model.py b/comfy_extras/nodes_upscale_model.py
index 2b5e49a55..4de5d7e37 100644
--- a/comfy_extras/nodes_upscale_model.py
+++ b/comfy_extras/nodes_upscale_model.py
@@ -37,9 +37,14 @@ class ImageUpscaleWithModel:
 
     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) * 256.0 #The 256.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)
+
         upscale_model.to(device)
         in_img = image.movedim(-1,-3).to(device)
-        free_memory = model_management.get_free_memory(device)
 
         tile = 512
         overlap = 32