import os import re import sys import time import math import inspect import itertools import torch import accelerate.hooks import accelerate.utils.modeling from modules.logger import log from modules import shared, devices, errors, model_quant, sd_models, sd_offload_aux from modules.timer import process as process_timer debug = os.environ.get('SD_MOVE_DEBUG', None) is not None verbose = os.environ.get('SD_MOVE_VERBOSE', None) is not None debug_move = log.trace if debug else lambda *args, **kwargs: None offload_allow_none = ['sd', 'sdxl'] offload_post = ['h1'] offload_hook_instance = None balanced_offload_exclude = ['CogView4Pipeline', 'MeissonicPipeline'] no_split_module_classes = [ "Linear", "Conv1d", "Conv2d", "Conv3d", "ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d", "Embedding", "SDNQLinear", "SDNQConv1d", "SDNQConv2d", "SDNQConv3d", "SDNQConvTranspose1d", "SDNQConvTranspose2d", "SDNQConvTranspose3d", "SDNQEmbedding", "WanTransformerBlock", "MiniMaxH3TransformerBlock", "MiniMaxH3TokenRefinerBlock", ] accelerate_dtype_byte_size = None move_stream = None def dtype_byte_size(dtype: torch.dtype): try: if dtype in [torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz]: dtype = accelerate.utils.modeling.CustomDtype.FP8 except Exception: # catch since older torch many not have defined dtypes pass return accelerate_dtype_byte_size(dtype) def get_signature(cls): signature = inspect.signature(cls.__init__, follow_wrapped=True) return signature.parameters def disable_offload(sd_model): remove_group_offload(sd_model) # group hooks block the meta move at unload, keeping component weights alive for as long as any reference to the pipe survives if not getattr(sd_model, 'has_accelerate', False): return for module_name in get_module_names(sd_model): module = getattr(sd_model, module_name, None) if isinstance(module, torch.nn.Module): network_layer_name = getattr(module, "network_layer_name", None) try: module = accelerate.hooks.remove_hook_from_module(module, recurse=True) except Exception as e: log.warning(f'Offload: remove hook module={module_name} {e}') if network_layer_name: module.network_layer_name = network_layer_name sd_model.has_accelerate = False def set_accelerate(sd_model): def set_accelerate_to_module(model): if hasattr(model, "pipe"): set_accelerate_to_module(model.pipe) for module_name in get_module_names(model): component = getattr(model, module_name, None) if isinstance(component, torch.nn.Module): component.has_accelerate = True sd_model.has_accelerate = True set_accelerate_to_module(sd_model) if hasattr(sd_model, "prior_pipe"): set_accelerate_to_module(sd_model.prior_pipe) if hasattr(sd_model, "decoder_pipe"): set_accelerate_to_module(sd_model.decoder_pipe) def group_offload_config(main: bool) -> dict: """Effective group offload settings for one component. Components that run once per generation take the leaf no-stream policy regardless of the main settings, so their weights are never held in pinned host memory.""" stream = shared.opts.group_offload_stream if main else False blocks = max(1, int(shared.opts.group_offload_blocks)) if stream and blocks != 1: blocks = 1 # streamed prefetch supports one block per group; upstream clamps with a warning otherwise return { 'offload_type': shared.opts.group_offload_type if main else 'leaf_level', 'num_blocks_per_group': blocks, 'non_blocking': shared.opts.diffusers_offload_nonblocking, 'use_stream': stream, 'record_stream': shared.opts.group_offload_record and stream, # record without streams is rejected upstream 'low_cpu_mem_usage': stream and not shared.opts.group_offload_pin, } def remove_group_offload_component(module) -> bool: if getattr(module, 'sdnext_group_offload_sig', None) is None: module = getattr(module, 'model', None) # wrapper components carry the hooks on the inner model if module is None or getattr(module, 'sdnext_group_offload_sig', None) is None: return False from diffusers.hooks.group_offloading import _GROUP_OFFLOADING, _LAYER_EXECUTION_TRACKER, _LAZY_PREFETCH_GROUP_OFFLOADING from diffusers.hooks.hooks import HookRegistry registry = HookRegistry.check_if_exists_or_initialize(module) registry.remove_hook(_GROUP_OFFLOADING, recurse=True) registry.remove_hook(_LAYER_EXECUTION_TRACKER, recurse=True) registry.remove_hook(_LAZY_PREFETCH_GROUP_OFFLOADING, recurse=True) module.sdnext_group_offload_sig = None return True def remove_group_offload(sd_model): removed = [] for module_name in get_module_names(sd_model): module = getattr(sd_model, module_name, None) if isinstance(module, torch.nn.Module) and remove_group_offload_component(module): removed.append(module_name) for module_name in getattr(sd_model, 'sdnext_ondemand_modules', None) or []: module = getattr(sd_model, module_name, None) if module is not None: module.sdnext_ondemand = False if hasattr(module, '_hf_hook'): module = accelerate.hooks.remove_hook_from_module(module, recurse=True) removed.append(f'{module_name}:ondemand') if getattr(sd_model, 'sdnext_ondemand_modules', None): sd_model.sdnext_ondemand_modules = [] if removed: log.debug(f'Offload: type=group op=remove modules={removed}') def apply_group_offload_component(module, module_name: str, main: bool) -> bool: """Apply group offload to one component. Re-application with unchanged settings is a no-op: the hooks silently keep their original config when re-applied and raise before the first forward, so a changed config must remove the old hooks first.""" from diffusers.hooks import apply_group_offloading cfg = group_offload_config(main) if cfg['use_stream'] and not cfg['low_cpu_mem_usage']: size_gb, _params = get_module_size(module) pin_ok = getattr(module, 'sdnext_group_offload_pin', None) if pin_ok is None: # decide once per module: a granted pin moves the weights into locked memory, so re-reading available on the next apply would see it lower by the pinned size and revoke its own grant from modules import memstats avail_gb = memstats.ram_stats().get('avail', 0) reserve_gb = max(8.0, 0.25 * shared.cpu_memory) # pinned pages cannot be reclaimed or swapped, so a quarter of the machine, floored at 8 GB, stays pageable for the process and page cache limit_gb = (avail_gb - reserve_gb) if avail_gb > 0 else (0.5 * shared.cpu_memory) # budget from memory free right now; total-derived ceiling only when psutil cannot say pin_ok = size_gb <= limit_gb module.sdnext_group_offload_pin = pin_ok module.sdnext_group_offload_pin_limit = limit_gb if not pin_ok: # unpinned streaming degrades to per-transfer staging and leaf groups make that a per-module cost, # so the whole leaf+stream shape goes with the pin: few large synchronous groups instead cfg['low_cpu_mem_usage'] = True cfg['use_stream'] = False cfg['record_stream'] = False cfg['offload_type'] = 'block_level' cfg['num_blocks_per_group'] = max(4, int(shared.opts.group_offload_blocks)) log.warning(f'Offload: type=group module={module_name} size={size_gb:.3f} limit={getattr(module, "sdnext_group_offload_pin_limit", 0):.3f} pin=denied type=block_level blocks={cfg["num_blocks_per_group"]} expect ~{size_gb:.0f} GB transferred per step') sig = f'{devices.device}:{main}:' + ':'.join(str(v) for v in cfg.values()) if getattr(module, 'sdnext_group_offload_sig', None) == sig: return False if hasattr(module, '_hf_hook'): # leftover accelerate hooks from a previous offload mode abort the group apply upstream module = accelerate.hooks.remove_hook_from_module(module, recurse=True) remove_group_offload_component(module) module.requires_grad_(False) log.debug(f'Offload: type=group op=apply type={shared.opts.group_offload_type} module={module_name} pin={cfg["use_stream"] and not cfg["low_cpu_mem_usage"]}') # before the apply: pinning large components takes a while and would otherwise run silently apply_group_offloading(module, onload_device=devices.device, offload_device=devices.cpu, **cfg) module.sdnext_group_offload_sig = sig return True def set_group_resident(module): """VAE-class components never take group hooks: the hooks are forward-scoped, while pipelines enter through encode/decode, and tiled calls re-enter per tile.""" if hasattr(module, '_hf_hook'): module = accelerate.hooks.remove_hook_from_module(module, recurse=True) remove_group_offload_component(module) module.requires_grad_(False) module.to(devices.device) def group_offload_role(module_name) -> str: if any(m in module_name for m in ['vae']): return 'vae' if any(m in module_name for m in ['text_encoder', 'image_encoder', 'safety_checker']): return 'aux' return 'main' def has_entry_bridge(module) -> bool: """Entry points decorated with diffusers' apply_forward_hook fire _hf_hook.pre_forward, which is what carries the on-demand onload for encode and decode calls that bypass forward.""" for name in ('decode', 'encode'): fn = getattr(module, name, None) if fn is not None and getattr(fn, '__qualname__', '').startswith('apply_forward_hook'): return True return False class OnDemandHook(accelerate.hooks.ModelHook): """Whole-module onload for components entered through decode or encode rather than forward. Tiled calls re-enter inside one entry point, so the module is on device before the first tile; the return to cpu happens at the processing seams once outputs are materialized.""" def pre_forward(self, module, *args, **kwargs): param = next(module.parameters(), None) if param is not None and not devices.same_device(param.device, devices.device): t0 = time.time() module.to(devices.device, non_blocking=shared.opts.diffusers_offload_nonblocking) t1 = time.time() process_timer.add('onload', t1 - t0) debug_move(f'Offload: type=ondemand op=onload module={module.__class__.__name__} nonblocking={shared.opts.diffusers_offload_nonblocking} time={t1 - t0:.3f}') # working so no need to log return args, kwargs def apply_group_offload_vae(sd_model, module, module_name: str) -> str: """Placement policy for vae-class components, which never take group hooks. Components onload whole when their decode or encode entry point fires.""" if not has_entry_bridge(module): log.warning(f'Offload: type=group module={module_name} class={module.__class__.__name__} no entry bridge') set_group_resident(module) # TODO group offload: this will fail as vae will end up on cpu and there will be nothing to pull it back on gpu module.sdnext_ondemand = False # a lingering stamp would let the seams offload a component with no onload hook names = getattr(sd_model, 'sdnext_ondemand_modules', None) or [] if module_name in names: sd_model.sdnext_ondemand_modules = [n for n in names if n != module_name] return False if not getattr(module, 'sdnext_ondemand', False) or not hasattr(module, '_hf_hook'): if hasattr(module, '_hf_hook'): module = accelerate.hooks.remove_hook_from_module(module, recurse=True) remove_group_offload_component(module) module.requires_grad_(False) accelerate.hooks.add_hook_to_module(module, OnDemandHook(), append=False) module.sdnext_ondemand = True module.to(devices.cpu) names = getattr(sd_model, 'sdnext_ondemand_modules', None) or [] if module_name not in names: sd_model.sdnext_ondemand_modules = names + [module_name] return True def offload_ondemand(sd_model, include=[], exclude=[], reason='', force=False): """Return on-demand components to cpu once their outputs are materialized.""" if sd_model is None: return if force: # if force, all loaded modules are candidates names = [name for name, component in sd_model.components.items() if isinstance(component, torch.nn.Module)] else: names = getattr(sd_model, 'sdnext_ondemand_modules', None) if not names and hasattr(sd_model, 'pipe'): sd_model = sd_model.pipe names = getattr(sd_model, 'sdnext_ondemand_modules', None) for module_name in names or []: if include and module_name not in include: continue if exclude and module_name in exclude: continue module = getattr(sd_model, module_name, None) param = next(module.parameters(), None) if module is not None else None if param is not None and not devices.same_device(param.device, devices.cpu): t0 = time.time() module.to(devices.cpu, non_blocking=shared.opts.diffusers_offload_nonblocking) dt = time.time() - t0 process_timer.add('offload', dt) debug_move(f'Offload: type=ondemand op=offload module={module_name} nonblocking={shared.opts.diffusers_offload_nonblocking} reason="{reason}" time={dt:.3f}') devices.torch_gc() def report_group_stats(sd_model, module_names): """Per-component stats block once per loaded model; balanced mode prints its own from the hook map.""" if getattr(sd_model, 'sdnext_group_stats_reported', False): return sd_model.sdnext_group_stats_reported = True total = 0.0 counted = [] for module_name in module_names: module = getattr(sd_model, module_name, None) if isinstance(module, torch.nn.Module): total += get_module_size(module)[0] counted.append(module_name) report_model_stats(module_name, module) log.info(f'Model class={sd_model.__class__.__name__} modules={len(counted)} size={total:.3f}') def apply_modular_group_offload(sd_model): """Per-component group offload for modular pipelines, which lack the pipeline-level enable_*_offload entry points. The model and sequential modes also route here.""" if shared.opts.diffusers_offload_mode != 'group' and not getattr(sd_model, 'sdnext_modular_offload_warned', False): sd_model.sdnext_modular_offload_warned = True log.warning(f'Offload: desired={shared.opts.diffusers_offload_mode} override=group reason="modular pipeline"') applied = [] loaded = [name for name, component in sd_model.components.items() if isinstance(component, torch.nn.Module)] for name in loaded: module = getattr(sd_model, name, None) if 'text_encoder' in name: if apply_group_offload_component(getattr(module, 'model', module), name, main=False): applied.append(name) if 'vae' in name: if apply_group_offload_vae(sd_model, module, name): applied.append(name) else: if apply_group_offload_component(module, name, main=True): applied.append(name) # has_accelerate stays unset: group hooks are not accelerate hooks, and the modular pipeline's own to() skips group-offloaded components when move_model runs log.info(f'Offload: type=group type={shared.opts.group_offload_type} modules={applied}') report_group_stats(sd_model, ('transformer', 'transformer_ref', 'text_encoder', 'vae', 'audio_vae')) def apply_group_offload(sd_model): applied, resident, ondemand = [], [], [] for module_name in get_module_names(sd_model): module = getattr(sd_model, module_name, None) if not isinstance(module, torch.nn.Module): continue try: role = group_offload_role(module_name) if role == 'vae': if apply_group_offload_vae(sd_model, module, module_name) == 'ondemand': ondemand.append(module_name) else: resident.append(module_name) elif apply_group_offload_component(module, module_name, main=role == 'main'): applied.append(module_name) except Exception as e: log.error(f'Offload: type=group module={module_name} {e}') set_accelerate(sd_model) if applied: log.info(f'Offload: type=group type={shared.opts.group_offload_type} modules={applied} resident={resident} ondemand={ondemand}') report_group_stats(sd_model, get_module_names(sd_model)) return sd_model def apply_model_offload(sd_model, quiet:bool=False): try: remove_group_offload(sd_model) log.quiet(quiet, f'Offload: type={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') 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: log.error(f'Offload: type={shared.opts.diffusers_offload_mode} {e}') def apply_sequential_offload(sd_model, op:str='model', quiet:bool=False): try: remove_group_offload(sd_model) log.quiet(quiet, f'Offload: type={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') 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: log.error(f'Offload: type={shared.opts.diffusers_offload_mode} {e}') def apply_none_offload(sd_model, quiet:bool=False): if shared.sd_model_type not in offload_allow_none: log.warning(f'Offload: type={shared.opts.diffusers_offload_mode} cls={shared.sd_model.__class__.__name__} large model') else: log.quiet(quiet, f'Offload: type={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') try: sd_model.has_accelerate = False remove_group_offload(sd_model) 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, force=True) 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() if sd_model is None: log.warning(f'{op} is not loaded') return if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate): sd_model.has_accelerate = False if accelerate_dtype_byte_size is None: accelerate_dtype_byte_size = accelerate.utils.modeling.dtype_byte_size accelerate.utils.modeling.dtype_byte_size = dtype_byte_size if sd_models.get_diffusers_task(sd_model) == sd_models.DiffusersTaskType.MODULAR and shared.opts.diffusers_offload_mode in {'model', 'sequential', 'group'}: apply_modular_group_offload(sd_model) process_timer.add('offload', time.time() - t0) return if shared.opts.diffusers_offload_mode == "none": log.warning('Offload: type=none "use balanced offload with model type set not to offload"') apply_none_offload(sd_model, quiet=quiet) if shared.opts.diffusers_offload_mode == "model" and hasattr(sd_model, "enable_model_cpu_offload"): log.warning('Offload: type=model "use balanced offload instead"') apply_model_offload(sd_model, quiet=quiet) if shared.opts.diffusers_offload_mode == "sequential" and hasattr(sd_model, "enable_sequential_cpu_offload"): apply_sequential_offload(sd_model, op=op, quiet=quiet) if shared.opts.diffusers_offload_mode == "group": sd_model = apply_group_offload(sd_model) if shared.opts.diffusers_offload_mode == "balanced": sd_model = apply_balanced_offload(sd_model, force=force) process_timer.add('offload', time.time() - t0) class OffloadHook(accelerate.hooks.ModelHook): def __init__(self, checkpoint_name): if shared.opts.diffusers_offload_max_gpu_memory > 1: shared.opts.diffusers_offload_max_gpu_memory = 0.75 if shared.opts.diffusers_offload_max_cpu_memory > 1: shared.opts.diffusers_offload_max_cpu_memory = 0.75 self.checkpoint_name = checkpoint_name self.min_watermark = shared.opts.diffusers_offload_min_gpu_memory self.max_watermark = shared.opts.diffusers_offload_max_gpu_memory self.cpu_watermark = shared.opts.diffusers_offload_max_cpu_memory self.offload_always = [m.strip() for m in re.split(';|,| ', shared.opts.diffusers_offload_always) if len(m.strip()) > 2] self.offload_never = [m.strip() for m in re.split(';|,| ', shared.opts.diffusers_offload_never) if len(m.strip()) > 2] self.gpu = int(shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory * 1024*1024*1024) self.cpu = int(shared.cpu_memory * shared.opts.diffusers_offload_max_cpu_memory * 1024*1024*1024) self.offload_map = {} self.param_map = {} self.last_pre = None self.last_post = None self.last_cls = None gpu = f'{(shared.gpu_memory * shared.opts.diffusers_offload_min_gpu_memory):.2f}-{(shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory):.2f}:{shared.gpu_memory:.2f}' log.info(f'Offload: type=balanced op=init watermark={self.min_watermark}-{self.max_watermark} gpu={gpu} cpu={shared.cpu_memory:.3f} limit={shared.opts.cuda_mem_fraction:.2f} always={self.offload_always} never={self.offload_never} pre={shared.opts.diffusers_offload_pre} streams={shared.opts.diffusers_offload_streams}') self.validate() super().__init__() def validate(self): if shared.opts.diffusers_offload_mode != 'balanced': return if shared.opts.diffusers_offload_min_gpu_memory < 0 or shared.opts.diffusers_offload_min_gpu_memory > 1: shared.opts.diffusers_offload_min_gpu_memory = 0.2 log.warning(f'Offload: type=balanced op=validate: watermark low={shared.opts.diffusers_offload_min_gpu_memory} invalid value') if shared.opts.diffusers_offload_max_gpu_memory < 0.1 or shared.opts.diffusers_offload_max_gpu_memory > 1: shared.opts.diffusers_offload_max_gpu_memory = 0.7 log.warning(f'Offload: type=balanced op=validate: watermark high={shared.opts.diffusers_offload_max_gpu_memory} invalid value') if shared.opts.diffusers_offload_min_gpu_memory > shared.opts.diffusers_offload_max_gpu_memory: shared.opts.diffusers_offload_min_gpu_memory = shared.opts.diffusers_offload_max_gpu_memory log.warning(f'Offload: type=balanced op=validate: watermark low={shared.opts.diffusers_offload_min_gpu_memory} reset') if shared.opts.diffusers_offload_max_gpu_memory * shared.gpu_memory < 3: log.warning(f'Offload: type=balanced op=validate: watermark high={shared.opts.diffusers_offload_max_gpu_memory} low memory') def model_size(self): return sum(self.offload_map.values()) def matches(self, module, names: list, module_name: str | None = None) -> bool: """Match against an always/never list by class name or by pipeline component name. Component entries such as `text_encoder` cover every architecture without listing each encoder class.""" if module.__class__.__name__ in names: return True module_name = module_name or getattr(module, 'module_name', None) return module_name is not None and module_name in names def init_hook(self, module): return module def offload_allowed(self, module): if hasattr(module, "offload_never"): return False if hasattr(module, 'nets') and any(hasattr(n, "offload_never") for n in module.nets): return False if shared.sd_model_type.lower() in [m.lower().strip() for m in re.split(r'[ ,]+', shared.opts.models_not_to_offload)]: return False return True def pre_forward(self, module, *args, **kwargs): _id = id(module) do_offload = (self.last_pre != _id) or (module.__class__.__name__ != self.last_cls) if do_offload and self.offload_allowed(module): # offload every other module first time when new module starts pre-forward if shared.opts.diffusers_offload_pre: t0 = time.time() debug_move(f'Offload: type=balanced op=pre module={module.__class__.__name__}') sd_offload_aux.evict_aux(reason=f'pre:{module.__class__.__name__}') for pipe in get_pipe_variants(): for module_name in get_module_names(pipe): module_instance = getattr(pipe, module_name, None) if (module_instance is not None) and (_id != id(module_instance)) and (not self.matches(module_instance, self.offload_never, module_name)) and (not devices.same_device(getattr(module_instance, "device", devices.cpu), devices.cpu)): apply_balanced_offload_to_module(module_instance, op='pre') self.last_cls = module.__class__.__name__ process_timer.add('offload', time.time() - t0) if not devices.same_device(getattr(module, "device", devices.cpu), devices.device): # move-to-device t0 = time.time() device_index = torch.device(devices.device).index if device_index is None: device_index = 0 max_memory = { device_index: self.gpu, "cpu": self.cpu } device_map = getattr(module, "balanced_offload_device_map", None) if (device_map is None) or (max_memory != getattr(module, "balanced_offload_max_memory", None)): device_map = accelerate.infer_auto_device_map(module, max_memory=max_memory, no_split_module_classes=no_split_module_classes, verbose=verbose, clean_result=False, ) offload_dir = getattr(module, "offload_dir", os.path.join(shared.opts.accelerate_offload_path, module.__class__.__name__)) if debug: log.trace(f'Offload: type=balanced op=dispatch map={device_map}') if device_map is not None: skip_keys = getattr(module, "_skip_keys", None) try: module = accelerate.dispatch_model(module, main_device=torch.device(devices.device), device_map=device_map, offload_dir=offload_dir, skip_keys=skip_keys, force_hooks=True, ) except Exception as e: # reapply hook log.warning(f'Offload: type=balanced op=dispatch module={module.__class__.__name__} {e}') module = accelerate.hooks.remove_hook_from_module(module, recurse=True) module.balanced_offload_device_map = None sd_models.move_model(module, devices.device, force=True) module = accelerate.hooks.add_hook_to_module(module, self, append=True) module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access module.balanced_offload_device_map = device_map module.balanced_offload_max_memory = max_memory process_timer.add('onload', time.time() - t0) if debug: for _i, pipe in enumerate(get_pipe_variants()): for module_name in get_module_names(pipe): module_instance = getattr(pipe, module_name, None) log.trace(f'Offload: type=balanced op=pre:status forward={module.__class__.__name__} module={module_name} class={module_instance.__class__.__name__} pipe={_i} device={getattr(module_instance, "device", devices.cpu)} dtype={module_instance.dtype}') self.last_pre = _id return args, kwargs def post_forward(self, module, output): if self.last_post != id(module): self.last_post = id(module) if getattr(module, "offload_post", False) and (module.device != devices.cpu): apply_balanced_offload_to_module(module, op='post') return output def detach_hook(self, module): return module def get_pipe_variants(pipe=None): if pipe is None: if shared.sd_loaded: pipe = shared.sd_model else: return [pipe] variants = [pipe] if hasattr(pipe, "pipe"): variants.append(pipe.pipe) if hasattr(pipe, "prior_pipe"): variants.append(pipe.prior_pipe) if hasattr(pipe, "decoder_pipe"): variants.append(pipe.decoder_pipe) return variants def get_module_names(pipe=None, exclude=None): def is_valid(module): if isinstance(getattr(pipe, module, None), torch.nn.ModuleDict): return True if isinstance(getattr(pipe, module, None), torch.nn.ModuleList): return True if isinstance(getattr(pipe, module, None), torch.nn.Module): return True return False if exclude is None: exclude = [] if pipe is None: if shared.sd_loaded: pipe = shared.sd_model else: return [] modules_names = [] try: dict_keys = pipe._internal_dict.keys() # pylint: disable=protected-access modules_names.extend(dict_keys) except Exception: pass try: dict_keys = get_signature(pipe).keys() modules_names.extend(dict_keys) except Exception: pass modules_names = [m for m in modules_names if m not in exclude and not m.startswith('_')] modules_names = [m for m in modules_names if is_valid(m)] modules_names = sorted(set(modules_names)) return modules_names def get_module_memory(module: torch.nn.Module) -> dict[str, float]: tensors = list(itertools.chain(module.parameters(), module.buffers())) logical_gib = sum(tensor.numel() * tensor.element_size() for tensor in tensors) / 1024**3 storages = {} for tensor in tensors: try: storage = tensor.untyped_storage() except (AttributeError, RuntimeError): continue storages[(storage.data_ptr(), storage.nbytes())] = storage.nbytes() storage_gib = sum(storages.values()) / 1024**3 return { "logical": round(logical_gib, 3), "storage": round(storage_gib, 3), "overhead": round(storage_gib - logical_gib, 3), "tensors": len(tensors), "storages": len(storages), } def get_module_size(module: torch.nn.Module) -> tuple[float, float]: module_size = 0 param_num = 0 if not isinstance(module, torch.nn.Module): return 0, 0 try: # module_size = sum(p.numel() * p.element_size() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024 tensors = set(itertools.chain(module.parameters(recurse=True), module.buffers(recurse=True))) module_size = sum(t.numel() * t.element_size() for t in tensors) / 1024**3 param_num = sum(p.numel() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024 except Exception as e: log.error(f'Offload: type=balanced op=calc module={module.__class__.__name__} {e}') module_size = 0 param_num = 0 return module_size, param_num def get_module_sizes(pipe=None, exclude=None): if exclude is None: exclude = [] modules = {} for module_name in get_module_names(pipe, exclude): module_size = offload_hook_instance.offload_map.get(module_name, None) if module_size is None: module = getattr(pipe, module_name, None) module_size, param_num = get_module_size(module) offload_hook_instance.offload_map[module_name] = module_size offload_hook_instance.param_map[module_name] = param_num modules[module_name] = module_size modules = sorted(modules.items(), key=lambda x: x[1], reverse=True) return modules def move_module_to_cpu(module, op='unk', force:bool=False): def do_move(module): if shared.opts.diffusers_offload_streams: global move_stream # pylint: disable=global-statement if move_stream is None: move_stream = torch.cuda.Stream(device=devices.device) with torch.cuda.stream(move_stream): module = module.to(devices.cpu) else: module = module.to(devices.cpu) return module try: module_name = getattr(module, "module_name", module.__class__.__name__) module_size = offload_hook_instance.offload_map.get(module_name, offload_hook_instance.model_size()) used_gpu, used_ram = devices.torch_gc(fast=True) perc_gpu = used_gpu / shared.gpu_memory prev_gpu = used_gpu module_cls = module.__class__.__name__ op = f'{op}:skip' if force: op = f'{op}:force' module = do_move(module) used_gpu -= module_size elif offload_hook_instance.matches(module, offload_hook_instance.offload_never, module_name): op = f'{op}:never' elif offload_hook_instance.matches(module, offload_hook_instance.offload_always, module_name): op = f'{op}:always' module = do_move(module) used_gpu -= module_size elif perc_gpu > shared.opts.diffusers_offload_min_gpu_memory: op = f'{op}:mem' module = do_move(module) used_gpu -= module_size if debug: quant = getattr(module, "quantization_method", None) debug_move(f'Offload: type=balanced op={op} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f}:{shared.opts.diffusers_offload_min_gpu_memory} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={module_cls} size={module_size:.3f}') except Exception as e: if 'out of memory' in str(e): devices.torch_gc(fast=True, force=True, reason='oom') elif 'bitsandbytes' in str(e): pass else: log.error(f'Offload: type=balanced op=apply module={getattr(module, "__name__", None)} cls={module.__class__ if inspect.isclass(module) else None} {e}') if os.environ.get('SD_MOVE_DEBUG', None): errors.display(e, f'Offload: type=balanced op=apply module={getattr(module, "__name__", None)}') def apply_balanced_offload_to_module(module, op="apply", force:bool=False): module_name = getattr(module, "module_name", module.__class__.__name__) 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) try: module = accelerate.hooks.remove_hook_from_module(module, recurse=True) except Exception as e: log.warning(f'Offload remove hook: module={module_name} {e}') move_module_to_cpu(module, op=op, force=force) try: module = accelerate.hooks.add_hook_to_module(module, offload_hook_instance, append=True) except Exception as e: log.warning(f'Offload add hook: module={module_name} {e}') 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 module.offload_post = shared.sd_model_type in offload_post and module_name.startswith("text_encoder") if shared.opts.layerwise_quantization or getattr(module, 'quantization_method', None) == 'LayerWise': model_quant.apply_layerwise(module, quiet=True) # need to reapply since hooks were removed/re-added devices.torch_gc(fast=True, force=True, reason='offload') def get_logical_param_count(module: torch.nn.Module) -> int: if hasattr(module, "sdnq_dequantizer"): original_shape = module.sdnq_dequantizer.original_shape count = math.prod(original_shape) if getattr(module, "bias", None) is not None: count += module.bias.numel() return int(count) count = sum(p.numel() for p in module.parameters(recurse=False)) for child in module.children(): count += get_logical_param_count(child) return count def report_model_stats(module_name, module): try: size = offload_hook_instance.offload_map.get(module_name, 0) if offload_hook_instance is not None else 0 if size == 0: size, _params = get_module_size(module) quant = getattr(module, "quantization_method", None) params = sum(p.numel() for p in module.parameters(recurse=True)) logical = get_logical_param_count(module) log.debug(f'Module: name={module_name} cls={module.__class__.__name__} size={size:.3f} params={params} logical={logical} quant={quant}') except Exception as e: log.error(f'Module stats: name={module_name} {e}') def apply_balanced_offload(sd_model=None, exclude: list[str] | None = None, force: bool = False, silent: bool = False): global offload_hook_instance # pylint: disable=global-statement if shared.opts.diffusers_offload_mode != "balanced": return sd_model if sd_model is None: if not shared.sd_loaded: return sd_model sd_model = shared.sd_model if sd_model is None: return sd_model if exclude is None: exclude = [] if sd_model.__class__.__name__ in balanced_offload_exclude: return sd_model remove_group_offload(sd_model) t0 = time.time() cached = True checkpoint_name = sd_model.sd_checkpoint_info.name if getattr(sd_model, "sd_checkpoint_info", None) is not None else sd_model.__class__.__name__ if force or (offload_hook_instance is None) or (offload_hook_instance.min_watermark != shared.opts.diffusers_offload_min_gpu_memory) or (offload_hook_instance.max_watermark != shared.opts.diffusers_offload_max_gpu_memory) or (checkpoint_name != offload_hook_instance.checkpoint_name): cached = False offload_hook_instance = OffloadHook(checkpoint_name) if cached and shared.opts.diffusers_offload_pre: debug_move('Offload: type=balanced op=apply skip') return sd_model for pipe in get_pipe_variants(sd_model): for module_name, _module_size in get_module_sizes(pipe, exclude): module = getattr(pipe, module_name, None) if module is None: continue module.module_name = module_name module.offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name) apply_balanced_offload_to_module(module, op='apply', force=force) if not silent: report_model_stats(module_name, module) set_accelerate(sd_model) t1 = time.time() process_timer.add('offload', t1 - t0) 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={t1 - t0:.2f} type=balanced fn={fn}') if not cached: log.info(f'Model class={sd_model.__class__.__name__} modules={len(offload_hook_instance.offload_map)} size={offload_hook_instance.model_size():.3f}') return sd_model