README.md

---
license: other
library_name: pytorch
pipeline_tag: feature-extraction
base_model:
- Alibaba-NLP/gte-Qwen2-1.5B-instruct
- google/siglip-so400m-patch14-384
- Qwen/Qwen-Audio
tags:
- multimodal
- embedding
- retrieval
- audio
- image
- video
datasets:
- chuonghm/ACM
- chuonghm/OmniRet
---
# 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](https://arxiv.org/abs/2603.02098). 
**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

```bash
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

```python
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:

```python
[
    [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:

```python
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](examples/audios/-1rZFviqTTQ_000003.wav) | ![Example frame](examples/images/-1rZFviqTTQ_000003.jpg) | [MP4](examples/videos/-1rZFviqTTQ_000003.mp4) |
| `-25e5qcELvw_000011` | [WAV](examples/audios/-25e5qcELvw_000011.wav) | ![Example frame](examples/images/-25e5qcELvw_000011.jpg) | [MP4](examples/videos/-25e5qcELvw_000011.mp4) |

## ACM performance

Results below are retrieval recall percentages on
[`chuonghm/ACM`](https://huggingface.co/datasets/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.

| Direction | R@1 | R@5 | R@10 |
| --- | ---: | ---: | ---: |
| A,T to A | 6.821924 | 22.347683 | 32.839332 |
| A to I | 7.585139 | 21.981424 | 32.739938 |
| I to A | 6.501548 | 21.749226 | 32.585139 |
| A to V | 17.569659 | 45.201238 | 59.597523 |
| V to A | 19.349845 | 43.034056 | 56.114551 |

## 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](https://huggingface.co/datasets/chuonghm/ACM) and
[VGG-Sound](https://www.robots.ox.ac.uk/~vgg/data/vggsound/).

## Citation

```bibtex
@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}
}
```