From fbd0020ad42e81a5ed531aa28bf1bce076d56820 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 16 Aug 2026 10:08:47 +0200 Subject: [PATCH] refactor offloading Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 22 +- modules/lora/networks.py | 3 +- modules/sd_models.py | 2 +- modules/sd_offload.py | 734 +-------------------------------- modules/sd_offload_balanced.py | 282 +++++++++++++ modules/sd_offload_group.py | 254 ++++++++++++ modules/sd_offload_state.py | 34 ++ modules/sd_offload_utils.py | 180 ++++++++ package.json | 2 +- webui.py | 2 +- 10 files changed, 775 insertions(+), 740 deletions(-) create mode 100644 modules/sd_offload_balanced.py create mode 100644 modules/sd_offload_group.py create mode 100644 modules/sd_offload_state.py create mode 100644 modules/sd_offload_utils.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e8ad40b02..ff378a84d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,7 @@ - nunchaku-lite support for `torch==2.13` - update handlers for all authenticated workflows - update handlers for all hf-based progress bars + - offload options take effect immediately without restart/reload - log long torch autotune operations - utilize `torch.accelerator` where available - add `SD_DIFFUSERS_DEBUG` and `SD_TRANSFORMERS_DEBUG` env variables to trace diffusers and transformers internal operations @@ -47,18 +48,17 @@ - remove DirectML support latest release was over 2 years ago and is not compatible with modern frameworks - **Fixes** - - init hf env variables before gradio load - - lora skip init and rebuild offload state - - lora keep network multiplier on change - - improve handling of hf auth - - improve pipeline detection for non-cached models - - cleanup alt offload codepaths + - hf: init hf env variables before gradio load + - lora: skip init and rebuild offload state + - lora: keep network multiplier on change + - auth: improve handling of hf auth + - load: improve pipeline detection for non-cached models + - offload: cleanup alt offload codepaths - offload: text encoders no longer take the denoiser profile on modular pipelines - offload: components entered through encode or decode are detected by structure rather than by name - offload: group offload honors the never-offload and model-type exclusion settings - - offload: offload options take effect when changed instead of waiting for a model reload - settings: offload settings grouped into shared overrides and per-mode sections - - hf progress bars + - log: hf progress bars - ltx: send the guidance stack and cross-timestep on every 2.x call path - ltx: distilled variants no longer force dynamic shifting on, which remapped their sigma schedule - ltx: sampler shift now reaches flow-match schedulers @@ -66,9 +66,9 @@ - video: take the audio sample rate from the loaded vocoder - video: keep the shared text encoder out of the registry rows - video: use generic loader methods - - processing stats reporting - - image metadata handle correct image index - - gguf transformer loader + - log: processing stats reporting + - metadata: image metadata handle correct image index + - gguf: transformer loader ## Update for 2026-08-07 diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 35a412f0a..2edb0c9f5 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -33,8 +33,9 @@ def group_offload_strip(sd_model, component_name: str, stripped: dict): lost on the next onload. With hooks removed the weights rest on cpu and the component reports its truthful device, so writes land in place; the offload reapply at the end of the pass snapshots the result into fresh groups.""" + from modules.sd_offload_group import remove_group_offload_component component = getattr(sd_model, component_name, None) - sd_models.remove_group_offload_component(component) + remove_group_offload_component(component) stripped[component_name] = component.device return stripped[component_name] diff --git a/modules/sd_models.py b/modules/sd_models.py index af3df0981..e3ed437ad 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -16,7 +16,7 @@ from modules.memstats import memory_stats, gpu_stats from modules.shared_helpers import walk_files from modules.modeldata import model_data from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoint_titles, get_closest_checkpoint_match, update_model_hashes, write_metadata, checkpoints_list # pylint: disable=unused-import -from modules.sd_offload import get_module_names, disable_offload, set_diffuser_offload, apply_balanced_offload, set_accelerate, remove_group_offload_component, offload_ondemand, reapply_offload, offload_reapply_options # pylint: disable=unused-import +from modules.sd_offload import get_module_names, disable_offload, set_diffuser_offload, apply_balanced_offload, set_accelerate, offload_ondemand, reapply_offload # pylint: disable=unused-import from modules.sd_models_utils import NoWatermark, get_signature, get_call, path_to_repo, apply_function_to_model, read_state_dict, get_state_dict_from_checkpoint # pylint: disable=unused-import diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 3d0754d2e..95e4fc261 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -1,56 +1,14 @@ -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 import shared, devices, sd_models 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'] -group_offload_main = [ # component names entered once per denoising step - "unet", "transformer", "transformer_2", "transformer_ref", "unconditional_transformer", - "prior", "prior_prior", "decoder", "dit_model", "model", "controlnet", -] # a denoiser registered under any other name takes the aux profile until listed here -offload_reapply_options = [ # settings that re-place loaded components when changed - "group_offload_type", "group_offload_stream", "group_offload_record", "group_offload_pin", "group_offload_blocks", - "diffusers_offload_nonblocking", "models_not_to_offload", "diffusers_offload_never", "diffusers_offload_always", -] -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 -group_stats_reported = set() -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 +import modules.sd_offload_state as s +from modules.sd_offload_utils import get_module_names, get_module_size, get_module_memory, offload_list, offload_matches, offload_model_types, dtype_byte_size, get_pipe_variants, set_accelerate # pylint: disable=unused-import +from modules.sd_offload_balanced import apply_balanced_offload +from modules.sd_offload_group import apply_group_offload, remove_group_offload, offload_ondemand # pylint: disable=unused-import def disable_offload(sd_model): @@ -70,301 +28,15 @@ def disable_offload(sd_model): 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) - module.sdnext_ondemand = False # group placement replaces any on-demand hook - 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 - module.sdnext_group_offload_sig = 'partial' # a raise below leaves hooks that only a non-empty signature will remove - apply_group_offloading(module, onload_device=devices.device, offload_device=devices.cpu, **cfg) - module.sdnext_group_offload_sig = sig - return True - - -def offload_list(opt: str) -> list: - return [m.strip() for m in re.split(';|,| ', opt) if len(m.strip()) > 2] - - -def offload_matches(module, module_name: str | None, names: list) -> 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 offload_model_types() -> list: - return [m.lower().strip() for m in re.split(r'[ ,]+', shared.opts.models_not_to_offload) if m.strip()] # type codes like sd and f1 are two characters, so only empty fragments are dropped - - -def offload_excluded(module_name: str, module) -> bool: - """Whether the offload exclusion settings keep this component on the accelerator.""" - if shared.sd_model_type.lower() in offload_model_types(): - return True - return offload_matches(module, module_name, offload_list(shared.opts.diffusers_offload_never)) - - -def set_group_resident(module) -> bool: - """Keep a component on the accelerator with no hooks of any kind.""" - changed = False - if hasattr(module, '_hf_hook'): - module = accelerate.hooks.remove_hook_from_module(module, recurse=True) - changed = True - if remove_group_offload_component(module): - changed = True - module.sdnext_ondemand = False - module.requires_grad_(False) - if any(not devices.same_device(t.device, devices.device) for t in itertools.chain(module.parameters(), module.buffers())): # an interrupted generation can leave a group-hooked module split across devices - module.to(devices.device) - changed = True - return changed - - -def group_offload_role(module_name: str, module) -> str: - """Placement role for one component: resident to stay put, ondemand for whole-module onload, main for per-step denoisers, aux for the rest.""" - if offload_excluded(module_name, module): - return 'resident' - if has_entry_bridge(module): - return 'ondemand' # encode and decode bypass the forward that group hooks scope to - if callable(getattr(module, 'encode', None)) or callable(getattr(module, 'decode', None)): - log.warning(f'Offload: type=group module={module_name} cls={module.__class__.__name__} bridge=missing role=resident') # decorate the entry points with apply_forward_hook to make the component offloadable - return 'resident' # nothing fires an onload for an undecorated entry point, so any hook placement strands the weights on cpu - if not getattr(module, '_supports_group_offloading', True): - return 'ondemand' # upstream marks modules that read submodule weights outside those submodules' forward - if module_name in group_offload_main: - return 'main' - return 'aux' - - -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_ondemand(module) -> bool: - """Placement for components that never take group hooks: they onload whole at their entry point.""" - if getattr(module, 'sdnext_ondemand', False) and hasattr(module, '_hf_hook'): - return False - 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) - 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 - moved = [] - for pipe in get_pipe_variants(sd_model): - names = get_module_names(pipe) if force else (getattr(pipe, 'sdnext_ondemand_modules', None) or []) # force enumerates the pipe rather than the list a load pass left on it - for module_name in names: - if include and module_name not in include: - continue - if exclude and module_name in exclude: - continue - module = getattr(pipe, module_name, None) - if not isinstance(module, torch.nn.Module) or not getattr(module, 'sdnext_ondemand', False): - continue # nothing else has an onload to bring it back - param = next(module.parameters(), None) - if param is None or devices.same_device(param.device, devices.cpu): - continue - try: - t0 = time.time() - module.to(devices.cpu, non_blocking=shared.opts.diffusers_offload_nonblocking) - dt = time.time() - t0 - process_timer.add('offload', dt) - moved.append(module_name) - debug_move(f'Offload: type=ondemand op=offload module={module_name} nonblocking={shared.opts.diffusers_offload_nonblocking} reason="{reason}" time={dt:.3f}') - except Exception as e: - log.warning(f'Offload: type=ondemand op=offload module={module_name} {e}') - if moved: - devices.torch_gc(reason='ondemand') - - -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.""" - 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 checkpoint_name in group_stats_reported: # keyed by checkpoint since a task switch rebuilds the pipe object - return - group_stats_reported.add(checkpoint_name) - 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_group_offload(sd_model): - """Per-component group offload for classic and modular pipelines.""" - changed = False - placements = [] - module_names = get_module_names(sd_model) - for module_name in module_names: - module = getattr(sd_model, module_name, None) - if not isinstance(module, torch.nn.Module): - continue - try: - role = group_offload_role(module_name, module) - placements.append(f'{module_name}:{role}') - if role == 'resident': - applied = set_group_resident(module) - elif role == 'ondemand': - applied = apply_group_offload_ondemand(module) - else: - applied = apply_group_offload_component(module, module_name, main=role == 'main') - changed = changed or applied - except Exception as e: - log.error(f'Offload: type=group module={module_name} {e}') - sd_model.sdnext_ondemand_modules = [name for name in module_names if getattr(getattr(sd_model, name, None), 'sdnext_ondemand', False)] - if sd_models.get_diffusers_task(sd_model) != sd_models.DiffusersTaskType.MODULAR: # group hooks are not accelerate hooks, so modular pipelines stay unstamped - set_accelerate(sd_model) - if changed: - log.info(f'Offload: type=group modules={placements}') - else: - log.debug(f'Offload: type=group modules={placements}') - report_group_stats(sd_model, module_names) - return sd_model - - def reapply_offload(): """Re-place loaded components after an offload setting changed.""" - global offload_hook_instance # pylint: disable=global-statement if not shared.sd_loaded: return modular = sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.MODULAR if shared.opts.diffusers_offload_mode == 'group' or (modular and shared.opts.diffusers_offload_mode in {'model', 'sequential'}): apply_group_offload(shared.sd_model) elif shared.opts.diffusers_offload_mode == 'balanced': - offload_hook_instance = None # the hook snapshots the exclusion lists when constructed + s.offload_hook_instance = None # the hook snapshots the exclusion lists when constructed apply_balanced_offload(shared.sd_model) @@ -400,7 +72,7 @@ def apply_sequential_offload(sd_model, op:str='model', quiet:bool=False): def apply_none_offload(sd_model, quiet:bool=False): - if shared.sd_model_type not in offload_allow_none: + if shared.sd_model_type not in s.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}') @@ -416,15 +88,14 @@ def apply_none_offload(sd_model, quiet:bool=False): 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 + if s.accelerate_dtype_byte_size is None: + s.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'}: @@ -453,390 +124,3 @@ def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False, force:bool= 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 = offload_list(shared.opts.diffusers_offload_always) - self.offload_never = offload_list(shared.opts.diffusers_offload_never) - 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: - return offload_matches(module, module_name, 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 offload_model_types(): - 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 = [] - if hasattr(pipe, '_component_specs'): # modular pipelines name their components in specs; the config dict also carries scalars - modules_names.extend(pipe.components) - else: - 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 diff --git a/modules/sd_offload_balanced.py b/modules/sd_offload_balanced.py new file mode 100644 index 000000000..589ff3522 --- /dev/null +++ b/modules/sd_offload_balanced.py @@ -0,0 +1,282 @@ +import os +import sys +import time +import inspect +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 +from modules.sd_offload_utils import offload_list, offload_matches, offload_model_types, get_pipe_variants, get_module_names, get_module_size, set_accelerate, report_model_stats +from modules.sd_offload_group import remove_group_offload +import modules.sd_offload_state as s + + +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 = offload_list(shared.opts.diffusers_offload_always) + self.offload_never = offload_list(shared.opts.diffusers_offload_never) + 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: + return offload_matches(module, module_name, 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 offload_model_types(): + 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() + s.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=s.no_split_module_classes, + verbose=s.verbose, + clean_result=False, + ) + offload_dir = getattr(module, "offload_dir", os.path.join(shared.opts.accelerate_offload_path, module.__class__.__name__)) + if s.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 s.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_module_sizes(pipe=None, exclude=None): + if exclude is None: + exclude = [] + modules = {} + for module_name in get_module_names(pipe, exclude): + module_size = s.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) + s.offload_hook_instance.offload_map[module_name] = module_size + s.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: + if s.move_stream is None: + s.move_stream = torch.cuda.Stream(device=devices.device) + with torch.cuda.stream(s.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 = s.offload_hook_instance.offload_map.get(module_name, s.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 s.offload_hook_instance.matches(module, s.offload_hook_instance.offload_never, module_name): + op = f'{op}:never' + elif s.offload_hook_instance.matches(module, s.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 s.debug: + quant = getattr(module, "quantization_method", None) + s.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, s.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 s.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 apply_balanced_offload(sd_model=None, exclude: list[str] | None = None, force: bool = False, silent: bool = False): + 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 s.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 (s.offload_hook_instance is None) or (s.offload_hook_instance.min_watermark != shared.opts.diffusers_offload_min_gpu_memory) or (s.offload_hook_instance.max_watermark != shared.opts.diffusers_offload_max_gpu_memory) or (checkpoint_name != s.offload_hook_instance.checkpoint_name): + cached = False + s.offload_hook_instance = OffloadHook(checkpoint_name) + + if cached and shared.opts.diffusers_offload_pre: + s.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 + s.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(s.offload_hook_instance.offload_map)} size={s.offload_hook_instance.model_size():.3f}') + return sd_model diff --git a/modules/sd_offload_group.py b/modules/sd_offload_group.py new file mode 100644 index 000000000..f42c13115 --- /dev/null +++ b/modules/sd_offload_group.py @@ -0,0 +1,254 @@ +import time +import itertools +import torch +import accelerate.hooks +import accelerate.utils.modeling +from modules.logger import log +from modules import shared, devices, sd_models +from modules.timer import process as process_timer +from modules.sd_offload_utils import get_pipe_variants, get_module_names, get_module_size, set_accelerate, offload_excluded, report_model_stats +import modules.sd_offload_state as s + + +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) + module.sdnext_ondemand = False # group placement replaces any on-demand hook + 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 + module.sdnext_group_offload_sig = 'partial' # a raise below leaves hooks that only a non-empty signature will remove + 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) -> bool: + """Keep a component on the accelerator with no hooks of any kind.""" + changed = False + if hasattr(module, '_hf_hook'): + module = accelerate.hooks.remove_hook_from_module(module, recurse=True) + changed = True + if remove_group_offload_component(module): + changed = True + module.sdnext_ondemand = False + module.requires_grad_(False) + if any(not devices.same_device(t.device, devices.device) for t in itertools.chain(module.parameters(), module.buffers())): # an interrupted generation can leave a group-hooked module split across devices + module.to(devices.device) + changed = True + return changed + + +def group_offload_role(module_name: str, module) -> str: + """Placement role for one component: resident to stay put, ondemand for whole-module onload, main for per-step denoisers, aux for the rest.""" + if offload_excluded(module_name, module): + return 'resident' + if has_entry_bridge(module): + return 'ondemand' # encode and decode bypass the forward that group hooks scope to + if callable(getattr(module, 'encode', None)) or callable(getattr(module, 'decode', None)): + log.warning(f'Offload: type=group module={module_name} cls={module.__class__.__name__} bridge=missing role=resident') # decorate the entry points with apply_forward_hook to make the component offloadable + return 'resident' # nothing fires an onload for an undecorated entry point, so any hook placement strands the weights on cpu + if not getattr(module, '_supports_group_offloading', True): + return 'ondemand' # upstream marks modules that read submodule weights outside those submodules' forward + if module_name in s.group_offload_main: + return 'main' + return 'aux' + + +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) + s.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_ondemand(module) -> bool: + """Placement for components that never take group hooks: they onload whole at their entry point.""" + if getattr(module, 'sdnext_ondemand', False) and hasattr(module, '_hf_hook'): + return False + 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) + 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 + moved = [] + for pipe in get_pipe_variants(sd_model): + names = get_module_names(pipe) if force else (getattr(pipe, 'sdnext_ondemand_modules', None) or []) # force enumerates the pipe rather than the list a load pass left on it + for module_name in names: + if include and module_name not in include: + continue + if exclude and module_name in exclude: + continue + module = getattr(pipe, module_name, None) + if not isinstance(module, torch.nn.Module) or not getattr(module, 'sdnext_ondemand', False): + continue # nothing else has an onload to bring it back + param = next(module.parameters(), None) + if param is None or devices.same_device(param.device, devices.cpu): + continue + try: + t0 = time.time() + module.to(devices.cpu, non_blocking=shared.opts.diffusers_offload_nonblocking) + dt = time.time() - t0 + process_timer.add('offload', dt) + moved.append(module_name) + s.debug_move(f'Offload: type=ondemand op=offload module={module_name} nonblocking={shared.opts.diffusers_offload_nonblocking} reason="{reason}" time={dt:.3f}') + except Exception as e: + log.warning(f'Offload: type=ondemand op=offload module={module_name} {e}') + if moved: + devices.torch_gc(reason='ondemand') + + +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.""" + 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 checkpoint_name in s.group_stats_reported: # keyed by checkpoint since a task switch rebuilds the pipe object + return + s.group_stats_reported.add(checkpoint_name) + 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_group_offload(sd_model): + """Per-component group offload for classic and modular pipelines.""" + changed = False + placements = [] + module_names = get_module_names(sd_model) + for module_name in module_names: + module = getattr(sd_model, module_name, None) + if not isinstance(module, torch.nn.Module): + continue + try: + role = group_offload_role(module_name, module) + placements.append(f'{module_name}:{role}') + if role == 'resident': + applied = set_group_resident(module) + elif role == 'ondemand': + applied = apply_group_offload_ondemand(module) + else: + applied = apply_group_offload_component(module, module_name, main=role == 'main') + changed = changed or applied + except Exception as e: + log.error(f'Offload: type=group module={module_name} {e}') + sd_model.sdnext_ondemand_modules = [name for name in module_names if getattr(getattr(sd_model, name, None), 'sdnext_ondemand', False)] + if sd_models.get_diffusers_task(sd_model) != sd_models.DiffusersTaskType.MODULAR: # group hooks are not accelerate hooks, so modular pipelines stay unstamped + set_accelerate(sd_model) + if changed: + log.info(f'Offload: type=group modules={placements}') + else: + log.debug(f'Offload: type=group modules={placements}') + report_group_stats(sd_model, module_names) + return sd_model diff --git a/modules/sd_offload_state.py b/modules/sd_offload_state.py new file mode 100644 index 000000000..2c1297675 --- /dev/null +++ b/modules/sd_offload_state.py @@ -0,0 +1,34 @@ +import os +from modules.logger import log + + +# logging +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'] # used to warn if offloading=none + + +offload_post = ['h1'] +offload_hook_instance = None # instance of sd_offload_balanced.OffloadHook +balanced_offload_exclude = ['CogView4Pipeline', 'MeissonicPipeline'] +group_offload_main = [ # component names entered once per denoising step + "unet", "transformer", "transformer_2", "transformer_ref", "unconditional_transformer", + "prior", "prior_prior", "decoder", "dit_model", "model", "controlnet", +] # a denoiser registered under any other name takes the aux profile until listed here +offload_reapply_options = [ # settings that re-place loaded components when changed + "group_offload_type", "group_offload_stream", "group_offload_record", "group_offload_pin", "group_offload_blocks", + "diffusers_offload_nonblocking", "models_not_to_offload", "diffusers_offload_never", "diffusers_offload_always", +] +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 # monkey-patch accelerate.utils.modeling.dtype_byte_size +group_stats_reported = set() +move_stream = None diff --git a/modules/sd_offload_utils.py b/modules/sd_offload_utils.py new file mode 100644 index 000000000..eb43dc781 --- /dev/null +++ b/modules/sd_offload_utils.py @@ -0,0 +1,180 @@ +import re +import math +import inspect +import itertools +import torch +import accelerate +from modules import shared +from modules.logger import log +import modules.sd_offload_state as s + + +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 s.accelerate_dtype_byte_size(dtype) + + +def get_signature(cls): + signature = inspect.signature(cls.__init__, follow_wrapped=True) + return signature.parameters + + +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 = [] + if hasattr(pipe, '_component_specs'): # modular pipelines name their components in specs; the config dict also carries scalars + modules_names.extend(pipe.components) + else: + 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 offload_list(opt: str) -> list: + return [m.strip() for m in re.split(';|,| ', opt) if len(m.strip()) > 2] + + +def offload_matches(module, module_name: str | None, names: list) -> 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 offload_model_types() -> list: + return [m.lower().strip() for m in re.split(r'[ ,]+', shared.opts.models_not_to_offload) if m.strip()] # type codes like sd and f1 are two characters, so only empty fragments are dropped + + +def offload_excluded(module_name: str, module) -> bool: + """Whether the offload exclusion settings keep this component on the accelerator.""" + if shared.sd_model_type.lower() in offload_model_types(): + return True + return offload_matches(module, module_name, offload_list(shared.opts.diffusers_offload_never)) + + +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 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 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 = s.offload_hook_instance.offload_map.get(module_name, 0) if s.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}') diff --git a/package.json b/package.json index f7bc02726..602188652 100644 --- a/package.json +++ b/package.json @@ -41,7 +41,7 @@ "dev:kanvas": "cd extensions-builtin/sdnext-kanvas && build --profile development", "dev:core": "build --profile development --config ui/.build.json", "ruff": ". venv/bin/activate && ruff check", - "pylint": ". venv/bin/activate && pylint *.py modules/ pipelines/ scripts/ extensions-builtin/ | grep -v '^*'", + "pylint": ". venv/bin/activate && pylint --disable fixme *.py modules/ pipelines/ scripts/ extensions-builtin/ | grep -v '^*'", "pyright": ". venv/bin/activate && pyright --threads 4", "ty": ". venv/bin/activate && ty check --force-exclude", "codespell": ". venv/bin/activate && codespell", diff --git a/webui.py b/webui.py index 1c5780d14..e76812064 100644 --- a/webui.py +++ b/webui.py @@ -201,7 +201,7 @@ def load_model(): shared.opts.onchange("sd_unet_secondary", wrap_queued_call(lambda: modules.sd_unet.load_unet_secondary(shared.sd_model)), call=False) shared.opts.onchange("sd_text_encoder", wrap_queued_call(lambda: modules.sd_models.reload_text_encoder()), call=False) shared.opts.onchange("temp_dir", modules.gr_tempdir.on_tmpdir_changed) - for opt in modules.sd_models.offload_reapply_options: + for opt in modules.sd_offload_state.offload_reapply_options: shared.opts.onchange(opt, wrap_queued_call(modules.sd_models.reapply_offload), call=False) timer.startup.record("onchange")