From e3b087b6c06b3b4775aa40bfd24b469b510b05a4 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 11 Aug 2024 17:27:30 +0300 Subject: [PATCH] Add balanced offload mode and make offload modes a single choice list --- modules/dml/__init__.py | 2 -- modules/model_pixart.py | 2 +- modules/model_t5.py | 4 ++-- modules/processing_vae.py | 4 ++-- modules/prompt_parser_diffusers.py | 4 ++-- modules/sd_models.py | 16 +++++++++------- modules/sd_models_compile.py | 4 ++-- modules/shared.py | 9 ++++++--- 8 files changed, 24 insertions(+), 21 deletions(-) diff --git a/modules/dml/__init__.py b/modules/dml/__init__.py index 8eed1aad1..b289e5d2d 100644 --- a/modules/dml/__init__.py +++ b/modules/dml/__init__.py @@ -92,8 +92,6 @@ class OverrideItem(NamedTuple): opts_override_table = { "diffusers_generator_device": OverrideItem("CPU", None, "DirectML does not support torch Generator API"), - "diffusers_model_cpu_offload": OverrideItem(False, None, "Diffusers model CPU offloading does not support DirectML devices"), - "diffusers_seq_cpu_offload": OverrideItem(False, lambda opts: opts.diffusers_pipeline != "Stable Diffusion XL", "Diffusers sequential CPU offloading is available only on StableDiffusionXLPipeline with DirectML devices"), } diff --git a/modules/model_pixart.py b/modules/model_pixart.py index d57079e4e..f7918c54b 100644 --- a/modules/model_pixart.py +++ b/modules/model_pixart.py @@ -5,7 +5,7 @@ def load_pixart(checkpoint_info, diffusers_load_config={}): from modules import shared, devices, modelloader, model_t5 modelloader.hf_login() # shared.opts.data['cuda_dtype'] = 'FP32' # override - # shared.opts.data['diffusers_model_cpu_offload'] = True # override + # shared.opts.data['diffusers_offload_mode}'] = "cpu" # override # devices.set_cuda_params() fn = checkpoint_info.path.replace('huggingface/', '') t5 = model_t5.load_t5(shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir) diff --git a/modules/model_t5.py b/modules/model_t5.py index 277ec92f4..c425d9d7a 100644 --- a/modules/model_t5.py +++ b/modules/model_t5.py @@ -79,11 +79,11 @@ def set_t5(pipe, module, t5=None, cache_dir=None): return pipe t5 = load_t5(t5=t5, cache_dir=cache_dir) setattr(pipe, module, t5) - if shared.cmd_opts.lowvram or shared.opts.diffusers_seq_cpu_offload: + if shared.opts.diffusers_offload_mode == "sequential": from accelerate import cpu_offload getattr(pipe, module).to("cpu") cpu_offload(getattr(pipe, module), devices.device, offload_buffers=len(getattr(pipe, module)._parameters) > 0) # pylint: disable=protected-access - elif shared.cmd_opts.medvram or shared.opts.diffusers_model_cpu_offload: + elif shared.opts.diffusers_offload_mode == "cpu": if not hasattr(pipe, "_all_hooks") or len(pipe._all_hooks) == 0: # pylint: disable=protected-access pipe.enable_model_cpu_offload(device=devices.device) else: diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 5d50f0411..80ce38acf 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -35,7 +35,7 @@ def full_vae_decode(latents, model): t0 = time.time() if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False): base_device = sd_models.move_base(model, devices.cpu) - if not shared.cmd_opts.lowvram and not shared.opts.diffusers_seq_cpu_offload and hasattr(model, 'vae'): + if not shared.opts.diffusers_offload_mode == "sequential" and hasattr(model, 'vae'): sd_models.move_model(model.vae, devices.device) latents.to(model.vae.device) @@ -80,7 +80,7 @@ def full_vae_encode(image, model): debug('Moving to CPU: model=UNet') unet_device = model.unet.device sd_models.move_model(model.unet, devices.cpu) - if not shared.cmd_opts.lowvram and not shared.opts.diffusers_seq_cpu_offload and hasattr(model, 'vae'): + if not shared.opts.diffusers_offload_mode == "sequential" and hasattr(model, 'vae'): sd_models.move_model(model.vae, devices.device) encoded = model.vae.encode(image.to(model.vae.device, model.vae.dtype)).latent_dist.sample() if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py index 3d50bd5e3..a4ba0597c 100644 --- a/modules/prompt_parser_diffusers.py +++ b/modules/prompt_parser_diffusers.py @@ -168,7 +168,7 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c p.negative_embeds = [] p.negative_pooleds = [] - if (shared.cmd_opts.medvram or shared.opts.diffusers_model_cpu_offload) and hasattr(pipe, "_all_hooks") and hasattr(pipe, "maybe_free_model_hooks"): + if shared.opts.diffusers_offload_mode in {"balanced", "cpu"} and hasattr(pipe, "_all_hooks") and hasattr(pipe, "maybe_free_model_hooks"): # if the last job is interrupted, model will stay in the vram and cause oom, send everything back to cpu before continuing pipe.maybe_free_model_hooks() devices.torch_gc() @@ -204,7 +204,7 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c if debug_enabled: get_tokens('positive', prompts[0]) get_tokens('negative', negative_prompts[0]) - if (shared.cmd_opts.medvram or shared.opts.diffusers_model_cpu_offload) and hasattr(pipe, "_all_hooks") and hasattr(pipe, "maybe_free_model_hooks"): + if shared.opts.diffusers_offload_mode in {"balanced", "cpu"} and hasattr(pipe, "_all_hooks") and hasattr(pipe, "maybe_free_model_hooks"): # text encoder will stay in the vram and cause oom, send everything back to cpu before continuing pipe.maybe_free_model_hooks() debug(f"Prompt encode: time={(time.time() - t0):.3f}") diff --git a/modules/sd_models.py b/modules/sd_models.py index e9e6beae9..bee25aef8 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -727,15 +727,10 @@ def set_diffuser_offload(sd_model, op: str = 'model'): if sd_model is None: shared.log.warning(f'{op} is not loaded') return - if (shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) and (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram): - shared.log.warning(f'Setting {op}: Model CPU offload and Sequential CPU offload are not compatible') - shared.log.debug(f'Setting {op}: disabling model CPU offload') - shared.opts.diffusers_model_cpu_offload=False - shared.cmd_opts.medvram=False if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate): sd_model.has_accelerate = False if hasattr(sd_model, "enable_model_cpu_offload"): - if shared.cmd_opts.medvram or shared.opts.diffusers_model_cpu_offload: + if shared.opts.diffusers_offload_mode == "cpu": shared.log.debug(f'Setting {op}: enable model CPU offload') if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: shared.opts.diffusers_move_base = False @@ -748,7 +743,7 @@ def set_diffuser_offload(sd_model, op: str = 'model'): sd_model.maybe_free_model_hooks() sd_model.has_accelerate = True if hasattr(sd_model, "enable_sequential_cpu_offload"): - if shared.cmd_opts.lowvram or shared.opts.diffusers_seq_cpu_offload: + if shared.opts.diffusers_offload_mode == "sequential": shared.log.debug(f'Setting {op}: enable sequential CPU offload') if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: shared.opts.diffusers_move_base = False @@ -921,6 +916,11 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No "requires_safety_checker": False, # "use_safetensors": True, } + if shared.opts.diffusers_offload_mode == "balanced": + diffusers_load_config['device_map'] = "balanced" + if shared.opts.diffusers_offload_max_memory != 0: + diffusers_load_config['max_memory'] = {0:f"{shared.opts.diffusers_offload_max_memory}GB"} + if shared.opts.diffusers_model_load_variant != 'default': diffusers_load_config['variant'] = shared.opts.diffusers_model_load_variant if shared.opts.diffusers_pipeline == 'Custom Diffusers Pipeline' and len(shared.opts.custom_diffusers_pipeline) > 0: @@ -1195,6 +1195,8 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No insert_parser_highjack(sd_model.__class__.__name__) set_diffuser_options(sd_model, vae, op, offload=False) + if shared.opts.diffusers_offload_mode == "balanced": + sd_model.has_accelerate = True if shared.opts.nncf_compress_weights and not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"): sd_model = sd_models_compile.nncf_compress_weights(sd_model) # run this before move model so it can be compressed in CPU if shared.opts.optimum_quanto_weights: diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index e972bfa8b..445700283 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -150,7 +150,7 @@ def nncf_compress_model(model, op=None, sd_model=None): else: getattr(sd_model, quant_last_model_name).to(quant_last_model_device) devices.torch_gc(force=True) - if shared.cmd_opts.medvram or shared.opts.diffusers_model_cpu_offload or shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram: + if shared.cmd_opts.medvram or shared.cmd_opts.lowvram or shared.opts.diffusers_offload_mode != "none": quant_last_model_name = op quant_last_model_device = model.device else: @@ -205,7 +205,7 @@ def optimum_quanto_model(model, op=None, sd_model=None, weights=None): else: getattr(sd_model, quant_last_model_name).to(quant_last_model_device) devices.torch_gc(force=True) - if shared.cmd_opts.medvram or shared.opts.diffusers_model_cpu_offload or shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram: + if shared.cmd_opts.medvram or shared.cmd_opts.lowvram or shared.opts.diffusers_offload_mode != "none": quant_last_model_name = op quant_last_model_device = model.device else: diff --git a/modules/shared.py b/modules/shared.py index b5ab3fac8..f316c4508 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -358,11 +358,14 @@ if not (cmd_opts.lowvram or cmd_opts.medvram): if "gpu" in mem_stat: if mem_stat['gpu']['total'] <= 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: cmd_opts.medvram = True + offload_mode_default = "cpu" log.info(f"VRAM: Detected={mem_stat['gpu']['total']} GB Optimization=medvram") else: + offload_mode_default = "none" log.info(f"VRAM: Detected={mem_stat['gpu']['total']} GB Optimization=none") @@ -425,7 +428,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "cross_attention_optimization": OptionInfo(cross_attention_optimization_default, "Attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention(native) }), "sdp_options": OptionInfo(sdp_options_default, "SDP options", gr.CheckboxGroup, {"choices": ['Flash attention', 'Memory attention', 'Math attention'] }), "xformers_options": OptionInfo(['Flash attention'], "xFormers options", gr.CheckboxGroup, {"choices": ['Flash attention'] }), - "dynamic_attention_slice_rate": OptionInfo(4, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.1, "maximum": 16, "step": 0.1, "visible": native}), + "dynamic_attention_slice_rate": OptionInfo(4, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.1, "maximum": 24, "step": 0.1, "visible": native}), "sub_quad_sep": OptionInfo("

Sub-quadratic options

", "", gr.HTML, {"visible": not native}), "sub_quad_q_chunk_size": OptionInfo(512, "Attention query chunk size", gr.Slider, {"minimum": 16, "maximum": 8192, "step": 8, "visible": not native}), "sub_quad_kv_chunk_size": OptionInfo(512, "Attention kv chunk size", gr.Slider, {"minimum": 0, "maximum": 8192, "step": 8, "visible": not native}), @@ -527,8 +530,8 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), { "diffusers_move_refiner": OptionInfo(False, "Move refiner model to CPU when not in use"), "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_model_cpu_offload": OptionInfo(cmd_opts.medvram, "Model CPU offload (--medvram)"), - "diffusers_seq_cpu_offload": OptionInfo(cmd_opts.lowvram, "Sequential CPU offload (--lowvram)"), + "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_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"),