"""Build a runnable DK-VLM from a released checkpoint and answer questions with it. A checkpoint directory contains three things: ``language_model_lora/`` LoRA adapter for the Qwen2.5-VL or Qwen3-VL text tower ``geo_branch_state.pt`` geometry-branch gates and projections ``dkvlm_config.json`` which layers carry the branch, and its slice widths (optional; sensible defaults are used when absent) The legacy filename ``p2vlm_config.json`` is also accepted for older checkpoints. Base Qwen2.5-VL or Qwen3-VL weights are downloaded separately from the upstream model hub, and VGGT-Omega weights come from its own release. Nothing here fetches them for you. """ from __future__ import annotations import json from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Union import torch from PIL import Image from transformers import AutoConfig, AutoProcessor, Qwen3VLForConditionalGeneration from .frames import ( SampledFrames, image_content_elements, load_image_frames, processor_visual_grid, sample_video_frames, ) from .geometry import VGGTOmegaConfig, VGGTOmegaEngine from .qwen3_geo import DKVLMGeoWrapper, install_geo_attention from .tokens import DEFAULT_METRIC_GLOBAL_SCALE, build_geo_coords DEFAULT_GEO_LAYERS = (6, 12, 18, 27, 33) DEFAULT_TOKEN_SE3_DIM = 48 DEFAULT_POSE_GTA_DIM = 24 DEFAULT_LAMBDA_MIN = 0.25 DEFAULT_LAMBDA_MAX = 4.0 DEFAULT_NFRAMES = 32 DEFAULT_MAX_PIXELS = 786432 DEFAULT_MAX_NEW_TOKENS = 128 def _load_vl_base(path: Path, attn_implementation: str, dtype: torch.dtype): """Load Qwen3-VL or Qwen2.5-VL from ``path`` according to its config.""" model_type = str(getattr(AutoConfig.from_pretrained(str(path)), "model_type", "")) if model_type == "qwen3_vl": cls = Qwen3VLForConditionalGeneration elif model_type == "qwen2_5_vl": from transformers import Qwen2_5_VLForConditionalGeneration cls = Qwen2_5_VLForConditionalGeneration else: raise ValueError( f"unsupported model_type={model_type!r} at {path}; " "DK-VLM supports qwen3_vl and qwen2_5_vl" ) return cls.from_pretrained(str(path), attn_implementation=attn_implementation, dtype=dtype) def _layers_from_geo_state(state: Dict[str, torch.Tensor]) -> List[int]: return sorted( { int(key.split(".")[1]) for key in state if key.startswith("geo_gate.") and key.count(".") == 1 } ) def resolve_attn_implementation(preferred: str) -> str: """Prefer Flash Attention 2, then SDPA. ``eager`` is only used if requested.""" wanted = (preferred or "flash_attention_2").strip().lower() if wanted == "eager": return "eager" if wanted == "flash_attention_2": try: import flash_attn # noqa: F401 return "flash_attention_2" except Exception: return "sdpa" if wanted == "sdpa": return "sdpa" raise ValueError(f"unknown attn_implementation={preferred!r}") PRE_PROMPT = "These are frames of a video." MCA_POST_PROMPT = "Answer with the option's letter from the given choices directly." NA_POST_PROMPT = "Please answer the question using a single word or phrase." def build_prompt(question: str, options: Optional[Sequence[str]] = None) -> str: """Compose the question the way the model was trained to read it.""" if options: return "\n".join( [PRE_PROMPT, question, "Options:\n" + "\n".join(options), MCA_POST_PROMPT] ) return PRE_PROMPT + "\n" + question + "\n" + NA_POST_PROMPT @dataclass class GeoBranchConfig: """Shape of the geometry branch. Must match the checkpoint being loaded.""" target_layers: List[int] = field(default_factory=lambda: list(DEFAULT_GEO_LAYERS)) geo_branch_mrope: bool = False token_se3_dim: int = DEFAULT_TOKEN_SE3_DIM pose_gta_dim: int = DEFAULT_POSE_GTA_DIM lambda_min: float = DEFAULT_LAMBDA_MIN lambda_max: float = DEFAULT_LAMBDA_MAX metric_global_scale: float = DEFAULT_METRIC_GLOBAL_SCALE normalize_mode: str = "metric" @classmethod def from_checkpoint(cls, checkpoint_dir: Path) -> "GeoBranchConfig": path = Path(checkpoint_dir) / "dkvlm_config.json" if not path.is_file(): path = Path(checkpoint_dir) / "p2vlm_config.json" if not path.is_file(): return cls() raw = json.loads(path.read_text(encoding="utf-8")) cfg = cls() if "target_layers" in raw: cfg.target_layers = [int(x) for x in raw["target_layers"]] if "geo_branch_mrope" in raw: cfg.geo_branch_mrope = bool(raw["geo_branch_mrope"]) slice_pair = raw.get("token_se3_slice") if slice_pair: cfg.token_se3_dim = int(slice_pair[1]) - int(slice_pair[0]) slice_pair = raw.get("pose_gta_slice") if slice_pair: cfg.pose_gta_dim = int(slice_pair[1]) - int(slice_pair[0]) lambdas = raw.get("token_se3_lambdas") if lambdas: cfg.lambda_min = float(min(lambdas)) cfg.lambda_max = float(max(lambdas)) for key in ("metric_global_scale", "normalize_mode"): if key in raw: setattr(cfg, key, type(getattr(cfg, key))(raw[key])) return cfg class DKVLM: """A loaded DK-VLM: Qwen-VL with the geometry branch, plus the geometry front-end.""" def __init__( self, model: DKVLMGeoWrapper, processor: AutoProcessor, geo_cfg: GeoBranchConfig, geometry_engine: Optional[VGGTOmegaEngine] = None, *, nframes: int = DEFAULT_NFRAMES, max_pixels: Optional[int] = DEFAULT_MAX_PIXELS, max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS, ): self.model = model self.processor = processor self.geo_cfg = geo_cfg self.geometry_engine = geometry_engine self.nframes = int(nframes) self.max_pixels = max_pixels self.max_new_tokens = int(max_new_tokens) self.patch_size, self.merge_size, self.cell_size = processor_visual_grid(processor) @property def device(self) -> torch.device: return next(self.model.base.parameters()).device def compute_geometry(self, sampled: SampledFrames, sample_id: str = ""): """Run VGGT-Omega over the sampled frames.""" if self.geometry_engine is None: raise RuntimeError( "no geometry engine configured; construct DKVLM with geometry_ckpt=... " "or pass use_geometry=False to run DK-VLM without the geometric branch" ) return self.geometry_engine.infer( pil_images=sampled.frames, source_frame_names=sampled.frame_names, sample_id=sample_id, ) def build_inputs( self, sampled: SampledFrames, prompt: str, *, geometry=None, sample_id: str = "", ) -> Dict[str, Any]: """Tokenize frames plus prompt, and attach per-token geometry features.""" from qwen_vl_utils import process_vision_info messages = [ { "role": "user", "content": image_content_elements(sampled.frames, max_pixels=self.max_pixels) + [{"type": "text", "text": prompt}], } ] prompt_text = self.processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) image_inputs, video_inputs = process_vision_info( messages, image_patch_size=self.patch_size ) inputs = self.processor( text=[prompt_text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt", ) device = self.device inputs = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()} if geometry is None: return inputs cfg = self.model.base.config image_grid_thw = inputs.get("image_grid_thw") if image_grid_thw is None or int(image_grid_thw.numel()) == 0: raise ValueError( "expected image_grid_thw; DK-VLM sends frames as independent images" ) if int(image_grid_thw.shape[0]) != len(sampled): raise ValueError( f"image grid rows {int(image_grid_thw.shape[0])} != sampled frames {len(sampled)}" ) geo_coords, geo_mask, stats = build_geo_coords( geometry, inputs["input_ids"][0].cpu(), image_grid_thw.cpu(), spatial_merge_size=self.merge_size, patch_size=self.patch_size, image_token_id=int(cfg.image_token_id), video_token_id=int(cfg.video_token_id), normalize_mode=self.geo_cfg.normalize_mode, metric_global_scale=self.geo_cfg.metric_global_scale, sample_id=sample_id, ) inputs = dict(inputs) inputs["geo_coords"] = geo_coords.unsqueeze(0).to(device=device, dtype=torch.float32) inputs["geo_coords_mask"] = geo_mask.unsqueeze(0).to(device=device, dtype=torch.bool) self.last_geometry_stats = stats return inputs @torch.inference_mode() def generate(self, inputs: Dict[str, Any], **gen_kwargs) -> str: """Greedy-decode an answer and return the newly generated text.""" kwargs = { "max_new_tokens": self.max_new_tokens, "do_sample": False, "num_beams": 1, } kwargs.update(gen_kwargs) generated = self.model.generate(**inputs, **kwargs) trimmed = generated[0][len(inputs["input_ids"][0]) :] return self.processor.decode(trimmed, skip_special_tokens=True) def answer( self, source: Union[str, Path, SampledFrames, Sequence[Image.Image]], question: str, *, options: Optional[Sequence[str]] = None, use_geometry: bool = True, sample_id: str = "", **gen_kwargs, ) -> str: """Answer a question about a video, a directory of frames, or decoded frames.""" sampled = self.resolve_frames(source) geometry = self.compute_geometry(sampled, sample_id) if use_geometry else None inputs = self.build_inputs( sampled, build_prompt(question, options), geometry=geometry, sample_id=sample_id ) return self.generate(inputs, **gen_kwargs) def resolve_frames( self, source: Union[str, Path, SampledFrames, Sequence[Image.Image]] ) -> SampledFrames: """Normalize any accepted input into :class:`SampledFrames`.""" if isinstance(source, SampledFrames): return source if isinstance(source, (str, Path)): path = Path(source) if path.is_dir(): files = sorted( p for p in path.iterdir() if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".bmp", ".webp"} ) if not files: raise FileNotFoundError(f"no image files in {path}") return load_image_frames(files) if not path.is_file(): raise FileNotFoundError(f"video not found: {path}") return sample_video_frames(path, self.nframes) frames = list(source) if not frames: raise ValueError("empty frame sequence") return SampledFrames( frames=[f.convert("RGB") for f in frames], frame_indices=list(range(len(frames))), frame_names=[f"frame_{i:06d}.png" for i in range(len(frames))], ) def load_dkvlm( checkpoint_dir: Union[str, Path], *, base_model_path: Union[str, Path], geometry_ckpt: str = "", device: Optional[Union[str, torch.device]] = None, torch_dtype: Optional[torch.dtype] = None, attn_implementation: str = "flash_attention_2", geo_attention_backend: str = "sdpa", nframes: int = DEFAULT_NFRAMES, max_pixels: Optional[int] = DEFAULT_MAX_PIXELS, use_geometry: bool = True, verbose: bool = True, ) -> DKVLM: """Assemble base model, LoRA adapter, geometry branch and geometry front-end.""" ckpt = Path(checkpoint_dir).resolve() if not ckpt.is_dir(): raise FileNotFoundError(f"checkpoint directory not found: {ckpt}") lora_dir = ckpt / "language_model_lora" if not (lora_dir / "adapter_config.json").is_file(): parent = ckpt.parent / "language_model_lora" lora_dir = parent if (parent / "adapter_config.json").is_file() else lora_dir if not (lora_dir / "adapter_config.json").is_file(): raise FileNotFoundError( f"missing LoRA adapter_config.json in {lora_dir}; " "DK-VLM requires its released adapter and cannot fall back to base weights" ) if not any((lora_dir / name).is_file() for name in ( "adapter_model.safetensors", "adapter_model.bin", "adapter_model.safetensors.index.json", "adapter_model.bin.index.json", )): raise FileNotFoundError(f"missing LoRA adapter weights in {lora_dir}") dev = torch.device(device) if device is not None else torch.device( "cuda" if torch.cuda.is_available() else "cpu" ) dtype = torch_dtype or (torch.bfloat16 if dev.type == "cuda" else torch.float32) attn_impl = resolve_attn_implementation(attn_implementation) if verbose and attn_impl != attn_implementation: print(f"[dkvlm] attn_implementation {attn_implementation} -> {attn_impl}") base = _load_vl_base(Path(base_model_path), attn_impl, dtype).to(dev) base.eval() from peft import PeftModel base.model.language_model = PeftModel.from_pretrained( base.model.language_model, str(lora_dir), is_trainable=False ) if verbose: print(f"[dkvlm] loaded LoRA adapter from {lora_dir}") proc_src = ckpt if (ckpt / "preprocessor_config.json").is_file() else Path(base_model_path) processor = AutoProcessor.from_pretrained(str(proc_src), use_fast=False) if verbose: print(f"[dkvlm] processor from {proc_src}") geo_state_path = ckpt / "geo_branch_state.pt" if not geo_state_path.is_file(): parent = ckpt.parent / "geo_branch_state.pt" if not parent.is_file(): raise FileNotFoundError(f"missing geo_branch_state.pt in {ckpt}") geo_state_path = parent geo_state = torch.load(str(geo_state_path), map_location="cpu", weights_only=True) geo_cfg = GeoBranchConfig.from_checkpoint(ckpt) if not any((ckpt / name).is_file() for name in ("dkvlm_config.json", "p2vlm_config.json")): inferred = _layers_from_geo_state(geo_state) if inferred: geo_cfg.target_layers = inferred wrapper = DKVLMGeoWrapper( base, target_layers=geo_cfg.target_layers, use_geo_branch=True, geo_attention_backend=geo_attention_backend, geo_branch_mrope=geo_cfg.geo_branch_mrope, geo_token_se3_dim=geo_cfg.token_se3_dim, geo_pose_gta_dim=geo_cfg.pose_gta_dim, geo_token_se3_lambda_min=geo_cfg.lambda_min, geo_token_se3_lambda_max=geo_cfg.lambda_max, ).to(dev) install_geo_attention(base, wrapper) report = wrapper.load_geo_state(geo_state) wrapper.eval() if verbose: print( f"[dkvlm] geometry branch on layers {geo_cfg.target_layers}: " f"{report['loaded_qkv_tensors']} projection tensors, " f"{len(report['loaded_gates'])} gates" ) if report["ignored_keys"]: print( f"[dkvlm] ignored {len(report['ignored_keys'])} non-geometry entries " f"(e.g. {report['ignored_keys'][0]})" ) engine = None if use_geometry: engine = VGGTOmegaEngine( VGGTOmegaConfig(ckpt=geometry_ckpt, device=str(dev), max_frames=max(nframes, 32)) ) return DKVLM( wrapper, processor, geo_cfg, engine, nframes=nframes, max_pixels=max_pixels, )