#!/usr/bin/env python3 """DuoVLM-40M 推理封装（包内自包含：给一张图 + 一个问题 → 回答） 路径全部可用环境变量覆盖，默认相对包根： DUOVLM_WEIGHTS 权重 .pth 默认 <包根>/weights/duovlm-s3-final.pth DUOVLM_CFG 模型配置 默认 <权重同目录>/model_config.yaml DUOVLM_TOKENIZER 词表目录 默认 <包根>/tokenizer DUOVLM_CLIP 视觉塔 HF id 默认 openai/clip-vit-base-patch16 DUOVLM_CLIP_DIR 视觉塔本地目录 设了就优先用它（离线场景） DUOVLM_DEVICE cuda / cpu 默认自动 依赖：torch、transformers、litgpt==0.5.13（见 requirements.txt）。 """ from __future__ import annotations import os import time from contextlib import nullcontext from pathlib import Path import numpy as np import torch from PIL import Image HERE = Path(__file__).resolve().parent PKG = HERE.parent DEFAULT_CLIP = "openai/clip-vit-base-patch16" from decoding import decode_safe # noqa: E402 from duovlm import (ASSISTANT, BOS, EOS, EOT, IMAGE, N_IMG, DuoVLM, # noqa: E402 load_duovlm) DESC_Q = "Render a clear and concise summary of the photo." def resolve_paths(weights: str | None = None) -> dict: w = Path(weights or os.environ.get("DUOVLM_WEIGHTS") or PKG / "weights" / "duovlm-s3-final.pth") return { "weights": w, "cfg": Path(os.environ.get("DUOVLM_CFG") or (w.parent / "model_config.yaml")), "tokenizer": Path(os.environ.get("DUOVLM_TOKENIZER") or PKG / "tokenizer"), } class DuoVLMInfer: """一次性载入（约 5~15s），之后每次问答 0.1~1.8s（取决于是否要现跑视觉塔）。""" def __init__(self, weights: str | None = None, device: str | None = None, clip: str | None = None, verbose: bool = True): from litgpt.config import Config from litgpt.tokenizer import Tokenizer self.paths = resolve_paths(weights) self.dev = device or os.environ.get("DUOVLM_DEVICE") or ("cuda" if torch.cuda.is_available() else "cpu") self.dtype = torch.bfloat16 if self.dev == "cuda" else torch.float32 self.tok = Tokenizer(self.paths["tokenizer"]) try: cfg = Config.from_file(str(self.paths["cfg"])) except Exception: import yaml cfg = Config(**yaml.safe_load(self.paths["cfg"].read_text())) self.model = DuoVLM(cfg) info = load_duovlm(self.paths["weights"], self.model) self.model = self.model.to(self.dev, dtype=self.dtype).eval() self.step = info.get("step") self.extra = {k: (float(v) if hasattr(v, "item") else v) for k, v in (info.get("extra") or {}).items()} if not self.extra: # load_duovlm 不一定回传 extra，这里从权重直接补读（只读元数据） try: _d = torch.load(self.paths["weights"], map_location="cpu") self.step = self.step or _d.get("step") self.extra = {k: (float(v) if hasattr(v, "item") else v) for k, v in (_d.get("extra") or {}).items()} except Exception: pass self.n_params = sum(x.numel() for x in self.model.parameters()) self.clip_source = clip or os.environ.get("DUOVLM_CLIP_DIR") or os.environ.get("DUOVLM_CLIP") or DEFAULT_CLIP self._clip = None if verbose: print(f"[duovlm] 权重 {self.paths['weights'].name}（step {self.step}）" f" 参数 {self.n_params:,} 设备 {self.dev} dtype {self.dtype}") print(f"[duovlm] 视觉塔 {self.clip_source}（首次用到时下载/载入）", flush=True) def param_report(self) -> str: llm = sum(p.numel() for p in self.model.llm.parameters()) con = sum(p.numel() for p in self.model.connector.parameters()) return (f"LLM {llm:,} + connector {con:,} = {llm + con:,} 可训练参数") # ---------- 视觉塔 ---------- def _load_clip(self): if self._clip is None: from transformers import CLIPImageProcessor, CLIPVisionModel self.proc = CLIPImageProcessor.from_pretrained(self.clip_source) self._clip = CLIPVisionModel.from_pretrained(self.clip_source, dtype=self.dtype).to(self.dev).eval() return self._clip @torch.no_grad() def embed(self, image) -> np.ndarray: """返回 (196, 768) fp16 特征。image 可以是路径或 PIL.Image。""" clip = self._load_clip() im = Image.open(image) if isinstance(image, (str, Path)) else image px = self.proc(im.convert("RGB"), return_tensors="pt")["pixel_values"] ctx = torch.autocast("cuda", dtype=torch.bfloat16) if self.dev == "cuda" else nullcontext() with torch.no_grad(), ctx: h = clip(pixel_values=px.to(self.dev, dtype=self.dtype), output_hidden_states=True).hidden_states[-2] h = clip.vision_model.post_layernorm(h)[:, 1:, :] return h[0].float().cpu().numpy().astype(np.float16) # ---------- 生成 ---------- @torch.no_grad() def ask(self, image=None, question: str = "", blind: bool = False, no_image: bool = False, max_new: int = 12, ngram: int = 3, rep_penalty: float = 1.0, temperature: float = 0.0, top_k: int = 0, loop_break: bool = True) -> dict: """看图问答。blind=True 表示保留模板但图像位换零（不看图对照）； no_image=True 表示连 196 个图像位都去掉（纯文本，分布外）。""" q = (question or DESC_Q).replace("", " ").strip() qids = self.tok.encode(q).tolist() feat, src = None, "no-image" if no_image: ids = [BOS] + qids[: 512 - 4] + [EOT, ASSISTANT] else: ids = [BOS] + [IMAGE] * N_IMG + [4] + qids[: 512 - (N_IMG + 4)] + [EOT, ASSISTANT] t_emb = time.time() feat = self.embed(image) src = "blind(zeros)" if blind else "image" if blind: feat = np.zeros_like(feat) emb_ms = (time.time() - t_emb) * 1000 t0 = time.time() ft = None if feat is None else torch.from_numpy(feat[None]).to(self.dev) out, st = decode_safe(self.model, torch.tensor([ids], dtype=torch.long).to(self.dev), ft, max_new=max_new, ngram=ngram, rep_penalty=rep_penalty, temperature=temperature, top_k=top_k, loop_break=loop_break) toks = out[0] ans = " ".join(self.tok.decode(torch.tensor(toks)).split()) if toks else "" return {"answer": ans, "tokens": len(toks), "ms": round((time.time() - t0) * 1000), "emb_ms": round(locals().get("emb_ms", 0)), "source": src, "stats": {k: v for k, v in st.items() if v}} @torch.no_grad() def continue_text(self, seed: str, max_new: int = 40, **kw) -> dict: """纯文本续写（无图像位，Stage 1 原生格式）。""" ids = ([BOS] + self.tok.encode(seed).tolist())[: 512 - max_new - 1] t0 = time.time() out, st = decode_safe(self.model, torch.tensor([ids], dtype=torch.long).to(self.dev), None, max_new=max_new, **kw) toks = out[0] return {"answer": " ".join(self.tok.decode(torch.tensor(toks)).split()) if toks else "", "tokens": len(toks), "ms": round((time.time() - t0) * 1000), "source": "text", "stats": {k: v for k, v in st.items() if v}}