| """DiffusionGemma one-denoise decision readout and decoder-only LoRA. | |
| This module is included in the export as the inference helper. It does not | |
| merge adapters, change base weight sharing, or interpret scores as confidence. | |
| """ | |
| import json | |
| import math | |
| from pathlib import Path | |
| import random | |
| import time | |
| import torch | |
| from PIL import Image | |
| from peft import PeftModel | |
| from transformers import AutoConfig, AutoProcessor, DiffusionGemmaForBlockDiffusion | |
| from huggingface_hub import snapshot_download | |
| from quantized import RECIPE, enable_sequential_vision, load_streamed_int8 | |
| from imaging import image_key, VisionFeatureCache | |
| from core import MODEL_ID, MODEL_REVISION, SEED, prompt_spec, stable_seed, public_identity | |
| MAX_INPUT = 8192 | |
| VISUAL_TOKENS = 1120 | |
| SCAFFOLD = "<|channel>thought\n<channel|>\n" | |
| def raw_model(model): | |
| return model.get_base_model() if isinstance(model, PeftModel) else model | |
| def tie_pointers(model): | |
| m = raw_model(model) | |
| def weight(layer): | |
| return layer.base_layer.weight if hasattr(layer, "base_layer") else layer.weight | |
| return [(weight(e.self_attn.q_proj).data_ptr(), weight(d.self_attn.q_proj).data_ptr()) | |
| for e,d in zip(m.model.encoder.language_model.layers, m.model.decoder.layers)] | |
| def load_model(adapter_dir=None, device="cuda:0", precision=None): | |
| if adapter_dir: | |
| settings_path = Path(adapter_dir)/"inference_config.json" | |
| if not settings_path.exists(): | |
| raise ValueError("Use the complete exported adapter package, including inference_config.json") | |
| settings = json.loads(settings_path.read_text()) | |
| if settings["model"] != MODEL_ID or settings["revision"] != MODEL_REVISION: | |
| raise ValueError("Adapter base model/revision mismatch") | |
| if precision is not None and precision != settings["weights"]: | |
| raise ValueError("Requested precision differs from the adapter configuration") | |
| precision = settings["weights"] | |
| if precision == "experts-int8" and settings.get("quantization_recipe") != RECIPE: | |
| raise ValueError("Adapter quantization recipe differs from this inference helper") | |
| precision = precision or "experts-int8" | |
| image_policy = settings.get("image_policy", "original1120") if adapter_dir else "original1120" | |
| if image_policy != "original1120": | |
| raise ValueError("Unsupported exported image policy") | |
| dtype = torch.bfloat16 if device.startswith("cuda") else torch.float32 | |
| if precision == "experts-int8": | |
| directory = snapshot_download(MODEL_ID, revision=MODEL_REVISION, | |
| allow_patterns=["config.json", "model*.safetensors", "model.safetensors.index.json"], max_workers=2) | |
| config = AutoConfig.from_pretrained(directory, trust_remote_code=False) | |
| config._name_or_path = MODEL_ID | |
| model = load_streamed_int8(directory, config, device=device, dtype=dtype) | |
| elif precision == "bf16": | |
| model = DiffusionGemmaForBlockDiffusion.from_pretrained( | |
| MODEL_ID, revision=MODEL_REVISION, dtype=dtype, | |
| device_map={"": device}, low_cpu_mem_usage=True, | |
| trust_remote_code=False, attn_implementation="sdpa") | |
| else: | |
| raise ValueError(f"Unsupported base precision: {precision}") | |
| model.requires_grad_(False) | |
| enable_sequential_vision(model) | |
| if any(a != b for a,b in tie_pointers(model)): | |
| raise RuntimeError("Expected pretrained encoder/decoder weight sharing is absent") | |
| processor = AutoProcessor.from_pretrained( | |
| str(adapter_dir) if adapter_dir else MODEL_ID, | |
| **({} if adapter_dir else {"revision": MODEL_REVISION}), trust_remote_code=False) | |
| processor.image_processor.max_soft_tokens = VISUAL_TOKENS | |
| if adapter_dir: | |
| model = PeftModel.from_pretrained(model, str(adapter_dir), is_trainable=False) | |
| runner = DecisionModel(model.eval(), processor, device) | |
| runner.precision = precision | |
| runner.image_policy = image_policy | |
| return runner | |
| class DecisionModel: | |
| def __init__(self, model, processor, device): | |
| self.model = model | |
| self.processor = processor | |
| self.tokenizer = processor.tokenizer | |
| self.device = torch.device(device) | |
| self.vocab = raw_model(model).config.text_config.vocab_size | |
| self.last_input_tokens = 0 | |
| self.last_image_tokens = 0 | |
| self.image_policy = "original1120" | |
| self.feature_cache = None | |
| self.last_image_geometry = [] | |
| def enable_feature_cache(self, directory, namespace, max_bytes=16*2**30): | |
| self.feature_cache = VisionFeatureCache(directory, namespace, max_bytes) | |
| def canvas(self, spec, noise_seed): | |
| tok = self.tokenizer | |
| head = tok.encode(SCAFFOLD, add_special_tokens=False) | |
| def answer(labels): | |
| return head + tok.encode("\n".join(f"q{i+1}: {label}" for i,label in enumerate(labels)), add_special_tokens=False) | |
| base = answer(["A"] * len(spec["questions"])) | |
| positions, allowed_ids = [], [] | |
| for qi, q in enumerate(spec["questions"]): | |
| ids, position = [], None | |
| for j in range(len(q["options"])): | |
| labels = ["A"] * len(spec["questions"]) | |
| labels[qi] = chr(65+j) | |
| variant = answer(labels) | |
| differences = [i for i,(a,b) in enumerate(zip(base, variant)) if a != b] | |
| if j: | |
| if len(variant) != len(base) or len(differences) != 1: | |
| raise RuntimeError("Answer label is not a single stable token slot") | |
| if position is not None and differences[0] != position: | |
| raise RuntimeError("Answer labels occupy inconsistent positions") | |
| position = differences[0] | |
| ids.append(variant) | |
| if position is None: | |
| raise RuntimeError("Need at least two answer options") | |
| label_ids = [v[position] for v in ids] | |
| if len(set(label_ids)) != len(label_ids): | |
| raise RuntimeError("Duplicate answer token IDs") | |
| positions.append(position); allowed_ids.append(label_ids) | |
| close = tok.convert_tokens_to_ids("<turn|>") | |
| if close != 106 or tok.pad_token_id != 0: | |
| raise RuntimeError("Pinned OpenJev token convention changed") | |
| width = math.ceil((len(base)+1)/16)*16 | |
| if width > 256: | |
| raise RuntimeError("Answer scaffold exceeds model canvas") | |
| canvas = base + [close] + [tok.pad_token_id]*(width-len(base)-1) | |
| rng = random.Random(noise_seed) | |
| for p in positions: | |
| canvas[p] = rng.randrange(self.vocab) | |
| return torch.tensor([canvas], device=self.device), positions, allowed_ids | |
| def prepare(self, public, root): | |
| spec = prompt_spec(public) | |
| images, keys, geometry = [], [], [] | |
| for im in public["images"]: | |
| path = Path(root)/im["path"] | |
| with Image.open(path) as source: | |
| rgb, transform = source.convert("RGB"), {"policy": "original1120", "original_wh": list(source.size)} | |
| images.append(rgb) | |
| geometry.append(transform) | |
| if self.feature_cache is not None: | |
| keys.append(image_key(path, transform)) | |
| self.last_image_geometry = geometry | |
| content = [{"type": "image"} for _ in images] + [{"type": "text", "text": spec["state"]}] | |
| messages = [{"role": "system", "content": spec["system"]}, {"role": "user", "content": content}] | |
| text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False) | |
| inputs = self.processor(text=[text], images=[images], return_tensors="pt", truncation=False, | |
| images_kwargs={"max_soft_tokens": VISUAL_TOKENS, "do_resize": self.image_policy == "original1120"}) | |
| for im in images: | |
| im.close() | |
| self.last_input_tokens = inputs["input_ids"].shape[-1] | |
| self.last_image_tokens = int((inputs["input_ids"] == raw_model(self.model).config.image_token_id).sum()) | |
| if self.last_input_tokens > MAX_INPUT: | |
| raise RuntimeError(f"Input has {self.last_input_tokens} tokens, limit {MAX_INPUT}; no truncation permitted") | |
| if not self.last_image_tokens or "pixel_values" not in inputs or "mm_token_type_ids" not in inputs: | |
| raise RuntimeError("Processor did not supply images and their vision-attention mask") | |
| dtype = raw_model(self.model).model.encoder.vision_tower.dtype | |
| inputs = {k: v.to(device=self.device, dtype=dtype if v.is_floating_point() else v.dtype) for k,v in inputs.items()} | |
| encoder = raw_model(self.model).model.encoder | |
| encoder._jev_feature_cache = self.feature_cache | |
| encoder._jev_image_keys = keys if self.feature_cache is not None else None | |
| try: | |
| with torch.no_grad(): | |
| encoded = encoder(**inputs, return_dict=True) | |
| finally: | |
| encoder._jev_image_keys = None | |
| cache = encoded.past_key_values | |
| del encoded, inputs | |
| canvas, positions, ids = self.canvas(spec, stable_seed(SEED, "eval-public-v3", public_identity(public, root))) | |
| return {"cache": cache, "canvas": canvas, "positions": positions, "allowed_ids": ids, "spec": spec} | |
| def hidden(self, prepared): | |
| return raw_model(self.model).model.decoder( | |
| decoder_input_ids=prepared["canvas"], past_key_values=prepared["cache"], | |
| self_conditioning_logits=None, return_dict=True).last_hidden_state | |
| def slot_logits(self, prepared): | |
| model = raw_model(self.model) | |
| h = self.hidden(prepared)[0] | |
| cap = model.final_logit_softcapping | |
| result = [] | |
| for p, ids in zip(prepared["positions"], prepared["allowed_ids"]): | |
| weights = model.lm_head.weight[torch.tensor(ids, device=h.device)] | |
| logits = torch.nn.functional.linear(h[p].to(weights.dtype), weights).float() | |
| logits = torch.tanh(logits/cap)*cap | |
| result.append(logits) | |
| return result | |
| def predict(self, public, root): | |
| self.model.eval() | |
| if self.device.type == "cuda": | |
| torch.cuda.synchronize(self.device) | |
| start = time.perf_counter() | |
| cache_before = self.feature_cache.stats() if self.feature_cache else None | |
| with torch.no_grad(): | |
| prepared = self.prepare(public, root) | |
| logits = self.slot_logits(prepared) | |
| choices = [q["options"][int(l.argmax())][0] for q,l in zip(prepared["spec"]["questions"], logits)] | |
| support = [{k: float(v) for (k,_),v in zip(q["options"], l.softmax(-1).cpu())} | |
| for q,l in zip(prepared["spec"]["questions"], logits)] | |
| if self.device.type == "cuda": | |
| torch.cuda.synchronize(self.device) | |
| result = {"id": public["id"], "latency_s": time.perf_counter()-start, | |
| "input_tokens": self.last_input_tokens, "image_tokens": self.last_image_tokens, | |
| "semantic_support": support, "image_policy": self.image_policy} | |
| if self.feature_cache: | |
| now = self.feature_cache.stats() | |
| result["vision_cache"] = {k: now[k]-cache_before[k] for k in ("hits", "misses")} | |
| if public["task"] == "known_ui_differences": | |
| result["categories"] = {q["tag"]: c for q,c in zip(prepared["spec"]["questions"], choices)} | |
| else: | |
| result["choice"] = choices[0] | |
| return result | |
| if __name__ == "__main__": | |
| import argparse | |
| parser = argparse.ArgumentParser(description="One-step Jev decisions; probabilities are semantic support, not calibrated correctness.") | |
| parser.add_argument("--adapter", required=True) | |
| parser.add_argument("--input", required=True, help="One public dataset-style input JSON") | |
| parser.add_argument("--image-root", required=True) | |
| args = parser.parse_args() | |
| runner = load_model(args.adapter) | |
| print(json.dumps(runner.predict(json.loads(Path(args.input).read_text()), args.image_root), indent=2)) | |