vllm-omni/examples/offline_inference/qwen2_5_omni/utils.py

import tempfile
from typing import Optional, Union
from urllib.request import urlopen

import librosa
import requests
import resampy
import soundfile as sf
import torch
import torchvision.io
from processing_omni import fetch_image, fetch_video
from transformers import AutoConfig, AutoProcessor
from vllm.inputs import TextPrompt

from vllm_omni.inputs.data import OmniTokensPrompt

# Simple caches to avoid repeated heavy HF loads per prompt
_PROCESSOR_CACHE: dict[str, "AutoProcessor"] = {}
_CONFIG_CACHE: dict[str, "AutoConfig"] = {}


def get_system_prompt():
    return {
        "role": "system",
        "content": [
            {
                "type": "text",
                "text": (
                    "You are Qwen, a virtual human developed by the Qwen Team, "
                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
                    "as well as generating text and speech."
                ),
            }
        ],
    }


def resample_wav_to_16khz(input_filepath):
    data, original_sample_rate = sf.read(input_filepath)
    # Only use the first channel
    if len(data.shape) > 1:
        data = data[:, 0]
    # resample to 16kHz
    data_resampled = resampy.resample(data, sr_orig=original_sample_rate, sr_new=16000)
    return data_resampled


def fetch_and_read_video(args, video_url: str, fps=2):
    def read_video_with_torchvision(video_file_name: str):
        video, audio, info = torchvision.io.read_video(
            video_file_name,
            start_pts=0.0,
            end_pts=None,
            pts_unit="sec",
            output_format="TCHW",
        )

        total_frames, video_fps = video.size(0), info["video_fps"]
        total_duration = round(total_frames / video_fps, 3)
        nframes = int(total_frames / video_fps * fps)

        frame_timestamps = total_duration * torch.arange(1, nframes + 1) / nframes
        grid_timestamps = frame_timestamps[::2]
        second_per_grid = grid_timestamps[1] - grid_timestamps[0]

        idx = torch.linspace(0, video.size(0) - 1, nframes).round().long()
        video = video[idx]

        if args.legacy_omni_video:
            return [video, total_duration, nframes, second_per_grid.item()]
        else:
            return video

    def read_video_with_transformers(video_file_name: Union[str, list[str]]):
        video, total_duration, nframes, second_per_grid = fetch_video({"video": video_file_name})
        if total_duration is None and nframes is None:
            nframes = len(video)
            total_duration = 0.5 * nframes
            second_per_grid = 1.0
        if args.legacy_omni_video:
            return [video, total_duration, nframes, second_per_grid]
        else:
            return video

    def read_video(video_file_name: str):
        if args.use_torchvision:
            return read_video_with_torchvision(video_file_name)
        else:
            return read_video_with_transformers(video_file_name)

    if isinstance(video_url, str) and video_url.startswith("http"):
        with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
            resp = requests.get(video_url)
            assert resp.status_code == requests.codes.ok, (
                f"Failed to fetch video from {video_url}, status_code:{resp.status_code}, resp:{resp}"
            )

            temp_video_file.write(urlopen(video_url).read())
            temp_video_file_path = temp_video_file.name
            video_file_name = temp_video_file_path
            return read_video(video_file_name)
    else:
        video_file_name = video_url
        return read_video(video_file_name)


