fix(lora): harden select stack modes on quantized and offloaded models

Select pairs now ride the svd side channel on any SDNQ linear, not
only sub-8-bit ones: quantized backups are packed tensors, so the
weight rewrite path cannot recompute a winner from them and left
layers stripped mid-requantize. The sub-8-bit gate stays for dense
hosting, where requantize retains the delta at 8 bits and above.

Weight selection now only serves unquantized modules: finalize
iterates a snapshot so dead entries drop cleanly, materializes
balanced-offload modules before rewriting weights and skips modules
with stripped or quantized weights instead of corrupting the layer.

modules/lora/lora_sdnq.py

@@ -258,8 +258,8 @@ def append_factors(self, ups, downs):
     return segments, deq.use_quantized_matmul
 
 
-def host_candidate(self, network_layer_name, wanted_names):
-    """True when a non-factorable set on this layer should be hosted as a truncated svd."""
+def select_candidate(self, network_layer_name, wanted_names):
+    """True when this layer can carry a set on the svd channel; select pairs ride it at any bit width."""
     if not enabled():
         return False
     if int(getattr(shared.opts, 'lora_sdnq_host_rank', 0) or 0) <= 0:
@@ -268,11 +268,17 @@ def host_candidate(self, network_layer_name, wanted_names):
         return False
     if wanted_names == ():
         return False
+    return any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks)
+
+
+def host_candidate(self, network_layer_name, wanted_names):
+    """True when a non-factorable set on this layer should be hosted as a truncated svd."""
+    if not select_candidate(self, network_layer_name, wanted_names):
+        return False
     from sdnq.common import dtype_dict
     if dtype_dict[self.sdnq_dequantizer.weights_dtype]['num_bits'] >= 8:
         return False # requantize retains most of the delta at 8 bits and above; truncation would lose more than it saves
-
-    return any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks)
+    return True
 
 
 def apply_cached(self, network_layer_name, wanted_names):

modules/lora/lora_stack.py

@@ -248,12 +248,21 @@ def layer_flip_step(scores, total_steps):
     return total_steps
 
 
+def materialize_model():
+    """Weight-kind selection rewrites module weights outside the activation walk; rebuild balanced-offload modules real first (mirrors network_activate)."""
+    from modules import sd_models
+    if getattr(shared.opts, 'diffusers_offload_mode', None) == 'balanced' and getattr(shared, 'sd_model', None) is not None:
+        sd_models.apply_balanced_offload(shared.sd_model, force=True)
+
+
 def finalize(total_steps):
     """Build the inverted flip map for the pass; select-mode layers start at their step-0 winner."""
     state['total_steps'] = int(total_steps)
     state['gamma'] = (state['gamma_num'] / state['gamma_den']) if state['gamma_den'] > 0 else 1.0
     state['flips'] = {}
-    for layer_name, entry in state['entries'].items():
+    if any(e['kind'] == 'weight' for e in state['entries'].values()):
+        materialize_model()
+    for layer_name, entry in list(state['entries'].items()): # snapshot: apply_selection drops entries whose module died
         flip_at = layer_flip_step(entry['scores'], state['total_steps'])
         initial = 1 if flip_at == 0 else 0
         apply_selection(layer_name, entry, initial)
@@ -298,6 +307,9 @@ def apply_selection(layer_name, entry, winner):
 def weight_selection(module, entry, winner):
     from modules.lora import lora_common as l
     from modules.lora.lora_apply import network_apply_weights
+    if getattr(module, 'sdnq_dequantizer', None) is not None:
+        warn_once('select-sdnq-weight', 'Network stack: flip=skipped layer=quantized') # quantized backups are packed tensors; only the segment path can flip them
+        return
     backup = getattr(module, 'network_weights_backup', None)
     if not isinstance(backup, torch.Tensor): # fuse mode keeps a bool sentinel, not a pristine copy
         warn_once('select-nobackup', 'Network stack: flip=skipped backup=none')
@@ -306,6 +318,10 @@ def weight_selection(module, entry, winner):
     net_module = net.modules.get(entry['layer'], None) if net is not None else None
     if net_module is None:
         return
