wzliu
33e03b44c5
ray + mooncake connector work on 2 nodes, but needes several updates\n1.start ray inside vllm-omni \n2.placement group optimize \n3.fix close
313 lines
11 KiB
Python
313 lines
11 KiB
Python
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
|