AURA Omni Native Pipeline¶

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""OpenAI-compatible client for the AURA Omni pipeline."""

from __future__ import annotations

import base64
import io
import os

import soundfile as sf
from openai import OpenAI
from vllm.assets.audio import AudioAsset
from vllm.utils.argparse_utils import FlexibleArgumentParser

from vllm_omni.model_executor.stage_input_processors.aura_omni import (
    DEFAULT_QWEN3_TTS_REF_TEXT,
    default_qwen3_tts_ref_audio_path,
)

SEED = 42
DEFAULT_MODEL = "aurateam/AURA"
DEFAULT_VIDEO_URL = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"


def _encode_file(path: str) -> str:
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def _data_url(path: str, default_mime: str) -> str:
    suffix = os.path.splitext(path)[1].lower()
    mime_by_suffix = {
        ".wav": "audio/wav",
        ".mp3": "audio/mpeg",
        ".ogg": "audio/ogg",
        ".flac": "audio/flac",
        ".m4a": "audio/mp4",
        ".mp4": "video/mp4",
        ".webm": "video/webm",
        ".mov": "video/quicktime",
        ".avi": "video/x-msvideo",
        ".mkv": "video/x-matroska",
    }
    return f"data:{mime_by_suffix.get(suffix, default_mime)};base64,{_encode_file(path)}"


def media_url(path_or_url: str | None, *, kind: str) -> str:
    if path_or_url:
        if path_or_url.startswith(("http://", "https://", "data:")):
            return path_or_url
        if not os.path.exists(path_or_url):
            raise FileNotFoundError(f"{kind} file not found: {path_or_url}")
        return _data_url(path_or_url, "audio/wav" if kind == "audio" else "video/mp4")
    if kind == "audio":
        return AudioAsset("mary_had_lamb").url
    return DEFAULT_VIDEO_URL


def sampling_params_list() -> list[dict]:
    return [
        {"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": SEED},
        {
            "temperature": 0.5,
            "top_p": 1.0,
            "top_k": -1,
            "max_tokens": 256,
            "seed": SEED,
            "repetition_penalty": 1.0,
        },
        {
            "temperature": 0.9,
            "top_k": 50,
            "max_tokens": 4096,
            "seed": SEED,
            "detokenize": False,
            "repetition_penalty": 1.05,
            "stop_token_ids": [2150],
        },
        {
            "temperature": 0.0,
            "top_p": 1.0,
            "top_k": -1,
            "max_tokens": 65536,
            "seed": SEED,
            "repetition_penalty": 1.0,
        },
    ]


def parse_modalities(value: str | None) -> list[str] | None:
    if not value:
        return None
    return [item.strip() for item in value.split(",") if item.strip()]


def save_response(response, output_dir: str) -> None:
    os.makedirs(output_dir, exist_ok=True)
    for idx, choice in enumerate(response.choices):
        message = choice.message
        if message.content:
            out_txt = os.path.join(output_dir, f"choice_{idx}.txt")
            with open(out_txt, "w", encoding="utf-8") as f:
                f.write(str(message.content).strip() + "\n")
            print(f"Text saved to {out_txt}")
            print(message.content)
        if getattr(message, "audio", None):
            audio_bytes = base64.b64decode(message.audio.data)
            audio_np, sample_rate = sf.read(io.BytesIO(audio_bytes))
            out_wav = os.path.join(output_dir, f"choice_{idx}.wav")
            sf.write(out_wav, audio_np, int(sample_rate), format="WAV")
            print(f"Audio saved to {out_wav}")


def main(args) -> None:
    client = OpenAI(base_url=f"http://{args.host}:{args.port}/v1", api_key="EMPTY")
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "audio_url", "audio_url": {"url": media_url(args.audio_path, kind="audio")}},
                {"type": "video_url", "video_url": {"url": media_url(args.video_path, kind="video")}},
                {"type": "text", "text": args.prompt},
            ],
        }
    ]
    response = client.chat.completions.create(
        model=args.model,
        messages=messages,
        modalities=parse_modalities(args.modalities),
        extra_body={
            "sampling_params_list": sampling_params_list(),
            "additional_information": {
                "aura_system_prompt": args.aura_system_prompt,
                "tts_task_type": args.tts_task_type,
                "tts_language": args.tts_language,
                "tts_speaker": args.tts_speaker,
                "tts_instruct": args.tts_instruct,
                "tts_ref_audio": args.tts_ref_audio,
                "tts_ref_text": args.tts_ref_text,
                "tts_x_vector_only_mode": args.tts_x_vector_only_mode,
                "tts_pass_token_ids": args.tts_pass_token_ids,
            },
        },
        timeout=args.timeout,
    )
    save_response(response, args.output_dir)


def parse_args():
    parser = FlexibleArgumentParser(description="AURA Omni online serving client")
    parser.add_argument("--host", default="localhost")
    parser.add_argument("--port", type=int, default=8091)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--audio-path", default=None, help="Audio file, URL, or data URL.")
    parser.add_argument("--video-path", default=None, help="Video file, URL, or data URL.")
    parser.add_argument(
        "--prompt",
        default="Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English.",
    )
    parser.add_argument("--modalities", default="text,audio")
    parser.add_argument("--output-dir", default="output_aura_omni_online")
    parser.add_argument(
        "--aura-system-prompt",
        default=(
            "You are receiving a live video stream where the final frame is the present moment. "
            "Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English."
        ),
    )
    parser.add_argument("--tts-task-type", default="Base", choices=["Base", "CustomVoice"])
    parser.add_argument("--tts-language", default="English")
    parser.add_argument("--tts-speaker", default="Vivian")
    parser.add_argument("--tts-instruct", default="")
    parser.add_argument(
        "--tts-ref-audio",
        default=default_qwen3_tts_ref_audio_path(),
        help="Base-mode reference audio path/URL visible to server.",
    )
    parser.add_argument(
        "--tts-ref-text",
        default=DEFAULT_QWEN3_TTS_REF_TEXT,
        help="Base-mode reference audio transcript.",
    )
    parser.add_argument(
        "--tts-x-vector-only-mode",
        action="store_true",
        help="Use speaker embedding only for Base mode (disable ICL ref_text conditioning).",
    )
    parser.add_argument(
        "--tts-pass-token-ids",
        action="store_true",
        help="Pass AURA-generated assistant token ids directly to Qwen3-TTS. Defaults to sending text.",
    )
    parser.add_argument("--timeout", type=float, default=600.0)
    return parser.parse_args()


if __name__ == "__main__":
    main(parse_args())