diff --git a/CHANGELOG.md b/CHANGELOG.md index c54cc52f0..3b4ec410d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 1e10e56c5..df17ba13c 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -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 diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 6b721640c..0d18e875d 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -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 diff --git a/modules/sd_models.py b/modules/sd_models.py index 87691e8e1..d1ca3efc6 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -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") diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 89eddae81..3d5c18a44 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -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) diff --git a/modules/shared_legacy.py b/modules/shared_legacy.py index 8b49a3871..c33f565c7 100644 --- a/modules/shared_legacy.py +++ b/modules/shared_legacy.py @@ -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}),