File size: 3,933 Bytes
255b4a8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """Differentiable loader for the Qwen3-VL vision tower (Qwen/Qwen3-VL-8B-Instruct).
Qwen3-VL's vision model consumes pre-patchified `hidden_states` (seq, patch_dim) plus
a `grid_thw` tensor, not a plain image. We reproduce the Qwen image-processor patchify
in pure torch so gradients flow back to pixels, then mean-pool the last hidden state.
The full checkpoint stores the vision tower under the `model.visual.` prefix, so we
load the shards manually and strip that prefix into a standalone Qwen3VLVisionModel
(a plain from_pretrained(REPO) leaves every vision weight randomly initialized).
"""
from __future__ import annotations
import glob
import json
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from transformers import AutoConfig, Qwen3VLVisionModel
REPO = "Qwen/Qwen3-VL-8B-Instruct"
PREFIX = "model.visual."
PATCH = 16
TEMPORAL = 2
MERGE = 2
DEFAULT = 384 # multiple of PATCH*MERGE = 32
def load_qwen3vl(dtype: torch.dtype = torch.float32) -> nn.Module:
# fp16/bf16 backward through this tower produces NaN grads (attention/norm overflow),
# so we pin fp32 regardless of the ensemble's default custom-tower dtype.
dtype = torch.float32
cfg = AutoConfig.from_pretrained(REPO)
vcfg = cfg.vision_config
# Real (not meta) init so rotary/position buffers get materialized; we only pay a
# one-time random-init of a 0.6B tower, then overwrite params from the checkpoint.
model = Qwen3VLVisionModel(vcfg).to(dtype)
root = Path(snapshot_download(REPO, allow_patterns=["*.safetensors", "*.json"]))
idx = root / "model.safetensors.index.json"
shards = (set(json.load(open(idx))["weight_map"].values())
if idx.exists() else {p.name for p in root.glob("*.safetensors")})
state: dict[str, torch.Tensor] = {}
for shard in shards:
for k, v in load_file(root / shard).items():
if k.startswith(PREFIX):
state[k[len(PREFIX):]] = v.to(dtype)
missing, unexpected = model.load_state_dict(state, strict=False, assign=False)
real_missing = [m for m in missing if "rotary" not in m and "inv_freq" not in m]
if real_missing:
raise RuntimeError(f"Qwen3-VL vision still missing {len(real_missing)} weights: "
f"{real_missing[:6]}")
model.eval()
model.requires_grad_(False)
model._veil_cfg = vcfg
return model.to(device="cuda", dtype=dtype)
def _patchify(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""x: (1,3,H,W) [0,1] -> (seq, 3*TEMPORAL*PATCH*PATCH) flattened patches + grid_thw.
Mirrors Qwen2/3-VL image-processor reshape so token order matches the ViT.
"""
_, c, h, w = x.shape
gh, gw = h // PATCH, w // PATCH
xt = x[0].unsqueeze(0).repeat(TEMPORAL, 1, 1, 1) # (TEMPORAL, C, H, W)
p = xt.reshape(1, TEMPORAL, c,
gh // MERGE, MERGE, PATCH,
gw // MERGE, MERGE, PATCH)
p = p.permute(0, 3, 6, 4, 7, 2, 1, 5, 8).contiguous()
flat = p.reshape(gh * gw, c * TEMPORAL * PATCH * PATCH)
grid = torch.tensor([[1, gh, gw]], device=x.device, dtype=torch.long)
return flat, grid
def image_feat_qwen3vl(model: nn.Module, x: torch.Tensor) -> torch.Tensor:
"""x: (1,3,H,W) in [0,1] on cuda, requires_grad=True. Returns (1, D) pooled feature."""
assert x.shape[0] == 1 and x.dim() == 4
if x.shape[-1] != DEFAULT or x.shape[-2] != DEFAULT:
x = F.interpolate(x, size=(DEFAULT, DEFAULT), mode="bicubic", align_corners=False)
dtype = next(model.parameters()).dtype
flat, grid = _patchify(x)
out = model(flat.to(dtype=dtype), grid_thw=grid)
hs = out.last_hidden_state if hasattr(out, "last_hidden_state") else (
out[0] if isinstance(out, (tuple, list)) else out)
return hs.float().mean(dim=0, keepdim=True)
|