-    device = module.weight.device
+    weight = getattr(module, 'weight', None)
+    if weight is None or weight.is_meta:
+        warn_once('select-offloaded', 'Network stack: flip=skipped weight=offloaded')
+        return
+    device = weight.device
     updown = net_module.calc_updown(backup.to(device))[0]
     network_apply_weights(module, updown, None, device=device) # recomputes from the pristine backup, requantizing where the layer needs it

modules/lora/networks.py

@@ -112,7 +112,7 @@ def network_activate(include=None, exclude=None):
                         device = group_offload_strip(sd_model, component, group_stripped)
                     calced = False # tracks whether this iteration assembled the delta, so the fallthrough reuses it instead of recomputing
                     if select_active and component_wanted and not network_layer_name.startswith('lora_te'):
-                        if lora_sdnq.host_candidate(module, network_layer_name, component_wanted): # sub-8-bit SDNQ pairs ride the channel as separate segments
+                        if lora_sdnq.select_candidate(module, network_layer_name, component_wanted): # SDNQ pairs ride the channel as separate segments at any bit width; weight rewrites cannot flip a quantized layer
                             weights_backup = getattr(module, "network_weights_backup", None)
                             if weights_backup is not None and not isinstance(weights_backup, bool):
                                 network_apply_weights(module, None, None, device=device)

test/test-sdnq-lora-factors.py

@@ -1749,6 +1749,38 @@ def test_select_gated_off_when_compiled():
     return True
 
 
+def test_select_finalize_drops_dead_module():
+    import weakref
+    layer = build_layer('uint4')
+    n1, n2, _D1, _D2 = select_pair(layer, seed0=61, seed1=62)
+    with mock_model(lin=layer), select_mode('klora'):
+        activate(n1, n2)
+        entry = lora_stack.state['entries'].get('lora_transformer_test')
+        assert entry is not None, 'pair must register before the module dies'
+        entry['module'] = weakref.ref(torch.nn.Linear(2, 2)) # referent dies immediately: simulates offload re-wraps replacing a registered module
+        assert entry['module']() is None
+        lora_stack.reset(12)
+        assert 'lora_transformer_test' not in lora_stack.state['entries'], 'a dead module must drop its entry without breaking finalize'
+        activate()
+    return True
+
+
+def test_select_int8_pair_rides_segments():
+    layer = build_layer('int8')
+    n1, n2, D1, D2 = select_pair(layer, seed0=63, seed1=64)
+    with mock_model(lin=layer), select_mode('klora'):
+        Wdq0 = dq(layer)
+        activate(n1, n2)
+        entry = lora_stack.state['entries'].get('lora_transformer_test')
+        assert entry is not None and entry['kind'] == 'factor', 'an int8 pair must ride svd segments, not weight rewrites'
+        lora_stack.reset(16)
+        eff = dq(layer) - Wdq0
+        assert max(rho_of(eff, D1), rho_of(eff, D2)) > 0.99, 'initial selection must deliver one exact per-net delta'
+        activate()
+        assert torch.equal(dq(layer), Wdq0), 'removal must restore bit-exact'
+    return True
+
+
 def test_select_gate_dormant_without_pair():
     with select_mode('klora'):
         assert lora_stack.select_possible(1) is False, 'a single network must leave the fuse gate alone'
@@ -2009,7 +2041,7 @@ def run_tests():
     for fn in [test_select_flip_schedule_end_to_end, test_select_initial_style_when_ramp_starts_won, test_select_flip_is_inplace_and_shape_stable,
                test_select_matmul_transposed_layout, test_select_per_net_hosted_pair, test_select_reset_restores_initial_state,
                test_select_deactivate_from_midflip, test_select_requires_exactly_two_nets, test_select_gated_off_when_compiled,
-               test_select_gate_dormant_without_pair, test_stale_schedule_dropped_on_reapply,
+               test_select_finalize_drops_dead_module, test_select_int8_pair_rides_segments, test_select_gate_dormant_without_pair, test_stale_schedule_dropped_on_reapply,
                test_est_energy_matches_full_frobenius, test_select_weight_kind_plain_layer]:
         run_test(CAT_SELECT, fn)
     log.warning('=== Compile ===')