Shenglei FuandHongsheng Liu
47b1893676
Signed-off-by: Shenglei Fu <sfu@confluent.io> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
848 lines
30 KiB
Python
848 lines
30 KiB
Python
"""Synthetic media generation and media/text utilities for tests."""
|
|
|
|
import atexit
|
|
import base64
|
|
import concurrent.futures
|
|
import hashlib
|
|
import io
|
|
import logging
|
|
import math
|
|
import multiprocessing
|
|
import os
|
|
import random
|
|
import re
|
|
import subprocess
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
from concurrent.futures.process import BrokenProcessPool
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
from PIL import Image
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
_synthetic_media_fallback_dir: Path | None = None
|
|
|
|
|
|
def _resolve_synthetic_media_cache_dir(cache_dir: Path | str | None) -> Path:
|
|
if cache_dir is not None:
|
|
return Path(cache_dir).expanduser().resolve()
|
|
|
|
default = Path(tempfile.gettempdir()) / "vllm_omni_test_synthetic_media"
|
|
try:
|
|
default.mkdir(parents=True, exist_ok=True)
|
|
# Verify write access: the directory may exist but belong to
|
|
# another user (e.g. a previous CI job), causing PermissionError
|
|
# later when individual files are saved.
|
|
canary = default / ".write_test"
|
|
canary.touch()
|
|
canary.unlink()
|
|
return default
|
|
except (PermissionError, OSError):
|
|
global _synthetic_media_fallback_dir
|
|
if _synthetic_media_fallback_dir is None:
|
|
_synthetic_media_fallback_dir = Path(tempfile.mkdtemp(prefix="vllm_omni_test_synthetic_media_"))
|
|
return _synthetic_media_fallback_dir
|
|
|
|
|
|
def _np_array_from_mp4_bytes(video_bytes: bytes) -> np.ndarray:
|
|
"""Decode MP4 bytes to a (T, H, W, 3) uint8 RGB stack (matches in-memory synthetic frames)."""
|
|
import cv2
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
|
|
tmp.write(video_bytes)
|
|
path = tmp.name
|
|
cap = None
|
|
try:
|
|
cap = cv2.VideoCapture(path)
|
|
if not cap.isOpened():
|
|
raise RuntimeError("Failed to open cached synthetic video for decode")
|
|
frames: list[np.ndarray] = []
|
|
while True:
|
|
ok, frame_bgr = cap.read()
|
|
if not ok:
|
|
break
|
|
frames.append(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
|
|
if not frames:
|
|
raise RuntimeError("Cached synthetic video has no decodable frames")
|
|
return np.stack(frames, axis=0)
|
|
finally:
|
|
if cap is not None:
|
|
cap.release()
|
|
try:
|
|
os.unlink(path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def generate_synthetic_audio(
|
|
duration: int,
|
|
num_channels: int,
|
|
sample_rate: int = 48000,
|
|
*,
|
|
phrase_text: str = "test",
|
|
force_regenerate: bool = False,
|
|
cache_dir: Path | str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Generate TTS speech with pyttsx3 and return base64 string.
|
|
|
|
Caches the WAV under ``cache_dir`` when given, else under the default temp
|
|
subdirectory. Reuses the file when the same
|
|
``duration`` / ``num_channels`` / ``sample_rate`` / ``phrase_text`` are
|
|
requested unless ``force_regenerate`` is true.
|
|
|
|
The cache filename includes a SHA-256 digest of ``phrase_text`` so different
|
|
phrases never share a WAV cache entry.
|
|
"""
|
|
root = _resolve_synthetic_media_cache_dir(cache_dir)
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
phrase_key = hashlib.sha256(phrase_text.encode("utf-8")).hexdigest()
|
|
cache_path = root / f"synth_audio_d{duration}_ch{num_channels}_sr{sample_rate}_pt{phrase_key}.wav"
|
|
|
|
if not force_regenerate and cache_path.is_file():
|
|
data, _sr = sf.read(str(cache_path), dtype="float32", always_2d=True)
|
|
audio_bytes = cache_path.read_bytes()
|
|
return {
|
|
"np_array": np.asarray(data, dtype=np.float32),
|
|
"base64": base64.b64encode(audio_bytes).decode("utf-8"),
|
|
"file_path": str(cache_path.resolve()),
|
|
}
|
|
|
|
import pyttsx3
|
|
|
|
def _pick_voice(engine: pyttsx3.Engine) -> str | None:
|
|
voices = engine.getProperty("voices")
|
|
if not voices:
|
|
return None
|
|
|
|
preferred_tokens = (
|
|
"natural",
|
|
"jenny",
|
|
"sonia",
|
|
"susan",
|
|
"zira",
|
|
"aria",
|
|
"hazel",
|
|
"samantha",
|
|
"ava",
|
|
"allison",
|
|
"female",
|
|
"woman",
|
|
"english-us",
|
|
"en-us",
|
|
"english",
|
|
)
|
|
discouraged_tokens = (
|
|
"espeak",
|
|
"robot",
|
|
"mbrola",
|
|
"microsoft david",
|
|
"male",
|
|
"man",
|
|
)
|
|
|
|
best_voice = voices[0]
|
|
best_score = float("-inf")
|
|
for voice in voices:
|
|
voice_text = f"{getattr(voice, 'id', '')} {getattr(voice, 'name', '')}".lower()
|
|
voice_languages = " ".join(
|
|
lang.decode(errors="ignore") if isinstance(lang, bytes) else str(lang)
|
|
for lang in getattr(voice, "languages", [])
|
|
).lower()
|
|
combined_text = f"{voice_text} {voice_languages}"
|
|
score = 0
|
|
for idx, token in enumerate(preferred_tokens):
|
|
if token in combined_text:
|
|
score += 20 - idx
|
|
for token in discouraged_tokens:
|
|
if token in combined_text:
|
|
score -= 10
|
|
if "english" in combined_text or "en_" in combined_text or "en-" in combined_text:
|
|
score += 4
|
|
if "en-us" in combined_text or "english-us" in combined_text:
|
|
score += 4
|
|
if score > best_score:
|
|
best_score = score
|
|
best_voice = voice
|
|
|
|
return best_voice.id
|
|
|
|
def _resample_audio(audio: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray:
|
|
if src_sr == dst_sr or len(audio) == 0:
|
|
return audio.astype(np.float32)
|
|
src_len = audio.shape[0]
|
|
dst_len = max(1, int(round(src_len * float(dst_sr) / float(src_sr))))
|
|
src_idx = np.arange(src_len, dtype=np.float32)
|
|
dst_idx = np.linspace(0, src_len - 1, dst_len, dtype=np.float32)
|
|
resampled_channels: list[np.ndarray] = []
|
|
for ch in range(audio.shape[1]):
|
|
resampled_channels.append(np.interp(dst_idx, src_idx, audio[:, ch]).astype(np.float32))
|
|
return np.stack(resampled_channels, axis=1)
|
|
|
|
def _match_channels(audio: np.ndarray, target_channels: int) -> np.ndarray:
|
|
current_channels = audio.shape[1]
|
|
if current_channels == target_channels:
|
|
return audio.astype(np.float32)
|
|
if target_channels == 1:
|
|
return np.mean(audio, axis=1, keepdims=True, dtype=np.float32)
|
|
if current_channels == 1:
|
|
return np.repeat(audio, target_channels, axis=1).astype(np.float32)
|
|
collapsed = np.mean(audio, axis=1, keepdims=True, dtype=np.float32)
|
|
return np.repeat(collapsed, target_channels, axis=1).astype(np.float32)
|
|
|
|
def _trim_silence(audio: np.ndarray, threshold: float = 0.01) -> np.ndarray:
|
|
if len(audio) == 0:
|
|
return audio
|
|
energy = np.max(np.abs(audio), axis=1)
|
|
voiced = np.where(energy > threshold)[0]
|
|
if len(voiced) == 0:
|
|
return audio
|
|
start = max(0, int(voiced[0]) - int(sample_rate * 0.02))
|
|
end = min(len(audio), int(voiced[-1]) + int(sample_rate * 0.04) + 1)
|
|
return audio[start:end]
|
|
|
|
def _enhance_speech(audio: np.ndarray) -> np.ndarray:
|
|
if len(audio) == 0:
|
|
return audio.astype(np.float32)
|
|
enhanced = audio.astype(np.float32).copy()
|
|
enhanced -= np.mean(enhanced, axis=0, keepdims=True, dtype=np.float32)
|
|
if len(enhanced) > 1:
|
|
preemphasis = enhanced.copy()
|
|
preemphasis[1:] = enhanced[1:] - 0.94 * enhanced[:-1]
|
|
enhanced = 0.7 * enhanced + 0.3 * preemphasis
|
|
enhanced = np.sign(enhanced) * np.sqrt(np.abs(enhanced))
|
|
fade = min(len(enhanced) // 4, max(1, int(sample_rate * 0.01)))
|
|
if fade > 1:
|
|
ramp_in = np.linspace(0.0, 1.0, fade, dtype=np.float32)
|
|
ramp_out = np.linspace(1.0, 0.0, fade, dtype=np.float32)
|
|
enhanced[:fade] *= ramp_in[:, None]
|
|
enhanced[-fade:] *= ramp_out[:, None]
|
|
peak = float(np.max(np.abs(enhanced)))
|
|
if peak > 1e-8:
|
|
enhanced = enhanced / peak * 0.95
|
|
return enhanced.astype(np.float32)
|
|
|
|
num_samples = int(sample_rate * max(1, duration))
|
|
audio_data = np.zeros((num_samples, num_channels), dtype=np.float32)
|
|
|
|
engine = pyttsx3.init()
|
|
engine.setProperty("rate", 112)
|
|
engine.setProperty("volume", 1.0)
|
|
selected_voice = _pick_voice(engine)
|
|
if selected_voice is not None:
|
|
engine.setProperty("voice", selected_voice)
|
|
|
|
temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
|
|
temp_wav.close()
|
|
try:
|
|
engine.save_to_file(phrase_text, temp_wav.name)
|
|
engine.runAndWait()
|
|
engine.stop()
|
|
|
|
ready = False
|
|
for _ in range(50):
|
|
if os.path.exists(temp_wav.name) and os.path.getsize(temp_wav.name) > 44:
|
|
ready = True
|
|
break
|
|
time.sleep(0.1)
|
|
if not ready:
|
|
raise RuntimeError("pyttsx3 did not produce a WAV file in time.")
|
|
|
|
tts_audio, tts_sr = sf.read(temp_wav.name, dtype="float32", always_2d=True)
|
|
finally:
|
|
if os.path.exists(temp_wav.name):
|
|
os.unlink(temp_wav.name)
|
|
|
|
if len(tts_audio) == 0:
|
|
raise RuntimeError("pyttsx3 produced an empty WAV file.")
|
|
|
|
tts_audio = _resample_audio(tts_audio, tts_sr, sample_rate)
|
|
tts_audio = _match_channels(tts_audio, num_channels)
|
|
tts_audio = _trim_silence(tts_audio, threshold=0.012)
|
|
tts_audio = _enhance_speech(tts_audio)
|
|
|
|
lead_silence = min(int(sample_rate * 0.02), num_samples // 8)
|
|
pause_samples = int(sample_rate * 0.18)
|
|
start = lead_silence
|
|
phrase_len = tts_audio.shape[0]
|
|
while start < num_samples:
|
|
take = min(phrase_len, num_samples - start)
|
|
audio_data[start : start + take] = tts_audio[:take]
|
|
start += phrase_len + pause_samples
|
|
|
|
max_amp = float(np.max(np.abs(audio_data)))
|
|
if max_amp > 0:
|
|
audio_data = audio_data / max_amp * 0.95
|
|
|
|
sf.write(str(cache_path), audio_data, sample_rate, format="WAV", subtype="PCM_16")
|
|
audio_bytes = cache_path.read_bytes()
|
|
|
|
return {
|
|
"np_array": audio_data.copy(),
|
|
"base64": base64.b64encode(audio_bytes).decode("utf-8"),
|
|
"file_path": str(cache_path.resolve()),
|
|
}
|
|
|
|
|
|
def _mux_mp4_bytes_with_synthetic_audio(
|
|
video_mp4_bytes: bytes,
|
|
*,
|
|
num_frames: int,
|
|
fps: float = 30.0,
|
|
sample_rate: int = 48000,
|
|
) -> bytes:
|
|
duration_sec = num_frames / fps if fps > 0 else 0.0
|
|
duration_int = max(1, int(math.ceil(duration_sec)))
|
|
|
|
try:
|
|
audio_result = generate_synthetic_audio(
|
|
duration=duration_int,
|
|
num_channels=1,
|
|
sample_rate=sample_rate,
|
|
)
|
|
audio_pcm = audio_result["np_array"]
|
|
except Exception as e:
|
|
logger.warning("Synthetic video: generate_synthetic_audio failed (%s); using video-only MP4.", e)
|
|
return video_mp4_bytes
|
|
|
|
try:
|
|
import imageio_ffmpeg
|
|
|
|
ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()
|
|
except Exception:
|
|
ffmpeg_exe = "ffmpeg"
|
|
|
|
try:
|
|
with tempfile.TemporaryDirectory(prefix="syn_vid_mux_") as tmp:
|
|
vid_path = os.path.join(tmp, "video.mp4")
|
|
wav_path = os.path.join(tmp, "audio.wav")
|
|
out_path = os.path.join(tmp, "out.mp4")
|
|
with open(vid_path, "wb") as f:
|
|
f.write(video_mp4_bytes)
|
|
sf.write(wav_path, audio_pcm, sample_rate, format="WAV", subtype="PCM_16")
|
|
cmd = [
|
|
ffmpeg_exe,
|
|
"-y",
|
|
"-nostdin",
|
|
"-hide_banner",
|
|
"-loglevel",
|
|
"error",
|
|
"-i",
|
|
vid_path,
|
|
"-i",
|
|
wav_path,
|
|
"-c:v",
|
|
"copy",
|
|
"-c:a",
|
|
"aac",
|
|
"-b:a",
|
|
"128k",
|
|
"-shortest",
|
|
"-movflags",
|
|
"+faststart",
|
|
out_path,
|
|
]
|
|
subprocess.run(cmd, check=True, stdin=subprocess.DEVNULL, timeout=300)
|
|
with open(out_path, "rb") as f:
|
|
return f.read()
|
|
except (
|
|
FileNotFoundError,
|
|
subprocess.CalledProcessError,
|
|
subprocess.TimeoutExpired,
|
|
OSError,
|
|
) as e:
|
|
logger.warning("Synthetic video: audio mux failed (%s); using video-only MP4.", e)
|
|
return video_mp4_bytes
|
|
|
|
|
|
def generate_synthetic_video(
|
|
width: int,
|
|
height: int,
|
|
num_frames: int,
|
|
*,
|
|
embed_audio: bool = False,
|
|
force_regenerate: bool = False,
|
|
cache_dir: Path | str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Generate synthetic MP4 (optional AAC audio). Caches final bytes by
|
|
``width`` / ``height`` / ``num_frames`` / ``embed_audio`` unless
|
|
``force_regenerate`` is true. Cache root: ``cache_dir`` if given, else the
|
|
default temp subdirectory.
|
|
"""
|
|
root = _resolve_synthetic_media_cache_dir(cache_dir)
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
cache_path = root / f"synth_video_w{width}_h{height}_nf{num_frames}_ea{int(embed_audio)}.mp4"
|
|
|
|
if not force_regenerate and cache_path.is_file():
|
|
video_bytes = cache_path.read_bytes()
|
|
return {
|
|
"np_array": _np_array_from_mp4_bytes(video_bytes),
|
|
"base64": base64.b64encode(video_bytes).decode("utf-8"),
|
|
"file_path": str(cache_path.resolve()),
|
|
}
|
|
|
|
import cv2
|
|
import imageio
|
|
|
|
num_balls = random.randint(3, 8)
|
|
balls = []
|
|
for _ in range(num_balls):
|
|
radius = min(width, height) // 8
|
|
if radius < 1:
|
|
raise ValueError(f"Video dimensions ({width}x{height}) too small")
|
|
x = random.randint(radius, width - radius)
|
|
y = random.randint(radius, height - radius)
|
|
speed = random.uniform(3.0, 8.0)
|
|
angle = random.uniform(0, 2 * math.pi)
|
|
vx = speed * math.cos(angle)
|
|
vy = speed * math.sin(angle)
|
|
color_bgr = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255))
|
|
balls.append({"x": x, "y": y, "vx": vx, "vy": vy, "radius": radius, "color_bgr": color_bgr})
|
|
|
|
video_frames = []
|
|
for _ in range(num_frames):
|
|
frame_bgr = np.zeros((height, width, 3), dtype=np.uint8)
|
|
for ball in balls:
|
|
ball["x"] += ball["vx"]
|
|
ball["y"] += ball["vy"]
|
|
if ball["x"] - ball["radius"] <= 0 or ball["x"] + ball["radius"] >= width:
|
|
ball["vx"] = -ball["vx"]
|
|
ball["x"] = max(ball["radius"], min(width - ball["radius"], ball["x"]))
|
|
if ball["y"] - ball["radius"] <= 0 or ball["y"] + ball["radius"] >= height:
|
|
ball["vy"] = -ball["vy"]
|
|
ball["y"] = max(ball["radius"], min(height - ball["radius"], ball["y"]))
|
|
x, y = int(ball["x"]), int(ball["y"])
|
|
radius = int(ball["radius"])
|
|
cv2.circle(frame_bgr, (x, y), radius, ball["color_bgr"], -1)
|
|
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
|
video_frames.append(frame_rgb)
|
|
|
|
fps = 30
|
|
buffer = io.BytesIO()
|
|
writer_kwargs = {
|
|
"format": "mp4",
|
|
"fps": fps,
|
|
"codec": "libx264",
|
|
"quality": 7,
|
|
"pixelformat": "yuv420p",
|
|
"macro_block_size": 16,
|
|
"ffmpeg_params": ["-preset", "medium", "-crf", "23", "-movflags", "+faststart", "-pix_fmt", "yuv420p"],
|
|
}
|
|
try:
|
|
with imageio.get_writer(buffer, **writer_kwargs) as writer:
|
|
for frame in video_frames:
|
|
writer.append_data(frame)
|
|
buffer.seek(0)
|
|
video_only_bytes = buffer.read()
|
|
except Exception as e:
|
|
print(f"Warning: Failed to encode synthetic video: {e}")
|
|
raise
|
|
video_bytes = (
|
|
_mux_mp4_bytes_with_synthetic_audio(video_only_bytes, num_frames=num_frames, fps=float(fps))
|
|
if embed_audio
|
|
else video_only_bytes
|
|
)
|
|
|
|
cache_path.write_bytes(video_bytes)
|
|
|
|
return {
|
|
"np_array": np.array(video_frames),
|
|
"base64": base64.b64encode(video_bytes).decode("utf-8"),
|
|
"file_path": str(cache_path.resolve()),
|
|
}
|
|
|
|
|
|
def generate_synthetic_image(
|
|
width: int,
|
|
height: int,
|
|
*,
|
|
force_regenerate: bool = False,
|
|
cache_dir: Path | str | None = None,
|
|
seed: int | None = None,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Random colored squares on white background. Caches JPEG by ``width`` /
|
|
``height`` unless ``force_regenerate`` is true. Cache root: ``cache_dir``
|
|
if given, else the default temp subdirectory.
|
|
"""
|
|
if seed is not None:
|
|
random.seed(seed)
|
|
|
|
root = _resolve_synthetic_media_cache_dir(cache_dir)
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
cache_path = root / f"synth_image_w{width}_h{height}.jpg"
|
|
|
|
if not force_regenerate and cache_path.is_file():
|
|
from PIL import Image as PILImage
|
|
|
|
image = PILImage.open(cache_path)
|
|
image.load()
|
|
image_bytes = cache_path.read_bytes()
|
|
return {
|
|
"np_array": np.array(image).copy(),
|
|
"base64": base64.b64encode(image_bytes).decode("utf-8"),
|
|
"file_path": str(cache_path.resolve()),
|
|
}
|
|
|
|
from PIL import ImageDraw
|
|
|
|
image = Image.new("RGB", (width, height), (255, 255, 255))
|
|
draw = ImageDraw.Draw(image)
|
|
num_squares = random.randint(3, 8)
|
|
for _ in range(num_squares):
|
|
square_size = random.randint(max(1, min(width, height) // 8), max(2, min(width, height) // 4))
|
|
x = random.randint(0, max(0, width - square_size - 1))
|
|
y = random.randint(0, max(0, height - square_size - 1))
|
|
color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
|
|
border_width = random.randint(1, 5)
|
|
draw.rectangle([x, y, x + square_size, y + square_size], fill=color, outline=(0, 0, 0), width=border_width)
|
|
|
|
image.save(str(cache_path), format="JPEG", quality=85, optimize=True)
|
|
image_bytes = cache_path.read_bytes()
|
|
|
|
return {
|
|
"np_array": np.array(image).copy(),
|
|
"base64": base64.b64encode(image_bytes).decode("utf-8"),
|
|
"file_path": str(cache_path.resolve()),
|
|
}
|
|
|
|
|
|
_TEST_ASSETS_ROOT = Path(__file__).resolve().parents[1] / "assets"
|
|
|
|
_AUDIO_MIME_BY_SUFFIX = {
|
|
".wav": "audio/wav",
|
|
".mp3": "audio/mpeg",
|
|
".flac": "audio/flac",
|
|
".ogg": "audio/ogg",
|
|
}
|
|
|
|
|
|
def get_asset_path(relative_path: str | os.PathLike) -> Path:
|
|
"""Resolve a path under ``tests/assets/`` to its absolute on-disk location."""
|
|
return _TEST_ASSETS_ROOT / Path(relative_path)
|
|
|
|
|
|
def load_test_audio_data_url(relative_path: str | os.PathLike) -> str:
|
|
"""Load a vendored test audio file under ``tests/assets/`` as a base64 data URL.
|
|
|
|
Used by tests that need real reference audio (e.g. voice cloning) without
|
|
relying on the server's ability to fetch external URLs at request time.
|
|
"""
|
|
path = get_asset_path(relative_path)
|
|
mime = _AUDIO_MIME_BY_SUFFIX.get(path.suffix.lower(), "application/octet-stream")
|
|
encoded = base64.b64encode(path.read_bytes()).decode("ascii")
|
|
return f"data:{mime};base64,{encoded}"
|
|
|
|
|
|
def decode_b64_image(b64: str):
|
|
img = Image.open(io.BytesIO(base64.b64decode(b64)))
|
|
img.load()
|
|
return img
|
|
|
|
|
|
def concat_audio(audio_val) -> np.ndarray:
|
|
"""Flatten a multimodal audio payload to mono float32 samples.
|
|
|
|
Engines return ``multimodal_output["audio"]`` as a tensor, a list of
|
|
per-chunk tensors (streaming decoders), or an array-like; concatenate
|
|
in order and return a 1-D ``np.float32`` array (empty when a list has
|
|
no tensors).
|
|
"""
|
|
import torch
|
|
|
|
if isinstance(audio_val, list):
|
|
tensors = [t.detach().cpu().float().reshape(-1) for t in audio_val if isinstance(t, torch.Tensor)]
|
|
if not tensors:
|
|
return np.zeros((0,), dtype=np.float32)
|
|
return torch.cat(tensors, dim=-1).numpy().astype(np.float32, copy=False)
|
|
if isinstance(audio_val, torch.Tensor):
|
|
return audio_val.detach().cpu().float().reshape(-1).numpy()
|
|
return np.asarray(audio_val, dtype=np.float32).reshape(-1)
|
|
|
|
|
|
def preprocess_text(text):
|
|
import opencc
|
|
|
|
word_to_num = {
|
|
"zero": "0",
|
|
"one": "1",
|
|
"two": "2",
|
|
"three": "3",
|
|
"four": "4",
|
|
"five": "5",
|
|
"six": "6",
|
|
"seven": "7",
|
|
"eight": "8",
|
|
"nine": "9",
|
|
"ten": "10",
|
|
}
|
|
for word, num in word_to_num.items():
|
|
pattern = r"\b" + re.escape(word) + r"\b"
|
|
text = re.sub(pattern, num, text, flags=re.IGNORECASE)
|
|
|
|
text = re.sub(r"[^\w\s]", "", text)
|
|
text = re.sub(r"\s+", " ", text)
|
|
cc = opencc.OpenCC("t2s")
|
|
text = cc.convert(text)
|
|
text = re.sub(r"(?<=[\u4e00-\u9fff])\s+(?=[\u4e00-\u9fff])", "", text)
|
|
return text.lower().strip()
|
|
|
|
|
|
def cosine_similarity_text(text1, text2, n: int = 3):
|
|
from collections import Counter
|
|
|
|
if not text1 or not text2:
|
|
return 0.0
|
|
|
|
text1 = preprocess_text(text1)
|
|
text2 = preprocess_text(text2)
|
|
print(f"cosine similarity text1 is: {text1}, text2 is: {text2}")
|
|
|
|
ngrams1 = [text1[i : i + n] for i in range(len(text1) - n + 1)]
|
|
ngrams2 = [text2[i : i + n] for i in range(len(text2) - n + 1)]
|
|
counter1 = Counter(ngrams1)
|
|
counter2 = Counter(ngrams2)
|
|
|
|
all_ngrams = set(counter1.keys()) | set(counter2.keys())
|
|
vec1 = [counter1.get(ng, 0) for ng in all_ngrams]
|
|
vec2 = [counter2.get(ng, 0) for ng in all_ngrams]
|
|
dot_product = sum(a * b for a, b in zip(vec1, vec2))
|
|
norm1 = sum(a * a for a in vec1) ** 0.5
|
|
norm2 = sum(b * b for b in vec2) ** 0.5
|
|
if norm1 == 0 or norm2 == 0:
|
|
return 0.0
|
|
cosine = dot_product / (norm1 * norm2)
|
|
# Down-weight when lengths differ: repeated/hallucinated transcripts stay
|
|
# high in bag-of-ngrams cosine (e.g. ABCABCABC vs ABC) but should score low.
|
|
len1, len2 = len(text1), len(text2)
|
|
length_harmony = (2.0 * min(len1, len2)) / (len1 + len2)
|
|
return cosine * length_harmony
|
|
|
|
|
|
class _AudioBuffer:
|
|
"""Minimal replacement for pydub.AudioSegment used by test helpers."""
|
|
|
|
def __init__(self, data: np.ndarray, sample_rate: int):
|
|
self.data = data
|
|
self.sample_rate = sample_rate
|
|
|
|
def export(self, buf: io.BytesIO, format: str = "wav"):
|
|
sf.write(buf, self.data, self.sample_rate, format=format.upper())
|
|
buf.seek(0)
|
|
|
|
|
|
def _merge_base64_audio_to_segment(base64_list: list[str]) -> _AudioBuffer:
|
|
from vllm.multimodal.media.audio import AudioMediaIO
|
|
|
|
io_ = AudioMediaIO()
|
|
chunks: list[np.ndarray] = []
|
|
sample_rate: int | None = None
|
|
for b64 in base64_list:
|
|
raw = base64.b64decode(b64.split(",", 1)[-1])
|
|
waveform, sr = io_.load_bytes(raw)
|
|
if sample_rate is None:
|
|
sample_rate = int(sr)
|
|
chunks.append(waveform)
|
|
merged = np.concatenate(chunks) if chunks else np.array([], dtype=np.float32)
|
|
return _AudioBuffer(merged, sample_rate or 16000)
|
|
|
|
|
|
@contextmanager
|
|
def _serialize_whisper_model_download(model_size: str = "small"):
|
|
"""Serialize Whisper cache writes across processes (Linux/Unix), per model."""
|
|
import fcntl
|
|
|
|
lock_path = Path.home() / ".cache" / "whisper" / f".{model_size}_model_download.lock"
|
|
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
|
f = open(lock_path, "a+b")
|
|
try:
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
|
|
yield
|
|
finally:
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|
|
f.close()
|
|
|
|
|
|
def _select_whisper_device() -> str:
|
|
device_index = None
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
if current_omni_platform.is_available():
|
|
n = current_omni_platform.get_device_count()
|
|
# Single-GPU runners (e.g. the L4 nightly): the model server already
|
|
# occupies device 0, and a Whisper model resident there competes with
|
|
# it for VRAM and OOMs. Only borrow an accelerator when a spare device
|
|
# exists; otherwise validate on CPU.
|
|
if n > 1:
|
|
device_index = n - 1
|
|
|
|
if device_index is None:
|
|
return "cpu"
|
|
|
|
torch_device = current_omni_platform.get_torch_device(device_index)
|
|
current_omni_platform.set_device(torch_device)
|
|
return str(torch_device)
|
|
|
|
|
|
# Populated in the transcription worker, not in the pytest process.
|
|
_WHISPER_MODELS: dict[str, Any] = {}
|
|
|
|
|
|
def _get_whisper_model(model_size: str) -> Any:
|
|
model = _WHISPER_MODELS.get(model_size)
|
|
if model is None:
|
|
import whisper
|
|
|
|
# The device is picked on first load and the model stays on it for the
|
|
# worker's lifetime: the current server or runner fixture instance, or
|
|
# the test module for callers that transcribe without those fixtures.
|
|
device = _select_whisper_device()
|
|
with _serialize_whisper_model_download(model_size):
|
|
model = whisper.load_model(model_size, device=device)
|
|
_WHISPER_MODELS[model_size] = model
|
|
return model
|
|
|
|
|
|
def _whisper_transcribe_in_current_process(
|
|
output_path: str, model_size: str = "small", language: str | None = None
|
|
) -> str:
|
|
model = _get_whisper_model(model_size)
|
|
text = model.transcribe(
|
|
output_path,
|
|
temperature=0.0,
|
|
word_timestamps=True,
|
|
condition_on_previous_text=False,
|
|
# None keeps whisper's auto-detection. Do not default this to a
|
|
# language: callers include non-English audio tests.
|
|
language=language,
|
|
)["text"]
|
|
return text or ""
|
|
|
|
|
|
# Serializes a whole submit->result->cleanup on the parent side, so at most one
|
|
# call is ever in flight on the single-worker pool. The child already runs
|
|
# transcriptions serially, so this costs no throughput; it lets a failed call
|
|
# discard the worker without disrupting another caller's in-flight future.
|
|
_TRANSCRIBER_CALL_LOCK = threading.Lock()
|
|
# Guards the _TRANSCRIBER pointer itself.
|
|
_TRANSCRIBER_LOCK = threading.Lock()
|
|
_TRANSCRIBER: concurrent.futures.ProcessPoolExecutor | None = None
|
|
|
|
|
|
def _get_transcriber() -> concurrent.futures.ProcessPoolExecutor:
|
|
global _TRANSCRIBER
|
|
with _TRANSCRIBER_LOCK:
|
|
if _TRANSCRIBER is None:
|
|
ctx = multiprocessing.get_context("spawn")
|
|
_TRANSCRIBER = concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx)
|
|
return _TRANSCRIBER
|
|
|
|
|
|
def _discard_transcriber(executor: concurrent.futures.ProcessPoolExecutor) -> None:
|
|
"""Drop ``executor``, but only while it is still the current one.
|
|
|
|
Identity-checked so a stale reference can never shut down a newer worker that
|
|
was installed after ``executor`` was replaced.
|
|
"""
|
|
global _TRANSCRIBER
|
|
with _TRANSCRIBER_LOCK:
|
|
if _TRANSCRIBER is not executor:
|
|
return
|
|
_TRANSCRIBER = None
|
|
# Joining the worker can block; do it outside the lock.
|
|
executor.shutdown(wait=True)
|
|
|
|
|
|
def release_audio_transcriber() -> None:
|
|
"""Shut the transcription worker down, freeing the device memory its models hold.
|
|
|
|
Called when a server or runner fixture instance tears down, so that the next
|
|
one -- including the next parametrization inside the same test module -- does
|
|
not initialize its model while Whisper still occupies the device. A module
|
|
teardown fixture repeats it for tests that transcribe without those fixtures.
|
|
|
|
Takes the call lock, so it waits for any in-flight transcription to finish
|
|
rather than shutting the worker down underneath it.
|
|
"""
|
|
global _TRANSCRIBER
|
|
with _TRANSCRIBER_CALL_LOCK:
|
|
with _TRANSCRIBER_LOCK:
|
|
executor, _TRANSCRIBER = _TRANSCRIBER, None
|
|
if executor is not None:
|
|
executor.shutdown(wait=True)
|
|
|
|
|
|
def convert_audio_file_to_text(output_path: str, model_size: str = "small", language: str | None = None) -> str:
|
|
"""Convert an audio file to text in a reused, isolated subprocess.
|
|
|
|
The worker outlives the call so its Whisper model is loaded once rather than
|
|
once per transcription. The call lock serializes callers onto the single
|
|
worker (the child already transcribes serially, so this costs no throughput),
|
|
which lets a failed call tear the worker down without racing another caller.
|
|
The worker caches one model per size it is asked for, so a run that escalates
|
|
to a stronger ASR keeps both models resident until release.
|
|
|
|
Any failure discards the worker, restoring the failure isolation of the old
|
|
one-process-per-call design: a failure (a ``torch`` OOM, or a
|
|
``KeyboardInterrupt``/``SystemExit`` transported through the future or
|
|
raised while ``result()`` blocks) propagates unchanged after the worker --
|
|
and its resident model -- is torn down, and a dead worker
|
|
(``BrokenProcessPool``) is additionally retried once.
|
|
"""
|
|
with _TRANSCRIBER_CALL_LOCK:
|
|
for attempt in range(2):
|
|
executor = _get_transcriber()
|
|
try:
|
|
return executor.submit(
|
|
_whisper_transcribe_in_current_process, output_path, model_size, language
|
|
).result()
|
|
except BrokenProcessPool:
|
|
_discard_transcriber(executor)
|
|
if attempt == 1:
|
|
raise
|
|
except BaseException:
|
|
# Any other failure -- a task exception (a torch OOM included), or
|
|
# a KeyboardInterrupt/SystemExit transported through the future or
|
|
# interrupting result() -- leaves the worker and its model
|
|
# resident; drop it so it cannot contaminate later calls, matching
|
|
# the old per-call teardown. Do not retry.
|
|
_discard_transcriber(executor)
|
|
raise
|
|
raise AssertionError("unreachable")
|
|
|
|
|
|
def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small", language: str | None = None) -> str:
|
|
output_fd, output_path = tempfile.mkstemp(prefix="test_", suffix=".wav")
|
|
os.close(output_fd)
|
|
if os.environ.get("VLLM_OMNI_KEEP_REQUEST_MEDIA", "").lower() not in ("1", "true", "yes"):
|
|
atexit.register(Path(output_path).unlink, missing_ok=True)
|
|
data, samplerate = sf.read(io.BytesIO(raw_bytes))
|
|
sf.write(output_path, data, samplerate, format="WAV", subtype="PCM_16")
|
|
print(f"audio data is saved: {output_path}")
|
|
return convert_audio_file_to_text(output_path, model_size, language)
|
|
|
|
|
|
__all__ = [
|
|
"_merge_base64_audio_to_segment",
|
|
"concat_audio",
|
|
"convert_audio_bytes_to_text",
|
|
"convert_audio_file_to_text",
|
|
"cosine_similarity_text",
|
|
"decode_b64_image",
|
|
"generate_synthetic_audio",
|
|
"generate_synthetic_image",
|
|
"generate_synthetic_video",
|
|
"get_asset_path",
|
|
"load_test_audio_data_url",
|
|
"preprocess_text",
|
|
"release_audio_transcriber",
|
|
]
|