add low/high threshold to balanced offload

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -42,6 +42,8 @@
   - Flux: do not recast quants  
 - **Offload** improvements:  
   - faster and more compatible *balanced* mode  
+  - balanced offload: units are now in percentage instead of bytes  
+  - balanced offload: add both high and low watermark  
 - **UI**:  
   - improved stats on generate completion  
   - improved live preview display and performance  

modules/devices.py

@@ -224,7 +224,7 @@ def torch_gc(force=False, fast=False):
         timer.process.records['gc'] = 0
     timer.process.records['gc'] += t1 - t0
     if not force or collected == 0:
-        return
+        return used_gpu
     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) }
@@ -233,6 +233,7 @@ def torch_gc(force=False, fast=False):
     results = { 'collected': collected, 'saved': saved }
     fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
     log.debug(f'GC: utilization={utilization} gc={results} before={before} after={after} device={torch.device(get_optimal_device_name())} fn={fn} time={round(t1 - t0, 2)}') # pylint: disable=protected-access
+    return used_gpu
 
 
 def set_cuda_sync_mode(mode):

modules/lora/networks.py

@@ -218,6 +218,7 @@ def maybe_recompile_model(names, te_multipliers):
 
 
 def load_networks(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None):
+    timer['list'] = 0
     global backup_size # pylint: disable=global-statement
     networks_on_disk: list[network.NetworkOnDisk] = [available_network_aliases.get(name, None) for name in names]
     if any(x is None for x in networks_on_disk):

modules/processing_helpers.py

@@ -368,14 +368,13 @@ def validate_sample(tensor):
     sample = 255.0 * np.moveaxis(sample, 0, 2) if not shared.native else 255.0 * sample
     with warnings.catch_warnings(record=True) as w:
         cast = sample.astype(np.uint8)
-    minimum, maximum, mean = np.min(cast), np.max(cast), np.mean(cast)
-    if len(w) > 0 or minimum == maximum:
+    if len(w) > 0:
         nans = np.isnan(sample).sum()
         cast = np.nan_to_num(sample)
         cast = cast.astype(np.uint8)
         vae = shared.sd_model.vae.dtype if hasattr(shared.sd_model, 'vae') else None
         upcast = getattr(shared.sd_model.vae.config, 'force_upcast', None) if hasattr(shared.sd_model, 'vae') and hasattr(shared.sd_model.vae, 'config') else None
-        shared.log.error(f'Decode: sample={sample.shape} invalid={nans} mean={mean} dtype={dtype} vae={vae} upcast={upcast} failed to validate')
+        shared.log.error(f'Decode: sample={sample.shape} invalid={nans} dtype={dtype} vae={vae} upcast={upcast} failed to validate')
         if upcast is not None and not upcast:
             setattr(shared.sd_model.vae.config, 'force_upcast', True) # noqa: B010
             shared.log.warning('Decode: upcast=True set, retry operation')

modules/sd_models.py

@@ -361,7 +361,7 @@ def set_diffuser_offload(sd_model, op: str = 'model'):
             shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}')
     if shared.opts.diffusers_offload_mode == "balanced":
         try:
-            shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} threshold={shared.opts.diffusers_offload_max_gpu_memory} limit={shared.opts.cuda_mem_fraction}')
+            shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} watermarks low={shared.opts.diffusers_offload_min_gpu_memory} high={shared.opts.diffusers_offload_max_gpu_memory} limit={shared.opts.cuda_mem_fraction:.2f}')
             sd_model = apply_balanced_offload(sd_model)
         except Exception as e:
             shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}')
@@ -369,6 +369,16 @@ def set_diffuser_offload(sd_model, op: str = 'model'):
 
 
 class OffloadHook(accelerate.hooks.ModelHook):
+    def __init__(self):
+        if shared.opts.diffusers_offload_max_gpu_memory > 1:
+            shared.opts.diffusers_offload_max_gpu_memory = 0.75
+        if shared.opts.diffusers_offload_max_cpu_memory > 1:
+            shared.opts.diffusers_offload_max_cpu_memory = 0.75
+        self.gpu = int(shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory * 1024*1024*1024)
+        self.cpu = int(shared.cpu_memory * shared.opts.diffusers_offload_max_cpu_memory * 1024*1024*1024)
+        shared.log.info(f'Init offload: type=balanced gpu={self.gpu} cpu={self.cpu}')
+        super().__init__()
+
     def init_hook(self, module):
         return module
 
