fix(lora): finish the activation pass when it aborts

The error limiter halts a pass by raising, and nothing between the raise
and the caller put the model back. A halted pass left group offload hooks
stripped from every component the walk had reached, left a sequential
model on the cpu with offload disabled, and left the counters other
modules read describing the pass before it.

The epilogue moves into finish_pass under a finally, so the model returns
to its offload mode and the counters describe the pass that just ran. The
abort still reaches the caller.

Pass state is reset in one place in lora_sdnq now. Two of the six
accumulators were not being cleared at the start of a pass, and a stale
routed layer suppresses the fallback count for that layer next time.

modules/lora/lora_sdnq.py

@@ -564,6 +564,16 @@ def note_fallback(self, network_layer_name):
         fallback_layers.append(network_layer_name)
 
 
+def reset_pass():
+    """Clear every per-pass accumulator, so a pass that raised leaves nothing behind for the next one."""
+    fallback_layers.clear()
+    hosted_layers.clear()
+    hosted_ranks.clear()
+    factor_layers.clear()
+    select_layers.clear()
+    routed_layers.clear() # note_fallback reads this to suppress double counting, so a stale entry silences a real fallback
+
+
 def report_fallbacks():
     hits, misses = lora_factor_cache.flush()
     if hits > 0 or misses > 0:

modules/lora/networks.py

@@ -276,46 +276,13 @@ def apply_generic(ctx, module, network_layer_name, batch):
     ctx.tick(f'networks={len(l.loaded_networks)} modules={ctx.active_components} layers={ctx.total} weights={ctx.applied_weight} bias={ctx.applied_bias} backup={bs} device={ctx.device}')
 
 
-def network_activate(include=None, exclude=None):
-    if exclude is None:
-        exclude = []
-    if include is None:
-        include = []
-    promote_pending()
-    t0 = time.time()
-    ctx = ActivationPass(lora_overrides.fuse_native()) # fuse resolved once: the backup, apply and restore paths must agree
-    with limit_errors("network_activate") as elimit:
-        ctx.elimit = elimit
-        ctx.sd_model = prepare_model_for_write(ctx.sd_model)
-        modules, components, ctx.active_components, ctx.total = collect_components(ctx.sd_model, include, exclude, default_components, restore_filtered=True)
-        ctx.pbar, ctx.task = pass_progress('activate', ctx.total, len(l.loaded_networks) > 0)
-        with devices.inference_context(), ctx.pbar:
-            applied_layers.clear()
-            lora_sdnq.fallback_layers.clear() # a raise mid-pass leaves stale entries behind
-            lora_sdnq.hosted_layers.clear()
-            lora_sdnq.factor_layers.clear()
-            lora_sdnq.select_layers.clear()
-            for component in modules.keys():
-                ctx.component_wanted = ctx.wanted_names if component in components else () # the pass tuple itself, never a copy
-                ctx.device = getattr(ctx.sd_model, component, None).device
-                for _, module in modules[component]:
-                    network_layer_name = getattr(module, 'network_layer_name', None)
-                    if should_skip(module, network_layer_name, ctx.component_wanted, ctx.stack_sig):
-                        ctx.tick()
-                        continue
-                    lora_stack.drop(network_layer_name) # re-application invalidates any live selection schedule; the select branch re-registers
-                    if ctx.group_offload and component not in ctx.group_stripped and group_will_mutate(module, network_layer_name, l.loaded_networks):
-                        ctx.device = group_offload_strip(ctx.sd_model, component, ctx.group_stripped)
-                    if try_select(ctx, module, network_layer_name):
-                        continue
-                    if try_factors(ctx, module, network_layer_name):
-                        continue
-                    hosted, batch = try_hosted(ctx, module, network_layer_name)
-                    if hosted:
-                        continue
-                    apply_generic(ctx, module, network_layer_name, batch)
-            if ctx.task is not None and len(applied_layers) == 0:
-                ctx.pbar.remove_task(ctx.task) # hide progress bar for no action
+def finish_pass(ctx, t0):
+    """Publish what the pass did and put the model back under its offload mode.
+
+    Runs even when the error limiter aborts the walk: the hooks it stripped
+    and the offload it disabled have to come back, and the counters other
+    modules read have to describe this pass.
+    """
     global native_active, refused_writes # pylint: disable=global-statement
     lora_sdnq.report_fallbacks()
     native_active = len(l.loaded_networks) > 0
@@ -326,11 +293,54 @@ def network_activate(include=None, exclude=None):
         log.error(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} weights={ctx.applied_weight} bias={ctx.applied_bias} refused={ctx.refused} network partially applied')
     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={ctx.active_components} layers={ctx.total} weights={ctx.applied_weight} bias={ctx.applied_bias} refused={ctx.refused} backup={round(ctx.backup_size/1024/1024/1024, 2)} fuse={ctx.fuse}:{shared.opts.lora_fuse_diffusers} device={ctx.device} time={l.timer.summary}')
