add cudaMallocAsync

CHANGELOG.md

@@ -120,6 +120,8 @@
   - Secondary sampler add option "same as primary"  
   - Change attention mechanism on-the-fly without model reload, thanks @Disty0  
   - Update stable-fast with support for torch 2.2.2 and 2.3.0, thanks @Aptronymist
+  - Add torch *cudaMallocAsync* in compute options  
+    Can improve memory utilization on compatible GPUs (RTX and newer)  
   - Support controlnet manually downloads models in both standalone and diffusers format  
     For standalone, simply copy safetensors file to `models/control/controlnet` folder  
     For diffusers format, create folder with model name in `models/control/controlnet/`  

installer.py

@@ -875,7 +875,6 @@ def set_environment():
     os.environ.setdefault('K_DIFFUSION_USE_COMPILE', '0')
     os.environ.setdefault('NUMEXPR_MAX_THREADS', '16')
     os.environ.setdefault('PYTHONHTTPSVERIFY', '0')
-    os.environ.setdefault('PYTORCH_CUDA_ALLOC_CONF', 'garbage_collection_threshold:0.8,max_split_size_mb:512')
     os.environ.setdefault('SAFETENSORS_FAST_GPU', '1')
     os.environ.setdefault('TF_CPP_MIN_LOG_LEVEL', '2')
     os.environ.setdefault('TF_ENABLE_ONEDNN_OPTS', '0')
@@ -885,6 +884,11 @@ def set_environment():
     os.environ.setdefault('DO_NOT_TRACK', '1')
     os.environ.setdefault('HF_HUB_CACHE', opts.get('hfcache_dir', os.path.join(os.path.expanduser('~'), '.cache', 'huggingface', 'hub')))
     log.debug(f'HF cache folder: {os.environ.get("HF_HUB_CACHE")}')
+    allocator = f'garbage_collection_threshold:{opts.get("torch_gc_threshold", 80)/100:0.2f},max_split_size_mb:512'
+    if opts.get("torch_malloc", "native") == 'cudaMallocAsync':
+        allocator += ',backend:cudaMallocAsync'
+    os.environ.setdefault('PYTORCH_CUDA_ALLOC_CONF', allocator)
+    log.debug(f'Torch allocator: "{allocator}"')
     if sys.platform == 'darwin':
         os.environ.setdefault('PYTORCH_ENABLE_MPS_FALLBACK', '1')
 

modules/devices.py

@@ -380,11 +380,8 @@ def randn_without_seed(shape):
         return torch.randn(shape, device=cpu).to(device)
     return torch.randn(shape, device=device)
 
-
 def autocast(disable=False):
-    if disable:
-        return contextlib.nullcontext()
-    if dtype == torch.float32 or shared.cmd_opts.precision == "Full":
+    if disable or dtype == torch.float32 or shared.cmd_opts.precision == "Full":
         return contextlib.nullcontext()
     if shared.cmd_opts.use_directml:
         return torch.dml.amp.autocast(dtype)

modules/extra_networks.py

@@ -1,6 +1,6 @@
 import re
 from collections import defaultdict
-from modules import errors, shared
+from modules import errors, shared, devices
 
 
 extra_network_registry = {}
@@ -82,24 +82,26 @@ def activate(p, extra_network_data, step=0):
         shared.log.warning("Composable LoRA not compatible with 'lora_force_diffusers'")
         stepwise = False
     shared.opts.data['lora_functional'] = stepwise or functional
-    for extra_network_name, extra_network_args in extra_network_data.items():
-        extra_network = extra_network_registry.get(extra_network_name, None)
-        if extra_network is None:
-            errors.log.warning(f"Skipping unknown extra network: {extra_network_name}")
-            continue
-        try:
-            extra_network.activate(p, extra_network_args, step=step)
-        except Exception as e:
-            errors.display(e, f"activating extra network: name={extra_network_name} args:{extra_network_args}")
+    with devices.autocast():
+        for extra_network_name, extra_network_args in extra_network_data.items():
+            extra_network = extra_network_registry.get(extra_network_name, None)
+            if extra_network is None:
+                errors.log.warning(f"Skipping unknown extra network: {extra_network_name}")
+                continue
+            try:
+                extra_network.activate(p, extra_network_args, step=step)
+            except Exception as e:
+                errors.display(e, f"activating extra network: name={extra_network_name} args:{extra_network_args}")
+
+        for extra_network_name, extra_network in extra_network_registry.items():
+            args = extra_network_data.get(extra_network_name, None)
+            if args is not None:
+                continue
+            try:
+                extra_network.activate(p, [])
+            except Exception as e:
+                errors.display(e, f"activating extra network: name={extra_network_name}")
 