@@ -377,10 +387,7 @@ class OffloadHook(accelerate.hooks.ModelHook):
             device_index = torch.device(devices.device).index
             if device_index is None:
                 device_index = 0
-            max_memory = {
-                device_index: int(shared.opts.diffusers_offload_max_gpu_memory * 1024*1024*1024),
-                "cpu": int(shared.opts.diffusers_offload_max_cpu_memory * 1024*1024*1024),
-            }
+            max_memory = { device_index: self.gpu, "cpu": self.cpu }
             device_map = getattr(module, "balanced_offload_device_map", None)
             if device_map is None or max_memory != getattr(module, "balanced_offload_max_memory", None):
                 device_map = accelerate.infer_auto_device_map(module, max_memory=max_memory)
@@ -399,10 +406,13 @@ class OffloadHook(accelerate.hooks.ModelHook):
         return module
 
 
-offload_hook_instance = OffloadHook()
+offload_hook_instance = None
 
 
 def apply_balanced_offload(sd_model):
+    global offload_hook_instance # pylint: disable=global-statement
+    if offload_hook_instance is None:
+        offload_hook_instance = OffloadHook()
     t0 = time.time()
     excluded = ['OmniGenPipeline']
     if sd_model.__class__.__name__ in excluded:
@@ -414,6 +424,7 @@ def apply_balanced_offload(sd_model):
         checkpoint_name = sd_model.__class__.__name__
 
     def apply_balanced_offload_to_module(pipe):
+        used_gpu = devices.torch_gc(fast=True)
         if hasattr(pipe, "pipe"):
             apply_balanced_offload_to_module(pipe.pipe)
         if hasattr(pipe, "_internal_dict"):
@@ -429,7 +440,9 @@ def apply_balanced_offload(sd_model):
                 max_memory = getattr(module, "balanced_offload_max_memory", None)
                 module = accelerate.hooks.remove_hook_from_module(module, recurse=True)
                 try:
-                    module = module.to(devices.cpu, non_blocking=True)
+                    if used_gpu > 100 * shared.opts.diffusers_offload_min_gpu_memory:
+                        module = module.to(devices.cpu, non_blocking=True)
+                        used_gpu = devices.torch_gc(fast=True)
                     module.offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name)
                     module = accelerate.hooks.add_hook_to_module(module, offload_hook_instance, append=True)
                     module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access

modules/shared.py

@@ -560,8 +560,9 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), {
     "diffusers_extract_ema": OptionInfo(False, "Use model EMA weights when possible"),
     "diffusers_generator_device": OptionInfo("GPU", "Generator device", gr.Radio, {"choices": ["GPU", "CPU", "Unset"]}),
     "diffusers_offload_mode": OptionInfo(startup_offload_mode, "Model offload mode", gr.Radio, {"choices": ['none', 'balanced', 'model', 'sequential']}),
-    "diffusers_offload_max_gpu_memory": OptionInfo(round(gpu_memory * 0.75, 1), "Max GPU memory before balanced offload", gr.Slider, {"minimum": 0, "maximum": gpu_memory, "step": 0.01, "visible": True }),
-    "diffusers_offload_max_cpu_memory": OptionInfo(round(cpu_memory * 0.75, 1), "Max CPU memory before balanced offload", gr.Slider, {"minimum": 0, "maximum": cpu_memory, "step": 0.01, "visible": False }),
+    "diffusers_offload_min_gpu_memory": OptionInfo(0.25, "Balanced offload GPU low watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }),
+    "diffusers_offload_max_gpu_memory": OptionInfo(0.75, "Balanced offload GPU high watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }),
+    "diffusers_offload_max_cpu_memory": OptionInfo(0.75, "Balanced offload CPU high watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }),
     "diffusers_vae_upcast": OptionInfo("default", "VAE upcasting", gr.Radio, {"choices": ['default', 'true', 'false']}),
     "diffusers_vae_slicing": OptionInfo(True, "VAE slicing"),
     "diffusers_vae_tiling": OptionInfo(cmd_opts.lowvram or cmd_opts.medvram, "VAE tiling"),