From b84b782ba474aca85d12ca6ef3105c442cb244a6 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 7 Aug 2026 23:16:53 +0100 Subject: [PATCH] fix(lora): apply native networks in place under group offload Group offload hooks report the onload device at module level while the weights rest on cpu, so every native apply took the parameter replacement branch in assign_weight and detached the written layers from the hook's group bookkeeping. The activation and deactivation walks now remove a component's group hooks before its first weight write and reapply offload at the end of the pass: writes land in place on the resting tensors and fresh groups snapshot the result. - hooks come off lazily, only for components with a covered layer or a pending backup or factor-stash restore; repeat activations with an unchanged set leave the hooks untouched - remove_group_offload_component follows wrapper components to the inner model that carries the hooks --- modules/lora/networks.py | 38 ++++++++++++++++++++++++++++++++++++-- modules/sd_models.py | 2 +- modules/sd_offload.py | 4 +++- 3 files changed, 40 insertions(+), 4 deletions(-) diff --git a/modules/lora/networks.py b/modules/lora/networks.py index bc3a896ae..35a412f0a 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -13,6 +13,32 @@ native_active: bool = False default_components = ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'text_encoder_4', 'unet', 'transformer', 'transformer_2', 'llm_adapter'] +def group_will_mutate(module, network_layer_name: str, loaded) -> bool: + """True when the pass will write to this module: a loaded network covers its layer, a + tensor backup awaits restore, or an svd factor stash awaits removal.""" + if any(net.modules.get(network_layer_name, None) is not None for net in loaded): + return True + weights_backup = getattr(module, 'network_weights_backup', None) + if weights_backup is not None and not isinstance(weights_backup, bool): + return True + bias_backup = getattr(module, 'network_bias_backup', None) + if bias_backup is not None and not isinstance(bias_backup, bool): + return True + return getattr(module, 'sdnq_lora_svd_stash', None) is not None + + +def group_offload_strip(sd_model, component_name: str, stripped: dict): + """Group offload hooks come off before the first weight write in a component: a write + under live hooks either replaces a parameter out of the hook's group bookkeeping or is + 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.""" + component = getattr(sd_model, component_name, None) + sd_models.remove_group_offload_component(component) + stripped[component_name] = component.device + return stripped[component_name] + + def network_activate(include=None, exclude=None): if exclude is None: exclude = [] @@ -32,6 +58,8 @@ def network_activate(include=None, exclude=None): sd_models.move_model(sd_model, device=devices.cpu) elif shared.opts.diffusers_offload_mode == "balanced": sd_model = sd_models.apply_balanced_offload(sd_model, force=True) # dispatched modules hold meta tensors backed by the offload map; rebuild them real on cpu with hooks intact before touching weights + group_offload = shared.opts.diffusers_offload_mode == "group" + group_stripped = {} device = None modules = {} components = include if len(include) > 0 else default_components @@ -67,6 +95,8 @@ def network_activate(include=None, exclude=None): if task is not None: pbar.update(task, advance=1) continue + if group_offload and component not in group_stripped and group_will_mutate(module, network_layer_name, l.loaded_networks): + device = group_offload_strip(sd_model, component, group_stripped) backup_size += network_backup_weights(module, network_layer_name, component_wanted) if not component_wanted: weights_backup = getattr(module, "network_weights_backup", None) @@ -100,7 +130,7 @@ def network_activate(include=None, exclude=None): if l.debug and len(l.loaded_networks) > 0: log.debug(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={round(backup_size/1024/1024/1024, 2)} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers} device={device} time={l.timer.summary}') modules.clear() - if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential": + if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential" or len(group_stripped) > 0: sd_models.set_diffuser_offload(sd_model, op="model") @@ -121,6 +151,8 @@ def network_deactivate(include=None, exclude=None): sd_models.move_model(sd_model, device=devices.cpu) elif shared.opts.diffusers_offload_mode == "balanced": sd_model = sd_models.apply_balanced_offload(sd_model, force=True) # dispatched modules hold meta tensors backed by the offload map; rebuild them real on cpu with hooks intact before touching weights + group_offload = shared.opts.diffusers_offload_mode == "group" + group_stripped = {} modules = {} components = include if len(include) > 0 else ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'unet', 'transformer', 'llm_adapter'] @@ -148,6 +180,8 @@ def network_deactivate(include=None, exclude=None): if task is not None: pbar.update(task, advance=1) continue + if group_offload and component not in group_stripped and group_will_mutate(module, network_layer_name, l.previously_loaded_networks): + device = group_offload_strip(sd_model, component, group_stripped) batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, use_previous=True, elimit=elimit) if shared.opts.lora_fuse_native: network_apply_direct(module, batch_updown, batch_ex_bias, device=device, deactivate=True) @@ -163,5 +197,5 @@ def network_deactivate(include=None, exclude=None): if l.debug and len(l.previously_loaded_networks) > 0: log.debug(f'Network deactivate: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} modules={active_components} layers={total} apply={len(applied_layers)} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers} time={l.timer.summary}') modules.clear() - if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential": + if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential" or len(group_stripped) > 0: sd_models.set_diffuser_offload(sd_model, op="model") diff --git a/modules/sd_models.py b/modules/sd_models.py index c4a829509..f050b7bfc 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -16,7 +16,7 @@ from modules.memstats import memory_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 # 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 # pylint: disable=unused-import from modules.sd_models_utils import NoWatermark, get_signature, get_call, path_to_repo, apply_function_to_model, read_state_dict, get_state_dict_from_checkpoint # pylint: disable=unused-import diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 58699d8b0..a4e15f406 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -96,7 +96,9 @@ def group_offload_config(main: bool) -> dict: def remove_group_offload_component(module): if getattr(module, 'sdnext_group_offload_sig', None) is None: - return + 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 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)