diff --git a/modules/sd_offload.py b/modules/sd_offload.py
index ca6c66570..58699d8b0 100644
--- a/modules/sd_offload.py
+++ b/modules/sd_offload.py
@@ -76,30 +76,142 @@ def set_accelerate(sd_model):
set_accelerate_to_module(sd_model.decoder_pipe)
-def apply_group_offload(sd_model, op:str='model'):
- offload_dct = {
- 'onload_device': devices.device,
- 'offload_device': devices.cpu,
- 'offload_type': shared.opts.group_offload_type,
- 'num_blocks_per_group': shared.opts.group_offload_blocks,
+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': shared.opts.group_offload_stream,
- 'record_stream': shared.opts.group_offload_record,
- 'low_cpu_mem_usage': False,
+ '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,
}
- if shared.opts.group_offload_type == 'block_level':
- offload_dct['exclude_modules'] = ['vae']
- log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} options={offload_dct}')
- if hasattr(sd_model, "enable_group_offload"):
- sd_model.enable_group_offload(**offload_dct)
- else:
- log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} not supported')
+
+
+def remove_group_offload_component(module):
+ if getattr(module, 'sdnext_group_offload_sig', None) is None:
+ return
+ 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
+
+
+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 getattr(module, 'sdnext_group_offload_sig', None) is not None:
+ remove_group_offload_component(module)
+ removed.append(module_name)
+ if removed:
+ log.debug(f'Offload: type=group op=remove modules={removed}')
+
+
+def apply_group_offload_component(module, module_name: str, main: bool, op: str = 'model') -> 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)
+ limit_gb = 0.5 * shared.cpu_memory # heuristic ceiling: pinned memory is non-pageable, and half of system memory must stay available to everything else
+ if size_gb > limit_gb:
+ cfg['low_cpu_mem_usage'] = True
+ log.warning(f'Setting {op}: offload=group module={module_name} size={size_gb:.3f} limit={limit_gb:.3f} pin=dynamic memory guard')
+ sig = f'{devices.device}:{main}:' + ':'.join(str(v) for v in cfg.values())
+ if getattr(module, 'sdnext_group_offload_sig', None) == sig:
+ return False
+ if hasattr(module, '_hf_hook'): # leftover accelerate hooks from a previous offload mode abort the group apply upstream
+ module = accelerate.hooks.remove_hook_from_module(module, recurse=True)
+ remove_group_offload_component(module)
+ module.requires_grad_(False)
+ apply_group_offloading(module, onload_device=devices.device, offload_device=devices.cpu, **cfg)
+ module.sdnext_group_offload_sig = sig
+ return True
+
+
+def set_group_resident(module):
+ """VAE-class components never take group hooks: the hooks are forward-scoped, while
+ pipelines enter through encode/decode, and tiled calls re-enter per tile."""
+ if hasattr(module, '_hf_hook'):
+ module = accelerate.hooks.remove_hook_from_module(module, recurse=True)
+ remove_group_offload_component(module)
+ module.requires_grad_(False)
+ module.to(devices.device)
+
+
+def group_offload_role(module_name: str, module) -> str:
+ cls = module.__class__.__name__
+ if 'vae' in module_name.lower() or cls.startswith(('Autoencoder', 'VQModel', 'AsymmetricAutoencoder', 'ConsistencyDecoder')):
+ return 'resident'
+ if module_name.startswith(('text_encoder', 'image_encoder', 'safety_checker')):
+ return 'aux'
+ return 'main'
+
+
+def apply_modular_group_offload(sd_model, op:str='model'):
+ """Per-component group offload for modular pipelines, which lack the pipeline-level
+ enable_*_offload entry points. The model and sequential modes also route here."""
+ if shared.opts.diffusers_offload_mode != 'group' and not getattr(sd_model, 'sdnext_modular_offload_warned', False):
+ sd_model.sdnext_modular_offload_warned = True
+ log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} not supported on modular pipelines: using group offload')
+ applied = []
+ for name in ('transformer', 'transformer_ref'):
+ transformer = getattr(sd_model, name, None)
+ if transformer is not None and apply_group_offload_component(transformer, name, main=True, op=op):
+ applied.append(name)
+ text_encoder = getattr(sd_model, 'text_encoder', None)
+ if text_encoder is not None:
+ # offload targets the inner model when present: conditioning may call it directly,
+ # and hooks on the wrapper forward would never fire
+ if apply_group_offload_component(getattr(text_encoder, 'model', text_encoder), 'text_encoder', main=False, op=op):
+ applied.append('text_encoder')
+ for name in ('vae', 'audio_vae'):
+ component = getattr(sd_model, name, None)
+ if component is not None:
+ set_group_resident(component)
+ applied.append(f'{name}:device')
+ # has_accelerate stays unset: group hooks are not accelerate hooks, and the modular
+ # pipeline's own to() skips group-offloaded components when move_model runs
+ if any(':' not in name for name in applied):
+ log.info(f'Setting {op}: offload=group type={shared.opts.group_offload_type} modules={applied}')
+
+
+def apply_group_offload(sd_model, op:str='model'):
+ applied, resident = [], []
+ for module_name in get_module_names(sd_model):
+ module = getattr(sd_model, module_name, None)
+ if not isinstance(module, torch.nn.Module):
+ continue
+ try:
+ role = group_offload_role(module_name, module)
+ if role == 'resident':
+ set_group_resident(module)
+ resident.append(module_name)
+ elif apply_group_offload_component(module, module_name, main=role == 'main', op=op):
+ applied.append(module_name)
+ except Exception as e:
+ log.error(f'Setting {op}: offload=group module={module_name} {e}')
set_accelerate(sd_model)
+ if applied:
+ log.info(f'Setting {op}: offload=group type={shared.opts.group_offload_type} modules={applied} resident={resident}')
return sd_model
def apply_model_offload(sd_model, op:str='model', quiet:bool=False):
try:
+ remove_group_offload(sd_model)
log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}')
if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner:
shared.opts.diffusers_move_base = False
@@ -117,6 +229,7 @@ def apply_model_offload(sd_model, op:str='model', quiet:bool=False):
def apply_sequential_offload(sd_model, op:str='model', quiet:bool=False):
try:
+ remove_group_offload(sd_model)
log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}')
if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner:
shared.opts.diffusers_move_base = False
@@ -144,6 +257,7 @@ def apply_none_offload(sd_model, op:str='model', quiet:bool=False):
log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}')
try:
sd_model.has_accelerate = False
+ remove_group_offload(sd_model)
if hasattr(sd_model, 'maybe_free_model_hooks'):
sd_model.maybe_free_model_hooks()
sd_model = accelerate.hooks.remove_hook_from_module(sd_model, recurse=True)
@@ -535,6 +649,7 @@ def apply_balanced_offload(sd_model=None, exclude: list[str] | None = None, forc
exclude = []
if sd_model.__class__.__name__ in balanced_offload_exclude:
return sd_model
+ remove_group_offload(sd_model)
t0 = time.time()
cached = True
diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py
index 5944dc426..4aa269a11 100644
--- a/modules/ui_definitions.py
+++ b/modules/ui_definitions.py
@@ -153,6 +153,7 @@ def create_settings(cmd_opts):
"group_offload_type": OptionInfo("leaf_level", "Group offload type", gr.Radio, {"choices": ['leaf_level', 'block_level']}),
"group_offload_stream": OptionInfo(False, "Use torch streams", gr.Checkbox),
'group_offload_record': OptionInfo(False, "Record torch streams", gr.Checkbox),
+ 'group_offload_pin': OptionInfo(True, "Pin offload memory", gr.Checkbox),
'group_offload_blocks': OptionInfo(1, "Offload blocks", gr.Number),
}))
diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json
index 149d6f381..710940442 100644
--- a/ui/locale/locale_en.json
+++ b/ui/locale/locale_en.json
@@ -620,7 +620,7 @@
{"id":"","label":"Generic","localized":"","hint":"","ui":"video"},
{"id":"","label":"Google GenAI","localized":"","hint":"","ui":"settings_model_options"},
{"id":"","label":"Group Offload","localized":"","hint":"","ui":"settings_offload"},
- {"id":"","label":"Group offload type","localized":"","hint":"Granularity used by group offload.
- leaf_level: offloads at the smallest module level; maximum memory savings, slower
- block_level: offloads groups of transformer blocks (size set by Offload blocks); faster with less savings, and keeps the VAE resident
Applies only when Model offload mode is group.
Default is leaf_level.","reload":"model","ui":"settings_offload"},
+ {"id":"","label":"Group offload type","localized":"","hint":"Granularity used by group offload.
- leaf_level: offloads at the smallest module level; maximum memory savings, slower
- block_level: offloads groups of transformer blocks (size set by Offload blocks); faster with less savings
The VAE stays resident on the GPU in both modes, and text encoders always offload at leaf level.
Applies only when Model offload mode is group.
Default is leaf_level.","reload":"model","ui":"settings_offload"},
{"id":"","label":"Grid Options","localized":"","hint":"","ui":"settings_saving-images"},
{"id":"","label":"Grids","localized":"","hint":"","ui":"settings_saving-paths"},
{"id":"","label":"Guider","localized":"","hint":"","ui":"txt2img"},
@@ -1067,6 +1067,7 @@
{"id":"","label":"OpenBody","localized":"","hint":"","ui":"control"}
],
"p": [
+ {"id":"","label":"Pin offload memory","localized":"","hint":"Keeps the CPU copy of every stream-offloaded weight in pinned non-pageable memory for the fastest transfers, at a host memory cost equal to the full module size.
When disabled, memory is pinned one group at a time during transfer: slower, but the weights stay pageable and use no extra memory at rest.
Modules larger than half of system memory fall back to per-group pinning automatically.
Applies only to group offload with Use torch streams enabled.
Enabled by default.","reload":"model","ui":"settings_offload"},
{"id":"extras_nav","label":"Process","localized":"","hint":"Process existing image
Can be used to upscale images, remove backgrounds, obfuscate NSFW content, apply various filters and effects"},
{"id":"txt2img_prompts","label":"Prompts","localized":"","hint":"Image prompt and negative prompt","ui":"txt2img"},
{"id":"txt2img_pause","label":"Pause","localized":"","hint":"Pause processing","ui":"txt2img"},
@@ -1237,7 +1238,7 @@
{"id":"","label":"Rebase","localized":"","hint":"","ui":"tab_update"},
{"id":"","label":"Repos","localized":"","hint":"","ui":"component-8779"},
{"id":"","label":"Refiner model","localized":"","hint":"Refiner model used for second-pass operations","ui":"settings_sd"},
- {"id":"","label":"Record torch streams","localized":"","hint":"Records CUDA stream usage during group offload so reused buffers stay correct when Use torch streams is enabled.
Applies only to group offload with streams.
Disabled by default.","reload":"model","ui":"settings_offload"},
+ {"id":"","label":"Record torch streams","localized":"","hint":"Skips a stream synchronization each time a group offloads, letting transfers and compute overlap more tightly.
Slightly faster at the cost of slightly higher VRAM use; correctness is maintained either way.
Has no effect unless Use torch streams is enabled.
Applies only to group offload with streams.
Disabled by default.","reload":"model","ui":"settings_offload"},
{"id":"","label":"Remote VAE image type","localized":"","hint":"","ui":"settings_vae_encoder"},
{"id":"","label":"Remote VAE for encode","localized":"","hint":"","ui":"settings_vae_encoder"},
{"id":"","label":"RAS enabled","localized":"","hint":"","ui":"settings_advanced"},
@@ -1562,7 +1563,7 @@
{"id":"","label":"Username","localized":"","hint":"","ui":"component-8823"},
{"id":"","label":"UNET model","localized":"","hint":"","ui":"settings_sd"},
{"id":"","label":"UNET model secondary","localized":"","hint":"Override for the second transformer of dual-transformer architectures:
- Ideogram 4: the unconditional transformer
- Wan combined stage: the second expert (transformer_2)
Shown only when the loaded model has a second transformer.
Default loads the component from the base model.","ui":"settings_sd"},
- {"id":"","label":"Use torch streams","localized":"","hint":"In group offload, uses CUDA streams to prefetch the next group while the current one runs, hiding transfer latency.
Faster, but uses more VRAM.
Applies only to group offload.
Disabled by default.","reload":"model","ui":"settings_offload"},
+ {"id":"","label":"Use torch streams","localized":"","hint":"In group offload, uses CUDA streams to prefetch the next group while the current one runs, hiding transfer latency.
Faster, but offloaded weights are then staged in pinned non-pageable host memory: the whole module when Pin offload memory is enabled, one group at a time otherwise.
Text encoders are exempt and always offload without streams.
Applies only to group offload.
Disabled by default.","reload":"model","ui":"settings_offload"},
{"id":"","label":"Use SVD quantization","localized":"","hint":"Adds a low-rank (SVDQuant) correction on top of SDNQ to recover accuracy lost at low bit widths, at the cost of extra size and compute.
Tuned by SVD rank size and SVD steps.
Disabled by default.","reload":"model","ui":"settings_quantization"},
{"id":"","label":"Use Dynamic quantization","localized":"","hint":"Picks a per-layer weight type automatically instead of one type everywhere. Each layer starts at the Quantization type (the minimum) and steps up to higher precision until its error meets the Dynamic loss threshold.
Protects error-sensitive layers at the cost of a larger model and slower load.
Disabled by default.","reload":"model","ui":"settings_quantization"},
{"id":"","label":"Use Hadamard rotations","localized":"","hint":"Applies a Hadamard rotation before quantizing to spread out weight outliers, which can improve accuracy at low bit widths.
Group size is set by Hadamard group size.
Disabled by default.","reload":"model","ui":"settings_quantization"},