From b036c2fc3b73b861147475e076aaed7bb1d657f8 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Fri, 21 Jun 2024 12:57:15 -0400
Subject: [PATCH] improve gc threshold

---
 CHANGELOG.md                  |  2 ++
 modules/control/processors.py |  2 +-
 modules/devices.py            | 15 +++++++++------
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 19dec1762..76dac5867 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,7 @@ But there's more than SD3:
 - support for **PixArt-Sigma** in small/medium/large variants  
 - support for **HunyuanDiT 1.1**  
 - (finally) new release of **Torch-DirectML**  
+- additional efficiencies for users with low vram gpus  
 - over 20 overall fixes  
 
 ### Model Improvements
@@ -47,6 +48,7 @@ But there's more than SD3:
   *note*: new directml is finally based on modern `torch` 2.3.1!  
 - extra networks: info display now contains link to source url if model if its known  
   works for civitai and huggingface models  
+- force gc for lowvram users and improve gc logging
 - improved google.colab support
 - css tweaks for standardui
 - css tweaks for modernui
diff --git a/modules/control/processors.py b/modules/control/processors.py
index eccaf1e30..19683f6ea 100644
--- a/modules/control/processors.py
+++ b/modules/control/processors.py
@@ -139,7 +139,7 @@ class Processor():
         self.model = None
         self.processor_id = processor_id
         # self.override = None
-        devices.torch_gc()
+        # devices.torch_gc()
         self.load_config = { 'cache_dir': cache_dir }
 
     def config(self, processor_id = None):
diff --git a/modules/devices.py b/modules/devices.py
index 9f762b544..d72555192 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -140,11 +140,8 @@ 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}')
+    threshold = 0 if (shared.cmd_opts.lowvram and not shared.cmd_opts.use_zluda) else shared.opts.torch_gc_threshold
+    if force or threshold == 0 or used_gpu >= threshold or used_ram >= threshold:
         force = True
     if oom > previous_oom:
         previous_oom = oom
@@ -163,7 +160,13 @@ def torch_gc(force=False):
         except Exception:
             pass
     t1 = time.time()
-    log.debug(f'GC: collected={collected} device={torch.device(get_optimal_device_name())} {memstats.memory_stats()} time={round(t1 - t0, 2)}')
+    mem = memstats.memory_stats()
+    saved = round(gpu.get('used', 0) - mem.get('gpu', {}).get('used', 0), 2)
+    before = { 'gpu': gpu.get('used', 0), 'ram': ram.get('used', 0) }
+    after = { 'gpu': mem.get('gpu', {}).get('used', 0), 'ram': mem.get('ram', {}).get('used', 0), 'retries': mem.get('retries', 0), 'oom': mem.get('oom', 0) }
+    utilization = { 'gpu': used_gpu, 'ram': used_ram, 'threshold': threshold }
+    results = { 'collected': collected, 'saved': saved }
+    log.debug(f'GC: utilization={utilization} gc={results} beofre={before} after={after} device={torch.device(get_optimal_device_name())} fn={sys._getframe(1).f_code.co_name} time={round(t1 - t0, 2)}') # pylint: disable=protected-access
 
 
 def set_cuda_sync_mode(mode):