mirror of
https://github.com/vladmandic/automatic
synced 2026-09-09 22:38:43 +02:00
CalamitousFelicitousness
44a4a8e2a5
The ramp read 1.5 when the option was absent while the option itself defaults to 0.0, so a config without the key ran a ramp the settings page said was off. Only reachable where the options registry is not loaded, which is where the offline suites run.
419 lines
19 KiB
Python
419 lines
19 KiB
Python
"""Stack modes for combining multiple LoRA networks beyond plain summation.
|
|
|
|
Dense modes (ties, dare_ties, dare_linear, magnitude_prune) combine the
|
|
networks' dense deltas elementwise; the result rides the normal apply tail
|
|
(side-channel hosting on sub-8-bit SDNQ, requantize at int8 and above,
|
|
direct add on unquantized layers). Select modes (klora, estlora) keep both
|
|
networks' contributions separate and choose a per-layer winner, shifting
|
|
from the first loaded network (subject) toward the second (style) across
|
|
the sampling steps. Selection scores depend only on the weights, so the
|
|
shift reduces to at most one flip per layer per generation, executed from
|
|
the step callback against a schedule finalized at apply time.
|
|
|
|
TIES arXiv:2306.01708, DARE arXiv:2311.03099, K-LoRA arXiv:2502.18461,
|
|
EST-LoRA arXiv:2508.02165 (its measured style-discrepancy estimate is
|
|
exposed as an option instead of being derived from probe generations).
|
|
"""
|
|
|
|
import time
|
|
import weakref
|
|
import hashlib
|
|
|
|
import torch
|
|
|
|
from modules import shared
|
|
from modules.logger import log
|
|
|
|
|
|
DENSE_MODES = ('ties', 'dare_ties', 'dare_linear', 'magnitude_prune')
|
|
SELECT_MODES = ('klora', 'estlora')
|
|
KLORA_BETA = 0.5 # the paper's fixed ramp offset; only the slope is user-tunable
|
|
ROW_CHUNK = 512 # fp32 interiors run in first-dim slices; also fixes the DARE draw sequence
|
|
SAMPLE_CAP = 1 << 22 # strided subsample bound for magnitude quantiles (full-size quantile exceeds torch limits)
|
|
|
|
state: dict = {'entries': {}, 'flips': {}, 'gamma': 1.0, 'gamma_e': 1.0, 'total_steps': 0, 'finalized': False, 'reported': None}
|
|
warned: set = set()
|
|
warned_context = None
|
|
|
|
|
|
def mode():
|
|
return getattr(shared.opts, 'lora_stack_mode', 'sum') or 'sum'
|
|
|
|
|
|
def density():
|
|
return float(getattr(shared.opts, 'lora_stack_density', 0.5))
|
|
|
|
|
|
def ramp_alpha():
|
|
return float(getattr(shared.opts, 'lora_stack_alpha', 0.0))
|
|
|
|
|
|
def manual_discrepancy():
|
|
return float(getattr(shared.opts, 'lora_stack_discrepancy', 0.5))
|
|
|
|
|
|
def signature():
|
|
m = mode()
|
|
if m in DENSE_MODES:
|
|
return f'{m}:{density():.2f}'
|
|
if m in SELECT_MODES:
|
|
return f'{m}:{ramp_alpha():.2f}:{manual_discrepancy():.2f}'
|
|
return 'sum'
|
|
|
|
|
|
def warn_context():
|
|
"""Settings the degradation warnings below speak about."""
|
|
return (signature(), getattr(shared.opts, 'diffusers_offload_mode', ''), int(getattr(shared.opts, 'lora_sdnq_host_rank', 0) or 0), getattr(shared.opts, 'sd_model_checkpoint', ''))
|
|
|
|
|
|
def warn_once(key, message):
|
|
global warned_context # pylint: disable=global-statement
|
|
context = warn_context()
|
|
if context != warned_context:
|
|
warned.clear() # what was reported under the old settings says nothing about the new ones
|
|
warned_context = context
|
|
if key not in warned:
|
|
warned.add(key)
|
|
log.warning(message)
|
|
|
|
|
|
def select_blocked():
|
|
return 'Model' in (getattr(shared.opts, 'cuda_compile', None) or [])
|
|
|
|
|
|
def active_dense(n_contrib):
|
|
return mode() in DENSE_MODES and n_contrib >= 2
|
|
|
|
|
|
def select_possible(n_loaded):
|
|
"""True when the loaded set could engage a select mode; silent, for the fuse gate."""
|
|
return mode() in SELECT_MODES and n_loaded == 2 and not select_blocked()
|
|
|
|
|
|
def select_engaged():
|
|
"""True while selection schedules are live on model layers."""
|
|
return bool(state['entries'])
|
|
|
|
|
|
def active_select(n_loaded):
|
|
m = mode()
|
|
if m not in SELECT_MODES:
|
|
return False
|
|
if n_loaded != 2:
|
|
log.warning(f'Network stack: mode={m} networks={n_loaded} required=2 fallback=sum')
|
|
return False
|
|
if select_blocked():
|
|
log.warning(f'Network stack: mode={m} compile=model fallback=sum')
|
|
return False
|
|
return True
|
|
|
|
|
|
def seed_for(layer_name, net_name):
|
|
payload = f'{layer_name}|{net_name}|{mode()}|{round(density(), 6)}'
|
|
return int.from_bytes(hashlib.sha256(payload.encode()).digest()[:8], 'little')
|
|
|
|
|
|
def magnitude_threshold(delta, dens):
|
|
flat = delta.abs().flatten()
|
|
step = max(1, flat.numel() // SAMPLE_CAP)
|
|
return torch.quantile(flat[::step].float(), 1.0 - dens)
|
|
|
|
|
|
def dare_generator(device, layer_name, net_name):
|
|
gen = torch.Generator(device=device)
|
|
gen.manual_seed(seed_for(layer_name, net_name))
|
|
return gen
|
|
|
|
|
|
def combine(named_deltas, layer_name):
|
|
"""Combine per-network dense deltas under the active dense mode; returns a tensor in the first delta's dtype."""
|
|
m = mode()
|
|
dens = density()
|
|
deltas = [d for _, d in named_deltas]
|
|
out_dtype = deltas[0].dtype
|
|
result = torch.zeros_like(deltas[0], dtype=torch.float32)
|
|
thresholds = [magnitude_threshold(d, dens) for d in deltas] if m in ('ties', 'magnitude_prune') else [None] * len(deltas)
|
|
gens = [dare_generator(deltas[0].device, layer_name, name) for name, _ in named_deltas] if m in ('dare_ties', 'dare_linear') else [None] * len(deltas)
|
|
for start in range(0, deltas[0].shape[0], ROW_CHUNK):
|
|
stop = min(start + ROW_CHUNK, deltas[0].shape[0])
|
|
chunks = []
|
|
for i, d in enumerate(deltas):
|
|
c = d[start:stop].to(torch.float32)
|
|
if thresholds[i] is not None:
|
|
c = c * (c.abs() >= thresholds[i])
|
|
if gens[i] is not None:
|
|
keep = torch.rand(c.shape, generator=gens[i], device=c.device, dtype=torch.float32) < dens
|
|
c = c * keep / dens
|
|
chunks.append(c)
|
|
if m in ('ties', 'dare_ties'):
|
|
total = torch.stack(chunks).sum(dim=0)
|
|
elected = torch.sign(total)
|
|
agree = [c * ((torch.sign(c) == elected) & (c != 0)) for c in chunks]
|
|
count = torch.stack([(a != 0).to(torch.float32) for a in agree]).sum(dim=0).clamp(min=1.0)
|
|
result[start:stop] = torch.stack(agree).sum(dim=0) / count
|
|
else: # dare_linear, magnitude_prune: independent per-delta edits, plain sum
|
|
result[start:stop] = torch.stack(chunks).sum(dim=0)
|
|
return result.to(out_dtype)
|
|
|
|
|
|
def score_pair(d0, d1, rank0, rank1):
|
|
"""Selection scores for a dense delta pair: klora top-K sums (K = rank product) or est energies; plus abs-sums for the global balance.
|
|
|
|
Row-chunked fp32 interiors with fp64 accumulators and one device sync for
|
|
all four reductions. Full-tensor staging (fp32 copy, abs copy, top-k
|
|
workspace) peaks hundreds of MB per large layer, which collides with block
|
|
swapping on offloaded denoisers; chunking bounds the transient to the
|
|
chunk. The global top-K over per-chunk top-K candidates selects the same
|
|
element set as a whole-tensor top-K.
|
|
"""
|
|
k = max(1, int(rank0) * int(rank1)) if mode() == 'klora' else 0
|
|
accs = []
|
|
for d in (d0, d1):
|
|
score = torch.zeros((), device=d.device, dtype=torch.float64)
|
|
abs_sum = torch.zeros((), device=d.device, dtype=torch.float64)
|
|
cands = []
|
|
for start in range(0, d.shape[0], ROW_CHUNK):
|
|
c = d[start:start + ROW_CHUNK].to(torch.float32).abs() # out-of-place abs: to() may alias a caller-owned fp32 tensor
|
|
abs_sum += c.sum(dtype=torch.float64)
|
|
if k:
|
|
flat = c.flatten()
|
|
cands.append(torch.topk(flat, min(k, flat.numel()), sorted=False).values)
|
|
else:
|
|
score += c.square().sum(dtype=torch.float64)
|
|
if k and cands:
|
|
allc = torch.cat(cands) if len(cands) > 1 else cands[0]
|
|
score = torch.topk(allc, min(k, allc.numel()), sorted=False).values.sum(dtype=torch.float64)
|
|
accs.append((score, abs_sum))
|
|
packed = torch.stack([accs[0][0], accs[0][1], accs[1][0], accs[1][1]]).cpu()
|
|
return (float(packed[0]), float(packed[2])), (float(packed[1]), float(packed[3]))
|
|
|
|
|
|
def register_weight_pair(layer_name, module, per_net, wanted_names=None):
|
|
"""Score and register a weight-kind selection pair; True when the layer is scheduled.
|
|
|
|
The scores persist in the factor cache when a pass identity is given, so a
|
|
later apply of the same configuration registers from the record alone.
|
|
"""
|
|
from modules.lora import lora_common as l
|
|
if per_net is None or len(per_net) != 2:
|
|
return False
|
|
ranks, names = [], []
|
|
for net_name, d in per_net:
|
|
if d is None:
|
|
return False
|
|
net = next((n for n in l.loaded_networks if n.name == net_name), None)
|
|
net_module = net.modules.get(layer_name, None) if net is not None else None
|
|
if net_module is None:
|
|
return False
|
|
names.append(net_name)
|
|
ranks.append(int(getattr(net_module, 'dim', 0) or 0) or 64)
|
|
scores, abs_sums = score_pair(per_net[0][1], per_net[1][1], ranks[0], ranks[1])
|
|
if wanted_names is not None:
|
|
from modules.lora import lora_factor_cache
|
|
lora_factor_cache.begin_pass(wanted_names)
|
|
lora_factor_cache.store_scores(layer_name, scores, abs_sums)
|
|
register(layer_name, module, 'weight', scores, nets=tuple(names), abs_sums=abs_sums)
|
|
return True
|
|
|
|
|
|
def register_weight_pair_cached(layer_name, module, wanted_names):
|
|
"""Register a weight-kind pair from its cached score record; True when served.
|
|
|
|
The record was stored under the same configuration signature, which pins
|
|
the loaded pair, multipliers and stack settings, so both networks are known
|
|
to target the layer and the prompt-order roles are unchanged.
|
|
"""
|
|
from modules.lora import lora_common as l
|
|
from modules.lora import lora_factor_cache
|
|
if len(l.loaded_networks) != 2:
|
|
return False
|
|
lora_factor_cache.begin_pass(wanted_names)
|
|
rec = lora_factor_cache.lookup_scores(layer_name)
|
|
if rec is None:
|
|
return False
|
|
scores, abs_sums = rec
|
|
register(layer_name, module, 'weight', scores, nets=tuple(n.name for n in l.loaded_networks), abs_sums=abs_sums)
|
|
return True
|
|
|
|
|
|
def drop(layer_name):
|
|
"""Forget a layer's selection entry (its factors were removed or restored)."""
|
|
if layer_name is not None and state['entries'].pop(layer_name, None) is not None:
|
|
state['finalized'] = False
|
|
|
|
|
|
def score_energy(up, down):
|
|
"""EST layer score: squared Frobenius norm of up@down via the Gram identity, no materialization."""
|
|
u = up.to(torch.float32)
|
|
dn = down.to(torch.float32)
|
|
return float(((u.t() @ u) * (dn @ dn.t())).sum())
|
|
|
|
|
|
def clear():
|
|
state['entries'] = {}
|
|
state['flips'] = {}
|
|
state['gamma'] = 1.0
|
|
state['gamma_e'] = 1.0
|
|
state['total_steps'] = 0
|
|
state['finalized'] = False
|
|
state['reported'] = None
|
|
|
|
|
|
def register(layer_name, module, kind, scores, segments: tuple[tuple[int, int], tuple[int, int], bool] | None = None, nets=None, abs_sums=None):
|
|
"""Record a select-mode layer for schedule finalization.
|
|
|
|
kind 'factor': segments = ((s0, s1), (t0, t1), transposed) column ranges on the svd
|
|
channel; both segments' pristine values are stashed for flips. kind 'weight': nets =
|
|
the two network names; the winner delta is recomputed from the layer backup at
|
|
selection time. abs_sums feeds the global magnitude balance (klora gamma).
|
|
"""
|
|
entry = {'layer': layer_name, 'module': weakref.ref(module), 'kind': kind, 'segments': segments, 'scores': scores, 'nets': nets, 'abs_sums': abs_sums, 'stash': None}
|
|
if kind == 'factor':
|
|
if segments is None:
|
|
raise ValueError("segments is required when kind='factor'")
|
|
(s0, s1), (t0, t1), transposed = segments
|
|
up = module.svd_up.data
|
|
entry['stash'] = (segment_view(up, s0, s1, transposed).clone(), segment_view(up, t0, t1, transposed).clone())
|
|
state['entries'][layer_name] = entry
|
|
state['finalized'] = False
|
|
|
|
|
|
def segment_view(up, start, stop, transposed):
|
|
return up[start:stop] if transposed else up[:, start:stop]
|
|
|
|
|
|
def layer_flip_step(scores, total_steps):
|
|
"""First step index at which the style side wins; total_steps when it never does, 0 when style wins from the start."""
|
|
m = mode()
|
|
sc, ss = scores
|
|
for step in range(total_steps):
|
|
t = step / max(1, total_steps - 1)
|
|
if m == 'klora':
|
|
ramp = state['gamma'] * (ramp_alpha() * t + KLORA_BETA)
|
|
if ss * ramp > sc:
|
|
return step
|
|
else: # estlora: content keeps the layer while sc >= gamma_t * ss
|
|
# est energies are ||dW||^2, so a magnitude gap enters squared; balance the style side by
|
|
# the total-energy ratio (mirrors klora's gamma) so the louder adapter cannot win on scale alone
|
|
ramp = ramp_alpha() * t + (1.0 - manual_discrepancy())
|
|
if sc < ramp * ss * state['gamma_e']:
|
|
return step
|
|
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, silent=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)
|
|
# both balances derive from the live entries every time, so drops and re-registrations stay consistent by construction
|
|
num = sum(e['abs_sums'][0] for e in state['entries'].values() if e['abs_sums'] is not None)
|
|
den = sum(e['abs_sums'][1] for e in state['entries'].values() if e['abs_sums'] is not None)
|
|
state['gamma'] = (num / den) if den > 0 else 1.0
|
|
e_num = sum(e['scores'][0] for e in state['entries'].values()) # est scores ARE the per-layer energies; their totals give the scale-invariant balance
|
|
e_den = sum(e['scores'][1] for e in state['entries'].values())
|
|
state['gamma_e'] = (e_num / e_den) if e_den > 0 else 1.0
|
|
state['flips'] = {}
|
|
stats = {'weight_n': 0, 'factor_n': 0, 'materialize': 0.0, 'select': 0.0, 'w_move': 0.0, 'w_calc': 0.0, 'w_apply': 0.0}
|
|
state['stats'] = stats
|
|
stats['weight_n'] = sum(1 for e in state['entries'].values() if e['kind'] == 'weight')
|
|
stats['factor_n'] = len(state['entries']) - stats['weight_n']
|
|
if stats['weight_n'] > 0:
|
|
t0 = time.time()
|
|
materialize_model()
|
|
stats['materialize'] = time.time() - t0
|
|
style_first = 0
|
|
t0 = time.time()
|
|
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
|
|
style_first += initial
|
|
apply_selection(layer_name, entry, initial)
|
|
if 0 < flip_at < state['total_steps']:
|
|
state['flips'].setdefault(flip_at - 1, []).append(layer_name) # step callbacks fire after the denoise, so the flip runs one step early to be live during the crossover step's forward
|
|
stats['select'] = time.time() - t0
|
|
state['finalized'] = True
|
|
if len(state['entries']) > 0: # only a built schedule can carry a flip count, so this is the line that shows selection is live rather than requested
|
|
gamma = state['gamma_e'] if mode() == 'estlora' else state['gamma']
|
|
report = (mode(), len(state['entries']), style_first, sum(len(v) for v in state['flips'].values()), state['total_steps'], round(gamma, 3))
|
|
if report != state['reported']: # rebuilt every pass, so a batch would otherwise repeat one line per image
|
|
state['reported'] = report
|
|
log.info(f'Network load: type=LoRA stack={report[0]} layers={report[1]} style={report[2]} flips={report[3]} steps={report[4]} gamma={report[5]:.3f}')
|
|
# logged every pass: the reset runs outside the activate walk, so its cost is invisible to the load timers
|
|
log.debug(f'Network select: type=LoRA reset weight={stats["weight_n"]} factor={stats["factor_n"]} time={{materialize: {stats["materialize"]:.2f}, select: {stats["select"]:.2f}, move: {stats["w_move"]:.2f}, calc: {stats["w_calc"]:.2f}, apply: {stats["w_apply"]:.2f}}}')
|
|
|
|
|
|
def reset(total_steps):
|
|
"""Per-pass reset from set_callbacks_p: restore initial selections and reschedule for this pass's step count."""
|
|
if mode() not in SELECT_MODES or not state['entries'] or int(total_steps) <= 0:
|
|
return
|
|
finalize(total_steps)
|
|
|
|
|
|
def on_step(step):
|
|
"""Flip the layers whose crossover is this step; non-flip steps are a dict miss."""
|
|
if not state['finalized']:
|
|
return
|
|
layers = state['flips'].get(int(step), ())
|
|
if not layers:
|
|
return
|
|
t0 = time.time()
|
|
for layer_name in layers:
|
|
entry = state['entries'].get(layer_name)
|
|
if entry is not None:
|
|
apply_selection(layer_name, entry, 1)
|
|
log.debug(f'Network select: type=LoRA flip step={int(step)} layers={len(layers)} time={time.time() - t0:.2f}')
|
|
|
|
|
|
def apply_selection(layer_name, entry, winner):
|
|
module = entry['module']()
|
|
if module is None:
|
|
state['entries'].pop(layer_name, None)
|
|
return
|
|
if entry['kind'] == 'factor':
|
|
(s0, s1), (t0, t1), transposed = entry['segments']
|
|
up = module.svd_up.data
|
|
keep_seg, drop_seg = ((t0, t1), (s0, s1)) if winner == 1 else ((s0, s1), (t0, t1))
|
|
stash = entry['stash'][winner]
|
|
segment_view(up, keep_seg[0], keep_seg[1], transposed).copy_(stash.to(device=up.device, dtype=up.dtype))
|
|
segment_view(up, drop_seg[0], drop_seg[1], transposed).zero_()
|
|
else:
|
|
weight_selection(module, 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')
|
|
return
|
|
net = next((n for n in l.loaded_networks if n.name == entry['nets'][winner]), None)
|
|
net_module = net.modules.get(entry['layer'], None) if net is not None else None
|
|
if net_module is None:
|
|
return
|
|
weight = getattr(module, 'weight', None)
|
|
if weight is None or weight.is_meta:
|
|
warn_once('select-offloaded', 'Network stack: flip=skipped weight=offloaded')
|
|
return
|
|
from modules import devices
|
|
stats = state.get('stats') or {}
|
|
device = weight.device
|
|
t0 = time.time()
|
|
base = backup.to(devices.device) # a swapped-out layer keeps its weight on cpu; the delta matmul belongs on the accelerator regardless
|
|
t1 = time.time()
|
|
updown = net_module.calc_updown(base)[0].to(device)
|
|
t2 = time.time()
|
|
network_apply_weights(module, updown, None, device=device) # recomputes from the pristine backup, requantizing where the layer needs it
|
|
stats['w_move'] = stats.get('w_move', 0.0) + (t1 - t0)
|
|
stats['w_calc'] = stats.get('w_calc', 0.0) + (t2 - t1)
|
|
stats['w_apply'] = stats.get('w_apply', 0.0) + (time.time() - t2)
|