mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 17:24:32 +02:00
refactor(lora): extract the shared activation pass plumbing
Both passes opened with the same four steps written twice: bring the model into a writable state, enumerate the components to walk, open a progress bar, and probe a weight backup before restoring it. Pull each into a helper and call it from both entry points. The component collector keeps the two lists it is given, so deactivate still walks its own shorter set, and promotion now clears the staged config it consumed.
This commit is contained in:
CalamitousFelicitousness
+83
-68
@@ -16,6 +16,7 @@ applied_layers: list[str] = []
|
||||
refused_writes: int = 0 # deltas the modules would not take on the last activate pass; infotext reports the network as partial
|
||||
native_active: bool = False
|
||||
default_components = ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'text_encoder_4', 'unet', 'transformer', 'transformer_2', 'llm_adapter']
|
||||
deactivate_components = ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'unet', 'transformer', 'llm_adapter']
|
||||
|
||||
|
||||
def group_will_mutate(module, network_layer_name: str, loaded) -> bool:
|
||||
@@ -45,48 +46,91 @@ def group_offload_strip(sd_model, component_name: str, stripped: dict):
|
||||
return stripped[component_name]
|
||||
|
||||
|
||||
def network_activate(include=None, exclude=None):
|
||||
if exclude is None:
|
||||
exclude = []
|
||||
if include is None:
|
||||
include = []
|
||||
for net in l.loaded_networks: # promote staged multipliers only now: the deactivate pass ran against the previous values, which fuse-mode removal recomputes with
|
||||
def promote_pending():
|
||||
"""Promote staged multipliers onto the loaded networks; the deactivate pass ran against the previous values, which fuse-mode removal recomputes with."""
|
||||
for net in l.loaded_networks:
|
||||
pending = getattr(net, 'pending_config', None)
|
||||
if pending is not None:
|
||||
net.te_multiplier = pending['te']
|
||||
net.unet_multiplier = pending['unet']
|
||||
net.dyn_dim = pending['dyn']
|
||||
net.block_spec = pending.get('blocks', None)
|
||||
net.pending_config = None # promotion is one-shot
|
||||
|
||||
|
||||
def prepare_model_for_write(sd_model):
|
||||
"""Bring the model into a state where weight writes land; balanced offload returns a rebuilt model."""
|
||||
if shared.opts.diffusers_offload_mode == "sequential":
|
||||
sd_models.disable_offload(sd_model)
|
||||
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, silent=True) # dispatched modules hold meta tensors backed by the offload map; rebuild them real on cpu with hooks intact before touching weights
|
||||
return sd_model
|
||||
|
||||
|
||||
def collect_components(sd_model, include, exclude, defaults, restore_filtered):
|
||||
"""Modules to walk, as (modules, wanted components, walked component names, module count).
|
||||
|
||||
With restore_filtered the walk also covers the components a filter left
|
||||
out, so they restore to backup instead of freezing with stale weights;
|
||||
those names stay out of the reported list because nothing applies to them.
|
||||
"""
|
||||
components = include if len(include) > 0 else defaults
|
||||
components = [x for x in components if x not in exclude]
|
||||
filtered = [x for x in defaults if x not in components] if restore_filtered else []
|
||||
modules = {}
|
||||
active_components = []
|
||||
for name in components + filtered:
|
||||
component = getattr(sd_model, name, None)
|
||||
if component is not None and hasattr(component, 'named_modules'):
|
||||
if name in components:
|
||||
active_components.append(name)
|
||||
modules[name] = list(component.named_modules())
|
||||
return modules, components, active_components, sum(len(x) for x in modules.values())
|
||||
|
||||
|
||||
def pass_progress(action, total, show):
|
||||
"""Progress bar for one pass, or a nullcontext with no task when there is nothing to show."""
|
||||
if not show:
|
||||
return nullcontext(), None
|
||||
pbar = rp.Progress(rp.TextColumn(f'[cyan]Network: type=LoRA action={action}'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=console)
|
||||
return pbar, pbar.add_task(description='', total=total)
|
||||
|
||||
|
||||
def tensor_backup(module):
|
||||
"""The module's weight backup when it holds real tensors; None in fuse mode, where the backup is a marker."""
|
||||
weights_backup = getattr(module, 'network_weights_backup', None)
|
||||
return None if isinstance(weights_backup, bool) else weights_backup
|
||||
|
||||
|
||||
def restore_pristine(module, device):
|
||||
"""Put a backed-up layer back on its checkpoint weights, so a mechanism sees the pristine base."""
|
||||
if tensor_backup(module) is not None:
|
||||
network_apply_weights(module, None, None, device=device)
|
||||
|
||||
|
||||
def should_skip(module, network_layer_name, wanted, stack_sig):
|
||||
"""True when the pass has nothing to do here: no weight, interrupted, unnamed, or already carrying this set under these settings."""
|
||||
if getattr(module, 'weight', None) is None or shared.state.interrupted or network_layer_name is None:
|
||||
return True
|
||||
return getattr(module, 'network_current_names', ()) == wanted and getattr(module, 'network_current_stack', 'sum') == stack_sig
|
||||
|
||||
|
||||
def network_activate(include=None, exclude=None):
|
||||
if exclude is None:
|
||||
exclude = []
|
||||
if include is None:
|
||||
include = []
|
||||
promote_pending()
|
||||
t0 = time.time()
|
||||
fuse = lora_overrides.fuse_native() # resolve once: backup and apply passes must agree
|
||||
with limit_errors("network_activate") as elimit:
|
||||
sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
|
||||
if shared.opts.diffusers_offload_mode == "sequential":
|
||||
sd_models.disable_offload(sd_model)
|
||||
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, silent=True) # dispatched modules hold meta tensors backed by the offload map; rebuild them real on cpu with hooks intact before touching weights
|
||||
sd_model = prepare_model_for_write(getattr(shared.sd_model, "pipe", shared.sd_model))
|
||||
group_offload = shared.opts.diffusers_offload_mode == "group"
|
||||
group_stripped = {}
|
||||
device = None
|
||||
modules = {}
|
||||
components = include if len(include) > 0 else default_components
|
||||
components = [x for x in components if x not in exclude]
|
||||
filtered_components = [x for x in default_components if x not in components] # filtered components restore to backup so a filter means detached, not frozen with stale weights
|
||||
active_components = []
|
||||
for name in components + filtered_components:
|
||||
component = getattr(sd_model, name, None)
|
||||
if component is not None and hasattr(component, 'named_modules'):
|
||||
if name in components:
|
||||
active_components.append(name)
|
||||
modules[name] = list(component.named_modules())
|
||||
total = sum(len(x) for x in modules.values())
|
||||
if len(l.loaded_networks) > 0:
|
||||
pbar = rp.Progress(rp.TextColumn('[cyan]Network: type=LoRA action=activate'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=console)
|
||||
task = pbar.add_task(description='' , total=total)
|
||||
else:
|
||||
task = None
|
||||
pbar = nullcontext()
|
||||
modules, components, active_components, total = collect_components(sd_model, include, exclude, default_components, restore_filtered=True)
|
||||
pbar, task = pass_progress('activate', total, len(l.loaded_networks) > 0)
|
||||
applied_weight = 0
|
||||
applied_bias = 0
|
||||
refused = 0
|
||||
@@ -105,8 +149,7 @@ def network_activate(include=None, exclude=None):
|
||||
device = getattr(sd_model, component, None).device
|
||||
for _, module in modules[component]:
|
||||
network_layer_name = getattr(module, 'network_layer_name', None)
|
||||
current_names = getattr(module, "network_current_names", ())
|
||||
if getattr(module, 'weight', None) is None or shared.state.interrupted or (network_layer_name is None) or (current_names == component_wanted and getattr(module, 'network_current_stack', 'sum') == stack_sig):
|
||||
if should_skip(module, network_layer_name, component_wanted, stack_sig):
|
||||
if task is not None:
|
||||
pbar.update(task, advance=1)
|
||||
continue
|
||||
@@ -116,9 +159,7 @@ def network_activate(include=None, exclude=None):
|
||||
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.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)
|
||||
restore_pristine(module, device)
|
||||
applied = lora_sdnq.apply_select_cached(module, network_layer_name, component_wanted) # a stored score record and factor pair serve before the deltas are assembled
|
||||
if applied is None:
|
||||
per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True)
|
||||
@@ -139,8 +180,7 @@ def network_activate(include=None, exclude=None):
|
||||
lora_stack.warn_once('select-host-disabled', f'Network stack: mode={lora_stack.mode()} quant=sdnq host=disabled fallback=sum')
|
||||
else: # other layers select by recomputing the winner from the pristine backup at schedule time
|
||||
sel_backup = network_backup_weights(module, network_layer_name, component_wanted, fuse)
|
||||
weights_backup = getattr(module, "network_weights_backup", None)
|
||||
if weights_backup is not None and not isinstance(weights_backup, bool):
|
||||
if tensor_backup(module) is not None: # a flip recomputes the winner from the pristine tensor, which fuse mode does not keep
|
||||
if lora_stack.register_weight_pair_cached(network_layer_name, module, component_wanted): # a stored score record registers without assembling the pair
|
||||
backup_size += sel_backup
|
||||
network_apply_weights(module, None, None, device=device) # pristine until the schedule applies the winner
|
||||
@@ -163,9 +203,7 @@ def network_activate(include=None, exclude=None):
|
||||
pbar.update(task, advance=1)
|
||||
continue
|
||||
if lora_sdnq.factor_candidate(module, network_layer_name, component_wanted):
|
||||
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) # an earlier non-factorable set requantized this layer, restore the pristine base before attaching factors
|
||||
restore_pristine(module, device) # an earlier non-factorable set may have requantized this layer
|
||||
applied = lora_sdnq.apply_factors(module, network_layer_name, component_wanted)
|
||||
if applied is not None: # exact path took the layer; None falls through to hosting or requantize
|
||||
if applied and component_wanted:
|
||||
@@ -177,9 +215,7 @@ def network_activate(include=None, exclude=None):
|
||||
pbar.update(task, advance=1)
|
||||
continue
|
||||
if lora_sdnq.host_candidate(module, network_layer_name, component_wanted):
|
||||
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) # the hosted delta is measured against the pristine base
|
||||
restore_pristine(module, device) # the hosted delta is measured against the pristine base
|
||||
hosted = lora_sdnq.apply_cached(module, network_layer_name, component_wanted) # a stored entry serves the layer before the delta is assembled
|
||||
if hosted is None:
|
||||
batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, elimit=elimit)
|
||||
@@ -208,8 +244,7 @@ def network_activate(include=None, exclude=None):
|
||||
backup_size += network_backup_weights(module, network_layer_name, component_wanted, fuse)
|
||||
if not component_wanted:
|
||||
lora_stack.drop(network_layer_name) # a restored layer must leave the selection schedule
|
||||
weights_backup = getattr(module, "network_weights_backup", None)
|
||||
if weights_backup is None or isinstance(weights_backup, bool): # fuse mode has no tensor backup, restore stays with network_deactivate
|
||||
if tensor_backup(module) is None: # fuse mode has no tensor backup, restore stays with network_deactivate
|
||||
if task is not None:
|
||||
pbar.update(task, advance=1)
|
||||
continue
|
||||
@@ -274,31 +309,11 @@ def network_deactivate(include=None, exclude=None):
|
||||
return
|
||||
t0 = time.time()
|
||||
with limit_errors("network_deactivate") as elimit:
|
||||
sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
|
||||
if shared.opts.diffusers_offload_mode == "sequential":
|
||||
sd_models.disable_offload(sd_model)
|
||||
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, silent=True) # dispatched modules hold meta tensors backed by the offload map; rebuild them real on cpu with hooks intact before touching weights
|
||||
sd_model = prepare_model_for_write(getattr(shared.sd_model, "pipe", shared.sd_model))
|
||||
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']
|
||||
components = [x for x in components if x not in exclude]
|
||||
active_components = []
|
||||
for name in components:
|
||||
component = getattr(sd_model, name, None)
|
||||
if component is not None and hasattr(component, 'named_modules'):
|
||||
modules[name] = list(component.named_modules())
|
||||
active_components.append(name)
|
||||
total = sum(len(x) for x in modules.values())
|
||||
if len(l.previously_loaded_networks) > 0 and l.debug:
|
||||
pbar = rp.Progress(rp.TextColumn('[cyan]Network: type=LoRA action=deactivate'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=console)
|
||||
task = pbar.add_task(description='', total=total)
|
||||
else:
|
||||
task = None
|
||||
pbar = nullcontext()
|
||||
modules, _components, active_components, total = collect_components(sd_model, include, exclude, deactivate_components, restore_filtered=False)
|
||||
pbar, task = pass_progress('deactivate', total, len(l.previously_loaded_networks) > 0 and l.debug)
|
||||
refused = 0
|
||||
with devices.inference_context(), pbar:
|
||||
applied_layers.clear()
|
||||
|
||||
Reference in New Issue
Block a user