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)