LCO-Embedding-Omni

Quick Start

Note: We are only using the thinker component of Qwen2.5 Omni and drops the talker component.

Using Sentence Transformers

Install Sentence Transformers with the multimodal extras (for image, audio, and video support):

pip install "sentence_transformers[image,audio,video]" "transformers>=5.6.0"

import torch
from sentence_transformers import SentenceTransformer

model = SentenceTransformer(
    "LCO-Embedding/LCO-Embedding-Omni-3B",
    model_kwargs={
        "torch_dtype": torch.bfloat16,
        "attn_implementation": "flash_attention_2",  # pip install kernels; recommended but not mandatory
    },
)

The same "Summarize the above in one word:" instruction used in the paper is baked into the chat template, so encode() takes plain text, file paths, URLs, or multimodal dicts directly.

Text Retrieval

query = "What is the tallest mountain in the world?"
documents = [
    "Mount Everest is Earth's highest mountain above sea level, located in the Mahalangur Himal sub-range of the Himalayas. Its elevation of 8,848.86 metres was established by a joint Chinese-Nepali survey in 2020.",
    "K2, at 8,611 metres above sea level, is the second-highest mountain on Earth, after Mount Everest. It lies in the Karakoram range on the China-Pakistan border.",
    "Mount Kilimanjaro is a dormant volcano in Tanzania. It is the highest mountain in Africa, with its summit about 5,895 metres above sea level.",
]

query_embedding = model.encode(query)
document_embeddings = model.encode(documents)
print(model.similarity(query_embedding, document_embeddings))
# tensor([[0.6199, 0.5585, 0.5233]])

Image Retrieval

query = "How many input modalities does Qwen2.5-Omni support?"
documents = [
    "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/qwen2.5omni_hgf.png",
    "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/llama4_hgf.png",
]

query_embedding = model.encode(query)
document_embeddings = model.encode(documents, batch_size=1)
print(model.similarity(query_embedding, document_embeddings))
# tensor([[0.4396, 0.3418]])

Audio Retrieval

query = "A light piano piece"
documents = [
    "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/joe_hisaishi_summer.mp3",
    "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/jay_chou_superman_cant_fly.mp3",
]

query_embedding = model.encode(query)
document_embeddings = model.encode(documents, batch_size=1)
print(model.similarity(query_embedding, document_embeddings))
# tensor([[0.3809, 0.0858]])

Video Retrieval

# For video on smaller GPUs, cap the processor up front:
model[0].processing_kwargs.update({
    "video": {"max_pixels": 64 * 28 * 28, "do_sample_frames": True, "fps": 1},
})

query = "How to cook Mapo Tofu?"
documents = [
    "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/mapo_tofu.mp4",
    "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/zhajiang_noodle.mp4",
]

query_embedding = model.encode(query)
document_embeddings = model.encode(documents, batch_size=1)
print(model.similarity(query_embedding, document_embeddings))
# tensor([[0.6406, 0.5033]])

Multimodal Inputs

To embed a document that combines multiple modalities, pass a dict with any combination of "text", "image", "audio", and "video" keys instead of a single path or string:

documents = [
    {
        "text": "A cooking tutorial for Mapo Tofu",
        "video": "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/mapo_tofu.mp4",
    },
    {
        "image": "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/qwen2.5omni_hgf.png",
        "audio": "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/joe_hisaishi_summer.mp3",
    },
]
document_embeddings = model.encode(documents, batch_size=1)

Using Transformers

from transformers import Qwen2_5OmniThinkerForConditionalGeneration, Qwen2_5OmniProcessor
from qwen_omni_utils import process_mm_info

processor = Qwen2_5OmniProcessor.from_pretrained("LCO-Embedding/LCO-Embedding-Omni-3B") # or add a `max_pixels = 1280*28*28' for efficient encoding
model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained("LCO-Embedding/LCO-Embedding-Omni-3B",
                                                                    torch_dtype=torch.bfloat16,
                                                                    device_map="auto")

Text Batch Encodings:

texts = ["some random text", "a second random text", "a third random text"] * 30
batch_size = 8
text_prompt =  "{}\nSummarize the above text in one word:"

all_text_embeddings = []

with torch.no_grad():
    for i in tqdm(range(0, len(texts), batch_size)):
        batch_texts = texts[i : i + batch_size]
        batch_texts = [text_prompt.format(text) for text in batch_texts]
        messages = [[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text":text},
                ],

            }
        ] for text in batch_texts]
        text_inputs = processor.apply_chat_template(messages, tokenize = False, add_generation_prompt = True)
        text_inputs = processor(
        text = text_inputs,
        padding = True,
        return_tensors = "pt",
        )
        text_inputs = text_inputs.to("cuda")
        text_outputs = model(
            **text_inpu