-    for extra_network_name, extra_network in extra_network_registry.items():
-        args = extra_network_data.get(extra_network_name, None)
-        if args is not None:
-            continue
-        try:
-            extra_network.activate(p, [])
-        except Exception as e:
-            errors.display(e, f"activating extra network: name={extra_network_name}")
     if stepwise:
         p.extra_network_data = extra_network_data
         shared.opts.data['lora_functional'] = functional

modules/processing.py

@@ -265,8 +265,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     with devices.inference_context(), ema_scope_context():
         t0 = time.time()
         if not hasattr(p, 'skip_init'):
-            with devices.autocast():
-                p.init(p.all_prompts, p.all_seeds, p.all_subseeds)
+            p.init(p.all_prompts, p.all_seeds, p.all_subseeds)
         extra_network_data = None
         debug(f'Processing inner: args={vars(p)}')
         for n in range(p.n_iter):
@@ -293,8 +292,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                 break
             p.prompts, extra_network_data = extra_networks.parse_prompts(p.prompts)
             if not p.disable_extra_networks:
-                with devices.autocast():
-                    extra_networks.activate(p, extra_network_data)
+                extra_networks.activate(p, extra_network_data)
             if p.scripts is not None and isinstance(p.scripts, scripts.ScriptRunner):
                 p.scripts.process_batch(p, batch_number=n, prompts=p.prompts, seeds=p.seeds, subseeds=p.subseeds)
 

modules/processing_helpers.py

@@ -501,6 +501,7 @@ def set_latents(p):
     shared.sd_model.prepare_latents = dummy_prepare_latents # stop diffusers processing latents again
     return latents
 
+
 def apply_circular(enable, model):
     try:
         for layer in [layer for layer in model.unet.modules() if type(layer) is torch.nn.Conv2d]:
@@ -510,6 +511,7 @@ def apply_circular(enable, model):
     except Exception as e:
         debug(f"Diffusers tiling failed: {e}")
 
+
 def save_intermediate(p, latents, suffix):
     for i in range(len(latents)):
         from modules.processing import create_infotext
@@ -518,6 +520,7 @@ def save_intermediate(p, latents, suffix):
         for j in range(len(decoded)):
             images.save_image(decoded[j], path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix=suffix)
 
+
 def update_sampler(p, sd_model, second_pass=False):
     sampler_selection = p.hr_sampler_name if second_pass else p.sampler_name
     if hasattr(sd_model, 'scheduler') and sampler_selection != 'Default':

modules/sd_models.py

@@ -680,7 +680,7 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model'):
             if shared.opts.no_half_vae:
                 devices.dtype_vae = torch.float32
                 sd_model.vae.to(devices.dtype_vae)
-            shared.log.debug(f'Setting {op} VAE: upcast={sd_model.vae.config.get("force_upcast", None)}')
+            shared.log.debug(f'Setting {op} VAE: upcast={sd_model.vae.config.force_upcast}')
     if hasattr(sd_model, "enable_vae_slicing"):
         if shared.opts.diffusers_vae_slicing:
             shared.log.debug(f'Setting {op}: enable VAE slicing')

modules/shared.py

@@ -430,7 +430,8 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "cudnn_benchmark": OptionInfo(False, "Full-depth cuDNN benchmark feature"),
     "cudnn_deterministic": OptionInfo(False, "Use deterministic options for cuDNN"),
     "diffusers_fuse_projections": OptionInfo(False, "Fused projections"),
-    "torch_gc_threshold": OptionInfo(80, "Memory usage threshold for GC", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}),
+    "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'] }),
 
     "cuda_compile_sep": OptionInfo("<h2>Model Compile</h2>", "", gr.HTML),
     "cuda_compile": OptionInfo([] if not cmd_opts.use_openvino else ["Model", "VAE"], "Compile Model", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "Upscaler"]}),
@@ -962,7 +963,7 @@ class Options:
             info = self.data_labels.get(k, None)
             if info is not None and not self.same_type(info.default, v):
                 log.error(f"Error: bad setting value: {k}: {v} ({type(v).__name__}; expected {type(info.default).__name__})")
-            if info is None and k not in compatibility_opts:
+            if info is None and k not in compatibility_opts and not k.startswith('uiux_'):
                 unknown_settings.append(k)
         if len(unknown_settings) > 0:
             log.debug(f"Unknown settings: {unknown_settings}")