-    modules.clear()
     if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential" or len(ctx.group_stripped) > 0:
         sd_models.set_diffuser_offload(ctx.sd_model, op="model")
 
 
+def network_activate(include=None, exclude=None):
+    if exclude is None:
+        exclude = []
+    if include is None:
+        include = []
+    promote_pending()
+    t0 = time.time()
+    ctx = ActivationPass(lora_overrides.fuse_native()) # fuse resolved once: the backup, apply and restore paths must agree
+    applied_layers.clear()
+    lora_sdnq.reset_pass()
+    modules = {}
+    try:
+        with limit_errors("network_activate") as elimit:
+            ctx.elimit = elimit
+            ctx.sd_model = prepare_model_for_write(ctx.sd_model)
+            modules, components, ctx.active_components, ctx.total = collect_components(ctx.sd_model, include, exclude, default_components, restore_filtered=True)
+            ctx.pbar, ctx.task = pass_progress('activate', ctx.total, len(l.loaded_networks) > 0)
+            with devices.inference_context(), ctx.pbar:
+                for component in modules.keys():
+                    ctx.component_wanted = ctx.wanted_names if component in components else () # the pass tuple itself, never a copy
+                    ctx.device = getattr(ctx.sd_model, component, None).device
+                    for _, module in modules[component]:
+                        network_layer_name = getattr(module, 'network_layer_name', None)
+                        if should_skip(module, network_layer_name, ctx.component_wanted, ctx.stack_sig):
+                            ctx.tick()
+                            continue
+                        lora_stack.drop(network_layer_name) # re-application invalidates any live selection schedule; the select branch re-registers
+                        if ctx.group_offload and component not in ctx.group_stripped and group_will_mutate(module, network_layer_name, l.loaded_networks):
+                            ctx.device = group_offload_strip(ctx.sd_model, component, ctx.group_stripped)
+                        if try_select(ctx, module, network_layer_name):
+                            continue
+                        if try_factors(ctx, module, network_layer_name):
+                            continue
+                        hosted, batch = try_hosted(ctx, module, network_layer_name)
+                        if hosted:
+                            continue
+                        apply_generic(ctx, module, network_layer_name, batch)
+                if ctx.task is not None and len(applied_layers) == 0:
+                    ctx.pbar.remove_task(ctx.task) # hide progress bar for no action
+    finally:
+        finish_pass(ctx, t0)
+        modules.clear()
+
+
 def effective_mode():
     """Weight-state label for load logs: backup and fuse say how touched weights restore, factor means the whole load rode the svd channel and unload just drops factors."""
     if getattr(l, 'last_backup_size', 0) > 0:

test/test-sdnq-lora-factors.py

@@ -2515,6 +2515,38 @@ def test_nunchaku_entries_carry_the_network_interface():
     return True
 
 
+def test_aborted_pass_still_publishes_its_state():
+    layer = build_layer('uint4')
+    _A, _B, D = make_delta()
+    net = make_dense_net('aborted', layer, D)
+    reported = {'n': 0}
+    real_prepare = networks.prepare_model_for_write
+    real_report = lora_sdnq.report_fallbacks
+
+    def exploding(_sd_model):
+        raise RuntimeError('offload rebuild failed') # a raise before the walk binds anything the epilogue reads
+
+    def counting_report():
+        reported['n'] += 1
+        real_report()
+
+    networks.prepare_model_for_write = exploding
+    lora_sdnq.report_fallbacks = counting_report
+    try:
+        with mock_model(lin=layer):
+            raised = None
+            try:
+                activate(net)
+            except RuntimeError as e:
+                raised = e
+            assert raised is not None and 'offload rebuild failed' in str(raised), f'the original failure must reach the caller, got {raised!r}'
+            assert reported['n'] == 1, 'an aborted pass must still publish its counters and put the model back under its offload mode'
+    finally:
+        networks.prepare_model_for_write = real_prepare
+        lora_sdnq.report_fallbacks = real_report
+    return True
+
+
 def test_stacked_shape_mismatch_falls_back():
     from types import SimpleNamespace
     layer = build_layer('uint4')
@@ -2986,7 +3018,7 @@ def run_tests():
         run_test(CAT_COMPILE, fn)
     log.warning('=== Robustness ===')
     for fn in [test_remove_factors_after_device_move, test_stacked_shape_mismatch_falls_back, test_nunchaku_entries_carry_the_network_interface,
-               test_four_dim_oft_blocks_load_as_boft]:
+               test_four_dim_oft_blocks_load_as_boft, test_aborted_pass_still_publishes_its_state]:
         run_test(CAT_ROBUST, fn)
     log.warning('=== Block weights ===')
     for fn in [test_block_index_sd_unet_layout, test_block_index_sdxl_unet_layout, test_block_index_flux_chains_concatenate,