Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 01:04:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor offloading

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2026-08-16 10:08:47 +02:00
parent 51a598bfc3
commit fbd0020ad4
10 changed files with 775 additions and 740 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+11 -11
View File
@@ -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
modules/lora/networks.py
+2 -1
View File
@@ -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]
modules/sd_models.py
+1 -1
View File
@@ -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
modules/sd_offload.py
+9 -725
View File
@@ -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
modules/sd_offload_balanced.py
+282
View File
@@ -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
modules/sd_offload_group.py
+254
View File
@@ -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
modules/sd_offload_state.py
+34
View File
@@ -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
modules/sd_offload_utils.py
+180
View File
@@ -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}')
package.json
+1 -1
View File
@@ -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",
webui.py
+1 -1
View File
@@ -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")