From fb89e26b548551d3e9860a0e23db574678a941be Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 11 Aug 2024 19:01:49 +0300 Subject: [PATCH] Auto detect memory size ffor balaced offload --- modules/sd_models.py | 2 +- modules/shared.py | 19 ++++++++++++------- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index dbbcf4e9e..217fcd844 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -918,7 +918,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No } if shared.opts.diffusers_offload_mode == "balanced": diffusers_load_config['device_map'] = "balanced" - if shared.opts.diffusers_offload_max_memory != 0: + if shared.opts.diffusers_offload_max_memory > 0: device_index = torch.device(devices.device).index if device_index is None: device_index = 0 diff --git a/modules/shared.py b/modules/shared.py index f316c4508..d747344d7 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -353,20 +353,25 @@ def temp_disable_extensions(): return disabled +mem_stat = memory_stats() +if "gpu" in mem_stat: + gpu_memory = mem_stat['gpu']['total'] +else: + gpu_memory = 0 + if not (cmd_opts.lowvram or cmd_opts.medvram): - mem_stat = memory_stats() if "gpu" in mem_stat: - if mem_stat['gpu']['total'] <= 4: + if gpu_memory <= 4: cmd_opts.lowvram = True offload_mode_default = "sequential" - log.info(f"VRAM: Detected={mem_stat['gpu']['total']} GB Optimization=lowvram") - elif mem_stat['gpu']['total'] <= 8: + log.info(f"VRAM: Detected={gpu_memory} GB Optimization=lowvram") + elif gpu_memory <= 8: cmd_opts.medvram = True offload_mode_default = "cpu" - log.info(f"VRAM: Detected={mem_stat['gpu']['total']} GB Optimization=medvram") + log.info(f"VRAM: Detected={gpu_memory} GB Optimization=medvram") else: offload_mode_default = "none" - log.info(f"VRAM: Detected={mem_stat['gpu']['total']} GB Optimization=none") + log.info(f"VRAM: Detected={gpu_memory} GB Optimization=none") if devices.backend == "directml": # Force BMM for DirectML instead of SDP @@ -531,7 +536,7 @@ 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(offload_mode_default, "Model offload mode", gr.Radio, {"choices": ['none', 'balanced', 'cpu', 'sequential']}), - "diffusers_offload_max_memory": OptionInfo(0, "Max memory for balanced offload mode in GB", gr.Slider, {"minimum": 0, "maximum": 24, "step": 0.1,}), + "diffusers_offload_max_memory": OptionInfo(gpu_memory * 0.8, "Max memory for balanced offload mode in GB", gr.Slider, {"minimum": 0, "maximum": gpu_memory, "step": 0.1,}), "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"),