0.2.0: amoe.diffusion subsystem (relay/multiband/StepGatedSampler, dtype law, align grounded-negative, conditioning law), safetensors I/O + amoe-convert, diffusion invariants; lineage corrected to the audited 19-package record

README.md

@@ -4,34 +4,39 @@ tags:
 - adapters
 - mixture-of-experts
 - aleph
+- diffusion
 ---
 
 # amoe-lora — aleph mixture-of-experts adapters
 
 Train, **attach**, align, **detach** small aleph-addressed adapters on
-a frozen transformer — with the honesty diagnostics that found this
+a frozen trunk — with the honesty diagnostics that found this
 architecture's real failure modes built in as first-class API.
 
-Productized from the geolip-aleph research line
-([research record](https://huggingface.co/AbstractPhil/geolip-aleph-qwen-3.5-0.8b-instruct) ·
+Productized from the geolip-aleph research lines
+([LM research record](https://huggingface.co/AbstractPhil/geolip-aleph-qwen-3.5-0.8b-instruct) ·
+[diffusion research record](https://huggingface.co/AbstractPhil/geolip-aleph-diffusion) ·
 [working caption adapter](https://huggingface.co/AbstractPhil/qwen3.5-0.8b-relay-caption)).
 "LoRA-style" refers to the attach/detach usage pattern, not the math:
 these are 16-slot patch heads reading the residual stream through a
 closed-form sinh/cosh address over a 64-atom S³ codebook (~261k
 params/block at d=1024).
 
-**Status: 0.1 — working-state framework.** The runtime verbs (attach /
+**Status: 0.2 — working-state framework.** The runtime verbs (attach /
 toggle / detach) and checkpoint I/O are implemented and invariant-
-tested (CPU CI, no downloads). The train/align recipes are reference-
-grade transplants of the certified campaign recipes — runnable,
-single-GPU-verified, not yet optimized or deep-tested. See
-*Maturity* below.
+tested (CPU CI, no downloads), now on BOTH substrate families: language
+trunks (`amoe`) and diffusion denoisers (`amoe.diffusion` — SD1.5-class
+UNets, SDXL, Cosmos-Predict2/Anima DiT). 0.2 also ends the safetensors
+non-goal: anchors save/load `.safetensors` with a ComfyUI-ready key
+layout. The train recipes are reference-grade transplants of the
+certified campaign recipes — runnable, single-GPU-verified, not yet
+optimized or deep-tested. See *Maturity* below.
 
 ## Install
 
 ```
 pip install torch          # >= 2.1
-pip install -e .           # from this repo; extras: pip install -e .[hf]
+pip install -e .           # extras: .[hf] (LM), .[diffusion] (diffusers+safetensors)
 ```
 
 ## The five verbs
@@ -86,6 +91,72 @@ adapter trained on more draws than the space holds distinct questions
 memorizes the space and posts fake generalization — measured in the
 research line, twice).
 
+## Diffusion (new in 0.2)
+
+The same verbs on a frozen denoiser — the productized r2 campaign
+([geolip-aleph-diffusion](https://huggingface.co/AbstractPhil/geolip-aleph-diffusion):
+relay certified relay-favorable 3-for-3 substrates, the multiband
+mechanism with surgical band lesions, the step-gated controller, and
+the **conditioning law**).
+
+```python
+# pip install amoe-lora[diffusion]
+import torch, amoe.diffusion as ad
+from diffusers import StableDiffusionPipeline
+from huggingface_hub import hf_hub_download
+
+pipe = StableDiffusionPipeline.from_pretrained(
+    "stable-diffusion-v1-5/stable-diffusion-v1-5",
+    torch_dtype=torch.float32).to("cuda")
+
+h = ad.attach(pipe.unet, hf_hub_download(
+    "AbstractPhil/aleph-diffusion-adapters", "sd15/mb3_s0.safetensors"))
+img = ad.sample(pipe, h, "a lighthouse at dusk", seed=7)   # step-gated
+
+with h.lesion_band(2):                     # generate without the HIGH band
+    img_no_high = ad.sample(pipe, h, "a lighthouse at dusk", seed=7)
+base = h.detach()                          # bit-exact or raises
+```
+
+Training on a folder of images is five lines:
+
+```python
+cache = ad.data.build_cache("photos/", pipe, captions="photos/captions.txt")
+cfg = ad.TrainConfig(name="mystyle", objective="flow",
+                     adapter="multiband3", blob=False)
+ck = ad.train(pipe.unet, cache, cfg)
+ck.save("mystyle.anchor.safetensors")
+```
+
+Three certified deviations from the LM defaults, each encoded as law
+with its receipt (`amoe.diffusion.laws`):
+
+- **Dtype law**: adapters attach at the DECLARED trunk dtype (bf16
+  relays on a bf16 DiT), never a sniffed one — first-param sniffing can
+  read a still-meta fp32 tensor. Judged gauges stay fp32.
+- **`align()` is a grounded negative**: comparative routing on
+  diffusion was falsified at 2 seeds (state+sigma keys route nothing;
+  address-as-key dies against the repeated-key null in both raw and
+  M-hat forms). `ad.align()` raises with the record. The certified
+  alternative is STRUCTURAL banding — cosine crossfade windows on the
+  sigma axis (`adapter="multiband3"`), whose band lesions are surgical
+  (own-band damage 50–200× cross-band, 2 seeds).
+- **Conditioning law**: blob/structural supervision pays where x0
+  recovery is linear (flow: `x0 = x_t − σ·v`; ~125–200× the eps
+  effect) and is inert on eps — `blob=True` with `objective="eps"`
+  refuses unless forced.
+
+### safetensors + ComfyUI
+
+`ck.save("x.safetensors")` writes the canonical flat layout
+(`blocks.{site}.{param}`) with the full meta as safetensors metadata;
+`amoe-convert path.pt --substrate sd15` converts every legacy campaign
+shape with bitwise verification. A `comfyui-amoe` custom-node package
+(loader / attach / band-toggle / detach, with step gating installed as
+a pre-forward hook so ANY stock KSampler becomes step-gated) is the
+next-turn deliverable; the node consumes the canonical layout and
+asserts the anchor's site-count + width signature before patching.
+
 ## House laws (encoded, not optional)
 
 Pure Adam only (`laws.make_optimizer`); dense signed dispatch with an
@@ -101,7 +172,14 @@ the research record.
 - **Training**: the trainer is DDP-aware (rank-sharded sampling,
   `find_unused_parameters=False`; launch with `torchrun`). Correct by
   construction; multi-GPU smoke is **deferred** (single-GPU-verified).
-- **FSDP / tensor parallel / DeepSpeed**: not implemented; see
+- **Diffusion — the declared split (the standard paradigm)**: native
+  `amoe.diffusion.train` = single-GPU + DDP at framework level;
+  the PRODUCTION multi-GPU path is the
+  [diffusion-pipe fork](https://github.com/AbstractEyes/diffusion-pipe)
+  (DeepSpeed pipeline engine, aleph relays baked in, proven end-to-end
+  on the Anima 2B DiT — r2 exp004). Use the fork when you have real
+  scale; use the native trainer for the certified single-card recipes.
+- **FSDP / tensor parallel**: not implemented; see
   `docs/distributed.md` for why and the sketch.
 
 ## Maturity
@@ -113,7 +191,13 @@ the research record.
 | checkpoint I/O + legacy import | implemented; round-trip verified against the shipped campaign checkpoints |
 | train / align | reference-grade (the certified campaign recipes, transplanted); single-GPU |
 | DDP | implemented, untested on >1 GPU |
-| FSDP/TP, safetensors, token-level gating, streaming data | documented non-goals for 0.1 |
+| diffusion core (relay/multiband/windows) | ported verbatim from the certified r2 beds; invariant-tested |
+| diffusion attach / lesions / detach | implemented, invariant-tested incl. dtype law (CPU CI) |
+| diffusion train (eps/flow/blob/roles) | reference-grade transplant; single-GPU-verified recipes |
+| StepGatedSampler | eps/DDIM = the exp010-proven path; flow/Euler = same warp as the beds |
+| safetensors I/O + amoe-convert | implemented in 0.2 (bitwise-verified round trips); ComfyUI node = next turn |
+| SDXL binding | enumeration implemented; site count PINNED AT FIRST REAL RUN (never guessed) |
+| FSDP/TP, token-level gating, streaming data | documented non-goals |
 
 ## Loading the campaign's shipped artifacts
 
@@ -128,6 +212,9 @@ ck = load_anchor(hf_hub_download(
 
 Architecture and every default in this package trace to the
 [geolip-aleph-qwen-3.5-0.8b-instruct](https://huggingface.co/AbstractPhil/geolip-aleph-qwen-3.5-0.8b-instruct)
-campaign (21 experiment packages, raw ledgers, adversarial audits,
-two published self-retractions) and its
-[Qwen2.5 predecessor](https://huggingface.co/AbstractPhil/geolip-aleph-qwen).
+campaign (19 experiment packages, raw ledgers, adversarial audits, one
+published self-retraction plus a corrected overclaim), its
+[Qwen2.5 predecessor](https://huggingface.co/AbstractPhil/geolip-aleph-qwen),
+and — for the 0.2 diffusion subsystem — the
+[geolip-aleph-diffusion](https://huggingface.co/AbstractPhil/geolip-aleph-diffusion)
+campaign (16 experiment packages, 2-seed program, the conditioning law).

docs/distributed.md

@@ -1,4 +1,25 @@
-# Distributed story (0.1)
+# Distributed story (0.2)
+
+## The diffusion split (the standard paradigm, 0.2)
+
+Diffusion training has TWO sanctioned paths, by design:
+
+1. **Native `amoe.diffusion.train`** — single-GPU + the same DDP
+   posture as the LM trainer below (rank-sharded cache sampling,
+   `find_unused_parameters=False`, rank-0 saves). This is the
+   framework-level path for the certified single-card recipes;
+   correct-by-construction, multi-GPU smoke deferred.
+2. **The diffusion-pipe fork**
+   (https://github.com/AbstractEyes/diffusion-pipe) — the PRODUCTION
+   multi-GPU trainer: DeepSpeed pipeline engine with the aleph relays
+   baked into the model integrations (attach at declared dtype,
+   freeze-by-lr-0 trunk, plain-Adam branch, adapter-only saves in the
+   amoe anchor format). Proven end-to-end on the Anima 2B DiT
+   (r2 exp004). Multiband w_bands plumbing is working at
+   `pipeline_stages=1`; stages>1 is documented-deferred.
+
+No FSDP on either path, same reasons as below.
+
 
 ## Implemented
 

examples/attach_diffusion.py

@@ -0,0 +1,32 @@
+# attach_diffusion.py — the diffusion quickstart, runnable as-is on a
+# CUDA machine with `pip install amoe-lora[diffusion]` (Colab: paste the
+# cell; no argparse, no __file__).
+import torch
+
+import amoe.diffusion as ad
+
+from diffusers import StableDiffusionPipeline
+from huggingface_hub import hf_hub_download
+
+pipe = StableDiffusionPipeline.from_pretrained(
+    "stable-diffusion-v1-5/stable-diffusion-v1-5",
+    torch_dtype=torch.float32).to("cuda")
+
+# The certified multiband stack (exp008 s0; 2-seed surgical band lesions).
+anchor = hf_hub_download("AbstractPhil/aleph-diffusion-adapters",
+                         "sd15/mb3_s0.safetensors")
+h = ad.attach(pipe.unet, anchor)
+print(h.gates())                       # gate health at load
+
+img = ad.sample(pipe, h, "a lighthouse at dusk, film photo", seed=7)
+img.save("lighthouse.png")
+
+# Band lesion: generate WITHOUT the HIGH-noise/structure band — the
+# exp010 image-space instrument (lesioning HIGH cut grounding most and
+# was 14x LP-dominated: it carries coarse structure).
+with h.lesion_band(2):
+    ad.sample(pipe, h, "a lighthouse at dusk, film photo",
+              seed=7).save("lighthouse_no_high.png")
+
+base = h.detach()                      # verified bit-exact, or it raises
+print("detached clean")

pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "amoe-lora"
-version = "0.1.0"
-description = "Aleph mixture-of-experts adapters: train, attach, align, detach — with the honesty diagnostics built in"
+version = "0.2.0"
+description = "Aleph mixture-of-experts adapters: train, attach, align, detach — with the honesty diagnostics built in. 0.2 adds the diffusion subsystem (amoe.diffusion) and safetensors I/O."
 readme = "README.md"
 requires-python = ">=3.10"
 license = { text = "Apache-2.0" }
@@ -14,6 +14,10 @@ dependencies = ["torch>=2.1"]
 
 [project.optional-dependencies]
 hf = ["transformers>=4.57", "huggingface_hub"]
+diffusion = ["diffusers>=0.30", "safetensors>=0.4"]
+
+[project.scripts]
+amoe-convert = "amoe.io.convert:main"
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/amoe/__init__.py

@@ -16,4 +16,8 @@ from .train.trainer import train
 from .train.aligner import align
 from . import laws
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
+
+# The diffusion subsystem (amoe.diffusion) is imported lazily — its verbs
+# live under `import amoe.diffusion as ad`. Core is pure torch; diffusers
+# is only touched by bindings/samplers/data that need it.

src/amoe/binding/diffusion.py

@@ -0,0 +1,197 @@
+"""Diffusion model binding: where do the adapters attach on a denoiser?
+
+Same substrate-shim discipline as resolver.py — everything
+architecture-specific lives here and only here. Diffusion sites are
+SCATTERED (heterogeneous widths across down/mid/up), so the binding
+speaks (qualified_name, module, width) triples and replaces modules by
+attribute walk, not by swapping one ModuleList.
+
+DTYPE LAW: `declared_dtype` is the trunk dtype AFTER materialization.
+It refuses meta-device models — first-param sniffing during init can
+read a still-meta fp32 tensor (the R0b catch, fork f4f1c46). Attach
+after weights are real, or pass dtype= explicitly.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Callable, Protocol
+
+import torch
+import torch.nn as nn
+
+from .resolver import REGISTRY
+
+
+class DiffusionBinding(Protocol):
+    name: str
+    expected_sites: int | None
+
+    def sites(self, model) -> list[tuple[str, nn.Module, int]]: ...
+    def replace(self, model, name: str, new: nn.Module) -> None: ...
+    def declared_dtype(self, model) -> torch.dtype: ...
+    def probe(self, model) -> "torch.Tensor | None": ...
+
+
+def walk_replace(root: nn.Module, name: str, new: nn.Module) -> None:
+    """setattr/index walk to a qualified module name (the certified
+    attach pattern from the r2 beds)."""
+    parent = root
+    parts = name.split(".")
+    for p in parts[:-1]:
+        parent = getattr(parent, p) if not p.isdigit() else parent[int(p)]
+    last = parts[-1]
+    if last.isdigit():
+        parent[int(last)] = new
+    else:
+        setattr(parent, last, new)
+
+
+def _declared_dtype(model) -> torch.dtype:
+    metas = [n for n, p in model.named_parameters() if p.is_meta]
+    if metas:
+        raise RuntimeError(
+            "dtype law: model has meta-device parameters "
+            f"(e.g. {metas[0]}) — a dtype read here can lie (the R0b "
+            "catch). Materialize the weights first, or pass dtype= "
+            "explicitly at attach.")
+    dt = getattr(model, "dtype", None)
+    if isinstance(dt, torch.dtype):
+        return dt
+    return next(model.parameters()).dtype
+
+
+@dataclass
+class UNetBinding:
+    """diffusers UNet2DConditionModel — every BasicTransformerBlock.
+    SD15 certifies at 16 sites (asserted). SDXL's count is PINNED AT THE
+    FIRST REAL RUN (deferred-by-design F1): expected_sites=None means
+    'record, do not assert a guess'."""
+    name: str = "sd15_unet"
+    expected_sites: "int | None" = 16
+
+    def sites(self, model):
+        from diffusers.models.attention import BasicTransformerBlock
+        out = []
+        for name, mod in model.named_modules():
+            if isinstance(mod, BasicTransformerBlock):
+                out.append((name, mod, mod.norm1.normalized_shape[0]))
+        if self.expected_sites is not None:
+            assert len(out) == self.expected_sites, (
+                f"{self.name}: enumerated {len(out)} BasicTransformerBlocks, "
+                f"expected {self.expected_sites}")
+        assert out, f"{self.name}: no BasicTransformerBlocks found"
+        return out
+
+    def replace(self, model, name, new):
+        walk_replace(model, name, new)
+
+    def declared_dtype(self, model):
+        return _declared_dtype(model)
+
+    def probe(self, model):
+        """Deterministic denoiser fingerprint on a tiny fixed input,
+        built from the LIVE config (never from memory — F4)."""
+        cfg = model.config
+        n_down = len(cfg.block_out_channels)
+        size = 2 ** (n_down - 1) * 4               # divisible through the UNet
+        g = torch.Generator().manual_seed(1400)
+        dev = next(model.parameters()).device
+        dt = self.declared_dtype(model)
+        x = torch.randn(1, cfg.in_channels, size, size, generator=g).to(dev, dt)
+        ehs = torch.randn(1, 8, cfg.cross_attention_dim, generator=g).to(dev, dt)
+        kwargs = {}
+        if getattr(cfg, "addition_embed_type", None) == "text_time":
+            # SDXL-style added conditioning, shapes from the live config
+            ta = torch.randn(1, cfg.projection_class_embeddings_input_dim
+                             - 6 * cfg.addition_time_embed_dim,
+                             generator=g).to(dev, dt)
+            tids = torch.tensor([[size * 8, size * 8, 0, 0, size * 8, size * 8]],
+                                device=dev)
+            kwargs["added_cond_kwargs"] = {"text_embeds": ta, "time_ids": tids}
+        with torch.no_grad():
+            out = model(x, 17, ehs, return_dict=False, **kwargs)[0]
+        return out.detach().float().cpu()
+
+
+@dataclass
+class DiTBinding:
+    """Cosmos-Predict2/Anima-style DiT: a `blocks` ModuleList of uniform
+    width. The production trainer for this family is the diffusion-pipe
+    fork (which carries its own P-INIT/P-TOGGLE gates); the in-package
+    probe is None — detach verification degrades to structural restore
+    with a warning (documented F4 deferral)."""
+    name: str = "cosmos_dit"
+    expected_sites: "int | None" = None
+
+    def sites(self, model):
+        blocks = getattr(model, "blocks", None)
+        assert isinstance(blocks, nn.ModuleList) and len(blocks) > 0, (
+            "cosmos_dit binding expects a `blocks` ModuleList")
+        d = None
+        for attr in ("model_channels", "hidden_size", "dim"):
+            v = getattr(getattr(model, "config", model), attr,
+                        getattr(model, attr, None))
+            if isinstance(v, int):
+                d = v
+                break
+        if d is None:
+            lin = next(m for m in blocks[0].modules()
+                       if isinstance(m, nn.Linear))
+            d = lin.in_features
+        return [(f"blocks.{i}", b, d) for i, b in enumerate(blocks)]
+
+    def replace(self, model, name, new):
+        walk_replace(model, name, new)
+
+    def declared_dtype(self, model):
+        return _declared_dtype(model)
+
+    def probe(self, model):
+        return None
+
+
+DIFF_REGISTRY: dict[str, Callable[[], "DiffusionBinding"]] = {}
+
+
+def register_diffusion(key: str):
+    def deco(fn):
+        DIFF_REGISTRY[key] = fn
+        REGISTRY[key] = fn          # discoverable from the one registry
+        return fn
+    return deco
+
+
+@register_diffusion("sd15_unet")
+def _sd15():
+    return UNetBinding(name="sd15_unet", expected_sites=16)
+
+
+@register_diffusion("sdxl_unet")
+def _sdxl():
+    return UNetBinding(name="sdxl_unet", expected_sites=None)   # F1
+
+
+@register_diffusion("cosmos_dit")
+def _cosmos():
+    return DiTBinding()
+
+
+def resolve_diffusion(model, binding=None) -> DiffusionBinding:
+    if binding is not None:
+        if isinstance(binding, str):
+            if binding in DIFF_REGISTRY:
+                return DIFF_REGISTRY[binding]()
+            raise ValueError(f"unknown diffusion binding '{binding}'; "
+                             f"known: {sorted(DIFF_REGISTRY)}")
+        return binding
+    cls = type(model).__name__
+    if cls == "UNet2DConditionModel":
+        cad = int(getattr(model.config, "cross_attention_dim", 0) or 0)
+        return DIFF_REGISTRY["sdxl_unet" if cad >= 2048 else "sd15_unet"]()
+    if isinstance(getattr(model, "blocks", None), nn.ModuleList):
+        return DIFF_REGISTRY["cosmos_dit"]()
+    cands = [f"{n} (len {len(m)})" for n, m in model.named_modules()
+             if isinstance(m, nn.ModuleList) and len(m) >= 4]
+    raise ValueError(
+        "amoe.diffusion could not resolve an attach site — pass "
+        f"binding=<DiffusionBinding|key>. ModuleList candidates: {cands}")

src/amoe/diagnostics/diffusion.py

@@ -0,0 +1,64 @@
+"""Diffusion diagnostics — the campaign's honesty instruments as API.
+
+lesion_report   : the exp008/exp010 band-lesion battery shape — run YOUR
+                  gauge under each single-band lesion and read whether
+                  specialization is surgical (own-band damage >> cross).
+foreground_gauge: the exp012 role-aligned payer instrument — HIGH-band
+                  foreground-LP-x0 error (the gauge that showed multiband
+                  beating the matched monolith ~10% when every
+                  aggregate-eps comparison was blind to it).
+gate_stats      : gate health (sigmoid means; the substrate-dependent
+                  dynamics finding — core grew, lune shrank).
+diagnose        : one-call summary with the loud flags.
+"""
+from __future__ import annotations
+
+import torch
+
+from ..diffusion.laws import N_BANDS
+from ..diffusion.train.objectives import blob_lp_err
+
+
+def lesion_report(handle, eval_fn) -> dict:
+    """eval_fn() -> float, evaluated under all_on and each single-band
+    lesion. Multiband anchors only. Returns {'all_on': x,
+    'lesion_band0': ..., 'surgical': bool} where surgical means every
+    lesion moved the gauge (the exp008 signature reads per-band gauges;
+    with a single gauge this reports the monotone lesion profile the
+    exp010 battery certified)."""
+    out = {"all_on": float(eval_fn())}
+    for b in range(N_BANDS):
+        with handle.lesion_band(b):
+            out[f"lesion_band{b}"] = float(eval_fn())
+    deltas = [out[f"lesion_band{b}"] - out["all_on"] for b in range(N_BANDS)]
+    out["lesion_deltas"] = [round(d, 6) for d in deltas]
+    return out
+
+
+def foreground_gauge(x0_hat: torch.Tensor, x0: torch.Tensor,
+                     blob: torch.Tensor) -> float:
+    """Mean foreground-LP-x0 error (lower is better). Judge in fp32 (law)."""
+    return float(blob_lp_err(x0_hat.float(), x0.float(),
+                             blob.float()).mean())
+
+
+@torch.no_grad()
+def gate_stats(handle) -> dict:
+    return handle.gates()
+
+
+@torch.no_grad()
+def diagnose(handle) -> dict:
+    """One-call health readout: gates, amplitude telemetry (if armed),
+    and the standing scope note on gate bands."""
+    rep = {"kind": handle.kind, "n_sites": len(handle.names),
+           **gate_stats(handle)}
+    amp = handle.amplitude()
+    if amp:
+        vals = list(amp.values())
+        rep["delta_ratio_mean"] = round(sum(vals) / len(vals), 6)
+    rep["note"] = ("gate dynamics are SUBSTRATE-DEPENDENT on diffusion "
+                   "(exp006: core grew 0.047->0.070; exp001: lune shrank) "
+                   "— read direction against your own frozen baseline, "
+                   "not against the LM band")
+    return rep

src/amoe/diffusion/__init__.py

@@ -0,0 +1,29 @@
+"""amoe.diffusion — the aleph adapter verbs on diffusion trunks (0.2).
+
+Same five-verb grammar as the text line, three certified deviations
+(see .laws): the dtype law, structural banding instead of comparative
+routing (align() is a grounded negative), and the declared multi-GPU
+split (native train = single-GPU/DDP; production multi-GPU = the
+diffusion-pipe fork).
+
+Quickstart:
+    import amoe.diffusion as ad
+    h = ad.attach(pipe.unet, "mb3_s0.safetensors")
+    img = ad.sample(pipe, h, "a lighthouse at dusk", seed=7)
+    with h.lesion_band(2): ...
+    base = h.detach()          # bit-exact or raises
+"""
+from ..io.checkpoint import (DiffusionAnchorCheckpoint,
+                             load_diffusion_anchor)
+from .core.cond import AlephCondAdapter
+from .core.multiband import (BandBlockWrap, MultibandDelta, band_of,
+                             band_weights)
+from .core.relay import BlockWithRelay, RelayPatch2D
+from .runtime.attach import DiffusionAttachHandle, attach, detach
+from .runtime.sampler import StepGatedSampler, sample
+from .train.aligner import align
+from .train.config import DiffusionTrainConfig, RelaySpec
+from .train.config import DiffusionTrainConfig as TrainConfig
+from .train.trainer import ConditioningLawWarning, train
+from . import laws
+from .train import data

src/amoe/diffusion/core/__init__.py

@@ -0,0 +1,3 @@
+from .cond import AlephCondAdapter
+from .multiband import BandBlockWrap, MultibandDelta, band_of, band_weights
+from .relay import BlockWithRelay, RelayPatch2D

src/amoe/diffusion/core/cond.py

@@ -0,0 +1,52 @@
+"""AlephCondAdapter — conditioning-stream address injection (research-grade
+export; verbatim port of pod2/aleph_diffusion_core.py).
+
+STATUS (read before using): exp002 certified this channel as a Law-2 honest
+negative IN CONTEXT — appended to full text conditioning the address is
+marginally inert (real-vs-deranged −0.0009). But addr-only steering is real
+(+0.0287 after 2k steps, gates GREW) — the channel is redundant-in-context,
+not dead. The open redesign target is COMPLEMENTARITY: inject where text is
+not (reduced-text regimes, sigma-gated, non-text streams). Shipped as a
+component for that line of work, not as a quickstart verb.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class AlephCondAdapter(nn.Module):
+    """Frozen [n_addr, addr_dim] address -> n_addr conditioning positions.
+    Everything zero-init => tokens are EXACT zeros at init and when
+    disabled (the position-presence offset is handled by the bed: SD15
+    masked-append/length-toggle, Anima write-into-padding)."""
+
+    def __init__(self, cond_dim: int, addr_dim: int = 128, n_addr: int = 32):
+        super().__init__()
+        self.cond_dim, self.addr_dim, self.n_addr = cond_dim, addr_dim, n_addr
+        self.addr_proj = nn.Linear(addr_dim, cond_dim)
+        nn.init.zeros_(self.addr_proj.weight)
+        nn.init.zeros_(self.addr_proj.bias)
+        self.pos_table = nn.Parameter(torch.zeros(n_addr, cond_dim))
+        self.gate = nn.Parameter(torch.tensor(-3.0))
+        self.enabled = True
+
+    def assert_zero_init(self):
+        assert self.addr_proj.weight.abs().max().item() == 0.0
+        assert self.addr_proj.bias.abs().max().item() == 0.0
+        assert self.pos_table.abs().max().item() == 0.0
+
+    def tokens(self, addr: torch.Tensor) -> torch.Tensor:
+        """addr [B, n_addr, addr_dim] -> [B, n_addr, cond_dim]."""
+        assert addr.shape[-2:] == (self.n_addr, self.addr_dim), addr.shape
+        if not self.enabled:
+            return addr.new_zeros(*addr.shape[:-1], self.cond_dim)
+        return torch.sigmoid(self.gate) * (self.addr_proj(addr) + self.pos_table)
+
+    @classmethod
+    def from_state_dict(cls, sd: dict) -> "AlephCondAdapter":
+        cond_dim, addr_dim = sd["addr_proj.weight"].shape
+        n_addr = sd["pos_table"].shape[0]
+        m = cls(cond_dim, addr_dim=addr_dim, n_addr=n_addr)
+        m.load_state_dict(sd, strict=True)
+        return m

src/amoe/diffusion/core/multiband.py

@@ -0,0 +1,109 @@
+"""MultibandDelta — the certified band-assignment mechanism (verbatim port
+of pod2/dexp008_multiband.py; exp008: band lesions surgical 3/3 at both
+seeds, own-band damage 50–200× cross-band; exp012's role-aligned gauge:
+multiband beats the matched monolith ~10% on HIGH-band foreground
+structure, 2-seed).
+
+Per site: N_BANDS rank-r zero-init deltas (W+b, own gates −3.0), combined
+by smooth COSINE CROSSFADE WINDOWS on s01 — STRUCTURAL/positional gating
+on the sigma axis, dense and differentiable, no selection event (law 2 of
+amoe.diffusion.laws). Per-band gradients arise from the windows: a sample
+at s01=0.9 trains (almost) only the HIGH expert.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+
+from ..laws import BAND_EDGES, N_BANDS, XFADE
+
+
+def band_weights(s01: torch.Tensor) -> torch.Tensor:
+    """Smooth structural band windows: (B, N_BANDS), sum to 1 everywhere.
+    Cosine crossfade of half-width XFADE around each edge — positional
+    gating on the sigma axis, never comparative."""
+    def ramp(x):                     # 0 below -XFADE, 1 above +XFADE, smooth
+        t = ((x / XFADE).clamp(-1, 1) + 1) / 2
+        return 0.5 - 0.5 * torch.cos(t * math.pi)
+    e1, e2 = BAND_EDGES
+    up1, up2 = ramp(s01 - e1), ramp(s01 - e2)
+    low = 1 - up1
+    mid = up1 * (1 - up2)
+    high = up1 * up2
+    return torch.stack([low, mid, high], dim=-1)
+
+
+def band_of(s01: float) -> int:
+    e1, e2 = BAND_EDGES
+    return 0 if s01 <= e1 else (1 if s01 <= e2 else 2)
+
+
+class MultibandDelta(nn.Module):
+    """Three band experts per site, window-combined; per-expert enable
+    flags (band lesions) + global enable (toggle law, code-path skip)."""
+
+    def __init__(self, d: int, r: int = 16):
+        super().__init__()
+        self.d, self.r = d, r
+        self.down = nn.ModuleList(nn.Linear(d, r, bias=False)
+                                  for _ in range(N_BANDS))
+        self.up = nn.ModuleList(nn.Linear(r, d) for _ in range(N_BANDS))
+        for dn, up in zip(self.down, self.up):
+            nn.init.orthogonal_(dn.weight)
+            nn.init.zeros_(up.weight)
+            nn.init.zeros_(up.bias)
+        self.gates = nn.Parameter(torch.full((N_BANDS,), -3.0))
+        self.enabled = True
+        self.band_enabled = [True] * N_BANDS
+
+    def assert_zero_init(self):
+        for up in self.up:
+            assert up.weight.abs().max().item() == 0.0
+            assert up.bias.abs().max().item() == 0.0
+
+    def forward(self, x, w_bands):                  # w_bands: (B, N_BANDS)
+        if not self.enabled:
+            return x
+        if w_bands is None:
+            raise RuntimeError(
+                "MultibandDelta needs band windows — set them per step "
+                "(handle.set_band_windows / StepGatedSampler) or per "
+                "training batch from the sampled sigmas")
+        g = torch.sigmoid(self.gates)
+        delta = 0
+        w = w_bands.view(w_bands.shape[0],
+                         *([1] * (x.ndim - 2)), N_BANDS)
+        for b in range(N_BANDS):
+            if not self.band_enabled[b]:
+                continue
+            delta = delta + g[b] * w[..., b:b + 1] * self.up[b](
+                self.down[b](x))
+        return x + delta if not isinstance(delta, int) else x
+
+    @classmethod
+    def from_state_dict(cls, sd: dict) -> "MultibandDelta":
+        d = sd["down.0.weight"].shape[1]
+        r = sd["down.0.weight"].shape[0]
+        m = cls(d, r=r)
+        m.load_state_dict(sd, strict=True)
+        return m
+
+
+class BandBlockWrap(nn.Module):
+    """Wrap one trunk block with a window-consuming module. `w_bands` is
+    set externally (per training batch, or per sampling step by the
+    controller). Tuple outputs pass through."""
+
+    def __init__(self, block: nn.Module, mod: nn.Module):
+        super().__init__()
+        self.block = block
+        self.mod = mod
+        self.w_bands = None
+
+    def forward(self, *args, **kwargs):
+        out = self.block(*args, **kwargs)
+        h = out[0] if isinstance(out, tuple) else out
+        h = self.mod(h, self.w_bands)
+        return (h,) + out[1:] if isinstance(out, tuple) else h

src/amoe/diffusion/core/relay.py

@@ -0,0 +1,90 @@
+"""RelayPatch2D — the certified per-block residual relay for diffusion
+token streams (verbatim port of the r2 bed, pod2/aleph_diffusion_core.py;
+certified relay-favorable 3-for-3 substrates: Lune flow / SD15-core eps /
+Anima DiT flow — exp001/exp006/exp004).
+
+Differences from the text-side RelayPatchwork (amoe.core.adapter), each a
+paid-for law:
+  * trailing-dim generalization: forward works on any [..., d] stream;
+  * zero-init output weight AND bias (exp006 bias-leak law) — a fresh
+    relay is EXACTLY inert, P-INIT bit-exact;
+  * `enabled=False` returns x via CODE-PATH SKIP (`return x`, `is x`) —
+    the toggle law is bit-exact at every dtype;
+  * dims-from-state-dict reconstruction (never trust constructor defaults).
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from ...core.address import AlephAddress
+from ...core.adapter import SquaredReLU
+
+
+class RelayPatch2D(nn.Module):
+    """forward(x[..., d]) = x + sigmoid(gate) * consume(m_hat(proj(x) slots))."""
+
+    def __init__(self, d: int, n_slots: int = 16, K: int = 64,
+                 tau: float = 0.1, hidden: int = 178):
+        super().__init__()
+        self.d, self.n_slots, self.hidden = d, n_slots, hidden
+        self.proj = nn.Linear(d, n_slots * 4, bias=False)
+        nn.init.orthogonal_(self.proj.weight)
+        self.addr = AlephAddress(K, 4, tau)
+        self.consume = nn.Sequential(
+            nn.Linear(n_slots * 4, hidden), SquaredReLU(),
+            nn.LayerNorm(hidden), nn.Linear(hidden, d))
+        nn.init.zeros_(self.consume[-1].weight)   # zero-init out: weight
+        nn.init.zeros_(self.consume[-1].bias)     # AND bias (exp006 law)
+        self.gate = nn.Parameter(torch.tensor(-3.0))
+        self.enabled = True                       # plain attr, not a buffer
+        self.telemetry = False                    # read-only ratio stash
+        self.last_delta_ratio = None
+
+    def assert_zero_init(self):
+        w, b = self.consume[-1].weight, self.consume[-1].bias
+        assert w.abs().max().item() == 0.0, "out weight not zero-init"
+        assert b.abs().max().item() == 0.0, "out BIAS not zero-init (leak law)"
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.enabled:
+            return x
+        lead = x.shape[:-1]
+        slots = self.proj(x).view(*lead, self.n_slots, 4)
+        feats = self.addr.m_hat(slots).reshape(*lead, self.n_slots * 4)
+        delta = torch.sigmoid(self.gate) * self.consume(feats)
+        if self.telemetry:
+            with torch.no_grad():
+                self.last_delta_ratio = (
+                    delta.norm() / x.norm().clamp_min(1e-12)).item()
+        return x + delta
+
+    @classmethod
+    def from_state_dict(cls, sd: dict) -> "RelayPatch2D":
+        """Rebuild dims FROM the state dict, then strict-load."""
+        d = sd["proj.weight"].shape[1]
+        n_slots = sd["proj.weight"].shape[0] // 4
+        hidden = sd["consume.0.weight"].shape[0]
+        K = sd["addr.codebook"].shape[0]
+        m = cls(d, n_slots=n_slots, K=K, hidden=hidden)
+        m.load_state_dict(sd, strict=True)
+        return m
+
+
+class BlockWithRelay(nn.Module):
+    """Wrap one trunk block; the relay reads its hidden-states output.
+    Tuple outputs pass through (hybrid-safe, as amoe.core BlockWithAdapter).
+    Carries a `w_bands` slot for wrap-uniformity with multiband sites —
+    the relay ignores it."""
+
+    def __init__(self, block: nn.Module, relay: nn.Module):
+        super().__init__()
+        self.block = block
+        self.mod = relay
+        self.w_bands = None            # unused by relays; uniform wrap API
+
+    def forward(self, *args, **kwargs):
+        out = self.block(*args, **kwargs)
+        if isinstance(out, tuple):
+            return (self.mod(out[0]),) + out[1:]
+        return self.mod(out)

src/amoe/diffusion/laws.py

@@ -0,0 +1,53 @@
+"""amoe.diffusion.laws — the diffusion-line laws, encoded with their receipts.
+
+Every rule was paid for on the r2 rented card; the receipts live in the
+research record (https://huggingface.co/AbstractPhil/geolip-aleph-diffusion).
+
+1. DTYPE LAW (supersedes amoe's fp32 default ON DIFFUSION TRUNKS) — adapters
+   attach at the DECLARED trunk dtype, never a sniffed one: fp32 adapters on
+   a bf16 trunk spin fp32 noise into the environment, and first-param
+   sniffing can read a still-meta fp32 tensor (the R0b integration-gate
+   catch, diffusion-pipe fork f4f1c46). Judged gauges stay fp32.
+2. STRUCTURAL BANDING, NOT COMPARATIVE ROUTING — the certified partition of
+   the sigma axis is cosine crossfade windows (positional, dense,
+   differentiable; no selection event). Comparative routing on diffusion was
+   falsified 2-seed: state+sigma dispatch keys found nothing to route on
+   (exp007), and address-as-key died against the repeated-key null in both
+   raw-flattened and M-hat-bottlenecked forms (exp014/exp015).
+3. CONDITIONING LAW — structural/blob supervision pays where x0 recovery is
+   LINEAR (flow: x0 = x_t − σ·v; ~125–200× the eps effect, 2-seed) and is
+   inert on eps, which divides by a vanishing sqrt(alpha_bar) exactly in
+   the supervised band (exp012 vs exp013). blob-on-eps therefore refuses
+   unless forced. Operating point λ≈1 (3-point dose curve, exp013).
+4. ZERO-INIT W **AND** B — the exp006 bias-leak law: output weight and bias
+   both zero so a fresh adapter is EXACTLY inert (P-INIT bit-exact).
+5. TOGGLE LAW — disabled adapters return x via CODE-PATH SKIP (`return x`,
+   not a zero-mul), bit-exact at every dtype. Held post-train in every
+   experiment of the line.
+6. PURE ADAM, NEVER ADAMW — inherited program-wide (amoe.laws law 1).
+"""
+from __future__ import annotations
+
+from ..laws import make_optimizer, pin_precision  # noqa: F401  (the constructors)
+
+# The certified band constants (exp008; edges/xfade are LAW, not config):
+N_BANDS = 3
+BAND_EDGES = (0.35, 0.75)   # LOW [0,.35] | MID (.35,.75] | HIGH (.75,1]
+XFADE = 0.06                # cosine crossfade half-width
+BAND_ROLES = ("fidelity/detail (LOW noise)",
+              "continuity/semantics (MID)",
+              "diversity/structure (HIGH noise)")
+
+BLOB_LAMBDA = 1.0           # the dose-curve operating point (exp013, 3 points)
+FLOW_SHIFT = 2.5            # the trained sigma warp of the flow substrate
+
+# The negatives record (why align()/dispatch is not a diffusion verb in 0.2):
+ROUTING_NEGATIVES = (
+    "exp007: A=4 bank with dense signed state+sigma dispatch ties/loses to a "
+    "matched monolith under uniform pressure; sigma-band usage flat.",
+    "exp014: text/address dispatch keys within 0.07% of the state-key control; "
+    "raw-flattened [32,128] address KILLS routing (disease geometry).",
+    "exp015: M-hat slot-bottlenecked address key — routing excess 2.5e-06 over "
+    "the REPEATED-key null, match advantage ~0, both seeds. The frozen "
+    "address-as-key form is CLOSED; the open form is TE+AMoE joint training.",
+)

src/amoe/diffusion/runtime/__init__.py

@@ -0,0 +1,2 @@
+from .attach import DiffusionAttachHandle, attach, detach
+from .sampler import StepGatedSampler, sample

src/amoe/diffusion/runtime/attach.py

@@ -0,0 +1,195 @@
+"""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)

src/amoe/diffusion/runtime/sampler.py

@@ -0,0 +1,120 @@
+"""StepGatedSampler — the controller (exp010): per-step band gating at
+inference. Each denoising step sets the crossfade windows from the
+CURRENT noise level, activating band experts coarse-to-fine across the
+trajectory exactly as trained. Image space sees what eps-MSE could not:
+on the certified mb3 stack the controller moved grounding +0.089 over
+frozen, with monotone band-lesion signatures (exp010, image-space).
+
+Two objectives:
+  "eps"  — DDIM on the stock SD15-style schedule (the exp010-proven path);
+  "flow" — Euler on the SHIFT-warped sigma grid (the trained warp of the
+           flow substrate, exp013).
+
+Works for relay anchors too — the windows are simply unused.
+Uncond convention follows the certified beds: ZEROED conditioning (the
+CFG-dropout convention the adapters were trained with), not an
+empty-prompt embedding.
+"""
+from __future__ import annotations
+
+import torch
+
+from ..laws import FLOW_SHIFT
+from .attach import DiffusionAttachHandle
+
+
+def _default_ddim():
+    from diffusers import DDIMScheduler
+    # The stock SD15 schedule, constructed explicitly (no download):
+    return DDIMScheduler(num_train_timesteps=1000, beta_start=0.00085,
+                         beta_end=0.012, beta_schedule="scaled_linear",
+                         clip_sample=False, set_alpha_to_one=False,
+                         steps_offset=1, prediction_type="epsilon")
+
+
+class StepGatedSampler:
+    def __init__(self, model, handle: DiffusionAttachHandle, *,
+                 scheduler=None, objective: str = "eps",
+                 shift: float = FLOW_SHIFT, device=None,
+                 latent_channels: int = 4):
+        assert objective in ("eps", "flow"), objective
+        self.model, self.handle = model, handle
+        self.objective, self.shift = objective, shift
+        self.latent_channels = latent_channels
+        self.device = device or next(model.parameters()).device
+        self.sched = None
+        if objective == "eps":
+            self.sched = scheduler or _default_ddim()
+
+    @torch.no_grad()
+    def sample_latents(self, ehs_cond: torch.Tensor,
+                       ehs_uncond: "torch.Tensor | None" = None, *,
+                       seed: int, steps: int = 30, guidance: float = 7.5,
+                       latent_size: tuple = (64, 64)) -> torch.Tensor:
+        B = ehs_cond.shape[0]
+        dev = self.device
+        dt = ehs_cond.dtype
+        g = torch.Generator(device=dev).manual_seed(seed)
+        x = torch.randn(B, self.latent_channels, *latent_size,
+                        generator=g, device=dev, dtype=dt)
+        unc = (torch.zeros_like(ehs_cond) if ehs_uncond is None
+               else ehs_uncond)
+        ehs_in = torch.cat([ehs_cond, unc], dim=0)
+
+        if self.objective == "eps":
+            self.sched.set_timesteps(steps, device=dev)
+            x = x * self.sched.init_noise_sigma
+            for t in self.sched.timesteps:
+                s01 = torch.full((2 * B,), float(t) / 1000.0, device=dev)
+                self.handle.set_band_windows(s01)         # the step gate
+                xin = self.sched.scale_model_input(torch.cat([x, x]), t)
+                pred = self.model(xin, t, ehs_in, return_dict=False)[0]
+                p_c, p_u = pred.chunk(2)
+                pred = p_u + guidance * (p_c - p_u)
+                x = self.sched.step(pred, t, x).prev_sample
+            return x
+
+        # flow: Euler from s=1 -> 0 on the SHIFT-warped grid (exp013 warp)
+        u = torch.linspace(1.0, 0.0, steps + 1, device=dev)
+        s_grid = (self.shift * u) / (1 + (self.shift - 1) * u)
+        for i in range(steps):
+            s, s_next = s_grid[i], s_grid[i + 1]
+            s01 = torch.full((2 * B,), float(s), device=dev)
+            self.handle.set_band_windows(s01)             # the step gate
+            v = self.model(torch.cat([x, x]), s01[:1].expand(2 * B) * 1000,
+                           ehs_in, return_dict=False)[0]
+            v_c, v_u = v.chunk(2)
+            v = v_u + guidance * (v_c - v_u)
+            x = x + (s_next - s) * v                      # dx/ds = v
+        return x
+
+
+@torch.no_grad()
+def sample(pipe, handle: DiffusionAttachHandle, prompt: str, *,
+           negative_prompt: "str | None" = None, seed: int = 0,
+           steps: int = 30, guidance: float = 7.5, objective: str = "eps"):
+    """Ease-of-use path for a diffusers StableDiffusionPipeline whose UNet
+    carries an attached anchor: encode -> step-gated sample -> decode."""
+    dev = pipe.device if pipe.device.type != "meta" else "cuda"
+    tok = pipe.tokenizer(prompt, padding="max_length", truncation=True,
+                         max_length=pipe.tokenizer.model_max_length,
+                         return_tensors="pt")
+    ehs = pipe.text_encoder(tok.input_ids.to(dev))[0]
+    unc = None
+    if negative_prompt is not None:
+        ntok = pipe.tokenizer(negative_prompt, padding="max_length",
+                              truncation=True,
+                              max_length=pipe.tokenizer.model_max_length,
+                              return_tensors="pt")
+        unc = pipe.text_encoder(ntok.input_ids.to(dev))[0]
+    sz = pipe.unet.config.sample_size
+    sampler = StepGatedSampler(pipe.unet, handle, scheduler=None,
+                               objective=objective, device=dev)
+    lat = sampler.sample_latents(ehs, unc, seed=seed, steps=steps,
+                                 guidance=guidance, latent_size=(sz, sz))
+    img = pipe.vae.decode(lat / pipe.vae.config.scaling_factor).sample
+    img = (img / 2 + 0.5).clamp(0, 1)
+    from PIL import Image
+    import numpy as np
+    arr = (img[0].permute(1, 2, 0).float().cpu().numpy() * 255).astype("uint8")
+    return Image.fromarray(np.asarray(arr))

src/amoe/diffusion/train/__init__.py

@@ -0,0 +1,4 @@
+from .aligner import align
+from .config import DiffusionTrainConfig, RelaySpec
+from .trainer import ConditioningLawWarning, train
+from . import data, objectives

src/amoe/diffusion/train/aligner.py

@@ -0,0 +1,26 @@
+"""align() — a grounded negative on diffusion, kept for API symmetry.
+
+The text line's align() trains dispatch keys over frozen anchors. On
+diffusion that entire family was falsified at 2 seeds — see
+amoe.diffusion.laws.ROUTING_NEGATIVES (exp007: state+sigma keys find
+nothing to route on; exp014/exp015: address-as-key dead against the
+repeated-key null in both raw and M-hat-bottlenecked forms).
+
+The certified replacement for comparative routing on the sigma axis is
+STRUCTURAL banding (adapter='multiband3' in train(); exp008 lesions
+surgical 3/3 both seeds). The open research form — TE+AMoE joint
+training, where the text encoder adapter learns WITH the dispatch — is
+queued science (history/open_questions.md), not shippable API.
+"""
+from __future__ import annotations
+
+from ..laws import ROUTING_NEGATIVES
+
+
+def align(*args, **kwargs):
+    raise NotImplementedError(
+        "amoe.diffusion.align: comparative dispatch on diffusion is a "
+        "2-seed falsified line — there is nothing honest to align here. "
+        "Use train(adapter='multiband3') for the certified structural "
+        "banding, or follow the open TE+AMoE joint-training question. "
+        "The record:\n  - " + "\n  - ".join(ROUTING_NEGATIVES))

src/amoe/diffusion/train/config.py

@@ -0,0 +1,44 @@
+"""DiffusionTrainConfig — the certified recipes as the only knobs.
+
+Deliberately ABSENT (laws, not options): optimizer choice (pure Adam,
+wd=0), top-k/load-balancing (no comparative selectors), band edges/xfade
+(law constants, exp008).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from ..laws import BLOB_LAMBDA, FLOW_SHIFT
+
+
+@dataclass
+class RelaySpec:
+    n_slots: int = 16
+    K: int = 64
+    tau: float = 0.1
+    hidden: int = 178
+
+
+@dataclass
+class DiffusionTrainConfig:
+    name: str = "anchor"
+    objective: str = "eps"            # "eps" | "flow"
+    shift: float = FLOW_SHIFT         # flow sigma warp (exp013)
+    adapter: str = "relay"            # "relay" | "multiband3"
+    relay: RelaySpec = field(default_factory=RelaySpec)
+    rank: int = 16                    # multiband3 per-band rank
+    band_roles: bool = False          # exp009 HP/LP role pressure (measured
+                                      # inert at gauge scale — kept honest)
+    blob: bool = False                # foreground-LP-x0 supervision
+    blob_lambda: float = BLOB_LAMBDA  # λ≈1: the 3-point dose-curve operating
+                                      # point (exp013); 0.5/2.0 are the
+                                      # mapped under/over-couple arms
+    force_blob_on_eps: bool = False   # conditioning law: eps refuses blob
+                                      # unless forced (exp012, 2-seed inert)
+    steps: int = 3000
+    batch_size: int = 16
+    lr: float = 1e-3
+    cfg_dropout: float = 0.1
+    seed: int = 0
+    log_every: int = 50
+    base_schedule_id: str = "stable-diffusion-v1-5/stable-diffusion-v1-5"

src/amoe/diffusion/train/data.py

@@ -0,0 +1,97 @@
+"""build_cache — images + captions -> the trainer's cache dict.
+
+The trainer is cache-fed (the pod-bed discipline): latents and text
+features are encoded ONCE, then training never touches the VAE/TE again.
+The campaign's caches load directly too — any dict with lat/ehs
+(+ optional blob, val_lat/val_ehs) works.
+
+Blob targets: the campaign rasterized entity polygons from the fused
+dataset (research record, exp011). The generic path here accepts a
+directory of binary masks (PNG, same stem as each image), resized to the
+latent grid. Empty masks are counted and reported LOUDLY (the exp011
+schema lesson: an empty blob target is a silent zero).
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+
+@torch.no_grad()
+def build_cache(image_dir, pipe, *, captions=None, n_val: int = 256,
+                size: int = 512, blob_masks=None, device: str = "cuda",
+                batch: int = 8) -> dict:
+    """image_dir: folder of images. pipe: a diffusers pipeline carrying
+    vae/text_encoder/tokenizer. captions: None | txt file (one caption per
+    image, sorted order) | dict stem->caption. blob_masks: None | folder
+    of binary PNG masks (same stems)."""
+    from PIL import Image
+    import numpy as np
+
+    paths = sorted(p for p in Path(image_dir).iterdir()
+                   if p.suffix.lower() in (".png", ".jpg", ".jpeg", ".webp"))
+    assert paths, f"no images under {image_dir}"
+    if isinstance(captions, (str, Path)):
+        lines = Path(captions).read_text(encoding="utf-8").splitlines()
+        assert len(lines) >= len(paths), \
+            f"{len(lines)} captions for {len(paths)} images"
+        cap_of = {p.stem: lines[i] for i, p in enumerate(paths)}
+    elif isinstance(captions, dict):
+        cap_of = captions
+    else:
+        cap_of = {p.stem: "" for p in paths}
+
+    vae, te, tok = pipe.vae, pipe.text_encoder, pipe.tokenizer
+    sf = vae.config.scaling_factor
+    lat_sz = size // (2 ** (len(vae.config.block_out_channels) - 1))
+
+    lats, ehss, blobs, empty = [], [], [], 0
+    for i in range(0, len(paths), batch):
+        chunk = paths[i:i + batch]
+        imgs = []
+        for p in chunk:
+            im = Image.open(p).convert("RGB").resize((size, size))
+            imgs.append(torch.from_numpy(
+                np.asarray(im)).float().permute(2, 0, 1) / 127.5 - 1)
+        x = torch.stack(imgs).to(device)
+        lats.append((vae.encode(x).latent_dist.sample(
+            generator=torch.Generator(device=device).manual_seed(1400 + i))
+            * sf).cpu())
+        t = tok([cap_of.get(p.stem, "") for p in chunk],
+                padding="max_length", truncation=True,
+                max_length=tok.model_max_length, return_tensors="pt")
+        ehss.append(te(t.input_ids.to(device))[0].cpu())
+        if blob_masks is not None:
+            for p in chunk:
+                mp = Path(blob_masks) / f"{p.stem}.png"
+                if mp.exists():
+                    m = Image.open(mp).convert("L").resize(
+                        (lat_sz, lat_sz), Image.NEAREST)
+                    mm = (torch.from_numpy(np.asarray(m)) > 127).float()
+                else:
+                    mm = torch.zeros(lat_sz, lat_sz)
+                if mm.sum() == 0:
+                    empty += 1
+                blobs.append(mm)
+
+    lat = torch.cat(lats)
+    ehs = torch.cat(ehss)
+    cache = {"lat": lat[:-n_val] if n_val else lat,
+             "ehs": ehs[:-n_val] if n_val else ehs}
+    if n_val:
+        cache["val_lat"] = lat[-n_val:]
+        cache["val_ehs"] = ehs[-n_val:]
+    if blob_masks is not None:
+        bl = torch.stack(blobs)
+        cache["blob"] = bl[:-n_val] if n_val else bl
+        if n_val:
+            cache["val_blob"] = bl[-n_val:]
+        cache["blob_empty_count"] = empty
+        if empty:
+            print(f"WARNING: {empty}/{len(paths)} blob masks are EMPTY — "
+                  "empty targets train nothing and read as silent zeros "
+                  "(the exp011 schema lesson)", flush=True)
+    print(f"cache: {cache['lat'].shape[0]} train rows"
+          + (f" + {n_val} val" if n_val else ""), flush=True)
+    return cache

src/amoe/diffusion/train/objectives.py

@@ -0,0 +1,80 @@
+"""Objectives — the certified loss pieces, ported verbatim.
+
+eps  : stock epsilon-MSE on the shipped schedule (exp006 recipe).
+flow : rectified-flow v-MSE with the SHIFT warp; x0 recovery is EXACT AND
+       LINEAR at all sigma (x0 = x_t − σ·v) — the mechanism behind the
+       conditioning law (exp013).
+blob : foreground-LP-x0 coupling, weighted by the HIGH-band window
+       (exp012/exp013; pays ~125–200× more on flow than eps).
+roles: HP/LP frequency-reweighted band pressure (exp009 — measured
+       directional-but-negligible; shipped for completeness, honest).
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+
+# ── frequency filters (dexp009) ─────────────────────────────────────────
+
+def hp(x):
+    """High-pass: x − avgpool3(x) (finer-detail component)."""
+    return x - F.avg_pool2d(x, 3, stride=1, padding=1)
+
+
+def lp(x):
+    """Low-pass: avgpool7 (coarse structure)."""
+    return F.avg_pool2d(x, 7, stride=1, padding=3)
+
+
+# ── eps path (dexp006) ──────────────────────────────────────────────────
+
+def make_schedule(base_schedule_id: str, device):
+    """Stock training schedule (alphas_cumprod) from the shipped
+    scheduler config."""
+    from diffusers import DDPMScheduler
+    sch = DDPMScheduler.from_pretrained(base_schedule_id,
+                                        subfolder="scheduler")
+    assert sch.config.prediction_type == "epsilon", \
+        sch.config.prediction_type
+    return sch.alphas_cumprod.to(device)
+
+
+def add_noise(lat, noise, t, acp):
+    a = acp[t].sqrt()[:, None, None, None]
+    s = (1 - acp[t]).sqrt()[:, None, None, None]
+    return a * lat + s * noise
+
+
+# ── flow path (dexp013) ─────────────────────────────────────────────────
+
+def warp_sigma(u: torch.Tensor, shift: float) -> torch.Tensor:
+    return (shift * u) / (1 + (shift - 1) * u)
+
+
+def flow_pieces(lat, s, noise):
+    """x_t and the v target; x0 = x_t − σ·v holds exactly."""
+    s4 = s[:, None, None, None]
+    return noise * s4 + lat * (1 - s4), noise - lat   # x_t, v
+
+
+# ── blob coupling (dexp012/013) ─────────────────────────────────────────
+
+def blob_lp_err(x0_hat, x0, blob):
+    """Foreground-masked LP-x0 error, per-sample (B,). blob: (B, H, W)
+    binary mask on the latent grid."""
+    d2 = (lp(x0_hat) - lp(x0)) ** 2
+    m = blob[:, None]
+    denom = m.sum(dim=(1, 2, 3)).clamp_min(1.0) * d2.shape[1]
+    return (d2 * m).sum(dim=(1, 2, 3)) / denom
+
+
+# ── role pressure (dexp009) ─────────────────────────────────────────────
+
+def role_losses(pred, target, lam: float = 0.5):
+    """Per-sample (B,) losses for each band role: LOW +HP, MID std,
+    HIGH +LP."""
+    base = ((pred - target) ** 2).mean(dim=(1, 2, 3))
+    low = base + lam * ((hp(pred) - hp(target)) ** 2).mean(dim=(1, 2, 3))
+    high = base + lam * ((lp(pred) - lp(target)) ** 2).mean(dim=(1, 2, 3))
+    return low, base, high

src/amoe/diffusion/train/trainer.py

@@ -0,0 +1,196 @@
+"""train() — the certified diffusion recipes as a verb.
+
+Reference-grade transplant of the campaign recipes: exp006 (eps on the
+stock schedule) and exp013 (flow with the SHIFT warp, optional blob
+coupling at the λ≈1 operating point), on a frozen trunk with fresh
+zero-init adapters, pure Adam wd=0.
+
+DDP posture mirrors amoe.train: rank-sharded batch sampling, rank-0
+logging, correct-by-construction — multi-GPU smoke deferred (F6). The
+production multi-GPU path for this family is the diffusion-pipe fork
+(DeepSpeed; proven end-to-end on Anima, exp004). No FSDP.
+
+CONDITIONING LAW (encoded): blob supervision on objective='eps' REFUSES —
+it was measured inert-to-negative at 2 seeds exactly where the eps
+parameterization divides by a vanishing sqrt(alpha_bar) (exp012). Pass
+force_blob_on_eps=True to run the control anyway (a ConditioningLawWarning
+fires).
+"""
+from __future__ import annotations
+
+import warnings
+
+import torch
+import torch.nn as nn
+
+from ...binding.diffusion import resolve_diffusion
+from ...io.checkpoint import DiffusionAnchorCheckpoint
+from ..core.multiband import BandBlockWrap, MultibandDelta, band_weights
+from ..core.relay import BlockWithRelay, RelayPatch2D
+from ..laws import make_optimizer, pin_precision
+from .config import DiffusionTrainConfig
+from .objectives import (add_noise, blob_lp_err, flow_pieces, make_schedule,
+                         role_losses, warp_sigma)
+
+
+class ConditioningLawWarning(UserWarning):
+    pass
+
+
+def _dist():
+    try:
+        import torch.distributed as dist
+        if dist.is_available() and dist.is_initialized():
+            return dist.get_rank(), dist.get_world_size()
+    except Exception:                                       # noqa: BLE001
+        pass
+    return 0, 1
+
+
+def train(model, cache: dict, config: "DiffusionTrainConfig | None" = None,
+          *, binding=None, device: str = "cuda"
+          ) -> DiffusionAnchorCheckpoint:
+    cfg = config or DiffusionTrainConfig()
+    assert cfg.objective in ("eps", "flow"), cfg.objective
+    assert cfg.adapter in ("relay", "multiband3"), cfg.adapter
+    if cfg.blob and cfg.objective == "eps":
+        if not cfg.force_blob_on_eps:
+            raise ValueError(
+                "conditioning law: blob supervision on the eps objective "
+                "was measured inert at 2 seeds (exp012) — the eps x0 "
+                "recovery divides by a vanishing sqrt(alpha_bar) exactly "
+                "in the supervised band. Use objective='flow' (or v-pred "
+                "when available), or set force_blob_on_eps=True to run "
+                "the control.")
+        warnings.warn("blob-on-eps forced: this reproduces the exp012 "
+                      "CONTROL arm, not a paying configuration",
+                      ConditioningLawWarning)
+    if cfg.blob and "blob" not in cache:
+        raise ValueError("cfg.blob=True but the cache has no 'blob' masks "
+                         "(see amoe.diffusion.data.build_cache)")
+
+    rank, world = _dist()
+    pin_precision()                       # judged gauges stay fp32 (law)
+    b = resolve_diffusion(model, binding)
+    dt = b.declared_dtype(model)          # dtype law
+    model.requires_grad_(False)
+    model.eval()
+    if hasattr(model, "enable_gradient_checkpointing"):
+        model.enable_gradient_checkpointing()
+
+    sites = b.sites(model)
+    originals, wraps, mods = [], [], nn.ModuleList()
+    for name, block, d in sites:
+        if cfg.adapter == "relay":
+            s = cfg.relay
+            m = RelayPatch2D(d, n_slots=s.n_slots, K=s.K, tau=s.tau,
+                             hidden=s.hidden)
+            m.assert_zero_init()
+            wrap = BlockWithRelay(block, m)
+        else:
+            m = MultibandDelta(d, r=cfg.rank)
+            m.assert_zero_init()
+            wrap = BandBlockWrap(block, m)
+        p0 = next(block.parameters(), None)
+        m.to(device=p0.device if p0 is not None else device, dtype=dt)
+        b.replace(model, name, wrap)
+        originals.append((name, block))
+        wraps.append(wrap)
+        mods.append(m)
+
+    ddp_model = model
+    if world > 1:
+        ddp_model = torch.nn.parallel.DistributedDataParallel(
+            model, find_unused_parameters=False)
+
+    opt = make_optimizer([p for p in mods.parameters()], lr=cfg.lr)
+    acp = (make_schedule(cfg.base_schedule_id, device)
+           if cfg.objective == "eps" else None)
+    lat_all, ehs_all = cache["lat"], cache["ehs"]
+    blob_all = cache.get("blob")
+    n = lat_all.shape[0]
+    g = torch.Generator().manual_seed(cfg.seed * 1000 + rank)
+    gd = torch.Generator(device=device).manual_seed(cfg.seed * 1000 + rank)
+
+    def set_w(s01):
+        w = band_weights(s01)
+        for wr in wraps:
+            wr.w_bands = w
+        return w
+
+    for step in range(cfg.steps):
+        idx = torch.randint(rank, n, (cfg.batch_size,), generator=g)
+        idx = idx - (idx % world) + rank if world > 1 else idx
+        idx = idx.clamp(0, n - 1)
+        lat = lat_all[idx].to(device, torch.float32)
+        ehs = ehs_all[idx].to(device, torch.float32)
+        bsz = lat.shape[0]
+        drop = torch.rand(bsz, generator=gd, device=device) < cfg.cfg_dropout
+        ehs = ehs.clone()
+        ehs[drop] = 0
+        noise = torch.randn(lat.shape, generator=gd, device=device)
+
+        if cfg.objective == "eps":
+            t = torch.randint(0, 1000, (bsz,), generator=gd, device=device)
+            s01 = 1 - acp[t]                       # noise level proxy in [0,1]
+            w = set_w(s01) if cfg.adapter == "multiband3" else None
+            x_t = add_noise(lat, noise, t, acp)
+            pred = ddp_model(x_t.to(dt), t, ehs.to(dt),
+                             return_dict=False)[0].float()
+            target = noise
+            x0_hat = None
+        else:
+            u = torch.rand(bsz, generator=gd, device=device)
+            s = warp_sigma(u, cfg.shift)
+            w = set_w(s) if cfg.adapter == "multiband3" else None
+            x_t, v = flow_pieces(lat, s, noise)
+            pred = ddp_model(x_t.to(dt), s * 1000, ehs.to(dt),
+                             return_dict=False)[0].float()
+            target = v
+            x0_hat = x_t - s[:, None, None, None] * pred   # EXACT linear
+
+        if cfg.band_roles and cfg.adapter == "multiband3":
+            low, base, high = role_losses(pred, target)
+            loss_vec = w[:, 0] * low + w[:, 1] * base + w[:, 2] * high
+        else:
+            loss_vec = ((pred - target) ** 2).mean(dim=(1, 2, 3))
+        if cfg.blob and x0_hat is not None:
+            blob = blob_all[idx].to(device, torch.float32)
+            wb = w[:, 2] if w is not None else torch.ones_like(loss_vec)
+            loss_vec = loss_vec + cfg.blob_lambda * wb * blob_lp_err(
+                x0_hat, lat, blob)
+        loss = loss_vec.mean()
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        opt.step()
+        if rank == 0 and (step % cfg.log_every == 0 or
+                          step == cfg.steps - 1):
+            print(f"[amoe.diffusion.train {cfg.name}] step {step} "
+                  f"loss {loss.item():.5f}", flush=True)
+
+    adapters = {}
+    for i, m in enumerate(mods):
+        for k, v in m.state_dict().items():
+            adapters[f"{i}.{k}"] = v.detach().cpu()
+    meta = {
+        "name": cfg.name,
+        "substrate": {"family": b.name, "n_sites": len(sites),
+                      "site_names": [s[0] for s in sites],
+                      "widths": [s[2] for s in sites]},
+        "adapter": ({"kind": "relay", **vars(cfg.relay)}
+                    if cfg.adapter == "relay"
+                    else {"kind": "multiband3", "rank": cfg.rank}),
+        "objective": ({"kind": "eps"} if cfg.objective == "eps"
+                      else {"kind": "flow", "shift": cfg.shift}),
+        "blob": ({"lambda": cfg.blob_lambda} if cfg.blob else None),
+        "dtype": str(dt).replace("torch.", ""),
+        "seed": cfg.seed,
+        "recipe": {"optimizer": "adam", "lr": cfg.lr, "weight_decay": 0.0,
+                   "steps": cfg.steps, "batch": cfg.batch_size,
+                   "cfg_dropout": cfg.cfg_dropout,
+                   "band_roles": cfg.band_roles},
+    }
+    # restore the unwrapped trunk before returning (amoe.train contract)
+    for name, block in originals:
+        b.replace(model, name, block)
+    return DiffusionAnchorCheckpoint(adapters, meta)

src/amoe/io/checkpoint.py

@@ -96,6 +96,128 @@ def load_dispatch(path: str) -> DispatchCheckpoint:
     raise ValueError(f"unrecognized dispatch checkpoint at {path}")
 
 
+# ── diffusion anchors (amoe 0.2) ─────────────────────────────────────────
+
+DIFF_ANCHOR_FORMAT = "amoe.diffusion.anchor"
+
+
+@dataclass
+class DiffusionAnchorCheckpoint:
+    """One diffusion adapter stack: adapters keyed "{site_index}.{param}".
+    meta.adapter.kind is 'relay' | 'multiband3' | load-only legacy kinds
+    ('mono', 'bank'). Relay kind requires per-site addr.home (drift
+    gauge); multiband has no address, so the requirement is kind-aware."""
+    adapters: dict[str, torch.Tensor]
+    meta: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def kind(self) -> str:
+        return self.meta.get("adapter", {}).get("kind", "unknown")
+
+    @property
+    def n_sites(self) -> int:
+        return len({k.split(".", 1)[0] for k in self.adapters})
+
+    def per_site(self, i: int) -> dict[str, torch.Tensor]:
+        pre = f"{i}."
+        return {k[len(pre):]: v for k, v in self.adapters.items()
+                if k.startswith(pre)}
+
+    def save(self, path: str) -> str:
+        if str(path).endswith(".safetensors"):
+            from .safetensors_io import save_anchor_safetensors
+            return save_anchor_safetensors(self, path)
+        meta = dict(self.meta)
+        meta.setdefault("created", _dt.datetime.now(
+            _dt.timezone.utc).isoformat())
+        meta["content_hash"] = _content_hash(self.adapters)
+        torch.save({"format": DIFF_ANCHOR_FORMAT, "version": VERSION,
+                    "meta": meta, "adapters": self.adapters}, path)
+        return meta["content_hash"]
+
+
+def _sd_kind(sd: dict) -> str:
+    if any(k.endswith("addr.codebook") for k in sd):
+        return "relay"
+    if "down.0.weight" in sd:
+        return "multiband3"
+    if "down.weight" in sd:
+        return "mono"
+    if "key_proj" in sd or "keys" in sd:
+        return "bank"
+    return "unknown"
+
+
+def _flatten_stack(stack, *, kind: str, imported_from: str,
+                   reconstruct_home: bool = False) -> "DiffusionAnchorCheckpoint":
+    adapters, recon = {}, False
+    for i, sd in enumerate(stack):
+        sd = dict(sd)
+        if kind == "relay" and "addr.home" not in sd:
+            if not reconstruct_home:
+                raise ValueError(
+                    f"site {i} lacks addr.home (drift gauge); pass a source "
+                    "that carries it or use the fork import path")
+            import torch.nn.functional as _F
+            sd["addr.home"] = _F.normalize(
+                sd["addr.codebook"].detach().float(), dim=-1)
+            recon = True
+        for k, v in sd.items():
+            adapters[f"{i}.{k}"] = v
+    meta: dict[str, Any] = {"adapter": {"kind": kind},
+                            "imported_from": imported_from}
+    if recon:
+        meta["home_reconstructed"] = True     # drift gauge void — never silent
+    return DiffusionAnchorCheckpoint(adapters, meta)
+
+
+def load_diffusion_anchor(path: str, *, substrate: "dict | None" = None
+                          ) -> DiffusionAnchorCheckpoint:
+    """Load a diffusion anchor: the versioned format, a .safetensors
+    companion, or any of the campaign's legacy stack shapes
+    ({'relays': [...]}, {'mods': [...]}, {'banks'/'monos': [...]},
+    diffusion-pipe {'relays': {idx: sd}})."""
+    if str(path).endswith(".safetensors"):
+        from .safetensors_io import load_anchor_safetensors
+        ck = load_anchor_safetensors(path)
+    else:
+        blob = torch.load(path, map_location="cpu", weights_only=True)
+        if blob.get("format") == DIFF_ANCHOR_FORMAT:
+            ck = DiffusionAnchorCheckpoint(blob["adapters"],
+                                           blob.get("meta", {}))
+        elif "relays" in blob:
+            r = blob["relays"]
+            if isinstance(r, dict):        # diffusion-pipe fork save
+                stack = [r[k] for k in sorted(r, key=int)]
+                ck = _flatten_stack(stack, kind="relay",
+                                    imported_from="diffusion_pipe",
+                                    reconstruct_home=True)
+            else:                          # exp001/exp006 relay stacks
+                ck = _flatten_stack(list(r), kind="relay",
+                                    imported_from="legacy_relays")
+        elif "mods" in blob:               # exp008/009/011/012/013 stacks
+            stack = list(blob["mods"])
+            ck = _flatten_stack(stack, kind=_sd_kind(stack[0]),
+                                imported_from="legacy_mods")
+        elif "banks" in blob or "monos" in blob:   # exp007 (load-only)
+            key = "banks" if "banks" in blob else "monos"
+            ck = _flatten_stack(list(blob[key]), kind=_sd_kind(blob[key][0]),
+                                imported_from=f"legacy_{key}")
+            ck.meta["routing_negative"] = (
+                "exp007/014/015: comparative dispatch on diffusion is a "
+                "2-seed falsified line — load for analysis, not deployment")
+        else:
+            raise ValueError(f"unrecognized diffusion anchor at {path}")
+    if ck.kind == "relay":
+        blocks = {k.split(".", 1)[0] for k in ck.adapters}
+        for b in blocks:
+            if f"{b}.addr.home" not in ck.adapters:
+                raise ValueError(f"relay anchor site {b} lacks addr.home")
+    if substrate:
+        ck.meta.setdefault("substrate", dict(substrate))
+    return ck
+
+
 def import_legacy_keys(path: str, *, seed: int, d: int,
                        emb: int = 64) -> DispatchCheckpoint:
     """Reconstruct a full dispatch state from a keys-only legacy file

src/amoe/io/convert.py

@@ -0,0 +1,86 @@
+"""amoe-convert — .pt anchor stacks -> .safetensors companions, verified.
+
+Converts the versioned format AND every legacy campaign shape
+(load_diffusion_anchor handles the zoo). Each conversion round-trips the
+result and asserts BITWISE tensor equality before reporting success.
+
+Colab-safe usage (no argparse required):
+    from amoe.io.convert import convert_one, convert_dir
+    convert_one("mb3_s0.pt", substrate={"family": "sd15"}, objective="eps")
+    convert_dir("ckpts/", substrate={"family": "sd15"})
+
+CLI: amoe-convert <path.pt|dir> [--substrate sd15] [--objective eps]
+                  [--layout amoe|comfy] [--out OUT]
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import torch
+
+from .checkpoint import load_diffusion_anchor
+from .safetensors_io import load_anchor_safetensors, save_anchor_safetensors
+
+
+def convert_one(path, *, substrate: "dict | None" = None,
+                objective: "str | None" = None, layout: str = "amoe",
+                out: "str | None" = None) -> str:
+    path = Path(path)
+    ck = load_diffusion_anchor(str(path), substrate=substrate)
+    if objective:
+        ck.meta.setdefault("objective", {"kind": objective})
+    dst = Path(out) if out else path.with_suffix(".safetensors")
+    save_anchor_safetensors(ck, str(dst), key_layout=layout)
+    back = load_anchor_safetensors(str(dst))
+    assert set(back.adapters) == set(ck.adapters), "key set changed"
+    for k in ck.adapters:
+        a, b = ck.adapters[k], back.adapters[k]
+        assert a.dtype == b.dtype and torch.equal(a, b), \
+            f"bitwise mismatch at {k}"
+    print(f"converted {path.name} -> {dst.name} "
+          f"({len(ck.adapters)} tensors, kind={ck.kind}, verified bitwise)")
+    return str(dst)
+
+
+def convert_dir(root, **kw) -> list[str]:
+    out = []
+    for p in sorted(Path(root).rglob("*.pt")):
+        try:
+            out.append(convert_one(p, **kw))
+        except Exception as e:          # noqa: BLE001 — report and continue
+            print(f"SKIP {p.name}: {e}")
+    return out
+
+
+def main(argv=None):
+    args = list(argv if argv is not None else sys.argv[1:])
+    if not args:
+        print(__doc__)
+        return 1
+    kw = {}
+    target = args.pop(0)
+    while args:
+        flag = args.pop(0)
+        if flag == "--substrate":
+            kw["substrate"] = {"family": args.pop(0)}
+        elif flag == "--objective":
+            kw["objective"] = args.pop(0)
+        elif flag == "--layout":
+            kw["layout"] = args.pop(0)
+        elif flag == "--out":
+            kw["out"] = args.pop(0)
+        else:
+            print(f"unknown flag {flag}")
+            return 1
+    p = Path(target)
+    if p.is_dir():
+        kw.pop("out", None)
+        convert_dir(p, **kw)
+    else:
+        convert_one(p, **kw)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/amoe/io/safetensors_io.py

@@ -0,0 +1,110 @@
+"""safetensors serialization for amoe anchors (new in 0.2 — the diffusion
+line needs ComfyUI-consumable artifacts; a documented 0.1 non-goal ends).
+
+Two key layouts:
+  "amoe"  (canonical, portable): blocks.{site_index}.{param_path}
+  "comfy" (export): diffusion_model.{site_name}.aleph_relay.{param_path}
+          — the cosmos_predict2 save_adapter precedent. Requires
+          meta.substrate.site_names. The ComfyUI node consumes the
+          canonical layout and asserts the width signature (F2).
+
+safetensors metadata is str:str only, so the full meta rides as one JSON
+blob ("amoe_meta") plus greppable duplicates. Content-hash v2 is defined
+over sorted(key)+shape+dtype+little-endian bytes — format-independent
+(the .pt v1 hash depends on torch.save serialization).
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:   # pragma: no cover
+    from .checkpoint import DiffusionAnchorCheckpoint
+
+
+def _require_safetensors():
+    try:
+        import safetensors.torch as st  # noqa: F401
+        return st
+    except ImportError as e:            # pragma: no cover
+        raise ImportError(
+            "safetensors is required for this path: "
+            "pip install amoe-lora[diffusion] (or pip install safetensors)"
+        ) from e
+
+
+def content_hash_v2(tensors: dict[str, torch.Tensor]) -> str:
+    h = hashlib.sha256()
+    for k in sorted(tensors):
+        t = tensors[k].detach().cpu().contiguous()
+        h.update(k.encode())
+        h.update(str(tuple(t.shape)).encode())
+        h.update(str(t.dtype).encode())
+        # byte view works for every dtype incl. bf16 (no numpy dtype needed)
+        h.update(t.flatten().view(torch.uint8).numpy().tobytes())
+    return "sha256v2:" + h.hexdigest()
+
+
+def save_anchor_safetensors(ck: "DiffusionAnchorCheckpoint", path: str,
+                            *, key_layout: str = "amoe") -> str:
+    st = _require_safetensors()
+    if key_layout == "amoe":
+        tensors = {f"blocks.{k}": v.detach().cpu().contiguous()
+                   for k, v in ck.adapters.items()}
+    elif key_layout == "comfy":
+        names = ck.meta.get("substrate", {}).get("site_names")
+        if not names:
+            raise ValueError(
+                "comfy layout needs meta.substrate.site_names (the trunk "
+                "module paths); save the canonical 'amoe' layout instead")
+        tensors = {}
+        for k, v in ck.adapters.items():
+            i, param = k.split(".", 1)
+            tensors[f"diffusion_model.{names[int(i)]}.aleph_relay.{param}"] = \
+                v.detach().cpu().contiguous()
+    else:
+        raise ValueError(f"unknown key_layout '{key_layout}'")
+    chash = content_hash_v2({k: v for k, v in ck.adapters.items()})
+    meta = dict(ck.meta)
+    meta["content_hash_v2"] = chash
+    md = {"format": "amoe.diffusion.anchor", "version": "1",
+          "key_layout": key_layout, "amoe_meta": json.dumps(meta),
+          "adapter_kind": ck.kind,
+          "substrate_family": str(meta.get("substrate", {}).get("family", "")),
+          "content_hash_v2": chash}
+    st.save_file(tensors, path, metadata=md)
+    return chash
+
+
+def load_anchor_safetensors(path: str) -> "DiffusionAnchorCheckpoint":
+    st = _require_safetensors()
+    from .checkpoint import DiffusionAnchorCheckpoint
+    from safetensors import safe_open
+    with safe_open(path, framework="pt", device="cpu") as f:
+        md = f.metadata() or {}
+        tensors = {k: f.get_tensor(k) for k in f.keys()}
+    if md.get("format") != "amoe.diffusion.anchor":
+        raise ValueError(f"not an amoe.diffusion.anchor safetensors: {path}")
+    meta = json.loads(md.get("amoe_meta", "{}"))
+    layout = md.get("key_layout", "amoe")
+    if layout == "amoe":
+        adapters = {k[len("blocks."):]: v for k, v in tensors.items()}
+    else:                               # comfy layout round-trip
+        names = meta.get("substrate", {}).get("site_names", [])
+        idx = {n: i for i, n in enumerate(names)}
+        adapters = {}
+        for k, v in tensors.items():
+            body = k[len("diffusion_model."):]
+            site, param = body.split(".aleph_relay.")
+            adapters[f"{idx[site]}.{param}"] = v
+    ck = DiffusionAnchorCheckpoint(adapters, meta)
+    want = md.get("content_hash_v2")
+    if want:
+        got = content_hash_v2(adapters)
+        if got != want:
+            raise ValueError(f"content hash mismatch in {path}: "
+                             f"{got} != {want}")
+    return ck

src/amoe/testing/diffusion_invariants.py

@@ -0,0 +1,334 @@
+"""CPU-testable diffusion invariants on tiny stubs (no downloads).
+
+The contracts, each tracing to a paid-for receipt:
+  toggle law (code-path skip, bit-exact)        — every r2 experiment
+  P-INIT (zero-init W+B exactly inert)          — exp006 bias-leak law
+  detach bit-exact via binding probe            — the text line's exp011 bar
+  band windows sum-to-1, smooth, positional     — exp008
+  flow x0 recovery exact and linear             — exp013 (the conditioning law)
+  dtype law (declared, refuses meta)            — R0b catch, fork f4f1c46
+  checkpoint round trips (.pt/.safetensors/legacy shapes) — burndown record
+  blob-on-eps refusal                           — exp012, 2-seed
+"""
+from __future__ import annotations
+
+import os
+import tempfile
+
+import torch
+import torch.nn as nn
+
+
+# ── tiny stubs ──────────────────────────────────────────────────────────
+
+class _TinySite(nn.Module):
+    """Marker block with a width witness; tensor output."""
+
+    def __init__(self, d):
+        super().__init__()
+        self.d = d
+        self.lin = nn.Linear(d, d)
+
+    def forward(self, x):
+        return x + torch.tanh(self.lin(x))
+
+
+class _TinyTupleSite(_TinySite):
+    def forward(self, x):
+        return (x + torch.tanh(self.lin(x)), None)
+
+
+class TinyUNetStub(nn.Module):
+    """Scattered heterogeneous-width sites (down 8/8, mid 16, up 8) —
+    the SD15 shape in miniature, one tuple-output site for the
+    hybrid-safety path."""
+
+    def __init__(self, tuple_site=False):
+        super().__init__()
+        torch.manual_seed(11)
+        cls_last = _TinyTupleSite if tuple_site else _TinySite
+        self.down = nn.ModuleList([_TinySite(8), _TinySite(8)])
+        self.mid = _TinySite(16)
+        self.up = nn.ModuleList([cls_last(8)])
+        self.enc = nn.Linear(8, 16)
+        self.dec = nn.Linear(16, 8)
+
+    def forward(self, x):                       # x: [B, T, 8]
+        for blk in self.down:
+            out = blk(x)
+            x = out[0] if isinstance(out, tuple) else out
+        h = self.enc(x)
+        out = self.mid(h)
+        h = out[0] if isinstance(out, tuple) else out
+        x = self.dec(h)
+        out = self.up[0](x)
+        return out[0] if isinstance(out, tuple) else out
+
+
+class TinyBinding:
+    name = "tiny_diffusion"
+    expected_sites = 4
+
+    def sites(self, model):
+        from ..binding.diffusion import DiffusionBinding  # noqa: F401
+        out = [(n, m, m.d) for n, m in model.named_modules()
+               if isinstance(m, _TinySite)]
+        assert len(out) == self.expected_sites, len(out)
+        return out
+
+    def replace(self, model, name, new):
+        from ..binding.diffusion import walk_replace
+        walk_replace(model, name, new)
+
+    def declared_dtype(self, model):
+        from ..binding.diffusion import _declared_dtype
+        return _declared_dtype(model)
+
+    def probe(self, model):
+        g = torch.Generator().manual_seed(1400)
+        dt = self.declared_dtype(model)
+        x = torch.randn(2, 5, 8, generator=g).to(dt)
+        with torch.no_grad():
+            return model(x).detach().float().cpu()
+
+
+# ── checkpoint builders ─────────────────────────────────────────────────
+
+def _relay_ckpt(model, *, trained=True, name="t"):
+    from ..diffusion.core.relay import RelayPatch2D
+    from ..io.checkpoint import DiffusionAnchorCheckpoint
+    b = TinyBinding()
+    torch.manual_seed(23)
+    adapters = {}
+    widths = []
+    for i, (_, _, d) in enumerate(b.sites(model)):
+        r = RelayPatch2D(d, n_slots=4, K=8, hidden=12)
+        if trained:
+            nn.init.normal_(r.consume[-1].weight, std=0.05)
+            nn.init.normal_(r.consume[-1].bias, std=0.05)
+        for k, v in r.state_dict().items():
+            adapters[f"{i}.{k}"] = v.clone()
+        widths.append(d)
+    return DiffusionAnchorCheckpoint(adapters, {
+        "name": name, "adapter": {"kind": "relay"},
+        "substrate": {"family": "tiny_diffusion", "n_sites": len(widths),
+                      "site_names": [s[0] for s in b.sites(model)],
+                      "widths": widths}})
+
+
+def _mb_ckpt(model, *, trained=True):
+    from ..diffusion.core.multiband import MultibandDelta
+    from ..io.checkpoint import DiffusionAnchorCheckpoint
+    b = TinyBinding()
+    torch.manual_seed(29)
+    adapters = {}
+    for i, (_, _, d) in enumerate(b.sites(model)):
+        m = MultibandDelta(d, r=3)
+        if trained:
+            for up in m.up:
+                nn.init.normal_(up.weight, std=0.05)
+        for k, v in m.state_dict().items():
+            adapters[f"{i}.{k}"] = v.clone()
+    return DiffusionAnchorCheckpoint(adapters, {
+        "name": "mb", "adapter": {"kind": "multiband3"},
+        "substrate": {"family": "tiny_diffusion", "n_sites": 4}})
+
+
+# ── invariants ──────────────────────────────────────────────────────────
+
+def assert_band_windows():
+    from ..diffusion.core.multiband import band_of, band_weights
+    s = torch.linspace(0, 1, 501)
+    w = band_weights(s)
+    assert (w.sum(-1) - 1).abs().max() < 1e-6, "windows do not sum to 1"
+    assert (w >= 0).all() and (w <= 1).all()
+    for v, b in ((0.1, 0), (0.5, 1), (0.9, 2)):
+        assert band_of(v) == b
+        assert int(band_weights(torch.tensor([v])).argmax()) == b
+    # max slope of the cosine ramp is pi/(4*XFADE) ~= 13.1 per unit s01;
+    # on a 1/500 grid that is ~0.0262 per step — assert the theory bound
+    d = (w[1:] - w[:-1]).abs().max()
+    assert d < 0.03, f"windows not smooth (max step {d} > theory ~0.0262)"
+
+
+def assert_flow_x0_recovery():
+    from ..diffusion.train.objectives import flow_pieces
+    g = torch.Generator().manual_seed(5)
+    lat = torch.randn(6, 4, 8, 8, generator=g)
+    s = torch.rand(6, generator=g)
+    noise = torch.randn(6, 4, 8, 8, generator=g)
+    x_t, v = flow_pieces(lat, s, noise)
+    x0 = x_t - s[:, None, None, None] * v
+    assert (x0 - lat).abs().max() < 1e-6, "flow x0 recovery not linear/exact"
+
+
+def assert_toggle_law_diffusion(tuple_site=False):
+    from ..diffusion.runtime.attach import attach
+    m = TinyUNetStub(tuple_site=tuple_site)
+    b = TinyBinding()
+    g = torch.Generator().manual_seed(31)
+    x = torch.randn(2, 5, 8, generator=g)
+    with torch.no_grad():
+        base = m(x).clone()
+
+    # relay: P-INIT (fresh zero-init anchors are EXACTLY inert)
+    h = attach(m, _relay_ckpt(m, trained=False), binding=b)
+    with torch.no_grad():
+        assert torch.equal(m(x), base), "P-INIT broken (zero-init leak)"
+    h.detach(verify=True)
+
+    # relay: trained anchor moves the output; all_off is bit-exact
+    h = attach(m, _relay_ckpt(m, trained=True), binding=b)
+    with torch.no_grad():
+        on = m(x)
+    assert not torch.equal(on, base), "trained relay produced no delta"
+    with h.all_off():
+        with torch.no_grad():
+            assert torch.equal(m(x), base), "toggle law violated (relay)"
+    h.detach(verify=True)
+    with torch.no_grad():
+        assert torch.equal(m(x), base), "model altered after detach"
+
+    # multiband: windows drive it; all_off + full lesion are bit-exact
+    h = attach(m, _mb_ckpt(m, trained=True), binding=b)
+    h.set_band_windows(torch.tensor([0.9, 0.1]))
+    with torch.no_grad():
+        on = m(x)
+    assert not torch.equal(on, base), "trained multiband produced no delta"
+    with h.all_off():
+        with torch.no_grad():
+            assert torch.equal(m(x), base), "toggle law violated (multiband)"
+    with h.lesion_band(0, 1, 2):
+        with torch.no_grad():
+            assert torch.equal(m(x), base), "full band lesion not bit-exact"
+    h.detach(verify=True)
+
+
+def assert_lesion_guard():
+    from ..diffusion.runtime.attach import attach
+    m = TinyUNetStub()
+    h = attach(m, _relay_ckpt(m), binding=TinyBinding())
+    try:
+        with h.lesion_band(0):
+            pass
+        raise AssertionError("lesion_band on a relay anchor must TypeError")
+    except TypeError:
+        pass
+    h.detach(verify=True)
+
+
+def assert_dtype_law():
+    from ..binding.diffusion import _declared_dtype
+    from ..diffusion.runtime.attach import attach
+    m = TinyUNetStub().to(torch.bfloat16)
+    b = TinyBinding()
+    g = torch.Generator().manual_seed(31)
+    x = torch.randn(2, 5, 8, generator=g).to(torch.bfloat16)
+    with torch.no_grad():
+        base = m(x).clone()
+    ck = _relay_ckpt(m, trained=True)
+    h = attach(m, ck, binding=b)                 # dtype=None -> declared
+    for mod in h.modules():
+        for p in mod.parameters():
+            assert p.dtype == torch.bfloat16, "dtype law: adapter != trunk"
+    with h.all_off():
+        with torch.no_grad():
+            assert torch.equal(m(x), base), \
+                "toggle law must be dtype-independent (code-path skip)"
+    h.detach(verify=True)
+
+    class _Meta(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.p = nn.Parameter(torch.empty(1, device="meta"))
+    try:
+        _declared_dtype(_Meta())
+        raise AssertionError("declared_dtype must refuse meta params")
+    except RuntimeError:
+        pass
+
+
+def assert_checkpoint_roundtrip():
+    from ..io.checkpoint import load_diffusion_anchor
+    from ..io.safetensors_io import (load_anchor_safetensors,
+                                     save_anchor_safetensors)
+    m = TinyUNetStub()
+    ck = _relay_ckpt(m, trained=True)
+    tmp = tempfile.mkdtemp()
+
+    # .pt round trip
+    p = os.path.join(tmp, "a.pt")
+    ck.save(p)
+    back = load_diffusion_anchor(p)
+    assert back.kind == "relay" and set(back.adapters) == set(ck.adapters)
+    for k in ck.adapters:
+        assert torch.equal(back.adapters[k], ck.adapters[k])
+
+    # safetensors round trip, canonical + comfy layouts, hash verified
+    for layout in ("amoe", "comfy"):
+        sp = os.path.join(tmp, f"a_{layout}.safetensors")
+        save_anchor_safetensors(ck, sp, key_layout=layout)
+        b2 = load_anchor_safetensors(sp)
+        assert set(b2.adapters) == set(ck.adapters), layout
+        for k in ck.adapters:
+            assert torch.equal(b2.adapters[k], ck.adapters[k]), (layout, k)
+
+    # legacy shapes: relays list / mods list / fork dict (home rebuilt)
+    per = [ck.per_site(i) for i in range(4)]
+    lp = os.path.join(tmp, "legacy_relays.pt")
+    torch.save({"relays": per}, lp)
+    assert load_diffusion_anchor(lp).kind == "relay"
+
+    mb = _mb_ckpt(m)
+    mp = os.path.join(tmp, "legacy_mods.pt")
+    torch.save({"mods": [mb.per_site(i) for i in range(4)]}, mp)
+    assert load_diffusion_anchor(mp).kind == "multiband3"
+
+    fork = {str(i): {k: v for k, v in per[i].items() if k != "addr.home"}
+            for i in range(4)}
+    fp = os.path.join(tmp, "fork.pt")
+    torch.save({"relays": fork}, fp)
+    fk = load_diffusion_anchor(fp)
+    assert fk.meta.get("home_reconstructed") is True, \
+        "fork import must STAMP the reconstructed home"
+    assert "0.addr.home" in fk.adapters
+
+
+def assert_blob_on_eps_refusal():
+    from ..diffusion.train.config import DiffusionTrainConfig
+    from ..diffusion.train.trainer import train
+    cfg = DiffusionTrainConfig(objective="eps", blob=True)
+    try:
+        train(None, {}, cfg)
+        raise AssertionError("blob-on-eps must refuse (conditioning law)")
+    except ValueError as e:
+        assert "conditioning law" in str(e)
+
+
+def assert_align_negative():
+    from ..diffusion.train.aligner import align
+    try:
+        align()
+        raise AssertionError("align must raise with the record")
+    except NotImplementedError as e:
+        assert "falsified" in str(e)
+
+
+def run_all() -> None:
+    assert_band_windows()
+    assert_flow_x0_recovery()
+    assert_toggle_law_diffusion(tuple_site=False)
+    assert_toggle_law_diffusion(tuple_site=True)
+    assert_lesion_guard()
+    assert_dtype_law()
+    assert_checkpoint_roundtrip()
+    assert_blob_on_eps_refusal()
+    assert_align_negative()
+    print("amoe.diffusion invariants: band windows, exact flow x0, toggle "
+          "law (relay+multiband, tensor+tuple), P-INIT, bit-exact detach, "
+          "dtype law, checkpoint round trips (.pt/.safetensors/legacy/"
+          "fork), blob-on-eps refusal, align negative — ALL GREEN")
+
+
+if __name__ == "__main__":
+    run_all()

tests/test_invariants.py

@@ -1,3 +1,5 @@
 import sys; sys.path.insert(0, "src")
 from amoe.testing.invariants import run_all
 run_all()
+from amoe.testing.diffusion_invariants import run_all as run_diffusion
+run_diffusion()