File size: 9,128 Bytes
e71f1b7 55b1c71 e71f1b7 55b1c71 e71f1b7 55b1c71 d46e9f5 55b1c71 d46e9f5 e71f1b7 | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 | """fusion-embedding-2-k3-vision — Kimi K3's vision encoder, made language-searchable.
A small trained projector maps Kimi K3's frozen MoonViT-V2 vision features into the
frozen fusion-embedding-2 (Qwen3-VL-Embedding-2B) text space, so K3's own visual
features become searchable in language and sit in the same space as fusion-embedding's
text, image, video, and audio. This ships the projector only; the ~401M vision tower is
pulled from moonshotai/Kimi-K3 on first use (never the 2.8T language model).
from inference import K3VisionEmbedder
m = K3VisionEmbedder.from_pretrained("EximiusLabs/fusion-embedding-2-k3-vision")
v = m.embed_image("photo.jpg") # 2048-d, in the FE2 text space
Query with text using fusion-embedding-2 itself:
from inference import FusionEmbedder # from the fusion-embedding-2 repo
fe = FusionEmbedder.from_pretrained("EximiusLabs/fusion-embedding-2-2b-preview")
q = fe.embed_text("a dog on a beach")
score = (q @ v) # cosine, both unit-norm
MoonViT-V2 is not published as a standalone checkpoint; it lives inside Kimi-K3 with a
`vision_tower.` prefix. We load only those tensors from the shard index. Loading the
vision tower's modeling code also imports K3's language stack, so we stub that import
(the vision path never uses it) rather than pulling in the LLM's kernels.
Requires: torch, transformers>=4.57, einops, pillow, safetensors, huggingface_hub.
The Kimi-K3 base is under the Kimi K3 License (permissive, notice required); this
projector is released under CC-BY-NC-4.0.
"""
from __future__ import annotations
import os
import sys
import json
import types
from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
K3_REPO = "moonshotai/Kimi-K3"
BASE = "Qwen/Qwen3-VL-Embedding-2B"
VPREFIX = "vision_tower."
# --------------------------------------------------------------------------- #
# MoonViT-V2 loader: only the vision tower, no language model, no fla/triton
# --------------------------------------------------------------------------- #
def _install_llm_stub() -> None:
"""Stub K3's language modeling submodule so importing the ViT never pulls the LLM."""
from importlib.abc import Loader, MetaPathFinder
from importlib.machinery import ModuleSpec
class _StubLoader(Loader):
def create_module(self, spec):
m = types.ModuleType(spec.name)
m.__getattr__ = lambda name: type(name, (), {})
return m
def exec_module(self, module):
pass
class _StubFinder(MetaPathFinder):
def find_spec(self, fullname, path=None, target=None):
if fullname.endswith("modeling_kimi_linear"):
return ModuleSpec(fullname, _StubLoader())
return None
if not any(type(f).__name__ == "_StubFinder" for f in sys.meta_path):
sys.meta_path.insert(0, _StubFinder())
def load_moonvit_tower(device: str, dtype=torch.bfloat16, token: Optional[str] = None):
"""Return the frozen MoonViT-V2 tower and its image processor."""
import safetensors.torch as st
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoImageProcessor
from transformers.dynamic_module_utils import get_class_from_dynamic_module
token = token or os.environ.get("HF_TOKEN")
cfg = AutoConfig.from_pretrained(K3_REPO, trust_remote_code=True, token=token)
vcfg = cfg.vision_config
# the LLM stack imports transformers-internal symbols this build may not expose;
# shim them (only referenced by the LLM, which the ViT never runs).
import transformers.utils.generic as _g
if not hasattr(_g, "OutputRecorder"):
_g.OutputRecorder = type("OutputRecorder", (), {"__init__": lambda self, *a, **k: None})
if not hasattr(_g, "check_model_inputs"):
_g.check_model_inputs = lambda func=None, *a, **k: (func if callable(func) else (lambda f: f))
_install_llm_stub()
MoonViT = get_class_from_dynamic_module(
"modeling_kimi_k3.MoonViT3dPretrainedModel", K3_REPO, token=token)
# KimiK3VisionConfig prefixes vision fields with vt_; alias to the standard names
# (hidden_size is the ViT width 1024, not the 7168 LLM width).
for _std, _vt in (("hidden_size", "vt_hidden_size"),
("num_hidden_layers", "vt_num_hidden_layers"),
("num_attention_heads", "vt_num_attention_heads"),
("intermediate_size", "vt_intermediate_size")):
if not hasattr(vcfg, _std) and hasattr(vcfg, _vt):
setattr(vcfg, _std, getattr(vcfg, _vt))
vcfg._attn_implementation = "eager" # the ViT attn table has no sdpa path
tower = MoonViT(vcfg)
# load only vision_tower.* tensors from the shard index (never the MoE experts)
idx = json.load(open(hf_hub_download(K3_REPO, "model.safetensors.index.json", token=token)))
wmap = idx["weight_map"]
vkeys = {k: v for k, v in wmap.items() if k.startswith(VPREFIX)}
sd = {}
for sh in sorted(set(vkeys.values())):
t = st.load_file(hf_hub_download(K3_REPO, sh, token=token))
for k in vkeys:
if wmap[k] == sh:
sd[k[len(VPREFIX):]] = t[k]
miss, unexp = tower.load_state_dict(sd, strict=False)
if miss or unexp:
raise RuntimeError(f"MoonViT load mismatch: missing={len(miss)} unexpected={len(unexp)}")
for m in tower.modules():
if hasattr(m, "attn_implementation"):
m.attn_implementation = "eager"
tower = tower.to(device, dtype).eval()
for p in tower.parameters():
p.requires_grad_(False)
proc = AutoImageProcessor.from_pretrained(K3_REPO, trust_remote_code=True, token=token)
return tower, proc
def _build_projector(in_dim: int, out_dim: int) -> nn.Sequential:
return nn.Sequential(nn.LayerNorm(in_dim), nn.Linear(in_dim, 1024), nn.GELU(),
nn.Dropout(0.2), nn.Linear(1024, out_dim))
class K3VisionEmbedder:
"""Frozen MoonViT-V2 + a trained projector into the fusion-embedding-2 text space."""
def __init__(self, ckpt_path: str, config: Optional[dict] = None,
device: str = "cuda", dtype=torch.bfloat16):
if ckpt_path.endswith(".safetensors"):
import safetensors.torch as st
proj_sd = st.load_file(ckpt_path)
if config is None:
cfgp = os.path.join(os.path.dirname(ckpt_path), "config.json")
config = json.load(open(cfgp))["projector"]
in_dim, out_dim = int(config["in_dim"]), int(config["out_dim"])
else: # legacy .pt checkpoint
ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
proj_sd = ck["proj"]; c = ck["config"]
in_dim, out_dim = int(c["in_dim"]), int(c["out_dim"])
self.cfg = {"in_dim": in_dim, "out_dim": out_dim}
self.device = device
self.tower, self.proc = load_moonvit_tower(device, dtype)
self.proj = _build_projector(in_dim, out_dim).to(device).float()
self.proj.load_state_dict({k: v.float() for k, v in proj_sd.items()})
self.proj.eval()
@classmethod
def from_pretrained(cls, repo_or_path: str, device: str = "cuda",
revision: Optional[str] = None, **kw) -> "K3VisionEmbedder":
if os.path.exists(repo_or_path):
return cls(repo_or_path, device=device, **kw)
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError
try:
hf_hub_download(repo_or_path, "config.json", revision=revision) # sibling config
except Exception: # noqa: BLE001
pass
path = None
for name in ("model.safetensors", "k3vit2fe2_proj.safetensors", "k3vit2fe2_proj.pt"):
try:
path = hf_hub_download(repo_or_path, name, revision=revision)
break
except EntryNotFoundError:
continue
if path is None:
raise FileNotFoundError(f"no projector checkpoint found in {repo_or_path}")
return cls(path, device=device, **kw)
@torch.no_grad()
def _tower_feat(self, image) -> torch.Tensor:
from PIL import Image
if isinstance(image, (str, os.PathLike)):
image = Image.open(str(image))
image = image.convert("RGB")
enc = self.proc.preprocess([{"type": "image", "image": image}], return_tensors="pt")
pv = enc["pixel_values"].to(self.device, torch.bfloat16)
grid = enc["grid_thws"].to(self.device)
f = self.tower(pv, grid)[0].float() # [N_merged, 4, 1024]
return F.normalize(f.mean(dim=(0, 1)), dim=-1) # [1024]
@torch.no_grad()
def embed_image(self, image) -> torch.Tensor:
"""Embed an image via MoonViT-V2 into the fusion-embedding-2 text space (2048-d)."""
feat = self._tower_feat(image)
return F.normalize(self.proj(feat), dim=-1).cpu()
|