fusion-embedding-2-k3-vision v0.1-preview: project Kimi K3's MoonViT-V2 vision encoder into the fusion-embedding-2 space

README.md

@@ -0,0 +1,96 @@
+---
+license: cc-by-nc-4.0
+base_model: moonshotai/Kimi-K3
+base_model_relation: adapter
+library_name: fusion-embedding
+pipeline_tag: feature-extraction
+tags:
+  - multimodal
+  - embeddings
+  - retrieval
+  - vision
+  - kimi-k3
+  - moonvit
+---
+
+# fusion-embedding-2-k3-vision
+
+Kimi K3's vision encoder, made language-searchable.
+
+This is a small trained projector that maps Kimi K3's frozen vision encoder
+(**MoonViT-V2**, about 401M parameters) into the frozen
+[fusion-embedding-2](https://huggingface.co/EximiusLabs/fusion-embedding-2-2b-preview) text
+space. K3's own visual features become searchable in plain language, and land in the same
+space as fusion-embedding's text, image, video, and audio, plus its thermal (Ember) and
+motion (Tremor) sensor packs.
+
+It builds on **only the ~401M vision tower of Kimi K3, never the 2.8T language model**. The
+tower runs on a single GPU. The projector is the only trained weight, about 3.15M
+parameters.
+
+## Results
+
+Text-to-image retrieval on Flickr30k (test), 1000-image gallery, querying in the
+fusion-embedding-2 text space:
+
+| Metric | fusion-embedding-2-k3-vision | same recipe, DINOv2 backbone |
+| --- | --- | --- |
+| R@1 | 0.432 | |
+| R@5 | 0.740 | |
+| **R@10** | **0.836** | 0.749 |
+
+MoonViT-V2 was trained alongside a language model, so its features align to text more
+readily than a self-supervised vision backbone. On the identical projector recipe it beats
+the DINOv2 version by nearly 9 points at R@10.
+
+## Usage
+
+The projector ships here; the vision tower is pulled from `moonshotai/Kimi-K3` on first use
+(only the `vision_tower.*` tensors, never the 2.8T language model).
+
+```python
+from inference import K3VisionEmbedder
+m = K3VisionEmbedder.from_pretrained("EximiusLabs/fusion-embedding-2-k3-vision")
+v = m.embed_image("photo.jpg")             # 2048-d, in the fusion-embedding-2 text space
+```
+
+Query it with text using fusion-embedding-2 itself:
+
+```python
+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 = float(q @ v)                       # cosine, both unit-norm
+```
+
+## How it works
+
+1. Kimi K3's MoonViT-V2 vision tower (frozen) encodes an image into patch tokens (1024-d).
+2. The tokens are pooled into one image descriptor.
+3. A trained projector (`LayerNorm -> Linear -> GELU -> Dropout -> Linear`) maps that
+   descriptor into the 2048-d fusion-embedding-2 text space.
+4. Trained with image-caption InfoNCE on Flickr30k. Same recipe as the Fusion Perception
+   DINO to fusion-embedding projector, with K3's vision tower as the backbone.
+
+## What this is and is not
+
+- It is a projector on a frozen vision encoder, not a new foundation model, and not a
+  fine-tune of Kimi K3's language model.
+- The `base_model` tag points at `moonshotai/Kimi-K3` because the vision tower is a component
+  of that model. We use only MoonViT-V2, not the 2.8T language model.
+- It is a research preview.
+
+## License and attribution
+
+The projector weights and code in this repository are released under **CC-BY-NC-4.0**.
+
+The vision tower is part of **Kimi K3** (Moonshot AI), used under the **Kimi K3 License**,
+which permits derivative works with its notice retained. This model is a derivative that
+uses K3's frozen MoonViT-V2 vision encoder. Copyright 2026 Moonshot AI.
+
+## Family
+
+- [fusion-embedding-2](https://huggingface.co/EximiusLabs/fusion-embedding-2-2b-preview) (the base multimodal embedding model)
+- [Ember](https://huggingface.co/EximiusLabs/fusion-embedding-2-ember) (thermal) and [Tremor](https://huggingface.co/EximiusLabs/fusion-embedding-2-tremor) (motion) sensor packs
+- Cookbook: [give a Kimi K3 agent a multimodal memory](https://github.com/Eximius-Labs/cookbook)
+- Code and paper: [github.com/Eximius-Labs](https://github.com/Eximius-Labs), [arXiv:2607.18666](https://arxiv.org/abs/2607.18666)

config.json

@@ -0,0 +1,21 @@
+{
+  "model_type": "fusion-embedding-2-k3-vision",
+  "description": "A trained projector mapping Kimi K3's frozen MoonViT-V2 vision encoder into the fusion-embedding-2 text space, so K3's visual features become language-searchable in the shared multimodal space.",
+  "backbone": "moonshotai/Kimi-K3",
+  "backbone_component": "vision_tower (MoonViT-V2, ~401M params)",
+  "base_text_model": "Qwen/Qwen3-VL-Embedding-2B",
+  "fusion_embedding_model": "EximiusLabs/fusion-embedding-2-2b-preview",
+  "projector": {
+    "in_dim": 1024,
+    "out_dim": 2048,
+    "arch": "LayerNorm(1024)->Linear(1024,1024)->GELU->Dropout(0.2)->Linear(1024,2048)",
+    "params": 3150848,
+    "readout": "mean over merged patch tokens and 2x2 subpatches",
+    "weights_file": "k3vit2fe2_proj.pt"
+  },
+  "vision_attn_implementation": "eager",
+  "trained_on": "nlphuji/flickr30k",
+  "training": {"objective": "image-caption InfoNCE", "temperature": 0.07, "steps": 4000, "n_train": 6000},
+  "license": "cc-by-nc-4.0",
+  "base_license": "Kimi K3 License"
+}

inference.py

@@ -0,0 +1,175 @@
+"""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, device: str = "cuda", dtype=torch.bfloat16):
+        ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        self.cfg = ck["config"]
+        self.device = device
+        self.tower, self.proc = load_moonvit_tower(device, dtype)
+        self.proj = _build_projector(self.cfg["in_dim"], self.cfg["out_dim"]).to(device).float()
+        self.proj.load_state_dict({k: v.float() for k, v in ck["proj"].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):
+            path = repo_or_path
+        else:
+            from huggingface_hub import hf_hub_download
+            path = hf_hub_download(repo_or_path, "k3vit2fe2_proj.pt", revision=revision)
+        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()

k3vit2fe2_proj.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:566e0192c390ac41e342a77fe6826688df00faf9dcbb152cf4de8f377213d1d2
+size 6304654

results.json

@@ -0,0 +1,14 @@
+{
+  "flickr30k_test_t2i": {
+    "R@1": 0.432,
+    "R@5": 0.74,
+    "R@10": 0.836,
+    "gallery": 1000,
+    "task": "text-to-image retrieval in the fusion-embedding-2 text space"
+  },
+  "comparison": {
+    "dinov2_backbone_same_recipe_R@10": 0.749,
+    "note": "Same projector recipe trained on a frozen DINOv2 backbone instead of MoonViT-V2."
+  },
+  "training": {"n_train": 6000, "steps": 4000, "runtime_s": 505.4, "gpu": "A100-80GB"}
+}