EximiusLabs
/

fusion-embedding-2-2b-preview

@@ -1,183 +1,183 @@
-"""fusion-embedding inference — one embedding space for text, images, and audio.
-Serves BOTH architecture generations: fusion-embedding-1 (frozen base + trained
-resampler) and fusion-embedding-2 (adds modality-gated deep adapters — in-layer audio
-capacity whose gate leaves every text/image/video forward bitwise identical to the
-frozen base). The checkpoint's own config selects the architecture; an adapter
-checkpoint refuses to load without its adapters.
-Loads the frozen Qwen3-VL-Embedding base (native paths for text and images), the frozen
-Qwen2.5-Omni audio tower, and this repository's trained connector checkpoint. All inputs
-use the base model's official chat-template format; embedding quality is sensitive to
-this formatting, so use the templates provided here rather than constructing your own.
-    from inference import FusionEmbedder
-    fe = FusionEmbedder.from_pretrained("EximiusLabs/fusion-embedding-2-2b-preview")
-    a, t, i = fe.embed_audio("dog.wav"), fe.embed_text("a dog barks"), fe.embed_image("dog.jpg")
-Requires: fusion_embedding (pip install git+https://github.com/Eximius-Labs/fusion-embedding-1),
-transformers>=4.46, torchvision, pillow, soundfile, librosa.
-"""
-from __future__ import annotations
-import dataclasses
-import os
-from typing import TYPE_CHECKING, Optional, Union
-if TYPE_CHECKING:
-    import numpy as np
-import torch
-BASE_MODEL = "Qwen/Qwen3-VL-Embedding-2B"
-AUDIO_MODEL = "Qwen/Qwen2.5-Omni-7B"
-DEFAULT_QUERY_INSTRUCTION = "Retrieve images or text relevant to the user's query."
-DOC_INSTRUCTION = "Represent the user's input."
-CKPT_FILES = ("fusion-embedding-2-2b-preview.pt", "fusion-embedding-1-2b-preview.pt")
-def _chat(instruction: str, user_content: str) -> str:
-    """The base's official embedding format: system-turn instruction, assistant opener."""
-    return (f"<|im_start|>system\n{instruction}<|im_end|>\n"
-            f"<|im_start|>user\n{user_content}<|im_end|>\n"
-            f"<|im_start|>assistant\n")
-class FusionEmbedder:
-    def __init__(self, ckpt_path: str, device: str = "cuda", dtype=torch.bfloat16):
-        from transformers import AutoFeatureExtractor, AutoModel, AutoProcessor
-        from fusion_embedding.config import FusionConfig
-        from fusion_embedding.hf_components import BaseLMAdapter, load_audio_tower
-        from fusion_embedding.model import FusionEmbeddingModel, last_token_pool
-        self.device = device
-        self._pool = last_token_pool
-        ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
-        flds = {f.name for f in dataclasses.fields(FusionConfig)}
-        self.cfg = FusionConfig(**{k: v for k, v in ck["config"].items() if k in flds})
-        self.full = AutoModel.from_pretrained(BASE_MODEL, trust_remote_code=True, dtype=dtype)
-        self.full = self.full.to(device).eval()
-        for p in self.full.parameters():
-            p.requires_grad_(False)
-        self.proc = AutoProcessor.from_pretrained(BASE_MODEL, trust_remote_code=True)
-        self.tok = self.proc.tokenizer
-        tower, _, _ = load_audio_tower(AUDIO_MODEL, device=device, dtype=dtype)
-        self.fe_audio = AutoFeatureExtractor.from_pretrained(AUDIO_MODEL, trust_remote_code=True)
-        self.model = FusionEmbeddingModel(self.cfg, self.full.get_input_embeddings(),
-                                          BaseLMAdapter(self.full.language_model),
-                                          audio_encoder=tower)
-        self.model.resampler.to(device).float()
-        self.model.resampler.load_state_dict(ck["resampler"])
-        # fusion-embedding-2: the gated adapters are part of the model — running an
-        # adapter checkpoint without them would silently produce the unadapted model,
-        # so any presence mismatch is a hard error.
-        if ("adapters" in ck) != (self.model.audio_adapters is not None):
-            raise RuntimeError(
-                f"adapter presence mismatch: checkpoint has_adapters={'adapters' in ck} "
-                f"but config adapter_rank={self.cfg.adapter_rank} — corrupted artifact?")
-        if self.model.audio_adapters is not None:
-            self.model.audio_adapters.to(device).float()
-            self.model.audio_adapters.load_state_dict(ck["adapters"])
-        self.model.text_whitening.load_state_dict(ck["text_whitening"])   # identity if unfitted
-        self.model.eval()
-    # ------------------------------------------------------------------ loading
-    @classmethod
-    def from_pretrained(cls, repo_or_path: str, device: str = "cuda",
-                        revision: Optional[str] = None, **kw) -> "FusionEmbedder":
-        """Load from a local checkpoint path or an HF repo. ``revision`` pins a repo
-        tag/commit (e.g. ``"v0.1-preview"``, ``"v0.2-preview"``); default is latest."""
-        if os.path.exists(repo_or_path):
-            path = repo_or_path
-        else:
-            from huggingface_hub import hf_hub_download
-            from huggingface_hub.utils import EntryNotFoundError
-            path = None
-            for name in CKPT_FILES:                       # repo generation decides the name
-                try:
-                    path = hf_hub_download(repo_or_path, name, revision=revision)
-                    break
-                except EntryNotFoundError:
-                    continue
-            if path is None:
-                raise FileNotFoundError(f"no known checkpoint file in {repo_or_path} "
-                                        f"(looked for {CKPT_FILES})")
-        return cls(path, device=device, **kw)
-    # ------------------------------------------------------------------ helpers
-    def _finish(self, pooled: torch.Tensor, dim: Optional[int]) -> torch.Tensor:
-        from fusion_embedding.model import mrl_truncate_normalize
-        return mrl_truncate_normalize(pooled.float(), dim or self.cfg.mrl_default).squeeze(0).cpu()
-    # ------------------------------------------------------------------ audio
-    @torch.no_grad()
-    def embed_audio(self, audio: Union[str, "np.ndarray"], sr: Optional[int] = None,
-                    dim: Optional[int] = None) -> torch.Tensor:
-        import librosa
-        import soundfile as sf
-        if isinstance(audio, (str, os.PathLike)):
-            wav, sr = sf.read(str(audio), dtype="float32")
-        else:
-            wav = audio
-            assert sr is not None, "pass sr= when embedding a raw array"
-        if getattr(wav, "ndim", 1) > 1:
-            wav = wav.mean(axis=1)
-        target_sr = self.fe_audio.sampling_rate
-        if sr != target_sr:
-            wav = librosa.resample(wav, orig_sr=sr, target_sr=target_sr)
-        feats = self.fe_audio(wav, sampling_rate=target_sr, return_tensors="pt",
-                              return_attention_mask=True, padding="max_length", truncation=True)
-        mel = feats["input_features"][0]
-        am = feats.get("attention_mask")
-        if am is not None:
-            mel = mel[:, : int(am[0].sum().item())]
-        audio_tok = self.model.audio_tokens(
-            mel.unsqueeze(0).to(self.device),
-            torch.ones(1, mel.shape[1], dtype=torch.bool, device=self.device))
-        ids = torch.tensor([[self.cfg.audio_pad_id] * self.cfg.n_query + [self.cfg.eos_id]],
-                           device=self.device)
-        pooled = self.model.encode_audio(ids, torch.ones_like(ids), audio_tok)
-        return self._finish(pooled, dim)
-    # ------------------------------------------------------------------ text
-    @torch.no_grad()
-    def embed_text(self, text: str, instruction: str = DEFAULT_QUERY_INSTRUCTION,
-                   dim: Optional[int] = None) -> torch.Tensor:
-        ids = self.tok.encode(_chat(instruction, text), add_special_tokens=False)[:512]
-        ids_t = torch.tensor([ids], device=self.device)
-        pooled = self.model.encode_text(ids_t, torch.ones_like(ids_t))
-        return self._finish(self.model.text_whitening(pooled), dim)
-    # ------------------------------------------------------------------ image
-    @torch.no_grad()
-    def embed_image(self, image, dim: Optional[int] = None) -> torch.Tensor:
-        from PIL import Image
-        gate = getattr(self.model, "_adapter_gate", None)
-        if gate is not None and gate.active:
-            # The vision path runs through the same (hook-carrying) decoder layers;
-            # non-audio inputs must execute with the gate closed so the adapter
-            # branch never runs. Mirrors the encode_text guard.
-            raise RuntimeError("adapter gate is open during an image embed — "
-                               "non-audio inputs must run with the gate closed")
-        if isinstance(image, (str, os.PathLike)):
-            image = Image.open(str(image))
-        image = image.convert("RGB")
-        text = _chat(DOC_INSTRUCTION, "<|vision_start|><|image_pad|><|vision_end|>")
-        inputs = self.proc(text=[text], images=[image], return_tensors="pt").to(self.device)
-        h = self.full(**inputs).last_hidden_state
-        pooled = self._pool(h, inputs["attention_mask"])
-        return self._finish(pooled, dim)
-    # ------------------------------------------------------------------ cross-modal readout
-    @staticmethod
-    def center(embs: torch.Tensor) -> torch.Tensor:
-        """Per-modality mean-centering followed by renormalization. Recommended when ranking
-        a gallery of one modality against queries of another; improves cross-modal R@1 by
-        roughly two points across modality pairs in our evaluation."""
-        c = embs - embs.mean(dim=0, keepdim=True)
-        return torch.nn.functional.normalize(c, dim=-1)

