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

update offloading codepaths for non-balanced

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2026-08-11 22:09:12 +02:00
parent 5123bfd96b
commit f92ef90d3d
6 changed files with 31 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+1
View File
@@ -31,6 +31,7 @@
- lora keep network multiplier on change
- improve handling of hf auth
- improve pipeline detection for non-cached models
- cleanup alt offload codepaths
## Update for 2026-08-07
modules/processing_diffusers.py
-12
View File
@@ -374,20 +374,11 @@ def process_refine(p: processing.StableDiffusionProcessing, output):
if is_refiner_enabled(p):
if shared.opts.samples_save and not p.do_not_save_samples and shared.opts.save_images_before_refiner and hasattr(shared.sd_model, 'vae'):
save_intermediate(p, latents=output.images, suffix="-before-refiner")
if shared.opts.diffusers_move_base:
log.debug('Moving to CPU: model=base')
sd_models.move_model(shared.sd_model, devices.cpu)
if shared.state.interrupted or shared.state.skipped:
shared.sd_model = orig_pipeline
return output
jobid = shared.state.begin('Refine')
shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
if shared.opts.diffusers_move_refiner:
sd_models.move_model(shared.sd_refiner, devices.device)
if hasattr(shared.sd_refiner, 'unet'):
sd_models.move_model(shared.sd_model.unet, devices.device)
if hasattr(shared.sd_refiner, 'transformer'):
sd_models.move_model(shared.sd_model.transformer, devices.device)
p.ops.append('refine')
p.is_refiner_pass = True
@@ -444,9 +435,6 @@ def process_refine(p: processing.StableDiffusionProcessing, output):
if shared.opts.diffusers_offload_mode == "balanced":
shared.sd_refiner = sd_models.apply_balanced_offload(shared.sd_refiner)
elif shared.opts.diffusers_move_refiner:
log.debug('Moving to CPU: model=refiner')
sd_models.move_model(shared.sd_refiner, devices.cpu)
shared.state.end(jobid)
shared.state.nextjob()
p.is_refiner_pass = False
modules/processing_vae.py
+1 -19
View File
@@ -43,10 +43,6 @@ def full_vqgan_decode(latents, model):
devices.torch_gc(force=True)
shared.mem_mon.reset()
base_device = None
if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False):
base_device = sd_models.move_base(model, devices.cpu)
if shared.opts.diffusers_offload_mode == "balanced":
shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
elif shared.opts.diffusers_offload_mode != "sequential":
@@ -76,8 +72,6 @@ def full_vqgan_decode(latents, model):
if shared.opts.diffusers_offload_mode == "balanced":
shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
elif shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and base_device is not None:
sd_models.move_base(model, base_device)
t1 = time.time()
if debug:
log_debug(f'VAE memory: {shared.mem_mon.read()}')
@@ -97,10 +91,7 @@ def full_vae_decode(latents, model):
devices.torch_gc(force=True)
shared.mem_mon.reset()
base_device = None
if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False):
base_device = sd_models.move_base(model, devices.cpu)
elif shared.opts.diffusers_offload_mode != "sequential":
if shared.opts.diffusers_offload_mode != "sequential":
sd_models.move_model(model.vae, devices.device)
sd_models.set_vae_options(model, vae=None, op='decode')
@@ -169,9 +160,6 @@ def full_vae_decode(latents, model):
model.vae.apply(sd_models_utils.convert_to_faketensors)
devices.torch_gc(force=True)
elif shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and base_device is not None:
sd_models.move_base(model, base_device)
t1 = time.time()
if debug:
log_debug(f'VAE memory: {shared.mem_mon.read()}')
@@ -183,10 +171,6 @@ def full_vae_decode(latents, model):
def full_vae_encode(image, model):
t0 = time.time()
if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'):
log_debug('Moving to CPU: model=UNet')
unet_device = model.unet.device
sd_models.move_model(model.unet, devices.cpu)
if shared.opts.diffusers_offload_mode != "sequential" and hasattr(model, 'vae'):
sd_models.move_model(model.vae, devices.device)
if getattr(model.vae, 'sdnext_ondemand', False):
@@ -209,8 +193,6 @@ def full_vae_encode(image, model):
model.vae = model.vae.to(dtype=model.vae.orig_dtype)
del model.vae.orig_dtype
if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'):
sd_models.move_model(model.unet, unet_device)
t1 = time.time()
log.debug(f'Encode: vae="{vae_name}" upcast={upcast} slicing={getattr(model.vae, "use_slicing", None)} tiling={getattr(model.vae, "use_tiling", None)} latents={encoded.shape}:{encoded.device}:{encoded.dtype} time={t1-t0:.3f}')
return encoded
modules/sd_models.py
+1 -4
View File
@@ -239,7 +239,7 @@ def move_model(model, device=None, force=False):
fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
if getattr(model, 'vae', None) is not None and get_diffusers_task(model) != DiffusersTaskType.TEXT_2_IMAGE:
if device == devices.device and model.vae.device.type != "meta" and not getattr(model.vae, 'sdnext_ondemand', False): # force vae back to gpu if not in txt2img mode; on-demand vaes onload at their entry point instead
if (device == devices.device) and (model.vae.device.type != "meta") and not getattr(model.vae, 'sdnext_ondemand', False): # force vae back to gpu if not in txt2img mode; on-demand vaes onload at their entry point instead
model.vae.to(device)
if hasattr(model.vae, '_hf_hook'):
debug_move(f'Model move: to={device} class={model.vae.__class__} fn={fn}') # pylint: disable=protected-access
@@ -1037,9 +1037,6 @@ def load_diffuser(checkpoint_info: CheckpointInfo | None = None, op='model', rev
if op == 'model' and not (os.path.isdir(checkpoint_info.path) or checkpoint_info.type == 'huggingface'):
if getattr(shared.sd_model, 'sd_checkpoint_info', None) is not None and vae_file is not None:
sd_vae.apply_vae_config(shared.sd_model.sd_checkpoint_info.filename, vae_file, sd_model)
if op == 'refiner' and shared.opts.diffusers_move_refiner:
log.debug('Moving refiner model to CPU')
move_model(sd_model, devices.cpu)
else:
move_model(sd_model, devices.device)
timer.load.record("move")
modules/sd_offload.py
+28 -35
View File
@@ -56,7 +56,7 @@ def disable_offload(sd_model):
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}')
log.warning(f'Offload: remove hook module={module_name} {e}')
if network_layer_name:
module.network_layer_name = network_layer_name
sd_model.has_accelerate = False
@@ -131,7 +131,7 @@ def remove_group_offload(sd_model):
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:
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."""
@@ -156,7 +156,7 @@ def apply_group_offload_component(module, module_name: str, main: bool, op: str
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'Setting {op}: offload=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')
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
@@ -164,7 +164,7 @@ def apply_group_offload_component(module, module_name: str, main: bool, op: str
module = accelerate.hooks.remove_hook_from_module(module, recurse=True)
remove_group_offload_component(module)
module.requires_grad_(False)
log.debug(f'Setting {op}: offload=group op=apply 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
log.debug(f'Offload: type=group op=apply module={module_name} pin={cfg["use_stream"] and not cfg["low_cpu_mem_usage"]}') # before the apply: pinning large components takes a while and would otherwise run silently
apply_group_offloading(module, onload_device=devices.device, offload_device=devices.cpu, **cfg)
module.sdnext_group_offload_sig = sig
return True
@@ -275,22 +275,22 @@ def report_group_stats(sd_model, module_names):
log.info(f'Model class={sd_model.__class__.__name__} modules={len(counted)} size={total:.3f}')
def apply_modular_group_offload(sd_model, op:str='model'):
def apply_modular_group_offload(sd_model):
"""Per-component group offload for modular pipelines, which lack the pipeline-level
enable_*_offload entry points. The model and sequential modes also route here."""
if shared.opts.diffusers_offload_mode != 'group' and not getattr(sd_model, 'sdnext_modular_offload_warned', False):
sd_model.sdnext_modular_offload_warned = True
log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} not supported on modular pipelines: using group offload')
log.warning(f'Offload: desired={shared.opts.diffusers_offload_mode} override=group reason="modular pipeline"')
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):
if transformer is not None and apply_group_offload_component(transformer, name, main=True):
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):
if apply_group_offload_component(getattr(text_encoder, 'model', text_encoder), 'text_encoder', main=False):
applied.append('text_encoder')
for name in ('vae', 'audio_vae'):
component = getattr(sd_model, name, None)
@@ -300,11 +300,11 @@ def apply_modular_group_offload(sd_model, op:str='model'):
# 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}')
log.info(f'Offload: type=group type={shared.opts.group_offload_type} modules={applied}')
report_group_stats(sd_model, ('transformer', 'transformer_ref', 'text_encoder', 'vae', 'audio_vae'))
def apply_group_offload(sd_model, op:str='model'):
def apply_group_offload(sd_model):
applied, resident, ondemand = [], [], []
for module_name in get_module_names(sd_model):
module = getattr(sd_model, module_name, None)
@@ -317,44 +317,34 @@ def apply_group_offload(sd_model, op:str='model'):
ondemand.append(module_name)
else:
resident.append(module_name)
elif apply_group_offload_component(module, module_name, main=role == 'main', op=op):
elif apply_group_offload_component(module, module_name, main=role == 'main'):
applied.append(module_name)
except Exception as e:
log.error(f'Setting {op}: offload=group module={module_name} {e}')
log.error(f'Offload: type=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} ondemand={ondemand}')
log.info(f'Offload: type=group type={shared.opts.group_offload_type} modules={applied} resident={resident} ondemand={ondemand}')
report_group_stats(sd_model, get_module_names(sd_model))
return sd_model
def apply_model_offload(sd_model, op:str='model', quiet:bool=False):
def apply_model_offload(sd_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
shared.opts.diffusers_move_unet = False
shared.opts.diffusers_move_refiner = False
log.warning(f'Disabling {op} "Move model to CPU" since "Model CPU offload" is enabled')
log.quiet(quiet, f'Offload: type={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}')
if not hasattr(sd_model, "_all_hooks") or len(sd_model._all_hooks) == 0: # pylint: disable=protected-access
sd_model.enable_model_cpu_offload(device=devices.device)
else:
sd_model.maybe_free_model_hooks()
set_accelerate(sd_model)
except Exception as e:
log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}')
log.error(f'Offload: type={shared.opts.diffusers_offload_mode} {e}')
def apply_sequential_offload(sd_model, op:str='model', quiet:bool=False):
try:
remove_group_offload(sd_model)
log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}')
if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner:
shared.opts.diffusers_move_base = False
shared.opts.diffusers_move_unet = False
shared.opts.diffusers_move_refiner = False
log.warning(f'Disabling {op} "Move model to CPU" since "Sequential CPU offload" is enabled')
log.quiet(quiet, f'Offload: type={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}')
if sd_model.has_accelerate:
if op == "vae": # reapply sequential offload to vae
from accelerate import cpu_offload
@@ -366,14 +356,14 @@ def apply_sequential_offload(sd_model, op:str='model', quiet:bool=False):
sd_model.enable_sequential_cpu_offload(device=devices.device)
set_accelerate(sd_model)
except Exception as e:
log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}')
log.error(f'Offload: type={shared.opts.diffusers_offload_mode} {e}')
def apply_none_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:
log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} type={shared.sd_model.__class__.__name__} large model')
log.warning(f'Offload: type={shared.opts.diffusers_offload_mode} cls={shared.sd_model.__class__.__name__} large model')
else:
log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}')
log.quiet(quiet, f'Offload: type={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}')
try:
sd_model.has_accelerate = False
remove_group_offload(sd_model)
@@ -398,21 +388,24 @@ def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False, force:bool=
accelerate.utils.modeling.dtype_byte_size = dtype_byte_size
if sd_models.get_diffusers_task(sd_model) == sd_models.DiffusersTaskType.MODULAR and shared.opts.diffusers_offload_mode in {'model', 'sequential', 'group'}:
apply_modular_group_offload(sd_model, op=op)
apply_modular_group_offload(sd_model)
process_timer.add('offload', time.time() - t0)
return
if shared.opts.diffusers_offload_mode == "none":
apply_none_offload(sd_model, op=op, quiet=quiet)
log.warning('Offload: type=none "use balanced offload with model type set not to offload"')
apply_none_offload(sd_model, quiet=quiet)
sd_models.move_model(sd_model, devices.device, force=True)
if shared.opts.diffusers_offload_mode == "model" and hasattr(sd_model, "enable_model_cpu_offload"):
apply_model_offload(sd_model, op=op, quiet=quiet)
log.warning('Offload: type=model "use balanced offload instead"')
apply_model_offload(sd_model, quiet=quiet)
if shared.opts.diffusers_offload_mode == "sequential" and hasattr(sd_model, "enable_sequential_cpu_offload"):
apply_sequential_offload(sd_model, op=op, quiet=quiet)
if shared.opts.diffusers_offload_mode == "group":
sd_model = apply_group_offload(sd_model, op=op)
sd_model = apply_group_offload(sd_model)
if shared.opts.diffusers_offload_mode == "balanced":
sd_model = apply_balanced_offload(sd_model, force=force)
modules/shared_legacy.py
-3
View File
@@ -26,9 +26,6 @@ legacy_options = options_section(('legacy_options', "Legacy options"), {
"sd_vae_sliced_encode": LegacyOption(False, "VAE sliced encode", gr.Checkbox, {"visible": False}),
"nan_skip": LegacyOption(False, "Skip Generation if NaN found in latents", gr.Checkbox, {"visible": False}),
"sd_model_dict": LegacyOption('None', "Use separate base dict", gr.Dropdown, lambda: {"choices": ['None'], "visible": False}),
"diffusers_move_base": LegacyOption(False, "Move base model to CPU when using refiner", gr.Checkbox, {"visible": False }),
"diffusers_move_unet": LegacyOption(False, "Move base model to CPU when using VAE", gr.Checkbox, {"visible": False }),
"diffusers_move_refiner": LegacyOption(False, "Move refiner model to CPU when not in use", gr.Checkbox, {"visible": False }),
"diffusers_extract_ema": LegacyOption(False, "Use model EMA weights when possible", gr.Checkbox, {"visible": False }),
"batch_cond_uncond": LegacyOption(True, "Do conditional and unconditional denoising in one batch", gr.Checkbox, {"visible": False}),
"CLIP_stop_at_last_layers": LegacyOption(1, "CLiP-skip", gr.Slider, {"minimum": 1, "maximum": 8, "step": 1, "visible": False}),