"""dexp011_fused_multiband.py — exp011a: multiband on REAL FUSED DATA (+ the blob-target cache for the 8am blob-loss design review). Two jobs (stage plan phase 4 prep, D2 from the reconstruction brief): 1. Does REAL data (qwen-deepfashion-fused: real photos, joycaption NL, multi-entity scenes) change the band economics that synthetic single-object data couldn't pay for? Same certified beds (multiband3 vs mono48, uniform objective), new substrate data, exp010 battery after. 2. Build the SEGMENTATION BLOB TARGETS: fused_json entities[].mask.polygon (NORM_0_1000) rasterized to the 64x64 latent grid, cached per row. The blob LOSS mechanism is DELIBERATELY NOT implemented tonight — it is a new mechanism and goes to Phil's 8am review with targets ready (pacing law: don't step into something unprepared). Data gates: rows filtered on the dataset's own audit columns when present (age_audit/audit approved-only; counts printed; empty result fails LOUDLY — never silently train on people data past an available gate). Pod: bash pod2/run_exp011.sh """ from __future__ import annotations import io import json import os import sys import time 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 dexp006_sd15core_relay import make_schedule, add_noise, encode_nl77 from dexp008_multiband import (MultibandDelta, MonoDelta, band_weights, band_of, attach, load_unet, N_BANDS) DATASET = "AbstractPhil/qwen-deepfashion-fused" SD_BASE = "stable-diffusion-v1-5/stable-diffusion-v1-5" VAE_SCALE = 0.18215 N_TRAIN, N_VAL = 4096, 256 BATCH = int(os.environ.get("DEXP11_BATCH", "16")) STEPS = int(os.environ.get("DEXP11_STEPS", "3000")) SEED = int(os.environ.get("DEXP11_SEED", "0")) LR, CFG_DROPOUT = 1e-3, 0.1 DATA_DIR = ("/workspace/data/dexp011" if os.path.isdir("/workspace") else os.path.join(os.environ.get("GEOLIP_DATA", "./data"), "dexp011")) CKPT_DIR = ("/workspace/ckpts2/dexp011" if os.path.isdir("/workspace") else DATA_DIR) def rasterize_blob(fused_json_str: str, size: int = 64) -> np.ndarray: """Entity mask polygons (NORM_0_1000) -> union occupancy on the latent grid. Returns (size, size) float32 in [0,1]. Rows without polygons return zeros (counted by the caller).""" from PIL import Image, ImageDraw try: fj = json.loads(fused_json_str) except Exception: return np.zeros((size, size), dtype=np.float32) img = Image.new("L", (size, size), 0) draw = ImageDraw.Draw(img) n = 0 for ent in (fj.get("entities") or []): mask = ent.get("mask") or {} poly = mask.get("polygon") or [] if len(poly) >= 3: pts = [(p[0] / 1000.0 * (size - 1), p[1] / 1000.0 * (size - 1)) for p in poly if isinstance(p, (list, tuple)) and len(p) >= 2] if len(pts) >= 3: draw.polygon(pts, fill=255) n += 1 arr = np.asarray(img, dtype=np.float32) / 255.0 return arr def _iter_fused(need: int): from huggingface_hub import HfApi, hf_hub_download import pyarrow.parquet as pq api = HfApi() files = sorted(f for f in api.list_repo_files(DATASET, repo_type="dataset") if f.endswith(".parquet")) got, kept, gated = 0, 0, 0 cols_printed = False for fname in files: path = hf_hub_download(DATASET, fname, repo_type="dataset") tbl = pq.read_table(path) if not cols_printed: print(f"[fused] columns: {tbl.column_names}", flush=True) cols_printed = True for r in tbl.to_pylist(): got += 1 gate_vals = [str(r.get(c)) for c in ("age_audit", "audit") if c in r and r.get(c) is not None] if gate_vals and not any(("approved" in v.lower() or "pass" in v.lower()) for v in gate_vals): gated += 1 continue kept += 1 yield r if kept >= need: print(f"[fused] kept {kept} / seen {got} (gated out {gated})", flush=True) return raise AssertionError( f"exhausted dataset with only {kept} gated rows of {got} seen — " f"inspect gate columns before proceeding") def build_cache(device="cuda"): cache_f = os.path.join(DATA_DIR, "cache.pt") if os.path.exists(cache_f): print(f"[cache] exists: {cache_f}") return cache_f os.makedirs(DATA_DIR, exist_ok=True) from PIL import Image from diffusers import AutoencoderKL from transformers import CLIPTextModel, CLIPTokenizer vae = AutoencoderKL.from_pretrained( SD_BASE, subfolder="vae", torch_dtype=torch.float32).to(device).eval() tok = CLIPTokenizer.from_pretrained(SD_BASE, subfolder="tokenizer") te = CLIPTextModel.from_pretrained( SD_BASE, subfolder="text_encoder", torch_dtype=torch.float32).to(device).eval() g = torch.Generator(device=device).manual_seed(SEED) lat, ehs, blobs, jpegs = [], [], [], [] buf_img, buf_txt = [], [] n_blob_empty = 0 @torch.no_grad() def flush(): if not buf_img: return px = torch.stack(buf_img).to(device) lat.append((vae.encode(px).latent_dist.sample(generator=g) * VAE_SCALE).cpu()) ehs.append(encode_nl77(buf_txt, tok, te, device).cpu()) buf_img.clear() buf_txt.clear() n = 0 for r in _iter_fused(N_TRAIN + N_VAL): img = r["image"] img = img.get("bytes") if isinstance(img, dict) else img im = Image.open(io.BytesIO(img)).convert("RGB").resize((512, 512)) buf_img.append(torch.from_numpy(np.asarray(im)).float() .permute(2, 0, 1) / 127.5 - 1.0) cap = r.get("caption_joycaption") or r.get("prompt_fused") or "" buf_txt.append(str(cap)) b = rasterize_blob(str(r.get("fused_json", ""))) if b.max() == 0: n_blob_empty += 1 blobs.append(torch.from_numpy(b)) if n < 24: jpegs.append(img) n += 1 if len(buf_img) == 16: flush() if n % 512 == 0: print(f"[cache] {n}/{N_TRAIN + N_VAL} " f"(blob-empty {n_blob_empty})", flush=True) flush() lat, ehs = torch.cat(lat), torch.cat(ehs) blobs = torch.stack(blobs) gv = torch.Generator().manual_seed(SEED + 1) torch.save({"lat": lat[:N_TRAIN], "ehs": ehs[:N_TRAIN], "blob": blobs[:N_TRAIN], "val_lat": lat[N_TRAIN:], "val_ehs": ehs[N_TRAIN:], "val_blob": blobs[N_TRAIN:], "val_noise": torch.randn(N_VAL, 4, 64, 64, generator=gv), "val_t": torch.randint(0, 1000, (N_VAL,), generator=gv), "judge_jpegs": jpegs, "blob_empty_count": n_blob_empty}, cache_f) print(f"[cache] built: {N_TRAIN}+{N_VAL} rows, blob-empty {n_blob_empty}", flush=True) del vae, te torch.cuda.empty_cache() return cache_f def run(device="cuda"): torch.cuda.set_per_process_memory_fraction(MEM_FRACTION, 0) os.makedirs(CKPT_DIR, exist_ok=True) acp = make_schedule(device) with ledger_run("dexp011 fused cache + blob targets", budget_h=1.5): build_cache(device) cache = torch.load(os.path.join(DATA_DIR, "cache.pt"), map_location="cpu", weights_only=True) def set_w(wraps, s01): w = band_weights(s01) for wr in wraps: wr.w_bands = w def loss_of(unet, wraps, lat, ehs, gen): bsz = lat.shape[0] drop = torch.rand(bsz, generator=gen, device=device) < CFG_DROPOUT ehs = ehs.clone() ehs[drop] = 0 t = torch.randint(0, 1000, (bsz,), generator=gen, device=device) set_w(wraps, t.float() / 1000.0) noise = torch.randn(lat.shape, generator=gen, device=device) pred = unet(add_noise(lat, noise, t, acp), t, ehs, return_dict=False)[0] return F.mse_loss(pred, noise) @torch.no_grad() def val(unet, wraps): tot, per_band = [], {0: [], 1: [], 2: []} for i in range(0, N_VAL, 32): lat = cache["val_lat"][i:i + 32].to(device) ehs = cache["val_ehs"][i:i + 32].to(device) noise = cache["val_noise"][i:i + 32].to(device) t = cache["val_t"][i:i + 32].to(device) set_w(wraps, t.float() / 1000.0) pred = unet(add_noise(lat, noise, t, acp), t, ehs, return_dict=False)[0] mse = ((pred - noise) ** 2).mean(dim=(1, 2, 3)) tot += mse.tolist() for j, tv in enumerate((t.float() / 1000.0).tolist()): per_band[band_of(tv)].append(mse[j].item()) return (sum(tot) / len(tot), {f"band{b}": round(sum(v) / max(len(v), 1), 6) for b, v in per_band.items()}) def train(unet, wraps, params, label): opt = torch.optim.Adam(params, lr=LR, weight_decay=0.0) gen = torch.Generator(device=device).manual_seed(SEED + 42) idx = torch.Generator().manual_seed(SEED + 7) t0 = time.time() for step in range(1, STEPS + 1): sel = torch.randint(0, N_TRAIN, (BATCH,), generator=idx) loss = loss_of(unet, wraps, cache["lat"][sel].to(device), cache["ehs"][sel].to(device), gen) loss.backward() opt.step() opt.zero_grad(set_to_none=True) if step == 50 or step % 500 == 0: print(f"[{label}] step {step}: loss {loss.item():.4f} | " f"{(time.time() - t0) / step:.2f}s/step", flush=True) return round((time.time() - t0) / STEPS, 3) results = {"config": {"dataset": DATASET, "steps": STEPS, "batch": BATCH, "seed": SEED, "cond": "caption_joycaption nl77", "blob_empty": int(cache["blob_empty_count"])}} with ledger_run("dexp011 frozen", budget_h=0.2) as h: unet = load_unet(device) v, pb = val(unet, []) results["frozen"] = {"val": v, "per_band": pb} del unet torch.cuda.empty_cache() h["verdict"] = f"val {v:.5f}" for label, mk, r_note in (("multiband3", lambda d: MultibandDelta(d), "3 band experts"), ("monolith", lambda d: MonoDelta(d), "r48")): with ledger_run(f"dexp011 {label} s{SEED}", budget_h=2.5) as h: unet = load_unet(device) mods, wraps = attach(unet, mk) n_params = sum(p.numel() for p in mods.parameters()) spd = train(unet, wraps, mods.parameters(), label) v, pb = val(unet, wraps) entry = {"n_params": n_params, "val": v, "per_band": pb, "s_per_step": spd} if label == "multiband3": lesions = {} for b in range(N_BANDS): for m in mods: m.band_enabled[b] = False _, pb_les = val(unet, wraps) lesions[f"lesion_band{b}"] = pb_les for m in mods: m.band_enabled[b] = True entry["lesions"] = lesions torch.save({"mods": [m.state_dict() for m in mods]}, os.path.join(CKPT_DIR, f"{label}_s{SEED}.pt")) results[label] = entry del unet, mods torch.cuda.empty_cache() h["verdict"] = f"val {v:.5f}" mb, mo = results["multiband3"], results["monolith"] lesion_sur = {} for b in range(N_BANDS): les = mb["lesions"][f"lesion_band{b}"] own = les[f"band{b}"] - mb["per_band"][f"band{b}"] oth = sum(les[f"band{o}"] - mb["per_band"][f"band{o}"] for o in range(N_BANDS) if o != b) / (N_BANDS - 1) lesion_sur[f"band{b}"] = {"own": round(own, 6), "other": round(oth, 6), "surgical": own > 3 * abs(oth) and own > 0} results["verdict"] = { "mb_vs_mono": "multiband" if mb["val"] <= mo["val"] else "monolith", "per_band_wins": {b: mb["per_band"][b] < mo["per_band"][b] for b in ("band0", "band1", "band2")}, "lesions": lesion_sur, "note": "REAL fused data (D2); blob targets cached for the 8am " "blob-loss design review; 1-seed CANDIDATE", } with open(os.path.join(DATA_DIR, "results.json" if SEED == 0 else f"results_s{SEED}.json"), "w") as f: json.dump(results, f, indent=2) note(f"dexp011: {json.dumps(results['verdict']['mb_vs_mono'])}") print(json.dumps(results["verdict"], indent=2)) burn_down() return results def smoke(): fj = json.dumps({"entities": [ {"mask": {"polygon": [[100, 100], [900, 100], [900, 900], [100, 900]]}}, {"mask": {"polygon": [[0, 0], [200, 0], [100, 200]]}}]}) b = rasterize_blob(fj) assert b.shape == (64, 64) and 0.5 < b.mean() < 0.9, b.mean() assert rasterize_blob("not json").max() == 0 assert rasterize_blob(json.dumps({"entities": []})).max() == 0 print("dexp011 smoke PASSED (blob rasterizer; GPU run is pod work)") if __name__ == "__main__": if "--run" in sys.argv: run() else: smoke()