Remove duplicated library code; the canonical copy is on PyPI

README.md

@@ -117,8 +117,12 @@ model = QuadEmbed.from_pretrained(device="cpu")                   # force CPU
 ```
 
 Audio takes mono float32 arrays at 16 kHz; video takes a file path
-(`model.embed_video_file("clip.mp4")`). See
-`examples/inference_example.py` for a complete runnable script.
+(`model.embed_video_file("clip.mp4")`). A complete runnable example lives in
+[the source repo](https://github.com/mithilai/QuadEmbed/blob/main/package/examples/inference_example.py).
+
+This repo hosts the trained weights only; the library itself is published on
+[PyPI](https://pypi.org/project/quadembed/) so there is one canonical copy of
+the code rather than a duplicate here that can drift out of sync.
 
 ## Important caveat: two different vision checkpoints
 

examples/inference_example.py

@@ -1,26 +0,0 @@
-"""Embed text, an image, and a video into the same 768-dim space and compare
-with cosine similarity (embeddings are already L2-normalized, so a plain dot
-product is cosine similarity).
-
-Run from the repo root:
-    pip install -r requirements.txt
-    python examples/inference_example.py
-"""
-from PIL import Image
-
-from quadembed import QuadEmbed
-
-model = QuadEmbed.from_pretrained("checkpoints")  # device auto-detects cuda/cpu
-
-texts = ["a photo of a dog running on the beach", "a bowl of ramen noodles"]
-text_embeds = model.embed_text(texts)
-
-image = Image.open("your_image.jpg").convert("RGB")  # replace with a real path
-image_embeds = model.embed_image([image])
-
-similarity = text_embeds @ image_embeds.T
-for text, score in zip(texts, similarity[:, 0].tolist()):
-    print(f"{score:.3f}  {text}")
-
-# video_embeds = model.embed_video_file("your_clip.mp4")  # uses the video-tuned projector, see README
-# audio_embeds = model.embed_audio([your_16khz_float32_array])

quadembed/__init__.py

@@ -1,84 +0,0 @@
-"""QuadEmbed: a local, from-scratch-trained multimodal embedding model
-covering all four modalities, text, image, audio, and video, in one shared
-768-dim space. Reproduces the architecture behind Jina AI's GELATO
-(jina-embeddings-v5-omni): frozen text/vision/audio encoders plus two small
-trained projectors.
-
-    from quadembed import QuadEmbed
-
-    model = QuadEmbed.from_pretrained("checkpoints", device="cuda")
-    text_embeds = model.embed_text(["a dog running on the beach"])
-    image_embeds = model.embed_image([Image.open("photo.jpg")])
-    similarity = text_embeds @ image_embeds.T
-"""
-import os
-
-import torch
-import torch.nn.functional as F
-
-from .encoders import AudioEncoder, TextEncoder, VisionEncoder
-from .projectors import AudioProjector, VisionProjector
-from .video import embed_video, sample_frames
-
-__all__ = [
-    "QuadEmbed",
-    "TextEncoder",
-    "VisionEncoder",
-    "AudioEncoder",
-    "VisionProjector",
-    "AudioProjector",
-    "embed_video",
-    "sample_frames",
-]
-
-
-class QuadEmbed:
-    """Loads all three frozen encoders plus the trained projectors, and
-    exposes one `embed_*` method per modality, each returning an L2-normalized
-    [batch, 768] tensor in the shared space -- so any two outputs from any
-    two modalities can be compared with a plain dot product / cosine
-    similarity."""
-
-    def __init__(self, device: str = "cuda" if torch.cuda.is_available() else "cpu"):
-        self.device = device
-        self.text_encoder = TextEncoder(device)
-        self.vision_encoder = VisionEncoder(device)
-        self.audio_encoder = AudioEncoder(device)
-        self.vision_projector = VisionProjector(patch_dim=768, out_dim=768).to(device).eval()
-        self.audio_projector = AudioProjector(in_dim=1280, out_dim=768).to(device).eval()
-        self.video_projector = VisionProjector(patch_dim=768, out_dim=768).to(device).eval()
-
-    @classmethod
-    def from_pretrained(cls, checkpoint_dir: str, device: str = "cuda" if torch.cuda.is_available() else "cpu"):
-        model = cls(device=device)
-        model.vision_projector.load_state_dict(
-            torch.load(os.path.join(checkpoint_dir, "vision_projector.pt"), map_location=device)
-        )
-        model.audio_projector.load_state_dict(
-            torch.load(os.path.join(checkpoint_dir, "audio_projector.pt"), map_location=device)
-        )
-        model.video_projector.load_state_dict(
-            torch.load(os.path.join(checkpoint_dir, "video_projector.pt"), map_location=device)
-        )
-        return model
-
-    @torch.no_grad()
-    def embed_text(self, texts: list[str]) -> torch.Tensor:
-        return F.normalize(self.text_encoder.embed(texts).float(), dim=-1)
-
-    @torch.no_grad()
-    def embed_image(self, images: list) -> torch.Tensor:
-        patch_tokens, mask, spatial_shapes = self.vision_encoder.patch_tokens(images)
-        out = self.vision_projector(patch_tokens.float(), mask, spatial_shapes)
-        return F.normalize(out, dim=-1)
-
-    @torch.no_grad()
-    def embed_audio(self, arrays: list, sampling_rate: int = 16000) -> torch.Tensor:
-        frame_tokens = self.audio_encoder.frame_tokens(arrays, sampling_rate=sampling_rate)
-        out = self.audio_projector(frame_tokens.float())
-        return F.normalize(out, dim=-1)
-
-    @torch.no_grad()
-    def embed_video_file(self, video_path: str, num_frames: int = 4) -> torch.Tensor:
-        out = embed_video(self.vision_encoder, self.video_projector, video_path, self.device, num_frames=num_frames)
-        return F.normalize(out.unsqueeze(0), dim=-1)

quadembed/encoders.py

@@ -1,103 +0,0 @@
-"""Frozen encoder loaders.
-
-Substitutions vs. the GELATO paper's exact checkpoints, since the paper's
-adapted encoders (Qwen3.5's vision tower, Qwen2.5-Omni's audio tower) aren't
-distributed as standalone checkpoints, and the paper itself says they were
-*adapted from* these two source models -- so we use the source models
-directly, which the paper's own dimension numbers (1280 audio, patch-merge
-producing 3072 for nano) line up with:
-  text  : jinaai/jina-embeddings-v5-text-nano (exact match, publicly released)
-  vision: google/siglip2-base-patch16-naflex   (source SigLIP2 GELATO adapted)
-  audio : openai/whisper-large-v3 encoder      (source Whisper GELATO adapted)
-"""
-import torch
-import torch.nn as nn
-from transformers import (
-    AutoModel,
-    AutoTokenizer,
-    AutoImageProcessor,
-    WhisperModel,
-    WhisperFeatureExtractor,
-)
-
-try:
-    from transformers import Siglip2VisionModel as _VisionModelClass
-except ImportError:
-    from transformers import AutoModel as _VisionModelClass
-
-TEXT_MODEL_ID = "jinaai/jina-embeddings-v5-text-nano"
-VISION_MODEL_ID = "google/siglip2-base-patch16-naflex"
-AUDIO_MODEL_ID = "openai/whisper-large-v3"
-
-
-def _freeze(module: nn.Module) -> nn.Module:
-    module.eval()
-    for p in module.parameters():
-        p.requires_grad_(False)
-    return module
-
-
-def mean_pool(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
-    mask = attention_mask.unsqueeze(-1).to(last_hidden_state.dtype)
-    summed = (last_hidden_state * mask).sum(dim=1)
-    counts = mask.sum(dim=1).clamp(min=1e-9)
-    return summed / counts
-
-
-class TextEncoder:
-    def __init__(self, device: str, dtype=torch.bfloat16, max_length: int = 512):
-        self.tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_ID, trust_remote_code=True)
-        self.model = _freeze(
-            AutoModel.from_pretrained(TEXT_MODEL_ID, trust_remote_code=True, dtype=dtype).to(device)
-        )
-        self.device = device
-        self.max_length = max_length
-
-    @torch.no_grad()
-    def embed(self, texts: list[str]) -> torch.Tensor:
-        inputs = self.tokenizer(
-            texts, padding=True, truncation=True, max_length=self.max_length, return_tensors="pt"
-        ).to(self.device)
-        out = self.model(**inputs)
-        return mean_pool(out.last_hidden_state, inputs["attention_mask"])
-
-
-class VisionEncoder:
-    def __init__(self, device: str, dtype=torch.bfloat16):
-        self.processor = AutoImageProcessor.from_pretrained(VISION_MODEL_ID)
-        self.model = _freeze(_VisionModelClass.from_pretrained(VISION_MODEL_ID, dtype=dtype).to(device))
-        self.device = device
-
-    @torch.no_grad()
-    def patch_tokens(self, images: list) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
-        """Returns (patch_tokens [B,N,D], attention_mask [B,N] or None,
-        spatial_shapes [B,2] or None -- per-image (H_patches, W_patches)
-        before NaFlex pads every image in the batch to the same N, needed
-        for a real spatial 2x2 merge instead of grouping arbitrary tokens."""
-        inputs = self.processor(images=images, return_tensors="pt").to(self.device)
-        vision_module = getattr(self.model, "vision_model", self.model)
-        out = vision_module(**inputs)
-        mask = inputs.get("pixel_attention_mask")
-        spatial_shapes = inputs.get("spatial_shapes")
-        return out.last_hidden_state, mask, spatial_shapes
-
-
-class AudioEncoder:
-    def __init__(self, device: str, dtype=torch.bfloat16):
-        self.feature_extractor = WhisperFeatureExtractor.from_pretrained(AUDIO_MODEL_ID)
-        full = WhisperModel.from_pretrained(AUDIO_MODEL_ID, dtype=dtype)
-        self.model = full.encoder
-        del full.decoder
-        del full
-        self.model = _freeze(self.model.to(device))
-        self.device = device
-        self.dtype = dtype
-
-    @torch.no_grad()
-    def frame_tokens(self, audio_arrays: list, sampling_rate: int = 16000) -> torch.Tensor:
-        inputs = self.feature_extractor(
-            audio_arrays, sampling_rate=sampling_rate, return_tensors="pt"
-        ).to(self.device)
-        features = inputs["input_features"].to(self.dtype)
-        out = self.model(features)
-        return out.last_hidden_state

quadembed/projectors.py

@@ -1,111 +0,0 @@
-"""Trainable projectors that map frozen vision/audio encoder outputs into the
-frozen text encoder's embedding space -- the only parameters GELATO trains.
-
-Everything upstream (SigLIP2, Whisper encoder, jina-v5-text-nano) stays frozen;
-only VisionProjector.fc_vision_2 and AudioProjector.fc_audio get gradients.
-"""
-import torch
-import torch.nn as nn
-
-
-class VisionProjector(nn.Module):
-    """LayerNorm -> real 2x2 spatial merge -> trainable linear, matching
-    GELATO's "LayerNorm, 2x2 spatial merge, fc_vision_2" description
-    (nano: 3072->768).
-
-    Uses NaFlex's per-image `spatial_shapes` (H_patches, W_patches) to group
-    each image's own patch grid into genuine 2x2 spatial blocks before
-    concatenating channel-wise -- earlier versions approximated this with
-    sequential 4-token grouping (mixing patches from unrelated rows), which
-    is architecturally wrong and is kept here only as a fallback for
-    callers that don't have spatial_shapes available.
-    """
-
-    def __init__(self, patch_dim: int, out_dim: int, merge: int = 4, hidden_dim: int | None = None):
-        super().__init__()
-        self.merge = merge
-        self.norm = nn.LayerNorm(patch_dim, elementwise_affine=False)
-        if hidden_dim:
-            # not part of GELATO's own recipe (which uses a single linear
-            # fc_vision_2) -- an experiment to test whether a linear
-            # projector's capacity, not data, is capping R@1. See README's
-            # "round 6" for the result.
-            self.fc_vision_2 = nn.Sequential(
-                nn.Linear(patch_dim * merge, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, out_dim)
-            )
-        else:
-            self.fc_vision_2 = nn.Linear(patch_dim * merge, out_dim)
-        self.modality_delim = nn.Parameter(torch.zeros(out_dim))
-
-    def _spatial_merge_one(self, tokens: torch.Tensor, h: int, w: int) -> torch.Tensor:
-        # tokens: [N, D] for one image, first h*w entries are the valid grid
-        d = tokens.shape[-1]
-        h2, w2 = h - (h % 2), w - (w % 2)  # drop a trailing odd row/col rather than fabricate one
-        grid = tokens[: h * w].view(h, w, d)[:h2, :w2]
-        blocks = grid.reshape(h2 // 2, 2, w2 // 2, 2, d).permute(0, 2, 1, 3, 4).reshape(-1, 4 * d)
-        return blocks
-
-    def forward(
-        self,
-        patch_tokens: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        spatial_shapes: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        # patch_tokens: [B, N, patch_dim]
-        b, n, d = patch_tokens.shape
-        patch_tokens = self.norm(patch_tokens)
-
-        if spatial_shapes is not None:
-            pooled = []
-            for i in range(b):
-                h, w = int(spatial_shapes[i, 0]), int(spatial_shapes[i, 1])
-                blocks = self._spatial_merge_one(patch_tokens[i], h, w)
-                projected = self.fc_vision_2(blocks)  # [num_blocks, out_dim]
-                pooled.append(projected.mean(dim=0))
-            pooled = torch.stack(pooled, dim=0)
-        else:
-            pad = (-n) % self.merge
-            pt = patch_tokens
-            if pad:
-                pt = nn.functional.pad(pt, (0, 0, 0, pad))
-                if attention_mask is not None:
-                    attention_mask = nn.functional.pad(attention_mask, (0, pad))
-            merged = pt.reshape(b, pt.shape[1] // self.merge, d * self.merge)
-            projected = self.fc_vision_2(merged)  # [B, N/merge, out_dim]
-            if attention_mask is not None:
-                merged_mask = attention_mask.reshape(b, -1, self.merge).amax(dim=-1).to(projected.dtype)
-                weights = merged_mask.unsqueeze(-1)
-                pooled = (projected * weights).sum(dim=1) / weights.sum(dim=1).clamp(min=1e-6)
-            else:
-                pooled = projected.mean(dim=1)
-
-        return pooled + self.modality_delim
-
-
-class AudioProjector(nn.Module):
-    """Trainable linear (1280->768 for nano), matching GELATO's fc_audio."""
-
-    def __init__(self, in_dim: int, out_dim: int):
-        super().__init__()
-        self.fc_audio = nn.Linear(in_dim, out_dim)
-        self.modality_delim = nn.Parameter(torch.zeros(out_dim))
-
-    def forward(self, frame_tokens: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
-        # frame_tokens: [B, T, in_dim]
-        projected = self.fc_audio(frame_tokens)
-        if attention_mask is not None:
-            weights = attention_mask.unsqueeze(-1).to(projected.dtype)
-            pooled = (projected * weights).sum(dim=1) / weights.sum(dim=1).clamp(min=1e-6)
-        else:
-            pooled = projected.mean(dim=1)
-        return pooled + self.modality_delim
-
-
-def embed_video_from_frames(vision_projector: VisionProjector, per_frame_patch_tokens: list[torch.Tensor]) -> torch.Tensor:
-    """Video = mean of per-frame vision-projector embeddings.
-
-    No dedicated video encoder, matching GELATO: sample frames, run each
-    through the (shared, already-trained) vision projector, pool over time.
-    """
-    frame_embeds = torch.stack([vision_projector(tokens) for tokens in per_frame_patch_tokens], dim=1)  # [B, F, out_dim]
-    return frame_embeds.mean(dim=1)

quadembed/video.py

@@ -1,36 +0,0 @@
-"""Frame sampling + video embedding. No video encoder exists here -- a video
-is just N sampled frames run through the already-trained VisionProjector and
-mean-pooled, exactly as GELATO does."""
-import cv2
-import numpy as np
-import torch
-from PIL import Image
-
-
-def sample_frames(video_path: str, num_frames: int = 4) -> list:
-    cap = cv2.VideoCapture(video_path)
-    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    if total <= 0:
-        cap.release()
-        raise ValueError(f"no frames read from {video_path}")
-    indices = np.linspace(0, total - 1, num=min(num_frames, total), dtype=int)
-    frames = []
-    for idx in indices:
-        cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
-        ok, frame_bgr = cap.read()
-        if not ok:
-            continue
-        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
-        frames.append(Image.fromarray(frame_rgb))
-    cap.release()
-    if not frames:
-        raise ValueError(f"no frames decoded from {video_path}")
-    return frames
-
-
-@torch.no_grad()
-def embed_video(vision_encoder, vision_projector, video_path: str, device: str, num_frames: int = 4) -> torch.Tensor:
-    frames = sample_frames(video_path, num_frames=num_frames)
-    patch_tokens, mask, spatial_shapes = vision_encoder.patch_tokens(frames)  # [num_frames, N, 768]
-    frame_embeds = vision_projector(patch_tokens.float(), mask, spatial_shapes)  # [num_frames, 768]
-    return frame_embeds.mean(dim=0)

requirements.txt

@@ -1,7 +0,0 @@
-torch>=2.1
-transformers>=4.45
-pillow
-opencv-python
-soundfile
-librosa
-numpy