def make_inputs_qwen2_omni(
    args,
    messages: list[dict[str, Union[str, list[dict[str, str]]]]],
    use_audio_in_video: Optional[bool] = False,
    tokenize: bool = False,
) -> Union[OmniTokensPrompt, TextPrompt]:
    from transformers import AutoConfig, AutoProcessor

    # Cached processor/config to prevent per-prompt reloading and repeated warnings
    if args.model not in _PROCESSOR_CACHE:
        _PROCESSOR_CACHE[args.model] = AutoProcessor.from_pretrained(args.model)
    processor = _PROCESSOR_CACHE[args.model]

    config = _CONFIG_CACHE.get(args.model)
    if config is None:
        try:
            config = AutoConfig.from_pretrained(args.model)
        except Exception:
            config = None
        _CONFIG_CACHE[args.model] = config  # cache even if None to avoid retry storms

    # Decide legacy flag only once based on config (default True if unknown)
    if getattr(args, "legacy_omni_video", None) is None:
        if config is not None and hasattr(config, "architectures"):
            args.legacy_omni_video = "Qwen2_5OmniModel" not in config.architectures
        else:
            args.legacy_omni_video = True

    audios, images, videos = [], [], []
    for message in messages:
        if not isinstance(message["content"], list):
            message["content"] = [
                {
                    "type": "text",
                    "text": message["content"],
                }
            ]
        index, num_contents = 0, len(message["content"])
        while index < num_contents:
            ele = message["content"][index]
            if "type" not in ele:
                if "text" in ele:
                    ele["type"] = "text"
                elif "audio" in ele:
                    ele["type"] = "audio"
                elif "audio_url" in ele:
                    ele["type"] = "audio_url"
                elif "image" in ele:
                    ele["type"] = "image"
                elif "image_url" in ele:
                    ele["type"] = "image_url"
                elif "video" in ele:
                    ele["type"] = "video"
                elif "video_url" in ele:
                    ele["type"] = "video_url"
                else:
                    raise ValueError(f"Unknown ele: {ele}")

            if ele["type"] == "audio" or ele["type"] == "audio_url":
                if "audio_url" in ele:
                    audio_key = "audio_url"
                    with tempfile.NamedTemporaryFile(delete=True) as temp_audio_file:
                        temp_audio_file.write(urlopen(ele[audio_key]).read())
                        temp_audio_file_path = temp_audio_file.name
                        audios.append(resample_wav_to_16khz(temp_audio_file_path))
                        ele["audio"] = temp_audio_file_path
                elif "audio" in ele:
                    audio_key = "audio"
                    audios.append(resample_wav_to_16khz(ele[audio_key]))
                else:
                    raise ValueError(f"Unknown ele {ele}")
            elif use_audio_in_video and (ele["type"] == "video" or ele["type"] == "video_url"):
                # use video as audio as well
                if "video_url" in ele:
                    audio_key = "video_url"
                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
                        temp_video_file.write(urlopen(ele[audio_key]).read())
                        temp_video_file_path = temp_video_file.name
                        ele[audio_key] = temp_video_file_path
                        audios.append(librosa.load(temp_video_file_path, sr=16000)[0])
                        videos.append(fetch_and_read_video(args, temp_video_file_path))
                        ele["video"] = temp_video_file_path
                elif "video" in ele:
                    audio_key = "video"
                    audios.append(librosa.load(ele[audio_key], sr=16000)[0])
                    videos.append(fetch_and_read_video(args, audio_key))
                else:
                    raise ValueError(f"Unknown ele {ele}")
                # insert a audio after the video
                message["content"].insert(
                    index + 1,
                    {
                        "type": "audio",
                        "audio": ele[audio_key],
                    },
                )
                # no need to load the added audio again
                index += 1
            elif ele["type"] == "video" or ele["type"] == "video_url":
                if "video_url" in ele:
                    video_key = "video_url"
                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
                        temp_video_file.write(urlopen(ele["video_url"]).read())
                        temp_video_file_path = temp_video_file.name
                        videos.append(fetch_and_read_video(args, temp_video_file))
                        ele["video"] = temp_video_file_path
                else:
                    video_key = "video"
                    videos.append(fetch_and_read_video(args, ele[video_key]))
            elif ele["type"] == "image" or ele["type"] == "image_url":
                images.append(fetch_image(ele))

            # move to the next content
            index += 1

    prompt = processor.apply_chat_template(
        messages,
        tokenize=tokenize,
        add_generation_prompt=True,
        add_vision_id=True,
    )

    audios = audios if len(audios) > 0 else None
    images = images if len(images) > 0 else None
    videos = videos if len(videos) > 0 else None

    multi_modal_data = {}
    if audios:
        multi_modal_data["audio"] = audios
    if images:
        multi_modal_data["image"] = images
    if videos:
        multi_modal_data["video"] = videos

    if isinstance(prompt, list) and isinstance(prompt[0], (list, str)):
        prompt = prompt[0]

    if tokenize:
        return OmniTokensPrompt(
            prompt_token_ids=prompt,
            multi_modal_data=multi_modal_data,
        )
    else:
        return TextPrompt(
            prompt=prompt,
            multi_modal_data=multi_modal_data,
        )


def make_text_prompt(args, prompt):
    messages = [
        get_system_prompt(),
        {
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
            ],
        },
    ]

    prompt = make_inputs_qwen2_omni(args, messages, tokenize=args.tokenize)
    return prompt


def make_audio_in_video_v2_prompt(args):
    messages = [
        {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": (
                        "You are Qwen, a virtual human developed by the Qwen Team, "
                        "Alibaba Group, capable of perceiving auditory and visual "
                        "inputs, as well as generating text and speech."
                    ),
                }
            ],
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "video_url",
                    "video_url": ("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4"),
                },
            ],
        },
    ]
    prompt = make_inputs_qwen2_omni(
        args,
        messages,
        use_audio_in_video=True,
        tokenize=args.tokenize,
    )
    return prompt


def make_omni_prompt(args, prompt=None) -> Union[OmniTokensPrompt, list[OmniTokensPrompt]]:
    if args.prompt_type == "text":
        prompt = make_text_prompt(args, prompt)
    elif args.prompt_type == "audio-in-video-v2":
        prompt = make_audio_in_video_v2_prompt(args)
    else:
        raise ValueError(f"Unsupported prompt type: {args.prompt_type}")
    return prompt