| """fusion-embedding-1 inference — one embedding space for text, images, and audio. |
| |
| 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-1-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_FILE = "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"]) |
| self.model.text_whitening.load_state_dict(ck["text_whitening"]) |
| self.model.eval() |
|
|
| |
| @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 |
| path = hf_hub_download(repo_or_path, CKPT_FILE, revision=revision) |
| return cls(path, device=device, **kw) |
|
|
| |
| 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() |
|
|
| |
| @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) |
|
|
| |
| @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) |
|
|
| |
| @torch.no_grad() |
| def embed_image(self, image, dim: Optional[int] = None) -> torch.Tensor: |
| from PIL import Image |
| 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) |
|
|
| |
| @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) |
|
|
