exp011_fused_multiband/dexp011_fused_multiband.py

"""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()