automatic/modules/lora/networks.py

from contextlib import nullcontext
import time
import rich.progress as rp
from modules.errorlimiter import limit_errors
from modules.lora import lora_common as l
from modules.lora import lora_overrides
from modules.lora import lora_sdnq
from modules.lora.lora_apply import network_apply_weights, network_apply_direct, network_backup_weights, network_calc_weights
from modules import shared, devices, sd_models
from modules.logger import log, console


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']


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."""
    from modules.sd_offload_group import remove_group_offload_component
    component = getattr(sd_model, component_name, None)
    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 = []
    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
        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']
    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) # 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
        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()
        applied_weight = 0
        applied_bias = 0
        refused = 0
        with devices.inference_context(), pbar:
            wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in l.loaded_networks) if len(l.loaded_networks) > 0 else ()
            stack_sig = lora_sdnq.signature() # apply-mechanism token tracked beside network_current_names so a settings-only flip re-applies
            applied_layers.clear()
            lora_sdnq.fallback_layers.clear() # a raise mid-pass leaves stale entries behind
            lora_sdnq.hosted_layers.clear()
            backup_size = 0
            for component in modules.keys():
                component_wanted = wanted_names if component in components else ()
                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', '') == stack_sig):
                        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)
                    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
                        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:
                                applied_layers.append(network_layer_name)
                                applied_weight += 1
                            module.network_current_names = component_wanted
                            module.network_current_stack = stack_sig
                            if task is not 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
                        batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, elimit=elimit)
                        if batch_ex_bias is None: # bias deltas need the plain path; weight-only sets ride the side-channel without a weight backup
                            hosted = lora_sdnq.apply_hosted(module, network_layer_name, batch_updown, component_wanted)
                            if hosted is not None:
                                if hosted and component_wanted:
                                    applied_layers.append(network_layer_name)
                                    applied_weight += 1
                                module.network_current_names = component_wanted
                                module.network_current_stack = stack_sig
                                batch_updown, batch_ex_bias = None, None
                                del batch_updown, batch_ex_bias
                                if task is not None:
                                    pbar.update(task, advance=1)
                                continue
                        batch_updown, batch_ex_bias = None, None
                        del batch_updown, batch_ex_bias
                    stripped = lora_sdnq.remove_factors(module) # the mechanism gate can decline a layer still carrying attached factors; the weight path must start from the pristine channel
                    if stripped and not component_wanted: # factor-mode layers have no tensor backup, dropping the factors is the whole restore
                        module.network_current_names = ()
                        module.network_current_stack = stack_sig
                        if task is not None:
                            pbar.update(task, advance=1)
                        continue
                    backup_size += network_backup_weights(module, network_layer_name, component_wanted, fuse)
                    if not component_wanted:
                        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 task is not None:
                                pbar.update(task, advance=1)
                            continue
                        batch_updown, batch_ex_bias = None, None # restore-only pass, apply with no weights reverts to backup
                    else:
                        batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, elimit=elimit)
                        if batch_updown is not None:
                            lora_sdnq.note_fallback(module, network_layer_name) # only layers whose quantized weight actually takes a delta
                    if fuse:
                        weight_written, bias_written = network_apply_direct(module, batch_updown, batch_ex_bias, device=device)
                    else:
                        weight_written, bias_written = network_apply_weights(module, batch_updown, batch_ex_bias, device=device)
                    if batch_updown is not None or batch_ex_bias is not None:
                        applied_layers.append(network_layer_name)
                        applied_weight += 1 if weight_written else 0
                        applied_bias += 1 if bias_written else 0
                        refused += (batch_updown is not None and not weight_written) + (batch_ex_bias is not None and not bias_written) # a delta the module would not take leaves that layer on its base value
                    batch_updown, batch_ex_bias = None, None
                    del batch_updown, batch_ex_bias
                    module.network_current_names = component_wanted
                    module.network_current_stack = stack_sig
                    if task is not None:
                        bs = round(backup_size/1024/1024/1024, 2) if backup_size > 0 else None
                        pbar.update(task, advance=1, description=f'networks={len(l.loaded_networks)} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={bs} device={device}')

            if task is not None and len(applied_layers) == 0:
                pbar.remove_task(task) # hide progress bar for no action
    global native_active, refused_writes # pylint: disable=global-statement
    lora_sdnq.report_fallbacks()
    native_active = len(l.loaded_networks) > 0
    refused_writes = refused
    l.timer.activate += time.time() - t0
    if refused > 0:
        log.error(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} weights={applied_weight} bias={applied_bias} refused={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={active_components} layers={total} weights={applied_weight} bias={applied_bias} refused={refused} backup={round(backup_size/1024/1024/1024, 2)} fuse={fuse}:{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" or len(group_stripped) > 0:
        sd_models.set_diffuser_offload(sd_model, op="model")


def network_deactivate(include=None, exclude=None):
    if exclude is None:
        exclude = []
    if include is None:
        include = []
    fuse = lora_overrides.fuse_native() # must match network_activate: backup mode restores in its restore-only pass instead
    if not fuse or shared.opts.lora_force_diffusers:
        return
    if len(l.previously_loaded_networks) == 0:
        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) # 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']
        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()
        refused = 0
        with devices.inference_context(), pbar:
            applied_layers.clear()
            for component in modules.keys():
                device = getattr(sd_model, component, None).device
                for _, module in modules[component]:
                    network_layer_name = getattr(module, 'network_layer_name', None)
                    if shared.state.interrupted or network_layer_name is 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)
                    if lora_sdnq.remove_factors(module): # exact inverse for factor-mode layers, weights were never touched
                        applied_layers.append(network_layer_name)
                        module.network_current_names = ()
                        if task is not None:
                            pbar.update(task, advance=1)
                        continue
                    batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, use_previous=True, elimit=elimit)
                    if fuse:
                        weight_written, bias_written = network_apply_direct(module, batch_updown, batch_ex_bias, device=device, deactivate=True)
                    else:
                        weight_written, bias_written = network_apply_weights(module, batch_updown, batch_ex_bias, device=device, deactivate=True)
                    if batch_updown is not None or batch_ex_bias is not None:
                        applied_layers.append(network_layer_name)
                        refused += (batch_updown is not None and not weight_written) + (batch_ex_bias is not None and not bias_written) # a delta the module would not take stays applied on that layer
                    del batch_updown, batch_ex_bias
                    module.network_current_names = ()
                    if task is not None:
                        pbar.update(task, advance=1, description=f'networks={len(l.previously_loaded_networks)} modules={active_components} layers={total} unapply={len(applied_layers)}')
    l.timer.deactivate = time.time() - t0
    if refused > 0:
        log.error(f'Network unload: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} unapply={len(applied_layers)} refused={refused} network partially removed')
    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)} refused={refused} fuse={fuse}:{shared.opts.lora_fuse_diffusers} time={l.timer.summary}')
    modules.clear()
    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")