| """Record a reference encoding + reconstruction of the bundled images, for the stability tests. |
| |
| What is being pinned is the PICTURE, not the codes. Codes are an argmin over 4096 near-tied |
| codebook entries; fp16 flips about 0.5% of them against fp32. Measured, that costs 0.47/255 mean |
| absolute pixel difference -- against a reconstruction error of 10.25/255 -- and the differences sit |
| on high-frequency edges, where a code flip moves an edge by a fraction of a pixel and a single |
| pixel can swing 159/255 without anything semantic changing. So the reference stores both, and the |
| tests judge on mean pixel difference. |
| |
| JSON plus PNG, no binary blobs: a reference is something a human should be able to open and diff |
| when a new device disagrees. |
| |
| python -m satlink.tools.make_reference |
| """ |
| import argparse |
| import glob |
| import hashlib |
| import json |
| import os |
| import sys |
|
|
| import torch |
|
|
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) |
| from satlink.frame import VQFrame |
| from satlink.protocol import LinkProfile |
| from satlink.sat.encoder import SatEncoder |
|
|
|
|
|
|
| def _pixel_grid(model, prof): |
| f = model.patch_size * model.merge |
| return (max(f, round(prof.height / f) * f) // model.patch_size, |
| max(f, round(prof.width / f) * f) // model.patch_size) |
|
|
|
|
| def _teacher_forced_logprobs(model, frame, ids, pos, answer, device): |
| """-> the logprob the model assigns to each answer token, given the prompt and all earlier |
| answer tokens. One forward pass, deterministic, no decoding.""" |
| import torch |
| full = torch.cat([ids, torch.tensor(answer, device=device)]) |
| nxt = int(pos.max()) + 1 |
| tail = torch.arange(nxt, nxt + len(answer), device=device).expand(3, -1) |
| fpos = torch.cat([pos, tail], dim=1) |
| with torch.no_grad(), model.reference_attention(): |
| h = model(full, fpos, frame=frame) |
| logits = model.compute_logits(h[len(ids) - 1:-1]).float() |
| lp = torch.log_softmax(logits, -1) |
| return lp[torch.arange(len(answer)), torch.tensor(answer, device=device)].tolist() |
|
|
|
|
| def _record_answers(ref, a): |
| """Record, per image: the generated answer, and the logprob the model assigns to each of its |
| tokens under teacher forcing. |
| |
| The test uses BOTH. Teacher forcing is the stable measurement -- one forward pass, no decoding, |
| a continuous number that moves smoothly with numerical drift. Free generation is the one that |
| tells you what a user would actually see, but it is brittle by nature: greedy decoding turns |
| any near-tie into a different trajectory, so it is reported as a common-prefix length rather |
| than an equality check. |
| """ |
| import torch |
| from tokenizers import Tokenizer |
| from satlink.frame import VQFrame |
| from satlink.ground.infer import greedy |
| from satlink.ground.modeling_vqvlm import mrope_positions |
|
|
| if a.latent_ground: |
| from satlink.ground.modeling_latentvlm import LatentGlm4vForConditionalGeneration as M |
| model = M.from_pretrained(a.model, a.latent_ground, lora=a.lora, device=a.device) |
| tail, path = "<|assistant|>", "latent" |
| else: |
| from satlink.ground.modeling_vqvlm import VQGlm4vForConditionalGeneration as M |
| model = M.from_pretrained(a.model, a.decoder, LinkProfile(512, 512), device=a.device) |
| tail, path = "<|assistant|>\n", "pixel" |
| tok = Tokenizer.from_file(os.path.join(a.model, "tokenizer.json")) |
| IMG = model.image_token_id |
| ref["answers"] = {"path": path, "question": a.question, "items": {}} |
|
|
| for name, meta in sorted(ref["images"].items()): |
| prof = LinkProfile(meta["width"], meta["height"], n_tok=a.n_tok) |
| frame = VQFrame(torch.tensor(meta["codes"], dtype=torch.long), prof).to(a.device) |
| n_img = model.n_image_tokens(prof) if a.latent_ground else model.n_image_tokens() |
| gh, gw = model.image_grid(prof) if a.latent_ground else _pixel_grid(model, prof) |
| ids = torch.tensor( |
| tok.encode("[gMASK]<sop><|user|>\n<|begin_of_image|>", add_special_tokens=False).ids |
| + [IMG] * n_img |
| + tok.encode(f"<|end_of_image|>{a.question}{tail}", add_special_tokens=False).ids, |
| device=a.device) |
| pos = mrope_positions(ids, IMG, gh, gw, model.merge).to(a.device) |
| with model.reference_attention(): |
| ans = greedy(model, ids, pos, frame, max_new=a.answer_tokens) |
| lp = _teacher_forced_logprobs(model, frame, ids, pos, ans, a.device) |
| ref["answers"]["items"][name] = { |
| "tokens": ans, "logprobs": [round(x, 4) for x in lp], |
| "text": tok.decode(ans, skip_special_tokens=False)} |
| print(f" answer {name:24} {len(ans)} tok, mean logprob {sum(lp)/len(lp):.4f}") |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--assets", default="satlink/assets") |
| ap.add_argument("--sat", default="satlink/weights/sat_encoder.pt") |
| ap.add_argument("--decoder", default="satlink/weights/ground_decoder.pt") |
| ap.add_argument("--n-tok", type=int, default=48) |
| ap.add_argument("--model", default=None, |
| help="also record a reference ANSWER and its teacher-forced logprobs") |
| ap.add_argument("--latent-ground", default=None) |
| ap.add_argument("--lora", default=None) |
| ap.add_argument("--question", default="where is the largest building in the image") |
| ap.add_argument("--answer-tokens", type=int, default=96) |
| ap.add_argument("--device", default="cuda") |
| a = ap.parse_args() |
|
|
| from PIL import Image |
| import torchvision.transforms.functional as TF |
| from satlink.ground.decoder import GroundDecoder |
| gnd = GroundDecoder(a.decoder).to(a.device) |
|
|
| ref = {"n_tok": a.n_tok, "dtype": "float32", |
| "device": torch.cuda.get_device_name(0) if a.device.startswith("cuda") else "cpu", |
| "note": "codes recorded in fp32; tests compare RECONSTRUCTED PIXELS, not codes", |
| "images": {}} |
| os.makedirs(os.path.join(a.assets, "reference"), exist_ok=True) |
| for path in sorted(glob.glob(os.path.join(a.assets, "*.png"))): |
| img = Image.open(path).convert("RGB") |
| prof = LinkProfile(*img.size, n_tok=a.n_tok) |
| sat = SatEncoder(a.sat, prof, device=a.device, dtype=torch.float32) |
| codes = sat.encode(img) |
| with torch.no_grad(): |
| rec = gnd(VQFrame(codes.tokens, prof).to(a.device)).cpu() |
| name = os.path.basename(path) |
| rec_path = os.path.join(a.assets, "reference", name) |
| TF.to_pil_image(rec).save(rec_path, "PNG", optimize=True) |
| ref["images"][name] = { |
| "sha256": hashlib.sha256(open(path, "rb").read()).hexdigest()[:16], |
| "width": img.size[0], "height": img.size[1], "grid": list(prof.grid), |
| "tile_origins": [prof.tile_origin(i) for i in range(prof.n_tiles)], |
| "recon": f"reference/{name}", |
| "codes": codes.tolist(), |
| } |
| print(f" {name:26} {img.size[0]}x{img.size[1]:<5} {prof.n_tiles} crops -> " |
| f"{prof.frame_bytes} B, recon {os.path.getsize(rec_path)/1e3:.0f} KB") |
| if a.model: |
| _record_answers(ref, a) |
|
|
| out = os.path.join(a.assets, "reference.json") |
| json.dump(ref, open(out, "w"), indent=1) |
| print(f"-> {out} ({os.path.getsize(out)/1e6:.2f} MB, recorded on {ref['device']})") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
