Harness: KV-cached decode with exact engram window, verified token-identical to full recompute (which was O(n^2))

c1b69d3 4 months ago

6.86 kB

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

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