AuroLA-Omni-3B

AuroLA-Omni-3B

AuroLA-Omni-3B is a BF16 omni-modal embedding model built from both Qwen2.5-Omni-3B and AuroLA-3B. It supports text, image, audio, and video inputs through a SentenceTransformers-compatible checkpoint layout.

MMEB-V3 Results

Usage

SentenceTransformers

pip install "sentence_transformers[image,audio,video]" "transformers>=5.6.0" qwen-omni-utils

import torch
from sentence_transformers import SentenceTransformer

model_id = "Jazzcharles/AuroLA-Omni-3B"

model = SentenceTransformer(
    model_id,
    model_kwargs={
        "torch_dtype": torch.bfloat16,
        "attn_implementation": "sdpa",  # use "flash_attention_2" if flash-attn is installed
    },
)

model[0].processing_kwargs.update({
    "video": {"max_pixels": 64 * 28 * 28, "do_sample_frames": True, "fps": 1},
})

query_embedding = model.encode_query("A person cooking in a kitchen")

image_embeddings = model.encode_document(["path/to/image.jpg"], batch_size=1)
video_embeddings = model.encode_document(["path/to/video.mp4"], batch_size=1)
audio_embeddings = model.encode_document(["path/to/audio.wav"], batch_size=1)

print("image:", model.similarity(query_embedding, image_embeddings))
print("video:", model.similarity(query_embedding, video_embeddings))
print("audio:", model.similarity(query_embedding, audio_embeddings))

Transformers

When using raw Transformers, use the sentence_transformers chat template explicitly and take the last non-padding hidden state.

import torch
import torch.nn.functional as F
from transformers import AutoProcessor, Qwen2_5OmniThinkerForConditionalGeneration
from qwen_omni_utils import process_mm_info

model_id = "Jazzcharles/AuroLA-Omni-3B"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",  # use "flash_attention_2" if flash-attn is installed
).to(device).eval()

processor.tokenizer.padding_side = "left"


def encode_message(message):
    text = processor.apply_chat_template(
        message,
        chat_template="sentence_transformers",
        tokenize=False,
        add_generation_prompt=True,
    )
    if isinstance(text, list):
        text = text[0]
    text = text + "<|endoftext|>"

    audio_inputs, image_inputs, video_inputs = process_mm_info(
        message,
        use_audio_in_video=False,
    )
    inputs = processor(
        text=text,
        audio=audio_inputs,
        images=image_inputs,
        videos=video_inputs,
        videos_kwargs={"do_sample_frames": False, "fps": 1},
        return_tensors="pt",
        padding="longest",
    ).to(device)

    with torch.inference_mode():
        outputs = model(
            **inputs,
            return_dict=True,
            output_hidden_states=True,
            use_cache=False,
        )

    hidden_states = outputs.hidden_states[-1]
    attention_mask = inputs["attention_mask"]
    last_token_indices = attention_mask.size(1) - 1 - attention_mask.flip(dims=[1]).argmax(dim=1)
    reps = hidden_states[torch.arange(hidden_states.size(0), device=device), last_token_indices]
    return F.normalize(reps, p=2, dim=-1)


query = [{"role": "user", "content": [{"type": "text", "text": "A person cooking in a kitchen"}]}]
image = [{"role": "user", "content": [{"type": "image", "image": "path/to/image.jpg"}]}]
video = [{
    "role": "user",
    "content": [{
        "type": "video",
        "video": "path/to/video.mp4",
        "max_pixels": 128 * 28 * 28,
        "fps": 1,
    }],
}]
audio = [{"role": "user", "content": [{"type": "audio", "audio": "path/to/audio.wav"}]}]

query_embedding = encode_message(query)
for name, document in {"image": image, "video": video, "audio": audio}.items():
    similarity = torch.cosine_similarity(query_embedding, encode_message(document))
    print(name, similarity.item())

Citation

If you use this model, please cite:

@misc{xu2026scalingaudiotextretrievalmultimodal,
  title={Scaling Audio-Text Retrieval with Multimodal Large Language Models},
  author={Jilan Xu and Carl Thom{\'e} and Danijela Horak and Weidi Xie and Andrew Zisserman},
  year={2026},
  eprint={2602.18010},
  archivePrefix={arXiv},
  primaryClass={cs.SD},
  url={https://arxiv.org/abs/2602.18010}
}