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