exp010 shipped: StepGatedSampler + generation-side lesion battery (controller +0.089 grounding; coarse-to-fine confirmed by lesion ratios; diversity = open training goal)
c5da886 verified | """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") | |
| 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 | |
| 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() | |