File size: 7,809 Bytes
9b91042 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | """attach / DiffusionAttachHandle / detach — the diffusion runtime verbs.
attach(model, anchor) wraps every bound site (RelayPatch2D or
MultibandDelta by anchor kind) at the DECLARED trunk dtype (dtype law).
detach(verify=True) restores the retained modules and asserts BIT-EXACT
equality with a pre-attach probe where the binding provides one
(UNet bindings do; the DiT binding defers to the diffusion-pipe gates —
documented F4 — and degrades to structural restore with a warning).
"""
from __future__ import annotations
import contextlib
import warnings
import torch
import torch.nn as nn
from ...binding.diffusion import resolve_diffusion
from ...io.checkpoint import DiffusionAnchorCheckpoint, load_diffusion_anchor
from ..core.multiband import BandBlockWrap, MultibandDelta, band_weights
from ..core.relay import BlockWithRelay, RelayPatch2D
from ..laws import N_BANDS
class DiffusionAttachHandle:
def __init__(self, model, binding, originals, names, wraps, fingerprint):
self.model = model
self.binding = binding
self._originals = originals # [(qualified_name, module)]
self.names = list(names)
self._wraps = wraps # BlockWithRelay | BandBlockWrap
self._fingerprint = fingerprint
@property
def kind(self) -> str:
return ("multiband3"
if isinstance(self._wraps[0].mod, MultibandDelta) else "relay")
def modules(self):
return [w.mod for w in self._wraps]
# -- masking / toggling (never renormalizes anything) ---------------
def set_mask(self, mask: dict[str, bool]) -> None:
for n, w in zip(self.names, self._wraps):
w.mod.enabled = mask.get(n, True)
@contextlib.contextmanager
def only(self, *names):
prev = [w.mod.enabled for w in self._wraps]
self.set_mask({n: n in names for n in self.names})
try:
yield self
finally:
for w, e in zip(self._wraps, prev):
w.mod.enabled = e
@contextlib.contextmanager
def all_off(self):
prev = [w.mod.enabled for w in self._wraps]
self.set_mask({n: False for n in self.names})
try:
yield self
finally:
for w, e in zip(self._wraps, prev):
w.mod.enabled = e
@contextlib.contextmanager
def lesion_band(self, *bands: int):
"""Disable band expert(s) at every site (the exp008/exp010 lesion
instrument). Multiband anchors only."""
if self.kind != "multiband3":
raise TypeError("lesion_band is a multiband instrument; this "
f"anchor is kind={self.kind}")
bad = [b for b in bands if not 0 <= b < N_BANDS]
if bad:
raise ValueError(f"bands out of range: {bad}")
prev = [list(w.mod.band_enabled) for w in self._wraps]
for w in self._wraps:
for b in bands:
w.mod.band_enabled[b] = False
try:
yield self
finally:
for w, p in zip(self._wraps, prev):
w.mod.band_enabled = p
# -- band windows (the step gate) ------------------------------------
def set_band_windows(self, s01: torch.Tensor) -> None:
"""s01: (B,) noise levels in [0,1] -> cosine crossfade windows on
every wrap. No-op for relay anchors (windows unused)."""
w = band_weights(s01)
for wr in self._wraps:
wr.w_bands = w
# -- telemetry --------------------------------------------------------
def telemetry(self, on: bool) -> None:
for w in self._wraps:
if hasattr(w.mod, "telemetry"):
w.mod.telemetry = on
def amplitude(self) -> dict[str, float]:
out = {}
for n, w in zip(self.names, self._wraps):
r = getattr(w.mod, "last_delta_ratio", None)
if r is not None:
out[n] = round(float(r), 6)
return out
def gates(self) -> dict:
vals = []
for w in self._wraps:
m = w.mod
if hasattr(m, "gate"):
vals.append(torch.sigmoid(m.gate).item())
elif hasattr(m, "gates"):
vals.extend(torch.sigmoid(m.gates).tolist())
return {"gate_mean": round(sum(vals) / max(len(vals), 1), 5),
"gate_min": round(min(vals), 5) if vals else None,
"gate_max": round(max(vals), 5) if vals else None}
# -- detach -----------------------------------------------------------
def detach(self, *, verify: bool = True) -> nn.Module:
for name, mod in self._originals:
self.binding.replace(self.model, name, mod)
if verify:
if self._fingerprint is None:
warnings.warn(
"detach: no pre-attach fingerprint for this binding "
"(probe deferred — F4); structural restore only. The "
"diffusion-pipe path carries its own parity gates.")
else:
post = self.binding.probe(self.model)
if not torch.equal(post, self._fingerprint):
raise RuntimeError(
"detach verification FAILED: post-detach output is "
"not bit-exact with the pre-attach fingerprint")
return self.model
def attach(model, anchor, *, binding=None, dtype: "torch.dtype | None" = None,
strict: bool = True) -> DiffusionAttachHandle:
b = resolve_diffusion(model, binding)
sites = b.sites(model)
ck = (anchor if isinstance(anchor, DiffusionAnchorCheckpoint)
else load_diffusion_anchor(anchor))
kind = ck.kind
if kind not in ("relay", "multiband3"):
raise ValueError(
f"anchor kind '{kind}' is not attachable — 'mono'/'bank' are "
"research-record controls (monolith falsifier / falsified "
"dispatch bank), loadable for analysis via "
"load_diffusion_anchor only. See "
"amoe.diffusion.laws.ROUTING_NEGATIVES.")
if ck.n_sites != len(sites):
raise ValueError(
f"anchor carries {ck.n_sites} sites, binding '{b.name}' "
f"enumerates {len(sites)} — wrong substrate?")
if strict:
fam = ck.meta.get("substrate", {}).get("family")
if fam and fam not in b.name:
raise ValueError(
f"anchor substrate family '{fam}' vs binding '{b.name}' "
"(pass strict=False to override)")
dt = dtype if dtype is not None else b.declared_dtype(model) # dtype law
fingerprint = b.probe(model)
originals, wraps, names = [], [], []
for i, (name, block, d) in enumerate(sites):
sd = ck.per_site(i)
if kind == "relay":
mod = RelayPatch2D.from_state_dict(sd)
assert mod.d == d, (f"site {i} ({name}): anchor width {mod.d} "
f"!= trunk width {d}")
wrap = BlockWithRelay(block, mod)
else:
mod = MultibandDelta.from_state_dict(sd)
assert mod.d == d, (f"site {i} ({name}): anchor width {mod.d} "
f"!= trunk width {d}")
wrap = BandBlockWrap(block, mod)
p0 = next(block.parameters(), None)
dev = p0.device if p0 is not None else torch.device("cpu")
mod.to(device=dev, dtype=dt)
b.replace(model, name, wrap)
originals.append((name, block))
wraps.append(wrap)
names.append(name)
return DiffusionAttachHandle(model, b, originals, names, wraps,
fingerprint)
def detach(handle: DiffusionAttachHandle, *,
verify: bool = True) -> nn.Module:
return handle.detach(verify=verify)
|