diff --git a/modules/sd_offload.py b/modules/sd_offload.py index a522cbc28..283633ad8 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -7,7 +7,7 @@ import torch import accelerate.hooks import accelerate.utils.modeling from installer import log -from modules import shared, devices, errors, model_quant +from modules import shared, devices, errors, model_quant, sd_models from modules.timer import process as process_timer @@ -92,6 +92,60 @@ def apply_group_offload(sd_model, op:str='model'): return sd_model +def apply_model_offload(sd_model, op:str='model', quiet:bool=False): + try: + shared.log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') + if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: + shared.opts.diffusers_move_base = False + shared.opts.diffusers_move_unet = False + shared.opts.diffusers_move_refiner = False + shared.log.warning(f'Disabling {op} "Move model to CPU" since "Model CPU offload" is enabled') + if not hasattr(sd_model, "_all_hooks") or len(sd_model._all_hooks) == 0: # pylint: disable=protected-access + sd_model.enable_model_cpu_offload(device=devices.device) + else: + sd_model.maybe_free_model_hooks() + set_accelerate(sd_model) + except Exception as e: + shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}') + + +def apply_sequential_offload(sd_model, op:str='model', quiet:bool=False): + try: + shared.log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') + if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: + shared.opts.diffusers_move_base = False + shared.opts.diffusers_move_unet = False + shared.opts.diffusers_move_refiner = False + shared.log.warning(f'Disabling {op} "Move model to CPU" since "Sequential CPU offload" is enabled') + if sd_model.has_accelerate: + if op == "vae": # reapply sequential offload to vae + from accelerate import cpu_offload + sd_model.vae.to(devices.cpu) + cpu_offload(sd_model.vae, devices.device, offload_buffers=len(sd_model.vae._parameters) > 0) # pylint: disable=protected-access + else: + pass # do nothing if offload is already applied + else: + sd_model.enable_sequential_cpu_offload(device=devices.device) + set_accelerate(sd_model) + except Exception as e: + shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}') + + +def disable_offload(sd_model, op:str='model', quiet:bool=False): + if shared.sd_model_type in offload_warn or 'video' in shared.sd_model_type: + shared.log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} type={shared.sd_model.__class__.__name__} large model') + else: + shared.log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') + try: + sd_model.has_accelerate = False + if hasattr(sd_model, 'maybe_free_model_hooks'): + sd_model.maybe_free_model_hooks() + sd_model = accelerate.hooks.remove_hook_from_module(sd_model, recurse=True) + except Exception: + pass + sd_models.move_model(sd_model, devices.device) + + def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False, force:bool=False): global accelerate_dtype_byte_size # pylint: disable=global-statement t0 = time.time() @@ -105,58 +159,13 @@ def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False, force:bool= accelerate.utils.modeling.dtype_byte_size = dtype_byte_size if shared.opts.diffusers_offload_mode == "none": - if shared.sd_model_type in offload_warn or 'video' in shared.sd_model_type: - shared.log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} type={shared.sd_model.__class__.__name__} large model') - else: - shared.log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') - try: - sd_model.has_accelerate = False - if hasattr(sd_model, 'maybe_free_model_hooks'): - sd_model.maybe_free_model_hooks() - sd_model = accelerate.hooks.remove_hook_from_module(sd_model, recurse=True) - except Exception: - pass - try: - sd_model = sd_model.to(devices.device) - except Exception as e: - shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} move device={devices.device} {e}') + disable_offload(sd_model, op=op, quiet=quiet) if shared.opts.diffusers_offload_mode == "model" and hasattr(sd_model, "enable_model_cpu_offload"): - try: - shared.log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') - if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: - shared.opts.diffusers_move_base = False - shared.opts.diffusers_move_unet = False - shared.opts.diffusers_move_refiner = False - shared.log.warning(f'Disabling {op} "Move model to CPU" since "Model CPU offload" is enabled') - if not hasattr(sd_model, "_all_hooks") or len(sd_model._all_hooks) == 0: # pylint: disable=protected-access - sd_model.enable_model_cpu_offload(device=devices.device) - else: - sd_model.maybe_free_model_hooks() - set_accelerate(sd_model) - except Exception as e: - shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}') + apply_model_offload(sd_model, op=op, quiet=quiet) if shared.opts.diffusers_offload_mode == "sequential" and hasattr(sd_model, "enable_sequential_cpu_offload"): - try: - shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') - if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: - shared.opts.diffusers_move_base = False - shared.opts.diffusers_move_unet = False - shared.opts.diffusers_move_refiner = False - shared.log.warning(f'Disabling {op} "Move model to CPU" since "Sequential CPU offload" is enabled') - if sd_model.has_accelerate: - if op == "vae": # reapply sequential offload to vae - from accelerate import cpu_offload - sd_model.vae.to(devices.cpu) - cpu_offload(sd_model.vae, devices.device, offload_buffers=len(sd_model.vae._parameters) > 0) # pylint: disable=protected-access - else: - pass # do nothing if offload is already applied - else: - sd_model.enable_sequential_cpu_offload(device=devices.device) - set_accelerate(sd_model) - except Exception as e: - shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}') + apply_sequential_offload(sd_model, op=op, quiet=quiet) if shared.opts.diffusers_offload_mode == "group": sd_model = apply_group_offload(sd_model, op=op)