+"""fusion-embedding inference — one embedding space for text, images, and audio.
+Serves BOTH architecture generations: fusion-embedding-1 (frozen base + trained
+resampler) and fusion-embedding-2 (adds modality-gated deep adapters — in-layer audio
+capacity whose gate leaves every text/image/video forward bitwise identical to the
+frozen base). The checkpoint's own config selects the architecture; an adapter
+checkpoint refuses to load without its adapters.
+Loads the frozen Qwen3-VL-Embedding base (native paths for text and images), the frozen
+Qwen2.5-Omni audio tower, and this repository's trained connector checkpoint. All inputs
+use the base model's official chat-template format; embedding quality is sensitive to
+this formatting, so use the templates provided here rather than constructing your own.
+    from inference import FusionEmbedder
+    fe = FusionEmbedder.from_pretrained("EximiusLabs/fusion-embedding-2-2b-preview")
+    a, t, i = fe.embed_audio("dog.wav"), fe.embed_text("a dog barks"), fe.embed_image("dog.jpg")
+Requires: fusion_embedding (pip install git+https://github.com/Eximius-Labs/fusion-embedding),
+transformers>=4.46, torchvision, pillow, soundfile, librosa.
+"""
+from __future__ import annotations
+import dataclasses
+import os
+from typing import TYPE_CHECKING, Optional, Union
+if TYPE_CHECKING:
+    import numpy as np
+import torch
+BASE_MODEL = "Qwen/Qwen3-VL-Embedding-2B"
+AUDIO_MODEL = "Qwen/Qwen2.5-Omni-7B"
+DEFAULT_QUERY_INSTRUCTION = "Retrieve images or text relevant to the user's query."
+DOC_INSTRUCTION = "Represent the user's input."
+CKPT_FILES = ("fusion-embedding-2-2b-preview.pt", "fusion-embedding-1-2b-preview.pt")
+def _chat(instruction: str, user_content: str) -> str:
+    """The base's official embedding format: system-turn instruction, assistant opener."""
+    return (f"<|im_start|>system\n{instruction}<|im_end|>\n"
+            f"<|im_start|>user\n{user_content}<|im_end|>\n"
+            f"<|im_start|>assistant\n")
+class FusionEmbedder:
+    def __init__(self, ckpt_path: str, device: str = "cuda", dtype=torch.bfloat16):
+        from transformers import AutoFeatureExtractor, AutoModel, AutoProcessor
+        from fusion_embedding.config import FusionConfig
+        from fusion_embedding.hf_components import BaseLMAdapter, load_audio_tower
+        from fusion_embedding.model import FusionEmbeddingModel, last_token_pool
+        self.device = device
+        self._pool = last_token_pool
+        ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        flds = {f.name for f in dataclasses.fields(FusionConfig)}
+        self.cfg = FusionConfig(**{k: v for k, v in ck["config"].items() if k in flds})
+        self.full = AutoModel.from_pretrained(BASE_MODEL, trust_remote_code=True, dtype=dtype)
+        self.full = self.full.to(device).eval()
+        for p in self.full.parameters():
+            p.requires_grad_(False)
+        self.proc = AutoProcessor.from_pretrained(BASE_MODEL, trust_remote_code=True)
+        self.tok = self.proc.tokenizer
+        tower, _, _ = load_audio_tower(AUDIO_MODEL, device=device, dtype=dtype)
+        self.fe_audio = AutoFeatureExtractor.from_pretrained(AUDIO_MODEL, trust_remote_code=True)
+        self.model = FusionEmbeddingModel(self.cfg, self.full.get_input_embeddings(),
+                                          BaseLMAdapter(self.full.language_model),
+                                          audio_encoder=tower)
+        self.model.resampler.to(device).float()
+        self.model.resampler.load_state_dict(ck["resampler"])
+        # fusion-embedding-2: the gated adapters are part of the model — running an
+        # adapter checkpoint without them would silently produce the unadapted model,
+        # so any presence mismatch is a hard error.
+        if ("adapters" in ck) != (self.model.audio_adapters is not None):
+            raise RuntimeError(
+                f"adapter presence mismatch: checkpoint has_adapters={'adapters' in ck} "
+                f"but config adapter_rank={self.cfg.adapter_rank} — corrupted artifact?")
+        if self.model.audio_adapters is not None:
+            self.model.audio_adapters.to(device).float()
+            self.model.audio_adapters.load_state_dict(ck["adapters"])
+        self.model.text_whitening.load_state_dict(ck["text_whitening"])   # identity if unfitted
+        self.model.eval()
+    # ------------------------------------------------------------------ loading
+    @classmethod
+    def from_pretrained(cls, repo_or_path: str, device: str = "cuda",
+                        revision: Optional[str] = None, **kw) -> "FusionEmbedder":
+        """Load from a local checkpoint path or an HF repo. ``revision`` pins a repo
+        tag/commit (e.g. ``"v0.1-preview"``, ``"v0.2-preview"``); default is latest."""
+        if os.path.exists(repo_or_path):
+            path = repo_or_path
+        else:
+            from huggingface_hub import hf_hub_download
+            from huggingface_hub.utils import EntryNotFoundError
+            path = None
+            for name in CKPT_FILES:                       # repo generation decides the name
+                try:
+                    path = hf_hub_download(repo_or_path, name, revision=revision)
+                    break
+                except EntryNotFoundError:
+                    continue
+            if path is None:
+                raise FileNotFoundError(f"no known checkpoint file in {repo_or_path} "
+                                        f"(looked for {CKPT_FILES})")
+        return cls(path, device=device, **kw)
+    # ------------------------------------------------------------------ helpers
+    def _finish(self, pooled: torch.Tensor, dim: Optional[int]) -> torch.Tensor:
+        from fusion_embedding.model import mrl_truncate_normalize
+        return mrl_truncate_normalize(pooled.float(), dim or self.cfg.mrl_default).squeeze(0).cpu()
+    # ------------------------------------------------------------------ audio
+    @torch.no_grad()
+    def embed_audio(self, audio: Union[str, "np.ndarray"], sr: Optional[int] = None,
+                    dim: Optional[int] = None) -> torch.Tensor:
+        import librosa
+        import soundfile as sf
+        if isinstance(audio, (str, os.PathLike)):
+            wav, sr = sf.read(str(audio), dtype="float32")
+        else:
+            wav = audio
+            assert sr is not None, "pass sr= when embedding a raw array"
+        if getattr(wav, "ndim", 1) > 1:
+            wav = wav.mean(axis=1)
+        target_sr = self.fe_audio.sampling_rate
+        if sr != target_sr:
+            wav = librosa.resample(wav, orig_sr=sr, target_sr=target_sr)
+        feats = self.fe_audio(wav, sampling_rate=target_sr, return_tensors="pt",
+                              return_attention_mask=True, padding="max_length", truncation=True)
+        mel = feats["input_features"][0]
+        am = feats.get("attention_mask")
+        if am is not None:
+            mel = mel[:, : int(am[0].sum().item())]
+        audio_tok = self.model.audio_tokens(
+            mel.unsqueeze(0).to(self.device),
+            torch.ones(1, mel.shape[1], dtype=torch.bool, device=self.device))
+        ids = torch.tensor([[self.cfg.audio_pad_id] * self.cfg.n_query + [self.cfg.eos_id]],
+                           device=self.device)
+        pooled = self.model.encode_audio(ids, torch.ones_like(ids), audio_tok)
+        return self._finish(pooled, dim)
+    # ------------------------------------------------------------------ text
+    @torch.no_grad()
+    def embed_text(self, text: str, instruction: str = DEFAULT_QUERY_INSTRUCTION,
+                   dim: Optional[int] = None) -> torch.Tensor:
+        ids = self.tok.encode(_chat(instruction, text), add_special_tokens=False)[:512]
+        ids_t = torch.tensor([ids], device=self.device)
+        pooled = self.model.encode_text(ids_t, torch.ones_like(ids_t))
+        return self._finish(self.model.text_whitening(pooled), dim)
+    # ------------------------------------------------------------------ image
+    @torch.no_grad()
+    def embed_image(self, image, dim: Optional[int] = None) -> torch.Tensor:
+        from PIL import Image
+        gate = getattr(self.model, "_adapter_gate", None)
+        if gate is not None and gate.active:
+            # The vision path runs through the same (hook-carrying) decoder layers;
+            # non-audio inputs must execute with the gate closed so the adapter
+            # branch never runs. Mirrors the encode_text guard.
+            raise RuntimeError("adapter gate is open during an image embed — "
+                               "non-audio inputs must run with the gate closed")
+        if isinstance(image, (str, os.PathLike)):
+            image = Image.open(str(image))
+        image = image.convert("RGB")
+        text = _chat(DOC_INSTRUCTION, "<|vision_start|><|image_pad|><|vision_end|>")
+        inputs = self.proc(text=[text], images=[image], return_tensors="pt").to(self.device)
+        h = self.full(**inputs).last_hidden_state
+        pooled = self._pool(h, inputs["attention_mask"])
+        return self._finish(pooled, dim)
+    # ------------------------------------------------------------------ cross-modal readout
+    @staticmethod
+    def center(embs: torch.Tensor) -> torch.Tensor:
+        """Per-modality mean-centering followed by renormalization. Recommended when ranking
+        a gallery of one modality against queries of another; improves cross-modal R@1 by
+        roughly two points across modality pairs in our evaluation."""
+        c = embs - embs.mean(dim=0, keepdim=True)
+        return torch.nn.functional.normalize(c, dim=-1)