OmniRet

OmniRet is an instruction-aware embedding model for unified text, audio, image, and video retrieval. It maps every supported input to a normalized 4,096-dimensional vector; matrix multiplication therefore computes cosine similarity.

This repository publishes the consolidated Stage-2 checkpoint from Efficient and High-Fidelity Omni Modality Retrieval. However, due to a cluster issue, the original checkpoint is lost. This published model is trained on a smaller setting and does not have the same exact quantitative results in the paper.

Highlights

• Text, mono WAV audio, images, and videos share one embedding space.
• Inputs may combine text with one media type.
• Instructions are supplied per query instead of being forced onto documents.
• Output scores are cosine similarities in [-1, 1], not probabilities.
• The repository includes raw ACM media and a checked end-to-end example.

Model details

OmniRet uses Alibaba-NLP/gte-Qwen2-1.5B-instruct for text, google/siglip-so400m-patch14-384 for vision, and the Qwen-Audio encoder for audio. Modality projectors, shared media resampling, and Attention Sliced Wasserstein Pooling produce 4,096-dimensional embeddings. Videos use at most eight uniformly sampled frames.

The first initialization downloads the three pinned base-model snapshots. The repository config.json records their exact revisions.

Installation

git clone https://huggingface.co/chuonghm/OmniRet
cd OmniRet
pip install -r requirements.txt

A CUDA GPU is recommended. CPU loading is supported but substantially slower.

Usage

from pathlib import Path

import torch

from scripts.omniret_embedding import OmniRetEmbedder

root = Path("examples")
ids = ("-1rZFviqTTQ_000003", "-25e5qcELvw_000011")

queries = [
    {
        "instruction": "Retrieve the video that aligns with the audio.",
        "audio": root / "audios" / f"{ids[0]}.wav",
    },
    {
        "instruction": "Retrieve the audio that matches the given image.",
        "image": root / "images" / f"{ids[1]}.jpg",
    },
    {
        "instruction": "Retrieve the audio that matches the given video.",
        "video": root / "videos" / f"{ids[0]}.mp4",
    },
]
documents = [
    {modality: root / f"{modality}s" / f"{media_id}.{extension}"}
    for modality, extension in (("audio", "wav"), ("image", "jpg"), ("video", "mp4"))
    for media_id in ids
]

model = OmniRetEmbedder("chuonghm/OmniRet", torch_dtype=torch.bfloat16)

embeddings = model.process(queries + documents)
similarity_scores = embeddings[: len(queries)] @ embeddings[len(queries) :].T
print(similarity_scores.tolist())

Expected scores, with documents ordered as audio-1, audio-2, image-1, image-2, video-1, video-2:

[
    [0.32468632, -0.03668483, 0.28203371, 0.04291208, 0.42771006, 0.06420071],
    [-0.03580582, 0.16362236, -0.03176199, 0.33771127, -0.00980220, 0.16253276],
    [0.41171762, 0.00378584, 0.22568333, 0.01310393, 0.43492085, 0.08068858],
]

Run the checked copy with python examples/compute_similarity.py.

Each input accepts instruction, text, one of audio/image/video, and max_frames from 1 through 8. Media may be a local path or HTTP(S) URL.

If FlashAttention 2 is installed, enable it when constructing the embedder:

model = OmniRetEmbedder(
    "chuonghm/OmniRet",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)

Raw examples

The release includes two raw files for each media type:

ACM ID	Audio	Image	Video
-1rZFviqTTQ_000003	WAV		MP4
-25e5qcELvw_000011	WAV		MP4

ACM performance

Results below are retrieval recall percentages on chuonghm/ACM. A,T to A uses 4,251 queries and 4,251 candidates. Every audio-visual direction uses 1,292 queries and 5,480 candidates.

Limitations

• Each item supports at most one media type plus optional text.
• Audio must be mono 16-bit 16 kHz PCM WAV.
• This release does not provide AutoModel, vLLM, quantized, or hosted API integration.
• Scores may differ by about 1e-3 across devices and kernels.

License and media terms

The checkpoint includes a Qwen-Audio-derived component and is distributed under the Tongyi Qianwen License Agreement in LICENSE. Required attribution and the Apache-2.0 base-model notices are in NOTICE.

ACM benchmark metadata is MIT-licensed. The bundled example media derives from VGG-Sound and remains subject to VGG-Sound and source-video terms. See the ACM dataset card and VGG-Sound.

Citation

@article{huynh2026omniret,
  title     = {Efficient and High-Fidelity Omni Modality Retrieval},
  author    = {Huynh, Chuong and Luong, Manh and Shrivastava},
  journal   = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year      = {2026}
}