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

@@ -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")

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
 
 

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)