"""Inference-only merged CVRR, preserving the reference strict decoding path. Native projections and dense learned projections are separate. No A/B adapter tensors or external backbone download is required after loading the package. """ from __future__ import annotations import contextlib import json import types from pathlib import Path from typing import TYPE_CHECKING import torch from torch import nn from safetensors.torch import load_file from transformers import PreTrainedModel from .configuration_cvrr_merged import CVRRMergedConfig # Transformers 4.57's local-directory loader copies direct relative imports; # enumerate transitive helpers without eagerly importing every backbone. if TYPE_CHECKING: from .source_splitting import SplitContext from .source_perceive import PerceiveDeliberateLayout from .source_spatial import SpatialVisualRecurrentCell from .source_helpers import _layer_hidden from .source_gemma import GemmaCVRR from .source_internvl import InternVLCVRR from .configuration_source_qwen25 import CloseQwen2_5_VLConfig from .configuration_source_qwen3 import CloseQwen3VLConfig from .modeling_source_qwen25 import CloseQwen2_5_VLForConditionalGeneration from .modeling_source_qwen3 import CloseQwen3VLForConditionalGeneration class NativeMergedLinear(nn.Module): def __init__(self, native, merged): super().__init__() if tuple(native.weight.shape) != tuple(merged.shape): raise ValueError('Native/merged projection shape mismatch') self.base = native self.register_buffer('merged_weight', merged.to(device=native.weight.device, dtype=torch.float32)) self.enabled = False self.dropout = nn.Identity() self.requires_grad_(False) def _apply(self, fn, recurse=True): # .to(device) is allowed; a caller's dtype cast must not truncate the # FP32 merged weights and then merely upcast already-lost precision. original = self.merged_weight super()._apply(fn, recurse=recurse) self.merged_weight = original.to(device=self.merged_weight.device, dtype=torch.float32) return self def forward(self, x): if not self.enabled: return self.base(x) bias = None if self.base.bias is None else self.base.bias.float() return torch.nn.functional.linear(x.float(), self.merged_weight, bias).to(x.dtype) @contextlib.contextmanager def merged_execution(model, enabled): modules = model._cvrr_merged_projections previous = [m.enabled for m in modules] for m in modules: m.enabled = bool(enabled) try: yield finally: for m, old in zip(modules, previous): m.enabled = old def set_merged_execution(model, enabled): for m in model._cvrr_merged_projections: m.enabled = bool(enabled) def install_merged(runtime, weights, *, generic): cell = runtime.layers[runtime.cell_index] if generic else runtime.text_model.layers[int(runtime.config.ell_star) + 1] modules = [] for key, weight in sorted(weights.items()): if not key.endswith('.weight'): raise ValueError(f'Unexpected merged state key: {key}') parts = key.removesuffix('.weight').split('.') parent = cell for part in parts[:-1]: parent = getattr(parent, part) old = getattr(parent, parts[-1]) native = old.base if generic else old.base_layer wrapped = NativeMergedLinear(native, weight) setattr(parent, parts[-1], wrapped) modules.append(wrapped) # Deliberately a tuple, not another ModuleList alias in state_dict. runtime._cvrr_merged_projections = tuple(modules) if generic: runtime.lora = {str(i): m for i, m in enumerate(modules)} runtime.adapters = types.MethodType(merged_execution, runtime) else: runtime._recurrence_adapter_execution = types.MethodType(merged_execution, runtime) runtime._set_adapters_enabled = types.MethodType(set_merged_execution, runtime) if any('lora_' in n for n, _ in runtime.named_parameters()): raise RuntimeError('Unmerged LoRA parameters remain') runtime.requires_grad_(False) return runtime.eval() class CVRRMergedModel(PreTrainedModel): config_class = CVRRMergedConfig main_input_name = 'input_ids' def __init__(self, config, runtime=None): super().__init__(config) if runtime is None: raise ValueError('Load a packaged release using from_pretrained') self.runtime = runtime self.is_generic = 'Qwen' not in config.release['name'] @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): if args: raise TypeError('Positional loader overrides are not supported') config = kwargs.pop('config', None) device = kwargs.pop('device_map', kwargs.pop('device', 'cpu')) if device is None: device = 'cpu' if isinstance(device, dict) or str(device) == 'auto': raise ValueError('Use one complete replica per device; pass device_map="cuda:0" or "cpu"') dtype = kwargs.pop('dtype', kwargs.pop('torch_dtype', torch.bfloat16)) if dtype not in (torch.bfloat16, 'bfloat16', 'auto', None): raise ValueError('This release preserves BF16 native weights and FP32 merged projections') offline = kwargs.pop('local_files_only', False) revision = kwargs.pop('revision', None) token = kwargs.pop('token', None) kwargs.pop('trust_remote_code', None) kwargs.pop('_from_auto', None) kwargs.pop('adapter_kwargs', None) kwargs.pop('name_or_path', None) kwargs.pop('_commit_hash', None) if kwargs: raise TypeError(f'Unsupported loader options: {sorted(kwargs)}') path = Path(pretrained_model_name_or_path) if not path.is_dir(): from huggingface_hub import snapshot_download path = Path(snapshot_download(str(pretrained_model_name_or_path), revision=revision, token=token, local_files_only=offline)) if config is None: config = CVRRMergedConfig.from_pretrained(path, local_files_only=True) settings = json.loads((path / 'cvrr_release_config.json').read_text()) if config.release != settings: raise ValueError('Root configuration and release metadata disagree') if settings['format'] != 'cvrr_native_plus_merged_transition_v1': raise ValueError('Unsupported release layout') native_path = str(path / 'native_backbone') generic = 'Qwen' not in settings['name'] common = dict(ell_star=settings['ell_star'], steps=settings['inference_T'], beta=settings['inference_beta'], rank=settings['lora_rank'], alpha=settings['lora_alpha'], dropout=settings['lora_dropout'], device=device, offline=True) if generic: if 'InternVL' in settings['name']: from .source_internvl import InternVLCVRR runtime = InternVLCVRR(native_path, **common) else: from .source_gemma import GemmaCVRR runtime = GemmaCVRR(native_path, **common) else: if 'Qwen2.5' in settings['name']: from transformers import Qwen2_5_VLForConditionalGeneration as Native from .configuration_source_qwen25 import CloseQwen2_5_VLConfig as SourceConfig from .modeling_source_qwen25 import CloseQwen2_5_VLForConditionalGeneration as Source else: from transformers import Qwen3VLForConditionalGeneration as Native from .configuration_source_qwen3 import CloseQwen3VLConfig as SourceConfig from .modeling_source_qwen3 import CloseQwen3VLForConditionalGeneration as Source source_config = SourceConfig.from_pretrained(path / 'source_config', local_files_only=True) source_config.num_workspace_steps = settings['inference_T'] source_config.counterfactual_beta = settings['inference_beta'] for flag in ('counterfactual_reliance_loss', 'counterfactual_step_retention_loss', 'functional_state_loss', 'interface_loss'): setattr(source_config, flag, False) native = Native.from_pretrained(native_path, dtype=torch.bfloat16, device_map=str(device), local_files_only=True, attn_implementation='sdpa') runtime = Source(source_config, backbone=native).to(device) weights = load_file(path / 'merged_transition.safetensors', device='cpu') runtime = install_merged(runtime, weights, generic=generic) model = cls(config, runtime).eval() model.release_path = path return model def prepare_inputs(self, image, question, *, max_visual_tokens=8192, max_tiles=12): """Single-image, batch-one preparation using the reference model prompts.""" if max_visual_tokens < 1: raise ValueError('max_visual_tokens must be positive') device = next(self.runtime.parameters()).device if 'InternVL' in self.config.release['name']: from .source_internvl import InternVLArrowCollator from .source_helpers import dynamic_tiles, _normalize_tiles helper = InternVLArrowCollator(self.runtime, max_tiles=max_tiles) tiles = dynamic_tiles(image, image_size=helper.image_size, max_tiles=max_tiles, thumbnail=helper.use_thumbnail) prompt = helper._query(question, '', len(tiles)) self.tokenizer = self.runtime.tokenizer encoded = dict(self.tokenizer(prompt, return_tensors='pt')) encoded.update(pixel_values=_normalize_tiles(tiles).to(torch.bfloat16), image_flags=torch.ones(len(tiles), 1, dtype=torch.long)) else: from transformers import AutoProcessor processor = AutoProcessor.from_pretrained(self.release_path / 'native_backbone', local_files_only=True, use_fast=('Qwen3' in self.config.release['name'])) self.tokenizer = processor.tokenizer if not self.is_generic: cfg = self.runtime.config patch = int(cfg.vision_config.patch_size) * int(cfg.vision_config.spatial_merge_size) pixels = max_visual_tokens * patch * patch ip = processor.image_processor if isinstance(getattr(ip, 'size', None), dict): ip.size = dict(ip.size, longest_edge=pixels) ip.max_pixels = pixels prompt = processor.apply_chat_template([{'role':'user', 'content':[ {'type':'image'}, {'type':'text','text':question}]}], tokenize=False, add_generation_prompt=True) encoded = dict(processor(text=[prompt], images=[image], return_tensors='pt')) if not self.is_generic: txt = processor.apply_chat_template([{'role':'user','content':[ {'type':'text','text':question}]}], tokenize=False, add_generation_prompt=True) q = processor.tokenizer(txt, return_tensors='pt', add_special_tokens=False) encoded['question_ids'] = q['input_ids'] encoded['question_attention_mask'] = q['attention_mask'] result = {k: v.to(device) for k,v in encoded.items() if isinstance(v, torch.Tensor)} if 'pixel_values' in result: result['pixel_values'] = result['pixel_values'].to(torch.bfloat16) if self.is_generic: count = int((self.runtime._modality(result).eq(1) & result['attention_mask'].bool()).sum()) else: count = int(result['input_ids'].eq(self.runtime.config.image_token_id).sum()) allowed = {'input_ids','attention_mask','pixel_values','image_grid_thw', 'question_ids','question_attention_mask'} result = {k:v for k,v in result.items() if k in allowed} if not 0 < count <= max_visual_tokens: raise ValueError(f'Visual token count {count} exceeds cap {max_visual_tokens} or is empty') return result def get_input_embeddings(self): if self.is_generic: return self.runtime.base_model.get_input_embeddings() return self.runtime.backbone.get_input_embeddings() def save_pretrained(self, *args, **kwargs): raise NotImplementedError('Keep the exported directory intact; generic HF reserialization would change the release layout') def forward(self, *args, **kwargs): raise NotImplementedError('Inference-only release: use next_token_logits or generate') @torch.inference_mode() def next_token_logits(self, **inputs): if self.is_generic: trace = self.runtime.extract(inputs) state = self.runtime.rollout(trace)[-1] return self.runtime.next_token_logits(trace, state) upper, _cache, _states, _evidence = self.runtime.prefill(**inputs) mask = inputs.get('question_attention_mask') last = (mask.long().sum(-1) - 1 if mask is not None else torch.full((upper.shape[0],), upper.shape[1]-1, device=upper.device)) hidden = upper[torch.arange(upper.shape[0], device=upper.device), last.long()] return self.runtime.backbone.lm_head(self.runtime.text_model.norm(hidden)).float() @torch.inference_mode() def generate(self, *, do_sample=False, max_new_tokens=32, **inputs): if do_sample: raise ValueError('Only deterministic greedy decoding is supported by this release') if self.is_generic: raise NotImplementedError('Gemma/InternVL release validation currently covers strict next-token readout only') return self.runtime.generate_answer(**inputs, max_new_tokens=max_new_tokens) CVRRMergedModel.register_for_auto_class('AutoModelForImageTextToText')