Add torch_gc to state.nextjob, vae and upscale

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
 

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
 
 

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

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 = {