diff --git a/CHANGELOG.md b/CHANGELOG.md index 8ade15e58..4f6b10909 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,8 +50,10 @@ - **Memory** improvements: - faster and more compatible *balanced offload* mode - balanced offload: units are now in percentage instead of bytes - - balanced offload: add both high and low watermark - default is 25% for low-watermark (skip offload if memory usage is below 25%) and 70% high-watermark (must offload if memory usage is above 70%) + - balanced offload: add both high and low watermark and pinned threshold, defaults as below + 25% for low-watermark: skip offload if memory usage is below 25% + 70% high-watermark: must offload if memory usage is above 70% + 15% pin-watermark: any model component smaller than 15% of total memory is pinned and not offloaded - change-in-behavior: low-end systems, triggered by either `lowvrwam` or by detection of <=4GB will use *sequential offload* all other systems use *balanced offload* by default (can be changed in settings) diff --git a/modules/memstats.py b/modules/memstats.py index fd5f152a0..d43e5bbfa 100644 --- a/modules/memstats.py +++ b/modules/memstats.py @@ -4,6 +4,8 @@ import torch from modules import shared, errors fail_once = False +mem = {} + def gb(val: float): return round(val / 1024 / 1024 / 1024, 2) @@ -11,7 +13,7 @@ def gb(val: float): def memory_stats(): global fail_once # pylint: disable=global-statement - mem = {} + mem.clear() try: process = psutil.Process(os.getpid()) res = process.memory_info() @@ -41,6 +43,10 @@ def memory_stats(): return mem +def memory_cache(): + return mem + + def ram_stats(): try: process = psutil.Process(os.getpid()) diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 627eb281f..9f12e44ab 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -349,7 +349,7 @@ def process_refine(p: processing.StableDiffusionProcessing, output): def process_decode(p: processing.StableDiffusionProcessing, output): - shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae']) if output is not None: if not hasattr(output, 'images') and hasattr(output, 'frames'): shared.log.debug(f'Generated: frames={len(output.frames[0])}') @@ -463,7 +463,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing): timer.process.record('decode') shared.sd_model = orig_pipeline - shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) + # shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) if p.state == '': global last_p # pylint: disable=global-statement diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 1c4a45f07..77a89c512 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -104,8 +104,6 @@ def full_vae_decode(latents, model): if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False): base_device = sd_models.move_base(model, devices.cpu) - if shared.opts.diffusers_offload_mode == "balanced": - shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) elif shared.opts.diffusers_offload_mode != "sequential": sd_models.move_model(model.vae, devices.device) @@ -159,8 +157,8 @@ def full_vae_decode(latents, model): model.vae.apply(sd_models.convert_to_faketensors) devices.torch_gc(force=True) - if shared.opts.diffusers_offload_mode == "balanced": - shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) + # if shared.opts.diffusers_offload_mode == "balanced": + # shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) elif shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and base_device is not None: sd_models.move_base(model, base_device) t1 = time.time() diff --git a/modules/sd_models.py b/modules/sd_models.py index 8853916e4..bd69ba45b 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -18,7 +18,7 @@ from omegaconf import OmegaConf from ldm.util import instantiate_from_config from modules import paths, shared, shared_state, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_config, sd_models_compile, sd_hijack_accelerate, sd_detect from modules.timer import Timer, process as process_timer -from modules.memstats import memory_stats +from modules.memstats import memory_stats, memory_cache from modules.modeldata import model_data from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoints_list, checkpoint_titles, get_closet_checkpoint_match, model_hash, update_model_hashes, setup_model, write_metadata, read_metadata_from_safetensors # pylint: disable=unused-import @@ -416,9 +416,10 @@ class OffloadHook(accelerate.hooks.ModelHook): offload_hook_instance = None +offload_component_map = {} -def apply_balanced_offload(sd_model): +def apply_balanced_offload(sd_model, exclude=[]): global offload_hook_instance # pylint: disable=global-statement if shared.opts.diffusers_offload_mode != "balanced": return sd_model @@ -428,8 +429,6 @@ def apply_balanced_offload(sd_model): excluded = ['OmniGenPipeline'] if sd_model.__class__.__name__ in excluded: return sd_model - fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access - debug_move(f'Apply offload: type=balanced fn={fn}') checkpoint_name = sd_model.sd_checkpoint_info.name if getattr(sd_model, "sd_checkpoint_info", None) is not None else None if checkpoint_name is None: checkpoint_name = sd_model.__class__.__name__ @@ -442,32 +441,38 @@ def apply_balanced_offload(sd_model): keys = pipe._internal_dict.keys() # pylint: disable=protected-access else: keys = get_signature(pipe).keys() + keys = [k for k in keys if k not in exclude and not k.startswith('_')] for module_name in keys: # pylint: disable=protected-access module = getattr(pipe, module_name, None) - if isinstance(module, torch.nn.Module): - network_layer_name = getattr(module, "network_layer_name", None) - device_map = getattr(module, "balanced_offload_device_map", None) - max_memory = getattr(module, "balanced_offload_max_memory", None) - module = accelerate.hooks.remove_hook_from_module(module, recurse=True) - try: - do_offload = used_gpu > 100 * shared.opts.diffusers_offload_min_gpu_memory - debug_move(f'Balanced offload: gpu={used_gpu} ram={used_ram} current={module.device} dtype={module.dtype} op={"move" if do_offload else "skip"} component={module.__class__.__name__}') - if do_offload: - module = module.to(devices.cpu) - used_gpu, used_ram = devices.torch_gc(fast=True, force=True) - except Exception as e: - if 'bitsandbytes' not in str(e): - shared.log.error(f'Balanced offload: module={module_name} {e}') - if os.environ.get('SD_MOVE_DEBUG', None): - errors.display(e, f'Balanced offload: module={module_name}') - module.offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name) - module = accelerate.hooks.add_hook_to_module(module, offload_hook_instance, append=True) - module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access - if network_layer_name: - module.network_layer_name = network_layer_name - if device_map and max_memory: - module.balanced_offload_device_map = device_map - module.balanced_offload_max_memory = max_memory + if not isinstance(module, torch.nn.Module): + continue + network_layer_name = getattr(module, "network_layer_name", None) + device_map = getattr(module, "balanced_offload_device_map", None) + max_memory = getattr(module, "balanced_offload_max_memory", None) + module = accelerate.hooks.remove_hook_from_module(module, recurse=True) + module_size = offload_component_map.get(module_name, None) + if module_size is None: + module_size = sum(p.numel()*p.element_size() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024 + offload_component_map[module_name] = module_size + do_offload = (used_gpu > 100 * shared.opts.diffusers_offload_min_gpu_memory) and (module_size > shared.gpu_memory * shared.opts.diffusers_offload_pin_gpu_memory) + try: + debug_move(f'Balanced offload: gpu={used_gpu} ram={used_ram} current={module.device} dtype={module.dtype} op={"move" if do_offload else "skip"} component={module.__class__.__name__} size={module_size:.3f}') + if do_offload and module.device != devices.cpu: + module = module.to(devices.cpu) + used_gpu, used_ram = devices.torch_gc(fast=True, force=True) + except Exception as e: + if 'bitsandbytes' not in str(e): + shared.log.error(f'Balanced offload: module={module_name} {e}') + if os.environ.get('SD_MOVE_DEBUG', None): + errors.display(e, f'Balanced offload: module={module_name}') + module.offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name) + module = accelerate.hooks.add_hook_to_module(module, offload_hook_instance, append=True) + module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access + if network_layer_name: + module.network_layer_name = network_layer_name + if device_map and max_memory: + module.balanced_offload_device_map = device_map + module.balanced_offload_max_memory = max_memory apply_balanced_offload_to_module(sd_model) if hasattr(sd_model, "pipe"): @@ -478,7 +483,10 @@ def apply_balanced_offload(sd_model): apply_balanced_offload_to_module(sd_model.decoder_pipe) set_accelerate(sd_model) devices.torch_gc(fast=True) - process_timer.add('offload', time.time() - t0) + t = time.time() - t0 + process_timer.add('offload', t) + fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access + debug_move(f'Apply offload: time={t:.2f} type=balanced fn={fn}') return sd_model diff --git a/modules/shared.py b/modules/shared.py index 256850d21..97c32bd84 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -483,6 +483,7 @@ options_templates.update(options_section(('sd', "Models & Loading"), { "diffusers_offload_mode": OptionInfo(startup_offload_mode, "Model offload mode", gr.Radio, {"choices": ['none', 'balanced', 'model', 'sequential']}), "diffusers_offload_min_gpu_memory": OptionInfo(0.25, "Balanced offload GPU low watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }), "diffusers_offload_max_gpu_memory": OptionInfo(0.70, "Balanced offload GPU high watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }), + "diffusers_offload_pin_gpu_memory": OptionInfo(0.15, "Balanced offload GPU pin watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }), "diffusers_offload_max_cpu_memory": OptionInfo(0.90, "Balanced offload CPU high watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }), "advanced_sep": OptionInfo("

Advanced Options

", "", gr.HTML), diff --git a/scripts/cogvideo.py b/scripts/cogvideo.py index 284109857..a5efcd3e6 100644 --- a/scripts/cogvideo.py +++ b/scripts/cogvideo.py @@ -180,7 +180,7 @@ class Script(scripts.Script): else: shared.sd_model = sd_models.switch_pipe(diffusers.CogVideoXPipeline, shared.sd_model) args['num_frames'] = p.frames # only txt2vid has num_frames - shared.log.info(f'CogVideoX: class={shared.sd_model.__class__.__name__} frames={p.frames} input={args.get('video', None) or args.get('image', None)}') + shared.log.info(f"CogVideoX: class={shared.sd_model.__class__.__name__} frames={p.frames} input={args.get('video', None) or args.get('image', None)}") if debug: shared.log.debug(f'CogVideoX args: {args}') frames = shared.sd_model(**args).frames[0]