cuda memory limits

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

CHANGELOG.md

@@ -1,8 +1,12 @@
 # Change Log for SD.Next
 
-## Update for 2024-10-25
+## Update for 2024-10-26
 
 Improvements:
+- Torch CUDA set device memory limit
+  in *settings -> compute settings -> torch memory limit*  
+  default=0 meaning no limit, if set torch will limit memory usage to specified fraction  
+  *note*: this is not a hard limit, torch will try to stay under this value  
 - Model selector:
   - change-in-behavior
   - when typing, it will auto-load model as soon as exactly one match is found

modules/devices.py

@@ -51,8 +51,8 @@ def has_zluda() -> bool:
     if not cuda_ok:
         return False
     try:
-        device = torch.device("cuda")
-        return torch.cuda.get_device_name(device).endswith("[ZLUDA]")
+        dev = torch.device("cuda")
+        return torch.cuda.get_device_name(dev).endswith("[ZLUDA]")
     except Exception:
         return False
 
@@ -207,7 +207,7 @@ def torch_gc(force=False, fast=False):
         force = True
     if oom > previous_oom:
         previous_oom = oom
-        log.warning(f'GPU out-of-memory error: {mem}')
+        log.warning(f'Torch GPU out-of-memory error: {mem}')
         force = True
     if force:
         # actual gc
@@ -247,13 +247,26 @@ def set_cuda_sync_mode(mode):
         return
     try:
         import ctypes
-        log.info(f'Set cuda sync: mode={mode}')
+        log.info(f'Torch CUDA sync: mode={mode}')
         torch.cuda.set_device(torch.device(get_optimal_device_name()))
         ctypes.CDLL('libcudart.so').cudaSetDeviceFlags({'auto': 0, 'spin': 1, 'yield': 2, 'block': 4}[mode])
     except Exception:
         pass
 
 
+def set_cuda_memory_limit():
+    if not cuda_ok or opts.cuda_mem_fraction == 0:
+        return
+    from modules.shared import cmd_opts
+    try:
+        torch_gc(force=True)
+        mem = torch.cuda.get_device_properties(device).total_memory
+        torch.cuda.set_per_process_memory_fraction(float(opts.cuda_mem_fraction), cmd_opts.device_id if cmd_opts.device_id is not None else 0)
+        log.info(f'Torch CUDA memory limit: fraction={opts.cuda_mem_fraction:.2f} limit={round(opts.cuda_mem_fraction * mem / 1024 / 1024)} total={round(mem / 1024 / 1024)}')
+    except Exception as e:
+        log.warning(f'Torch CUDA memory limit: fraction={opts.cuda_mem_fraction:.2f} {e}')
+
+
 def test_fp16():
     global fp16_ok # pylint: disable=global-statement
     if fp16_ok is not None:
@@ -449,6 +462,7 @@ def set_dtype():
 
 def set_cuda_params():
     override_ipex_math()
+    set_cuda_memory_limit()
     set_cudnn_params()
     set_sdpa_params()
     set_dtype()

modules/shared.py

@@ -393,7 +393,7 @@ def get_default_modes():
             elif gpu_memory <= 8:
                 cmd_opts.medvram = True
                 default_offload_mode = "model"
-                log.info(f"Device detect: memory={gpu_memory:.1f} ptimization=medvram")
+                log.info(f"Device detect: memory={gpu_memory:.1f} optimization=medvram")
             else:
                 default_offload_mode = "none"
                 log.info(f"Device detect: memory={gpu_memory:.1f} optimization=none")
@@ -479,6 +479,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "cudnn_benchmark": OptionInfo(False, "Full-depth cuDNN benchmark feature"),
     "diffusers_fuse_projections": OptionInfo(False, "Fused projections"),
     "torch_expandable_segments": OptionInfo(False, "Torch expandable segments"),
+    "cuda_mem_fraction": OptionInfo(0.0, "Torch memory limit", gr.Slider, {"minimum": 0, "maximum": 2.0, "step": 0.05}),
     "torch_gc_threshold": OptionInfo(80, "Torch memory threshold for GC", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}),
     "torch_malloc": OptionInfo("native", "Torch memory allocator", gr.Radio, {"choices": ['native', 'cudaMallocAsync'] }),