| """d1_exp000_baselines.py — exp000: the SD15-Lune zero-shot baseline wall.
|
|
|
| Every later adapter claim reads against these rows (claims discipline). For
|
| each frozen model {json-prompt ckpt-2500, base lune 18765}: generate from GT
|
| json_prompt CONDITIONAL vs SHUFFLED-json (derangement — no fixed points), CLIP-L
|
| image-image cosine original-vs-regen per arm. The cond-minus-shuffled gap is
|
| the scene-information gauge (pod/v35_expP1_sd15.py judge, contract-checked).
|
|
|
| Data: AbstractPhil/synthetic-object-relations-json (the json trainer's own
|
| dataset), parquet-direct, held-out tail rows, columns AUTODETECTED and printed
|
| (never assumed). Sampler: pod2/d1_lune_sampler.py (SHIFT=2.5 grid — the
|
| trainer-matched sampler of record). All judge features fp32.
|
|
|
| Local: python pod2/d1_exp000_baselines.py --smoke (parse only)
|
| Pod: python pod2/d1_exp000_baselines.py --run (GPU; ledgered)
|
| """
|
| from __future__ import annotations
|
|
|
| import io
|
| import json
|
| import os
|
| import sys
|
|
|
| sys.path[:0] = ["pod2", "."]
|
|
|
| import torch
|
|
|
| from pod_ledger import ledger_run, note
|
| from aleph_diffusion_core import derangement
|
| from d1_lune_sampler import encode_clip_225, flow_sample, decode
|
|
|
| DATASET = "AbstractPhil/synthetic-object-relations-json"
|
| SD_BASE = "stable-diffusion-v1-5/stable-diffusion-v1-5"
|
| MODELS = [
|
| ("json_ckpt2500", "AbstractPhil/sd15-flow-lune-json-prompt",
|
| "checkpoint-00002500/unet", "json_prompt"),
|
| ("base_lune", "AbstractPhil/sd15-flow-lune-flux",
|
| "flux_t2_6_pose_t4_6_port_t1_4/checkpoint-00018765/unet", "json_prompt"),
|
| ("json_vit", "AbstractPhil/sd15-flow-lune-json-vit",
|
| "auto", "vit_json_prompt"),
|
| ]
|
| N, STEPS, GUIDANCE, SEED = 24, 30, 6.0, 1234
|
| OUT_DIR = ("/workspace/data/dexp000" if os.path.isdir("/workspace")
|
| else os.path.join(os.environ.get("GEOLIP_DATA", "./data"),
|
| "dexp000"))
|
|
|
|
|
| def load_rows(n=N):
|
| """Parquet-direct held-out tail; autodetect image + json columns."""
|
| 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"))
|
| assert files, f"no parquet files in {DATASET}"
|
| path = hf_hub_download(DATASET, files[-1], repo_type="dataset")
|
| tbl = pq.read_table(path)
|
| cols = tbl.column_names
|
| print(f"[exp000] columns: {cols}", flush=True)
|
| img_col = next((c for c in cols if c in
|
| ("image", "img", "jpeg", "png", "image_bytes")), None)
|
| cond_cols = sorted({m[3] for m in MODELS})
|
| assert img_col and all(c in cols for c in cond_cols), \
|
| f"column autodetect failed (img={img_col}, need {cond_cols}) — inspect"
|
| rows = tbl.slice(max(0, tbl.num_rows - n), n).to_pylist()
|
| out = []
|
| for r in rows:
|
| img = r[img_col]
|
| if isinstance(img, dict):
|
| img = img.get("bytes")
|
| row = {"image_bytes": img}
|
| for c in cond_cols:
|
| v = r[c]
|
| row[c] = v if isinstance(v, str) else json.dumps(v)
|
| out.append(row)
|
| print(f"[exp000] {len(out)} held-out rows (tail of {files[-1]})",
|
| flush=True)
|
| return out
|
|
|
|
|
| def judge_selftest(clip_model, clip_proc, device):
|
| """Known-answer scorer check (silent-zero law): img vs itself == 1,
|
| img vs noise well below."""
|
| import numpy as np
|
| from PIL import Image
|
| rng = np.random.default_rng(0)
|
| a = Image.fromarray(rng.integers(0, 255, (256, 256, 3), dtype=np.uint8))
|
| b = Image.fromarray(rng.integers(0, 255, (256, 256, 3), dtype=np.uint8))
|
|
|
| def feat(im):
|
| with torch.no_grad():
|
| inp = clip_proc(images=im, return_tensors="pt").to(device)
|
| out = clip_model.get_image_features(**inp)
|
| if not torch.is_tensor(out):
|
| out = getattr(out, "image_embeds", None) \
|
| if getattr(out, "image_embeds", None) is not None \
|
| else out.pooler_output
|
| f = out.float()
|
| return torch.nn.functional.normalize(f, dim=-1)
|
|
|
| same = (feat(a) @ feat(a).T).item()
|
| diff = (feat(a) @ feat(b).T).item()
|
| assert same > 0.999, f"judge self-test: self-cos {same}"
|
| assert diff < same, "judge self-test: noise pair >= self pair"
|
| print(f"[exp000] judge self-test ok (self {same:.4f}, noise {diff:.4f})",
|
| flush=True)
|
| return feat
|
|
|
|
|
| def run(device="cuda"):
|
| os.makedirs(OUT_DIR, exist_ok=True)
|
| import numpy as np
|
| from PIL import Image
|
| from diffusers import UNet2DConditionModel, AutoencoderKL
|
| from transformers import (CLIPTextModel, CLIPTokenizer, CLIPModel,
|
| CLIPProcessor)
|
|
|
| rows = load_rows()
|
| perm = derangement(len(rows), seed=SEED)
|
|
|
| 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()
|
| 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)
|
|
|
| orig_feats = torch.cat([
|
| feat(Image.open(io.BytesIO(r["image_bytes"])).convert("RGB"))
|
| for r in rows])
|
|
|
| def resolve_sub(repo, sub):
|
| if sub != "auto":
|
| return sub
|
| from huggingface_hub import HfApi
|
| cks = sorted({f.split("/")[0] for f in
|
| HfApi().list_repo_files(repo)
|
| if f.startswith("checkpoint-")})
|
| assert cks, f"no checkpoints in {repo}"
|
| return f"{cks[-1]}/unet"
|
|
|
| results = {}
|
| for label, repo, sub, cond_col in MODELS:
|
| prompts = [r[cond_col] for r in rows]
|
| shuffled = [prompts[i] for i in perm.tolist()]
|
| with torch.no_grad():
|
| ehs_cond = encode_clip_225(prompts, tok, te, device)
|
| ehs_shuf = encode_clip_225(shuffled, tok, te, device)
|
| with ledger_run(f"dexp000 baseline {label}", budget_h=1.0) as h:
|
| unet = UNet2DConditionModel.from_pretrained(
|
| repo, subfolder=resolve_sub(repo, sub),
|
| torch_dtype=torch.float16).to(device)
|
| unet.eval()
|
| arms = {}
|
| for arm, ehs in (("cond", ehs_cond), ("shuffled", ehs_shuf)):
|
| imgs = []
|
| for i in range(0, len(prompts), 6):
|
| lat = flow_sample(unet, ehs[i:i + 6].half(),
|
| n_steps=STEPS, guidance=GUIDANCE,
|
| seed=SEED + i, device=device)
|
| imgs.extend(decode(vae, lat.float()))
|
| pil = [Image.fromarray((im * 255).astype(np.uint8))
|
| for im in imgs]
|
| f = torch.cat([feat(p) for p in pil])
|
| cos = (f * orig_feats).sum(-1)
|
| arms[arm] = {"mean_cos": round(cos.mean().item(), 4),
|
| "per_row": [round(c, 4) for c in cos.tolist()]}
|
| for j, p in enumerate(pil[:6]):
|
| p.save(os.path.join(OUT_DIR, f"{label}_{arm}_{j}.png"))
|
| gap = arms["cond"]["mean_cos"] - arms["shuffled"]["mean_cos"]
|
| results[label] = {**arms, "cond_minus_shuffled": round(gap, 4)}
|
| h["verdict"] = f"cond-shuf gap {gap:+.4f}"
|
| del unet
|
| torch.cuda.empty_cache()
|
| note(f"exp000 {label}: {json.dumps(results[label]['cond_minus_shuffled'])}")
|
|
|
| with open(os.path.join(OUT_DIR, "results.json"), "w") as f:
|
| json.dump({"n": N, "steps": STEPS, "guidance": GUIDANCE, "seed": SEED,
|
| "sampler": "shifted SHIFT=2.5 (trainer-matched)",
|
| "models": results}, f, indent=2)
|
| print(json.dumps(results, indent=2))
|
| return results
|
|
|
|
|
| def smoke():
|
| p = derangement(N, seed=SEED)
|
| assert not (p == torch.arange(N)).any()
|
| assert len(MODELS) == 3 and all(len(m) == 4 for m in MODELS) and OUT_DIR
|
| print("d1_exp000_baselines smoke PASSED (parse + derangement; GPU run is "
|
| "pod work)")
|
|
|
|
|
| if __name__ == "__main__":
|
| if "--run" in sys.argv:
|
| run()
|
| else:
|
| smoke()
|
|
|
