"""dexp010_controller.py — exp010: THE STEP-GATED DIFFUSION CONTROLLER at inference + the generation-side band-lesion battery (the instrument upgrade). The controller (stage plan phase 3): during sampling, each denoising step sets the band windows from the CURRENT timestep — band experts activate coarse-to-fine across the trajectory exactly as trained. This file ships the controller as a usable module (StepGatedSampler) and judges what the bands DO IN IMAGE SPACE, where per-step eps-MSE may be gauge-blind (exp009 verdict). Arms (trained mb3 stack, exp008 s0): all_on | lesion_band0 (LOW-noise/detail) | lesion_band1 (MID) | lesion_band2 (HIGH-noise/structure) | frozen. Image-space gauges (deterministic, paired seeds): G1 round-trip: CLIP-L cond-vs-shuffled per arm (does lesioning band b change grounding?). G2 structure-vs-detail decomposition of the lesion effect: for the SAME (prompt, seed), LP-MSE(lesioned img, all_on img) = coarse-structure change; HP-MSE = fine-detail change. PREREG: lesion_band2 (HIGH) moves LP most (coarse/blob role); lesion_band0 (LOW) moves HP most (detail role); band1 intermediate. This is the image-space test of the coarse-to-fine thesis. G3 within-prompt diversity: 6 prompts x 4 seeds, mean pairwise CLIP distance per arm — PREREG: lesion_band2 reduces diversity (the HIGH band's diversity role). Pod: bash pod2/run_exp010.sh """ from __future__ import annotations import io import json import os import sys sys.path[:0] = ["pod2", "."] import numpy as np import torch import torch.nn.functional as F from pod_ledger import ledger_run, note, burn_down from d1_substrate import MEM_FRACTION from aleph_diffusion_core import derangement from dexp008_multiband import (MultibandDelta, band_weights, attach, load_unet, N_BANDS) SD_BASE = "stable-diffusion-v1-5/stable-diffusion-v1-5" DEXP8_CKPT = ("/workspace/ckpts2/dexp008" if os.path.isdir("/workspace") else "./data/dexp008") N_JUDGE, STEPS_GEN, GUIDANCE = 24, 30, 7.5 CKPT_FILE = os.environ.get("DEXP10_CKPT", "mb3_s0.pt") CKPT_DIR_ENV = os.environ.get("DEXP10_CKPT_DIR", "") TAG = os.environ.get("DEXP10_TAG", "mb3_s0") N_DIV_PROMPTS, N_DIV_SEEDS = 6, 4 SEED = 1234 DATA_DIR = ("/workspace/data/dexp010" if os.path.isdir("/workspace") else os.path.join(os.environ.get("GEOLIP_DATA", "./data"), "dexp010")) class StepGatedSampler: """THE CONTROLLER: DDIM sampling with per-step band gating — each step sets the crossfade windows from the current timestep, activating band experts coarse-to-fine across the trajectory.""" def __init__(self, unet, wraps, device): from diffusers import DDIMScheduler self.unet, self.wraps, self.device = unet, wraps, device self.sched = DDIMScheduler.from_pretrained(SD_BASE, subfolder="scheduler") @torch.no_grad() def sample(self, ehs_cond, seed, steps=STEPS_GEN, guidance=GUIDANCE): B = ehs_cond.shape[0] self.sched.set_timesteps(steps, device=self.device) g = torch.Generator(device=self.device).manual_seed(seed) x = torch.randn(B, 4, 64, 64, generator=g, device=self.device) x = x * self.sched.init_noise_sigma ehs_in = torch.cat([ehs_cond, torch.zeros_like(ehs_cond)], dim=0) for t in self.sched.timesteps: s01 = torch.full((2 * B,), float(t) / 1000.0, device=self.device) w = band_weights(s01) for wr in self.wraps: wr.w_bands = w # the step gate xin = self.sched.scale_model_input(torch.cat([x, x]), t) eps = self.unet(xin, t, ehs_in, return_dict=False)[0] e_c, e_u = eps.chunk(2) eps = e_u + guidance * (e_c - e_u) x = self.sched.step(eps, t, x).prev_sample return x def hp_img(x): return x - F.avg_pool2d(x, 3, stride=1, padding=1) def lp_img(x): return F.avg_pool2d(x, 7, stride=1, padding=3) def run(device="cuda"): torch.cuda.set_per_process_memory_fraction(MEM_FRACTION, 0) os.makedirs(DATA_DIR, exist_ok=True) from PIL import Image from diffusers import AutoencoderKL from transformers import (CLIPTextModel, CLIPTokenizer, CLIPModel, CLIPProcessor) from d1_exp000b_natural import load_rows, encode_nl77 from d1_exp000_baselines import judge_selftest rows = load_rows(N_JUDGE) prompts = [r["prompt"] for r in rows] perm = derangement(N_JUDGE, seed=SEED) tok = CLIPTokenizer.from_pretrained(SD_BASE, subfolder="tokenizer") te = CLIPTextModel.from_pretrained( SD_BASE, subfolder="text_encoder", torch_dtype=torch.float32).to(device).eval() vae = AutoencoderKL.from_pretrained( SD_BASE, subfolder="vae", torch_dtype=torch.float32).to(device).eval() clip = CLIPModel.from_pretrained( "openai/clip-vit-large-patch14", torch_dtype=torch.float32).to(device).eval() cproc = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14") feat = judge_selftest(clip, cproc, device) with torch.no_grad(): ehs_c = encode_nl77(prompts, tok, te, device) ehs_s = encode_nl77([prompts[i] for i in perm.tolist()], tok, te, device) orig = torch.cat([ feat(Image.open(io.BytesIO(r["image_bytes"])).convert("RGB")) for r in rows]) unet = load_unet(device) mods, wraps = attach(unet, lambda d: MultibandDelta(d)) ck_dir = CKPT_DIR_ENV or DEXP8_CKPT ck = torch.load(os.path.join(ck_dir, CKPT_FILE), map_location="cpu", weights_only=True) for m, sd in zip(mods, ck["mods"]): m.load_state_dict(sd, strict=True) sampler = StepGatedSampler(unet, wraps, device) def set_arm(arm): for m in mods: m.enabled = arm != "frozen" m.band_enabled = [True] * N_BANDS if arm.startswith("lesion_band"): m.band_enabled[int(arm[-1])] = False @torch.no_grad() def decode_imgs(lat): img = vae.decode(lat / 0.18215).sample return ((img / 2 + 0.5).clamp(0, 1)) arms = ["all_on", "lesion_band0", "lesion_band1", "lesion_band2", "frozen"] results = {"config": {"steps": STEPS_GEN, "guidance": GUIDANCE, "seed": SEED, "ckpt": TAG, "sampler": "DDIM step-gated"}} gen_store = {} with ledger_run("dexp010 controller battery", budget_h=1.5) as h: for arm in arms: set_arm(arm) imgs_all, cos = [], [] for i in range(0, N_JUDGE, 6): lat = sampler.sample(ehs_c[i:i + 6], seed=SEED + i) px = decode_imgs(lat) imgs_all.append(px.cpu()) pil = [Image.fromarray( (p.permute(1, 2, 0).numpy() * 255).astype(np.uint8)) for p in px.cpu()] f = torch.cat([feat(p) for p in pil]) cos += (f * orig[i:i + 6]).sum(-1).tolist() cond_mean = sum(cos) / len(cos) cos_s = [] for i in range(0, N_JUDGE, 6): lat = sampler.sample(ehs_s[i:i + 6], seed=SEED + i) px = decode_imgs(lat) pil = [Image.fromarray( (p.permute(1, 2, 0).numpy() * 255).astype(np.uint8)) for p in px.cpu()] f = torch.cat([feat(p) for p in pil]) cos_s += (f * orig[i:i + 6]).sum(-1).tolist() gen_store[arm] = torch.cat(imgs_all) results[arm] = { "round_trip": {"cond": round(cond_mean, 4), "shuffled": round(sum(cos_s) / len(cos_s), 4), "gap": round(cond_mean - sum(cos_s) / len(cos_s), 4)}} print(f"[exp010] {arm}: rt gap " f"{results[arm]['round_trip']['gap']:+.4f}", flush=True) # G2: structure-vs-detail decomposition of each lesion (vs all_on) base_imgs = gen_store["all_on"] for arm in ("lesion_band0", "lesion_band1", "lesion_band2", "frozen"): d = gen_store[arm] - base_imgs results[arm]["vs_all_on"] = { "lp_change": round(float((lp_img(gen_store[arm]) - lp_img(base_imgs)).pow(2).mean()), 6), "hp_change": round(float((hp_img(gen_store[arm]) - hp_img(base_imgs)).pow(2).mean()), 6), "total_change": round(float(d.pow(2).mean()), 6)} # G3: within-prompt diversity (6 prompts x 4 seeds) for arm in arms: set_arm(arm) div = [] for pi in range(N_DIV_PROMPTS): fs = [] for si in range(N_DIV_SEEDS): lat = sampler.sample(ehs_c[pi:pi + 1], seed=SEED + 1000 + 97 * si) px = decode_imgs(lat).cpu()[0] pil = Image.fromarray( (px.permute(1, 2, 0).numpy() * 255).astype(np.uint8)) fs.append(feat(pil)) fs = torch.cat(fs) sim = fs @ fs.T n = fs.shape[0] div.append(1 - ((sim.sum() - n) / (n * (n - 1))).item()) results[arm]["diversity"] = round(sum(div) / len(div), 5) h["verdict"] = json.dumps( {a: results[a]["round_trip"]["gap"] for a in arms}) # prereg checks lb0 = results["lesion_band0"]["vs_all_on"] lb2 = results["lesion_band2"]["vs_all_on"] results["prereg"] = { "G2_high_band_moves_LP_most": lb2["lp_change"] > lb0["lp_change"], "G2_low_band_moves_HP_most": lb0["hp_change"] > lb2["hp_change"], "G3_high_lesion_cuts_diversity": results["lesion_band2"]["diversity"] < results["all_on"]["diversity"], "note": "image-space coarse-to-fine test; 1-seed battery", } with open(os.path.join(DATA_DIR, "results.json" if TAG == "mb3_s0" else f"results_{TAG}.json"), "w") as f: json.dump(results, f, indent=2) note(f"dexp010: {json.dumps(results['prereg'])}") print(json.dumps(results["prereg"], indent=2)) burn_down() return results def smoke(): x = torch.randn(2, 3, 32, 32) assert hp_img(x).shape == x.shape and lp_img(x).shape == x.shape w = band_weights(torch.tensor([0.02, 0.5, 0.98])) assert w[0].argmax() == 0 and w[2].argmax() == 2 print("dexp010 smoke PASSED (filters + windows; GPU run is pod work)") if __name__ == "__main__": if "--run" in sys.argv: run() else: smoke()