From d63f35e298463d8d2c573176abd22df067dcc2ca Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 16 May 2024 18:02:55 -0400 Subject: [PATCH] add cudaMallocAsync --- CHANGELOG.md | 2 ++ installer.py | 6 +++++- modules/devices.py | 5 +---- modules/extra_networks.py | 38 ++++++++++++++++++----------------- modules/processing.py | 6 ++---- modules/processing_helpers.py | 3 +++ modules/sd_models.py | 2 +- modules/shared.py | 5 +++-- 8 files changed, 37 insertions(+), 30 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 350ce7fab..9c8b34f18 100644 --- a/CHANGELOG.md +++ b/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/` diff --git a/installer.py b/installer.py index 2c98cee36..14bb64a80 100644 --- a/installer.py +++ b/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') diff --git a/modules/devices.py b/modules/devices.py index 74ca9ded2..36425ff39 100644 --- a/modules/devices.py +++ b/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) diff --git a/modules/extra_networks.py b/modules/extra_networks.py index 63fd2b1fe..9ee2ece46 100644 --- a/modules/extra_networks.py +++ b/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 diff --git a/modules/processing.py b/modules/processing.py index 3bc852cf5..11bcc5040 100644 --- a/modules/processing.py +++ b/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) diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index 284e9a63d..30d23d311 100644 --- a/modules/processing_helpers.py +++ b/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': diff --git a/modules/sd_models.py b/modules/sd_models.py index e70ff4a47..3b476d4f6 100644 --- a/modules/sd_models.py +++ b/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') diff --git a/modules/shared.py b/modules/shared.py index f7944b5b5..b26025894 100644 --- a/modules/shared.py +++ b/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("

Model Compile

", "", 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}")