From 092a326c09e8ba81e5d913f34a4499ad7bcf92e1 Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Thu, 20 Jun 2024 14:47:30 +0300
Subject: [PATCH] Add torch_gc to state.nextjob, vae and upscale

---
 modules/devices.py            | 9 ++++++---
 modules/processing_helpers.py | 1 +
 modules/processing_vae.py     | 2 ++
 modules/shared_state.py       | 2 ++
 4 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/modules/devices.py b/modules/devices.py
index f3bdbaffe..9f762b544 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -140,13 +140,16 @@ def torch_gc(force=False):
         used_gpu = round(100 * gpu.get('used', 0) / gpu.get('total', 1)) if gpu.get('total', 1) > 1 else 0
     used_ram = round(100 * ram.get('used', 0) / ram.get('total', 1)) if ram.get('total', 1) > 1 else 0
     global previous_oom # pylint: disable=global-statement
+    if force or shared.opts.torch_gc_threshold == 0:
+        log.debug(f'Forced Torch GC: GPU={used_gpu}% RAM={used_ram}% {mem}')
+        force = True
+    elif used_gpu >= shared.opts.torch_gc_threshold or used_ram >= shared.opts.torch_gc_threshold:
+        log.info(f'High memory utilization: GPU={used_gpu}% RAM={used_ram}% {mem}')
+        force = True
     if oom > previous_oom:
         previous_oom = oom
         log.warning(f'GPU out-of-memory error: {mem}')
         force = True
-    if used_gpu >= shared.opts.torch_gc_threshold or used_ram >= shared.opts.torch_gc_threshold:
-        log.info(f'High memory utilization: GPU={used_gpu}% RAM={used_ram}% {mem}')
-        force = True
     if not force:
         return
 
diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py
index 5baf0193f..486dc09a5 100644
--- a/modules/processing_helpers.py
+++ b/modules/processing_helpers.py
@@ -400,6 +400,7 @@ def resize_hires(p, latents): # input=latents output=pil if not latent_upscaler
         else:
             resized_image = img
         resized_images.append(resized_image)
+    devices.torch_gc()
     return resized_images
 
 
diff --git a/modules/processing_vae.py b/modules/processing_vae.py
index 5db51a176..9b295e39c 100644
--- a/modules/processing_vae.py
+++ b/modules/processing_vae.py
@@ -140,6 +140,7 @@ def vae_decode(latents, model, output_type='np', full_quality=True):
     if shared.cmd_opts.profile:
         t1 = time.time()
         shared.log.debug(f'Profile: VAE decode: {t1-t0:.2f}')
+    devices.torch_gc()
     return imgs
 
 
@@ -155,4 +156,5 @@ def vae_encode(image, model, full_quality=True): # pylint: disable=unused-variab
         latents = full_vae_encode(image=tensor, model=shared.sd_model)
     else:
         latents = taesd_vae_encode(image=tensor)
+    devices.torch_gc()
     return latents
diff --git a/modules/shared_state.py b/modules/shared_state.py
index 470ee19c7..79ee20f19 100644
--- a/modules/shared_state.py
+++ b/modules/shared_state.py
@@ -41,10 +41,12 @@ class State:
         log.debug(f'Requested {"pause" if self.paused else "continue"}')
 
     def nextjob(self):
+        import modules.devices
         self.do_set_current_image()
         self.job_no += 1
         self.sampling_step = 0
         self.current_image_sampling_step = 0
+        modules.devices.torch_gc()
 
     def dict(self):
         obj = {