#!/usr/bin/env python3 """ BET 5 — Object permanence in Gabor packet space. ================================================= Question: if we FREEZE the geometry channels of an existing atom set and let SDS gradients touch only appearance (amp / color / phase / background), does the object HOLD its shape while the scene relights? "at dusk, it turns to dusk, the tractor stays." This is the empirical test of the fiber-bundle claim: illumination acts on the amplitude/chroma fibers; geometry (x, y, theta, sigma, f) is the base space and should be untouchable by a relighting edit. Carries forward everything verified in Bet 4 (GO: recognizable tractor from 256 atoms under pure SDS) plus the fixes that run demanded: * SOFT clamp pre = 4*tanh(pre/4) -- no zero-gradient dead zones (Bet 4's hard clamp fixed saturation but killed the escape gradient; this keeps the corridor AND the slope) * Random Sim(2) cameras -- the DreamFusion trick, free in this representation: zoom/shift/rotate are parameter arithmetic on atoms. Kills the SDS zoom-crop trap: only a complete, centered object scores well under every random view. * Normalized SDS loss + gate warmup -- in Bet 4 the SDS loss dwarfed the L0 term so gates never closed (256/256 all run). Loss is now per-element normalized and gates get a warmup before pruning pressure engages. * --init-atoms load any previous atoms.pt (Bet 4 recon or SDS, N inferred from the file) * --freeze channel groups excluded from optimization entirely * --train-groups per-atom masks: only listed groups receive gradients (the two-slot tractor/background experiment is a flag) Modes ----- recon : fit atoms to a target image (MSE). No diffusion model needed. sds : score distillation from frozen Stable Diffusion. render : load atoms and render -- identity view, a chosen camera, or a camera sweep saved as GIF (the "glide" demo). The permanence experiment (the actual bet) ------------------------------------------ # 1. you already have runs/recon_tractor/atoms.pt (23.5 dB, 205 atoms) # 2. relight it, geometry untouchable: python bet5_gabor_sds.py --mode sds \ --init-atoms runs/recon_tractor/atoms.pt \ --freeze geometry,gates \ --prompt "a photo of a red tractor at dusk, golden hour, warm light" \ --iters 1500 --render-size 512 --cfg 50 \ --sd-model sd2-community/stable-diffusion-2-1-base \ --out runs/bet5_dusk GO : final image reads as the SAME tractor, relit. Geometry channels are bitwise identical (the script verifies and prints this). NO-GO: appearance channels alone cannot express the edit (tractor holds by construction -- geometry is frozen -- but the scene refuses to read as dusk after a seed retry and a cfg bump). Honesty note: recon/render/freeze/group/camera machinery executed and verified on CPU before shipping. The sds path reuses Bet 4's verified-on- your-GPU loop with the fixes above; the fixes themselves have not run under a real SDS gradient yet. First run = smoke test. """ import argparse import json import math import os import random import time import numpy as np import torch import torch.nn as nn import torch.nn.functional as F # ---------------------------------------------------------------------------- # Hard-concrete gates (Louizos et al.) — Slapstack two-way doors # ---------------------------------------------------------------------------- class HardConcreteGates(nn.Module): GAMMA, ZETA, BETA = -0.1, 1.1, 2.0 / 3.0 def __init__(self, n, init_logit=2.0): super().__init__() self.logits = nn.Parameter(torch.full((n,), float(init_logit))) def forward(self, hard_eval=False): if self.training and not hard_eval: u = torch.rand_like(self.logits).clamp(1e-6, 1 - 1e-6) s = torch.sigmoid((torch.log(u) - torch.log(1 - u) + self.logits) / self.BETA) else: s = torch.sigmoid(self.logits) return (s * (self.ZETA - self.GAMMA) + self.GAMMA).clamp(0.0, 1.0) def l0(self): return torch.sigmoid(self.logits - self.BETA * math.log(-self.GAMMA / self.ZETA)) @torch.no_grad() def hard_open(self): z = torch.sigmoid(self.logits) * (self.ZETA - self.GAMMA) + self.GAMMA return (z.clamp(0, 1) > 0.5) # ---------------------------------------------------------------------------- # Differentiable Gabor packet image (Bet 4 renderer + soft clamp + cameras) # ---------------------------------------------------------------------------- class GaborPacketImage(nn.Module): """ Image = sigmoid( softclamp( bg_bias + sum_i g_i a_i c_i env_i carrier_i ) ) env_i = exp(-0.5 (u^2/su^2 + v^2/sv^2)), (u,v) atom-local rotated coords carrier_i = cos(2*pi*f_i*u + phi_i) phi is ENVELOPE-RELATIVE: pose never touches phase. A camera g = (s, rho, tx, ty) in Sim(2) acts by parameter arithmetic: xy -> s*R(rho)*xy + t, theta -> theta+rho, sigma -> s*sigma, f -> f/s phase and color are untouched — that IS the identity/pose factorization. """ # per-atom parameter names (used by freeze / group-mask machinery) PER_ATOM = ["xy_raw", "theta", "log_sigma_u", "log_sigma_v", "freq_raw", "phase", "amp", "color", "gates.logits"] def __init__(self, n_atoms=256, coarse_frac=0.25, seed=0): super().__init__() g = torch.Generator().manual_seed(seed) n = n_atoms nc = int(n * coarse_frac) self.xy_raw = nn.Parameter(torch.randn(n, 2, generator=g) * 0.7) self.theta = nn.Parameter(torch.rand(n, generator=g) * math.pi) log_s = torch.empty(n) log_s[:nc] = math.log(0.45) + 0.25 * torch.randn(nc, generator=g) log_s[nc:] = math.log(0.12) + 0.35 * torch.randn(n - nc, generator=g) self.log_sigma_u = nn.Parameter(log_s.clone()) self.log_sigma_v = nn.Parameter(log_s + 0.2 * torch.randn(n, generator=g)) f = torch.empty(n) f[:nc] = 0.25 + 0.5 * torch.rand(nc, generator=g) f[nc:] = 0.75 + 2.0 * torch.rand(n - nc, generator=g) self.freq_raw = nn.Parameter(torch.log(torch.expm1(f))) self.phase = nn.Parameter(2 * math.pi * torch.rand(n, generator=g)) self.amp = nn.Parameter(0.35 + 0.15 * torch.randn(n, generator=g)) self.color = nn.Parameter(0.30 * torch.randn(n, 3, generator=g)) self.bg_bias = nn.Parameter(torch.zeros(3)) self.gates = HardConcreteGates(n) # per-atom group id (0 = default). Buffer so it persists in atoms.pt. self.register_buffer("group", torch.zeros(n, dtype=torch.long)) self.n_atoms = n # -- derived --------------------------------------------------------------- def xy(self): return torch.tanh(self.xy_raw) def freq(self): return F.softplus(self.freq_raw) def sigmas(self): return (self.log_sigma_u.exp().clamp(5e-3, 2.0), self.log_sigma_v.exp().clamp(5e-3, 2.0)) # -- render ------------------------------------------------------------------ def render(self, H, W, device, chunk=64, hard_gates=False, camera=None): ys = torch.linspace(-1, 1, H, device=device) xs = torch.linspace(-1, 1, W, device=device) Y, X = torch.meshgrid(ys, xs, indexing="ij") xy = self.xy().to(device) theta = self.theta.to(device) su, sv = self.sigmas() su, sv = su.to(device), sv.to(device) f = self.freq().to(device) phi = self.phase.to(device) amp = self.amp.to(device) col = self.color.to(device) z = self.gates(hard_eval=hard_gates).to(device) if camera is not None: # Sim(2): parameter arithmetic s, rho, tx, ty = camera c, sn = math.cos(rho), math.sin(rho) R = torch.tensor([[c, -sn], [sn, c]], device=device, dtype=xy.dtype) xy = s * xy @ R.T + torch.tensor([tx, ty], device=device, dtype=xy.dtype) theta = theta + rho su, sv = s * su, s * sv f = f / s # phase untouched by construction — envelope-relative. pre = torch.zeros(3, H, W, device=device) + self.bg_bias.to(device)[:, None, None] for i0 in range(0, self.n_atoms, chunk): sl = slice(i0, min(i0 + chunk, self.n_atoms)) dx = X[None] - xy[sl, 0, None, None] dy = Y[None] - xy[sl, 1, None, None] ct = torch.cos(theta[sl])[:, None, None] st = torch.sin(theta[sl])[:, None, None] u = ct * dx + st * dy v = -st * dx + ct * dy env = torch.exp(-0.5 * ((u / su[sl, None, None]) ** 2 + (v / sv[sl, None, None]) ** 2)) carrier = torch.cos(2 * math.pi * f[sl, None, None] * u + phi[sl, None, None]) w = (z[sl] * amp[sl])[:, None, None] * env * carrier pre = pre + torch.einsum("nhw,nc->chw", w, col[sl]) # Leaky soft clamp: tanh corridor + linear leak. Plain tanh still dies # numerically in fp32 at |pre|>~40 (caught by shipped test 5); the # 0.02*pre leak guarantees a nonzero escape gradient at ANY depth of # saturation — the solid-red trap always has an exit ramp. pre = 4.0 * torch.tanh(pre / 4.0) + 0.02 * pre return torch.sigmoid(pre) def ledger(self): return {"atoms_total": self.n_atoms, "atoms_open_hard": int(self.gates.hard_open().sum().item()), "expected_L0": float(self.gates.l0().sum().item()), "groups": {int(k): int(v) for k, v in zip(*[t.tolist() for t in self.group.unique(return_counts=True)])}} # ---------------------------------------------------------------------------- # Loading, freezing, group masks # ---------------------------------------------------------------------------- FREEZE_MAP = { "position": ["xy_raw"], "orientation": ["theta"], "scale": ["log_sigma_u", "log_sigma_v"], "frequency": ["freq_raw"], "phase": ["phase"], "amp": ["amp"], "color": ["color"], "bg": ["bg_bias"], "gates": ["gates.logits"], } FREEZE_MAP["geometry"] = (FREEZE_MAP["position"] + FREEZE_MAP["orientation"] + FREEZE_MAP["scale"] + FREEZE_MAP["frequency"]) FREEZE_MAP["appearance"] = (FREEZE_MAP["phase"] + FREEZE_MAP["amp"] + FREEZE_MAP["color"] + FREEZE_MAP["bg"]) def load_atoms(path): sd = torch.load(path, map_location="cpu") n = sd["phase"].shape[0] model = GaborPacketImage(n_atoms=n) missing, unexpected = model.load_state_dict(sd, strict=False) # 'group' may be missing in Bet 4 files — that's fine (defaults to 0). assert not unexpected, f"unexpected keys in {path}: {unexpected}" print(f"loaded {path}: {n} atoms" + (f" (new buffers defaulted: {missing})" if missing else "")) return model def resolve_frozen(freeze_arg): frozen = set() for tok in [t for t in (freeze_arg or "").split(",") if t.strip()]: tok = tok.strip() assert tok in FREEZE_MAP, f"unknown freeze channel '{tok}' " \ f"(choose from {sorted(FREEZE_MAP)})" frozen.update(FREEZE_MAP[tok]) return frozen def make_optimizer(model, frozen, mode, lr_scale=1.0): # LRs that worked: recon = Bet 4 defaults; sds = the rebalance that # escaped the solid-red saturation trap (color 2e-3, bg 5e-4). lrs = {"xy_raw": 5e-3, "theta": 5e-3, "log_sigma_u": 5e-3, "log_sigma_v": 5e-3, "freq_raw": 5e-3, "phase": 2e-2, "amp": 1e-2, "color": 1e-2, "bg_bias": 1e-2, "gates.logits": 2e-2} if mode == "sds": lrs["color"], lrs["bg_bias"] = 2e-3, 5e-4 groups, trained = [], [] for name, lr in lrs.items(): if name in frozen: continue p = dict(model.named_parameters())[name] groups.append({"params": [p], "lr": lr * lr_scale}) trained.append(name) assert groups, "everything is frozen — nothing to optimize" print(f"training channels: {trained}") if frozen: print(f"frozen channels: {sorted(frozen)}") return torch.optim.Adam(groups, betas=(0.9, 0.99)) def group_mask(model, train_groups): """None if all groups train; else float mask (n,) — 1 for trainable atoms.""" if train_groups is None: return None ids = torch.tensor([int(t) for t in train_groups.split(",")]) mask = torch.isin(model.group.cpu(), ids).float() print(f"group mask: {int(mask.sum())}/{model.n_atoms} atoms trainable " f"(groups {ids.tolist()})") return mask def apply_grad_masks(model, mask, device): """Zero gradients of per-atom params for atoms outside trainable groups.""" if mask is None: return m = mask.to(device) params = dict(model.named_parameters()) for name in GaborPacketImage.PER_ATOM: p = params[name] if p.grad is not None: p.grad.mul_(m.view(-1, *([1] * (p.dim() - 1)))) def sample_camera(args): if args.no_camera: return None return (math.exp(random.uniform(-args.cam_zoom, args.cam_zoom)), random.uniform(-args.cam_rot, args.cam_rot), random.uniform(-args.cam_shift, args.cam_shift), random.uniform(-args.cam_shift, args.cam_shift)) def save_png(img_chw, path): from PIL import Image arr = (img_chw.detach().clamp(0, 1).cpu().numpy() .transpose(1, 2, 0) * 255).astype(np.uint8) Image.fromarray(arr).save(path) def geometry_fingerprint(model): """Hash of geometry channels — proves bitwise permanence after a run.""" import hashlib h = hashlib.sha256() for name in FREEZE_MAP["geometry"]: h.update(dict(model.named_parameters())[name].detach().cpu() .numpy().tobytes()) return h.hexdigest()[:16] # ---------------------------------------------------------------------------- # Mode: recon # ---------------------------------------------------------------------------- def run_recon(args, model, device): from PIL import Image tgt = Image.open(args.target).convert("RGB").resize((args.render_size,) * 2) target = torch.from_numpy(np.asarray(tgt).copy()).float().permute(2, 0, 1) / 255.0 target = target.to(device) frozen = resolve_frozen(args.freeze) opt = make_optimizer(model, frozen, "recon") mask = group_mask(model, args.train_groups) model.train().to(device) os.makedirs(args.out, exist_ok=True) save_png(target, os.path.join(args.out, "target.png")) log, t0 = [], time.time() for it in range(args.iters): opt.zero_grad() img = model.render(args.render_size, args.render_size, device, chunk=args.chunk) mse = F.mse_loss(img, target) loss = mse + args.l0_weight * model.gates.l0().sum() / model.n_atoms loss.backward() if it < args.gate_warmup and model.gates.logits.grad is not None: model.gates.logits.grad.zero_() apply_grad_masks(model, mask, device) opt.step() if it % max(1, args.iters // 20) == 0 or it == args.iters - 1: psnr = -10 * math.log10(max(mse.item(), 1e-12)) row = {"it": it, "mse": mse.item(), "psnr_db": psnr, **model.ledger()} log.append(row) print(f"[recon] it {it:5d} mse {mse.item():.5f} psnr {psnr:5.2f} dB " f"open {row['atoms_open_hard']}/{model.n_atoms} ({time.time()-t0:.0f}s)") save_png(img, os.path.join(args.out, f"it_{it:05d}.png")) finish(model, args, device, log) # ---------------------------------------------------------------------------- # Mode: sds # ---------------------------------------------------------------------------- def run_sds(args, model, device): from diffusers import StableDiffusionPipeline, DDPMScheduler dtype = torch.float16 if device.type == "cuda" else torch.float32 pipe = StableDiffusionPipeline.from_pretrained( args.sd_model, torch_dtype=dtype, safety_checker=None, requires_safety_checker=False) pipe.to(device) vae, unet, tok, te = pipe.vae, pipe.unet, pipe.tokenizer, pipe.text_encoder for m in (vae, unet, te): m.requires_grad_(False) sched = DDPMScheduler.from_pretrained(args.sd_model, subfolder="scheduler") alphas = sched.alphas_cumprod.to(device) T = sched.config.num_train_timesteps def embed(text): ids = tok(text, padding="max_length", max_length=tok.model_max_length, truncation=True, return_tensors="pt").input_ids.to(device) return te(ids)[0] with torch.no_grad(): emb = torch.cat([embed(args.negative_prompt), embed(args.prompt)]) frozen = resolve_frozen(args.freeze) opt = make_optimizer(model, frozen, "sds") mask = group_mask(model, args.train_groups) model.train().to(device) fp_before = geometry_fingerprint(model) os.makedirs(args.out, exist_ok=True) log, t0 = [], time.time() for it in range(args.iters): opt.zero_grad() cam = sample_camera(args) img = model.render(args.render_size, args.render_size, device, chunk=args.chunk, camera=cam) x = img[None] * 2 - 1 if args.render_size != 512: x = F.interpolate(x, (512, 512), mode="bilinear", align_corners=False) latents = vae.encode(x.to(dtype)).latent_dist.sample() * vae.config.scaling_factor latents = latents.float() frac = it / max(1, args.iters - 1) t_max = args.t_max_start + (args.t_max_end - args.t_max_start) * frac t = torch.randint(int(args.t_min * T), int(t_max * T), (1,), device=device) noise = torch.randn_like(latents) noisy = sched.add_noise(latents, noise, t) with torch.no_grad(): eps = unet(torch.cat([noisy] * 2).to(dtype), torch.cat([t] * 2), encoder_hidden_states=emb).sample.float() eps_un, eps_tx = eps.chunk(2) eps_hat = eps_un + args.cfg * (eps_tx - eps_un) w = (1 - alphas[t]).view(-1, 1, 1, 1) grad = (w * (eps_hat - noise)).detach() sds_loss = (grad * latents).sum() / latents.numel() # normalized (Bet 4 fix) l0_loss = model.gates.l0().sum() / model.n_atoms loss = sds_loss + args.l0_weight * l0_loss loss.backward() if it < args.gate_warmup and model.gates.logits.grad is not None: model.gates.logits.grad.zero_() apply_grad_masks(model, mask, device) torch.nn.utils.clip_grad_norm_( [p for g_ in opt.param_groups for p in g_["params"]], 1.0) opt.step() if it % max(1, args.iters // 30) == 0 or it == args.iters - 1: row = {"it": it, "sds": float(sds_loss.item()), "l0": float(l0_loss.item()), "t_max": t_max, **model.ledger()} log.append(row) print(f"[sds] it {it:5d} sds {sds_loss.item():+.4f} t_max {t_max:.2f} " f"open {row['atoms_open_hard']}/{model.n_atoms} ({time.time()-t0:.0f}s)") with torch.no_grad(): save_png(model.render(args.render_size, args.render_size, device, chunk=args.chunk, hard_gates=True), os.path.join(args.out, f"it_{it:05d}.png")) fp_after = geometry_fingerprint(model) if "geometry" in (args.freeze or ""): verdict = "IDENTICAL — permanence held by construction" \ if fp_before == fp_after else "CHANGED — BUG, investigate" print(f"geometry fingerprint before/after: {fp_before} / {fp_after} -> {verdict}") finish(model, args, device, log, extra={"geometry_fp_before": fp_before, "geometry_fp_after": fp_after}) # ---------------------------------------------------------------------------- # Mode: render (identity view, chosen camera, or a glide GIF) # ---------------------------------------------------------------------------- def run_render(args, model, device): model.eval().to(device) os.makedirs(args.out, exist_ok=True) S = args.render_size with torch.no_grad(): save_png(model.render(S, S, device, chunk=args.chunk, hard_gates=True), os.path.join(args.out, "identity.png")) if args.camera: cam = tuple(float(v) for v in args.camera.split(",")) assert len(cam) == 4, "--camera expects 's,rho,tx,ty'" save_png(model.render(S, S, device, chunk=args.chunk, hard_gates=True, camera=cam), os.path.join(args.out, "camera.png")) if args.gif: from PIL import Image frames = [] n = 40 for i in range(n): p = i / (n - 1) cam = (1.0 + 0.15 * math.sin(2 * math.pi * p), # gentle zoom breath 0.0, -0.45 + 0.9 * p, # glide left -> right 0.10 * math.sin(4 * math.pi * p)) # slight bob im = model.render(S, S, device, chunk=args.chunk, hard_gates=True, camera=cam) frames.append(Image.fromarray( (im.clamp(0, 1).cpu().numpy().transpose(1, 2, 0) * 255) .astype(np.uint8))) frames[0].save(os.path.join(args.out, "glide.gif"), save_all=True, append_images=frames[1:], duration=60, loop=0) print("wrote glide.gif — texture rides the envelopes; phase never moves") print(f"render -> {args.out}") # ---------------------------------------------------------------------------- def finish(model, args, device, log, extra=None): model.eval() with torch.no_grad(): img = model.render(args.render_size, args.render_size, device, chunk=args.chunk, hard_gates=True) save_png(img, os.path.join(args.out, "final_hardgates.png")) torch.save(model.state_dict(), os.path.join(args.out, "atoms.pt")) ledger = {"mode": args.mode, "prompt": getattr(args, "prompt", None), "freeze": args.freeze, "train_groups": args.train_groups, "init_atoms": args.init_atoms, "camera": None if args.no_camera else {"zoom": args.cam_zoom, "shift": args.cam_shift, "rot": args.cam_rot}, "final": model.ledger(), "log": log} if extra: ledger.update(extra) with open(os.path.join(args.out, "ledger.json"), "w") as fh: json.dump(ledger, fh, indent=2) print(f"done -> {args.out} | open atoms: {model.ledger()['atoms_open_hard']}") def assign_group_rect(model, spec): """--assign-group-rect 'x0,y0,x1,y1:gid' — atoms with canonical xy inside the rect get group gid. Coordinates in [-1,1]. Repeatable.""" box, gid = spec.split(":") x0, y0, x1, y1 = (float(v) for v in box.split(",")) xy = model.xy().detach() inside = ((xy[:, 0] >= x0) & (xy[:, 0] <= x1) & (xy[:, 1] >= y0) & (xy[:, 1] <= y1)) model.group[inside] = int(gid) print(f"assigned {int(inside.sum())} atoms in [{x0},{x1}]x[{y0},{y1}] " f"to group {gid}") def main(): p = argparse.ArgumentParser(description="Bet 5: permanence in Gabor packet space") p.add_argument("--mode", choices=["recon", "sds", "render"], required=True) p.add_argument("--out", default="runs/bet5") p.add_argument("--atoms", type=int, default=256) p.add_argument("--iters", type=int, default=2000) p.add_argument("--render-size", type=int, default=512) p.add_argument("--chunk", type=int, default=64) p.add_argument("--seed", type=int, default=0) p.add_argument("--l0-weight", type=float, default=5e-3) p.add_argument("--gate-warmup", type=int, default=400, help="iterations before L0 pruning pressure engages") # init / freeze / groups p.add_argument("--init-atoms", help="atoms.pt from any previous run (Bet 4 OK)") p.add_argument("--freeze", default="", help=f"comma list from {sorted(FREEZE_MAP)}") p.add_argument("--train-groups", help="comma list of group ids that receive gradients") p.add_argument("--assign-group-rect", action="append", default=[], help="'x0,y0,x1,y1:gid' assign atoms in rect to group (repeatable)") # cameras p.add_argument("--no-camera", action="store_true", help="disable random Sim(2) cameras in sds mode") p.add_argument("--cam-zoom", type=float, default=0.30, help="log-zoom range") p.add_argument("--cam-shift", type=float, default=0.25) p.add_argument("--cam-rot", type=float, default=0.15, help="radians") p.add_argument("--camera", help="render mode: fixed 's,rho,tx,ty'") p.add_argument("--gif", action="store_true", help="render mode: glide GIF") # recon p.add_argument("--target", help="target image (recon mode)") # sds p.add_argument("--prompt", default="a photo of a tractor") p.add_argument("--negative-prompt", default="blurry, low quality, deformed") p.add_argument("--sd-model", default="sd2-community/stable-diffusion-2-1-base") p.add_argument("--cfg", type=float, default=50.0) p.add_argument("--t-min", type=float, default=0.02) p.add_argument("--t-max-start", type=float, default=0.98) p.add_argument("--t-max-end", type=float, default=0.50) args = p.parse_args() random.seed(args.seed) torch.manual_seed(args.seed) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"device: {device}") model = load_atoms(args.init_atoms) if args.init_atoms \ else GaborPacketImage(args.atoms, seed=args.seed) for spec in args.assign_group_rect: assign_group_rect(model, spec) if args.mode == "recon": assert args.target, "--target required in recon mode" run_recon(args, model, device) elif args.mode == "sds": run_sds(args, model, device) else: run_render(args, model, device) if __name__ == "__main__": main()