File size: 7,515 Bytes
d8bf653 a815bcd d8bf653 a815bcd d8bf653 | 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 | """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"]) # 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
path = hf_hub_download(repo_or_path, CKPT_FILE, revision=revision)
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
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)
|