"""Caption an image with the trained Byrne-VLM (vision encoder + connector + LLM[+LoRA]). python generate.py --image photo.jpg --ckpt vlm_checkpoints/vlm_stage2_last.pt python generate.py --image photo.jpg --ckpt vlm_checkpoints/connector_last.pt # stage-1 only Loads the connector (and the Family/HRM-LoRA if the checkpoint has one), encodes the image to 196 tokens, and autoregressively decodes a caption. """ from __future__ import annotations import os os.environ.setdefault("USE_TF", "0") os.environ.setdefault("USE_FLAX", "0") import argparse import torch import torch.nn.functional as F from vlm_model import ByrneVLM, DEFAULT_IMAGE_TOKEN_ID from vision.dataset import load_image from anyres import anyres_tiles from spike_tokenizer import SpikeTokenizer from hrm_lora import apply_hrm_lora from family_lora import apply_family_lora def load_vlm(ckpt, lm_dir, vision_ckpt, device, anyres_grid=None): ck = torch.load(ckpt, map_location=device) cfg = ck.get("lora_cfg") new_vocab = (cfg or {}).get("new_vocab") vlm = ByrneVLM.from_pretrained(lm_dir, vision_ckpt, device=device, anyres_grid=anyres_grid, new_vocab=new_vocab) if cfg: # stage-2 checkpoint: rebuild the adapter, then load its weights if cfg.get("type", "hrm") == "family": apply_family_lora(vlm.lm, r=cfg["r"], steps=cfg["steps"], alpha=cfg["alpha"], n_routed=cfg["moe_routed"], n_shared=cfg["moe_shared"], top_k=cfg["moe_topk"]) else: apply_hrm_lora(vlm.lm, r=cfg["r"], steps=cfg["steps"], alpha=cfg["alpha"]) vlm.to(device) lora_sd = ck["hrm_lora"] own = dict(vlm.lm.named_parameters()) for n, p in lora_sd.items(): own[n].data.copy_(p.to(device)) vlm.connector.load_state_dict(ck["connector"]) vlm.eval() return vlm def _banned_by_no_repeat(out_ids, n): """Return the set of token ids that would complete a repeated n-gram.""" if n <= 0 or len(out_ids) < n - 1: return set() prefix = tuple(out_ids[-(n - 1):]) if n > 1 else () banned = set() for i in range(len(out_ids) - n + 1): if tuple(out_ids[i:i + n - 1]) == prefix: banned.add(out_ids[i + n - 1]) return banned @torch.no_grad() def caption(vlm, tok, image_path, device, max_new=48, temperature=0.0, top_k=20, repetition_penalty=1.0, no_repeat_ngram=0, letterbox=False): if vlm.anyres_grid: img = anyres_tiles(__import__("PIL.Image", fromlist=["Image"]).open(image_path), vlm.vision.cfg.image_size, vlm.anyres_grid).unsqueeze(0).to(device) else: img = load_image(image_path, vlm.vision.cfg.image_size, letterbox=letterbox).unsqueeze(0).to(device) img_embeds = vlm.encode_images(img) # [1, n_tokens, 640] (encode once) # KV-cached decode: a full recompute per token is O(n^2) -- minutes on CPU. # The Engram hashes the last up-to-3 RAW input embeds, which a single-token # step cannot see: it runs normally for the prefill, is disabled during # decode, and its per-token contribution is computed from a 3-wide window # of raw embeds -- verified token-identical to the full-recompute loop. embed = vlm.lm.model.embed_tokens eng_mod = vlm.lm.model.engram out = vlm.lm(inputs_embeds=img_embeds, use_cache=True) # prefill: Engram ON past, logits = out.past_key_values, out.logits[0, -1] recent = img_embeds[:, -2:, :] # n-gram window (raw embeds) out_ids = [] if eng_mod is not None: vlm.lm.model.engram = None try: for _ in range(max_new): # repetition penalty: down-weight tokens already produced (CTRL-style) if repetition_penalty != 1.0 and out_ids: idx = torch.tensor(list(set(out_ids)), device=device) sel = logits[idx] logits[idx] = torch.where(sel > 0, sel / repetition_penalty, sel * repetition_penalty) # no-repeat n-gram: hard-block tokens that would repeat an n-gram for b in _banned_by_no_repeat(out_ids, no_repeat_ngram): logits[b] = float("-inf") if temperature and temperature > 0: probs = F.softmax(logits / temperature, dim=-1) if top_k: v, _ = torch.topk(probs, top_k); probs[probs < v[-1]] = 0 nxt = torch.multinomial(probs, 1).item() else: nxt = int(logits.argmax()) if nxt == tok.eos_token_id: break out_ids.append(nxt) raw = embed(torch.tensor([[nxt]], device=device)) x = raw if eng_mod is not None: # exact Engram for the new token x = raw + eng_mod(torch.cat([recent, raw], dim=1))[:, -1:, :] out = vlm.lm(inputs_embeds=x, past_key_values=past, use_cache=True) past, logits = out.past_key_values, out.logits[0, -1] recent = torch.cat([recent, raw], dim=1)[:, -2:, :] finally: vlm.lm.model.engram = eng_mod # restore return tok.decode(out_ids) def main(): ap = argparse.ArgumentParser() ap.add_argument("--image", required=True) ap.add_argument("--ckpt", default="vlm_checkpoints/vlm_stage2_last.pt") ap.add_argument("--lm-dir", default=r"C:/Users/quaz/Desktop/benchmark_models/opd_v2") ap.add_argument("--vision-ckpt", default=r"C:/Users/quaz/Desktop/checkpoints/vision_self_best.pt") ap.add_argument("--max-new", type=int, default=48) ap.add_argument("--temperature", type=float, default=0.0) ap.add_argument("--repetition-penalty", type=float, default=1.0) ap.add_argument("--no-repeat-ngram", type=int, default=0) ap.add_argument("--letterbox", action="store_true", help="aspect-preserving pad (match training)") ap.add_argument("--tokenizer", default="tokenizer.json", help="use tokenizer_doctags.json for DocTags checkpoints") ap.add_argument("--anyres-grid", default=None, help="enable AnyRes tiling, e.g. '2x2' (must match how the ckpt was trained)") args = ap.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" grid = tuple(int(x) for x in args.anyres_grid.lower().split("x")) if args.anyres_grid else None tok = SpikeTokenizer(vocab_file=args.tokenizer) vlm = load_vlm(args.ckpt, args.lm_dir, args.vision_ckpt, device, anyres_grid=grid) cap = caption(vlm, tok, args.image, device, args.max_new, args.temperature, repetition_penalty=args.repetition_penalty, no_repeat_ngram=args.no_repeat_ngram, letterbox=args.letterbox) print(f"\n[{args.image}]\n {cap!r}") if __name__ == "__main__": main()