Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
1632 lines
68 KiB
Python
1632 lines
68 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import time
|
|
import wave
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Callable
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
|
|
# Prefer the bundled CUDA libraries in the active conda env over any
|
|
# incompatible system cuDNN injected through LD_LIBRARY_PATH.
|
|
os.environ.pop("LD_LIBRARY_PATH", None)
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from safetensors import safe_open
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
DEFAULT_PERF_CASES = REPO_ROOT / "tools" / "perf" / "model_perf_cases.json"
|
|
DEFAULT_REQUEST_CASES = REPO_ROOT / "tools" / "perf" / "model_perf_request_cases.json"
|
|
|
|
|
|
@dataclass
|
|
class PerfCase:
|
|
id: str
|
|
family: str
|
|
case_id: str
|
|
request_id: str
|
|
warmup_request_id: str | None = None
|
|
request_overrides: dict[str, Any] | None = None
|
|
warmup_request_overrides: dict[str, Any] | None = None
|
|
request_remove_keys: list[str] | None = None
|
|
warmup_request_remove_keys: list[str] | None = None
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Serial Python performance harness for audio.cpp model references.")
|
|
parser.add_argument("--perf-cases", type=Path, default=DEFAULT_PERF_CASES)
|
|
parser.add_argument("--request-cases", type=Path, default=DEFAULT_REQUEST_CASES)
|
|
parser.add_argument("--id", default="")
|
|
parser.add_argument("--backend", choices=("cpu", "cuda"), default="cuda")
|
|
parser.add_argument("--device", type=int, default=0)
|
|
parser.add_argument("--threads", type=int, default=1)
|
|
parser.add_argument("--summary-out", type=Path, default=None)
|
|
return parser.parse_args()
|
|
|
|
|
|
def read_json(path: Path) -> Any:
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def ensure_absolute_path(value: str) -> Path:
|
|
path = Path(value)
|
|
return path if path.is_absolute() else REPO_ROOT / path
|
|
|
|
|
|
def shorten_text(text: str, limit: int) -> str:
|
|
if len(text) <= limit:
|
|
return text
|
|
cut = text.rfind(" ", 0, limit)
|
|
if cut < max(8, limit // 2):
|
|
cut = limit
|
|
return text[:cut].rstrip()
|
|
|
|
|
|
def read_wav_mono_f32(path: Path) -> tuple[np.ndarray, int]:
|
|
pcm, sample_rate = sf.read(str(path), dtype="float32", always_2d=False)
|
|
pcm = np.asarray(pcm, dtype=np.float32)
|
|
if pcm.ndim == 2:
|
|
pcm = pcm.mean(axis=1)
|
|
return pcm, sample_rate
|
|
|
|
|
|
def write_wav_mono_f32(path: Path, audio: np.ndarray, sample_rate: int) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
audio = np.asarray(audio, dtype=np.float32).reshape(-1)
|
|
pcm = np.clip(audio, -1.0, 1.0)
|
|
pcm16 = np.round(pcm * 32767.0).astype(np.int16, copy=False)
|
|
with wave.open(str(path), "wb") as wav:
|
|
wav.setnchannels(1)
|
|
wav.setsampwidth(2)
|
|
wav.setframerate(sample_rate)
|
|
wav.writeframes(pcm16.tobytes())
|
|
|
|
|
|
def write_clipped_wav(source: Path, temp_dir: Path, seconds: float) -> Path:
|
|
samples, sample_rate = read_wav_mono_f32(source)
|
|
limit = max(1, min(len(samples), int(round(sample_rate * seconds))))
|
|
target = temp_dir / source.name
|
|
write_wav_mono_f32(target, samples[:limit], sample_rate)
|
|
return target
|
|
|
|
|
|
def resolve_request_paths(request: dict[str, Any]) -> dict[str, Any]:
|
|
out = dict(request)
|
|
for key in ("audio", "voice_ref", "source_audio", "target_voice", "prosody_ref", "style_ref"):
|
|
value = out.get(key)
|
|
if isinstance(value, str) and value:
|
|
out[key] = str(ensure_absolute_path(value))
|
|
return out
|
|
|
|
|
|
def derive_warmup_request(request: dict[str, Any], temp_dir: Path) -> dict[str, Any]:
|
|
out = resolve_request_paths(request)
|
|
for key, limit in (
|
|
("text", 80),
|
|
("target_text", 80),
|
|
("style_ref_text", 80),
|
|
("reference_text", 80),
|
|
("lyrics", 80),
|
|
("instruct", 64),
|
|
):
|
|
value = out.get(key)
|
|
if isinstance(value, str) and value:
|
|
out[key] = shorten_text(value, limit)
|
|
for key in ("audio", "voice_ref", "source_audio", "target_voice", "prosody_ref", "style_ref"):
|
|
value = out.get(key)
|
|
if isinstance(value, str) and value:
|
|
out[key] = str(write_clipped_wav(Path(value), temp_dir, 2.0))
|
|
if isinstance(out.get("duration_seconds"), (int, float)):
|
|
out["duration_seconds"] = min(float(out["duration_seconds"]), 4.0)
|
|
if isinstance(out.get("max_tokens"), (int, float)):
|
|
out["max_tokens"] = min(int(out["max_tokens"]), 64)
|
|
if isinstance(out.get("num_inference_steps"), (int, float)):
|
|
out["num_inference_steps"] = min(int(out["num_inference_steps"]), 8)
|
|
return out
|
|
|
|
|
|
def audio_duration_sec(audio: np.ndarray, sample_rate: int) -> float:
|
|
return float(len(np.asarray(audio).reshape(-1))) / float(sample_rate) if sample_rate > 0 else 0.0
|
|
|
|
|
|
def input_audio_duration_sec(request: dict[str, Any]) -> float:
|
|
for key in ("audio", "source_audio"):
|
|
value = request.get(key)
|
|
if isinstance(value, str) and value:
|
|
samples, sample_rate = read_wav_mono_f32(Path(value))
|
|
return audio_duration_sec(samples, sample_rate)
|
|
return 0.0
|
|
|
|
|
|
def load_perf_cases(path: Path) -> list[PerfCase]:
|
|
root = read_json(path)
|
|
return [
|
|
PerfCase(
|
|
id=item["id"],
|
|
family=item["family"],
|
|
case_id=item["case_id"],
|
|
request_id=item["request_id"],
|
|
warmup_request_id=item.get("warmup_request_id"),
|
|
request_overrides=item.get("request_overrides"),
|
|
warmup_request_overrides=item.get("warmup_request_overrides"),
|
|
request_remove_keys=item.get("request_remove_keys"),
|
|
warmup_request_remove_keys=item.get("warmup_request_remove_keys"),
|
|
)
|
|
for item in root["cases"]
|
|
]
|
|
|
|
|
|
def find_case(request_cases: dict[str, Any], case_id: str) -> dict[str, Any]:
|
|
for item in request_cases["cases"]:
|
|
if item["id"] == case_id:
|
|
return item
|
|
raise KeyError(case_id)
|
|
|
|
|
|
def find_request(case: dict[str, Any], request_id: str) -> dict[str, Any]:
|
|
for item in case["requests"]:
|
|
if item["id"] == request_id:
|
|
return item
|
|
raise KeyError(f"{case['id']}:{request_id}")
|
|
|
|
|
|
def materialize_perf_request(
|
|
request: dict[str, Any],
|
|
overrides: dict[str, Any] | None,
|
|
remove_keys: list[str] | None,
|
|
) -> dict[str, Any]:
|
|
out = dict(request)
|
|
for key in remove_keys or []:
|
|
out.pop(key, None)
|
|
if overrides:
|
|
out.update(overrides)
|
|
return out
|
|
|
|
|
|
def set_threads(threads: int) -> None:
|
|
os.environ["OMP_NUM_THREADS"] = str(max(1, threads))
|
|
os.environ["MKL_NUM_THREADS"] = str(max(1, threads))
|
|
|
|
|
|
def load_module(name: str, path: Path):
|
|
spec = importlib.util.spec_from_file_location(name, path)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError(f"failed to load module: {path}")
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def use_transformers457() -> None:
|
|
import transformers457 as _transformers457
|
|
import transformers457.dynamic_module_utils as _dynamic_module_utils
|
|
|
|
sys.modules["transformers"] = _transformers457
|
|
_dynamic_module_utils.transformers = _transformers457
|
|
|
|
|
|
def enable_legacy_llama_config_positional_args() -> None:
|
|
import transformers
|
|
|
|
original = transformers.LlamaConfig
|
|
if getattr(original, "_audiocpp_legacy_positional_enabled", False):
|
|
return
|
|
|
|
class LegacyLlamaConfig(original):
|
|
_audiocpp_legacy_positional_enabled = True
|
|
|
|
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
|
if args:
|
|
if len(args) != 5:
|
|
raise TypeError(f"Legacy LlamaConfig positional compatibility expects 5 args, got {len(args)}")
|
|
vocab_size, hidden_size, intermediate_size, num_hidden_layers, num_attention_heads = args
|
|
kwargs = {
|
|
"vocab_size": vocab_size,
|
|
"hidden_size": hidden_size,
|
|
"intermediate_size": intermediate_size,
|
|
"num_hidden_layers": num_hidden_layers,
|
|
"num_attention_heads": num_attention_heads,
|
|
**kwargs,
|
|
}
|
|
super().__init__(**kwargs)
|
|
|
|
transformers.LlamaConfig = LegacyLlamaConfig
|
|
|
|
|
|
def force_attention_implementation(module: object, value: str) -> None:
|
|
if hasattr(module, "attn_implementation"):
|
|
setattr(module, "attn_implementation", value)
|
|
if hasattr(module, "config"):
|
|
config = getattr(module, "config")
|
|
if hasattr(config, "attn_implementation"):
|
|
setattr(config, "attn_implementation", value)
|
|
if hasattr(config, "local_transformer_attn_implementation"):
|
|
setattr(config, "local_transformer_attn_implementation", value)
|
|
if hasattr(config, "gpt2_config") and getattr(config, "gpt2_config") is not None:
|
|
gpt2_config = getattr(config, "gpt2_config")
|
|
if hasattr(gpt2_config, "_attn_implementation"):
|
|
setattr(gpt2_config, "_attn_implementation", value)
|
|
if hasattr(gpt2_config, "attn_implementation"):
|
|
setattr(gpt2_config, "attn_implementation", value)
|
|
named_children = getattr(module, "named_children", None)
|
|
if callable(named_children):
|
|
for _, child in named_children():
|
|
force_attention_implementation(child, value)
|
|
|
|
|
|
def resample_linear_mono(waveform: np.ndarray, input_sample_rate: int, output_sample_rate: int) -> np.ndarray:
|
|
if input_sample_rate == output_sample_rate:
|
|
return waveform
|
|
scale = output_sample_rate / input_sample_rate
|
|
output_samples = int(round(len(waveform) * scale))
|
|
pos = np.arange(output_samples, dtype=np.float64) / scale
|
|
left = np.floor(pos).astype(np.int64)
|
|
right = np.minimum(left + 1, len(waveform) - 1)
|
|
frac = (pos - left).astype(np.float32)
|
|
return waveform[left] * (1.0 - frac) + waveform[right] * frac
|
|
|
|
|
|
class RefCitrinet:
|
|
LOG_EPS = float(np.ldexp(1.0, -24))
|
|
BN_EPS = 1.0e-3
|
|
NORM_EPS = 1.0e-5
|
|
|
|
def __init__(self, model_dir: Path, device: torch.device):
|
|
model_path = model_dir / "citrinet_256.safetensors"
|
|
self.device = device
|
|
self.tensors: dict[str, torch.Tensor] = {}
|
|
with safe_open(str(model_path), framework="pt", device="cpu") as f:
|
|
self.metadata = {
|
|
key: json.loads(value)
|
|
if value.startswith(("[", "{", "\"")) or value in {"true", "false", "null"} or value.lstrip("-").isdigit()
|
|
else value
|
|
for key, value in f.metadata().items()
|
|
}
|
|
for key in f.keys():
|
|
self.tensors[key] = f.get_tensor(key).float().contiguous().to(device)
|
|
self.sample_rate = int(self.metadata["sample_rate"])
|
|
self.n_mels = int(self.metadata["n_mels"])
|
|
self.n_fft = int(self.metadata["n_fft"])
|
|
self.hop = int(round(float(self.metadata["window_stride"]) * self.sample_rate))
|
|
self.win = int(round(float(self.metadata["window_size"]) * self.sample_rate))
|
|
self.pad_to = int(self.metadata["pad_to"])
|
|
self.blank_id = int(self.metadata["blank_id"])
|
|
self.jasper = list(self.metadata["jasper"])
|
|
self.window = self.tensors["preprocessor.featurizer.window"]
|
|
self.fb = self.tensors["preprocessor.featurizer.fb"].view(self.n_mels, self.n_fft // 2 + 1)
|
|
vocab_path = model_dir / str(self.metadata["vocab_file"])
|
|
self.vocab = [piece[:-1] if piece.endswith("\r") else piece for piece in vocab_path.read_text().split("\n")]
|
|
|
|
def bn(self, x: torch.Tensor, prefix: str) -> torch.Tensor:
|
|
w = self.tensors[f"{prefix}.weight"].view(1, -1, 1)
|
|
b = self.tensors[f"{prefix}.bias"].view(1, -1, 1)
|
|
mean = self.tensors[f"{prefix}.running_mean"].view(1, -1, 1)
|
|
var = self.tensors[f"{prefix}.running_var"].view(1, -1, 1)
|
|
return (x - mean) / torch.sqrt(var + self.BN_EPS) * w + b
|
|
|
|
def conv1d(self, x: torch.Tensor, prefix: str, stride: int, dilation: int, padding: int, groups: int, bias: bool) -> torch.Tensor:
|
|
weight = self.tensors[f"{prefix}.weight"]
|
|
if weight.ndim == 2:
|
|
weight = weight.unsqueeze(-1)
|
|
bias_t = self.tensors[f"{prefix}.bias"] if bias else None
|
|
return F.conv1d(x, weight, bias=bias_t, stride=stride, padding=padding, dilation=dilation, groups=groups)
|
|
|
|
def squeeze_excite(self, x: torch.Tensor, block: int, repeat: int) -> torch.Tensor:
|
|
se_index = 3 if repeat == 1 else repeat * 5 - 2
|
|
y = x.mean(dim=-1, keepdim=True)
|
|
y = self.conv1d(y, f"encoder.encoder.{block}.mconv.{se_index}.fc.0", 1, 1, 0, 1, False)
|
|
y = F.relu(y)
|
|
y = self.conv1d(y, f"encoder.encoder.{block}.mconv.{se_index}.fc.2", 1, 1, 0, 1, False)
|
|
return x * torch.sigmoid(y)
|
|
|
|
def block(self, x: torch.Tensor, block_index: int, cfg: dict) -> torch.Tensor:
|
|
residual_x = x
|
|
repeat = int(cfg["repeat"])
|
|
kernel = int(cfg["kernel"])
|
|
dilation = int(cfg["dilation"])
|
|
for r in range(repeat):
|
|
base = r * 5 if cfg["separable"] else r * 3
|
|
stride = int(cfg["stride"]) if r + 1 == repeat else 1
|
|
padding = dilation * (kernel - 1) // 2
|
|
if cfg["separable"]:
|
|
x = self.conv1d(x, f"encoder.encoder.{block_index}.mconv.{base}.conv", stride, dilation, padding, x.shape[1], False)
|
|
x = self.conv1d(x, f"encoder.encoder.{block_index}.mconv.{base + 1}.conv", 1, 1, 0, 1, False)
|
|
x = self.bn(x, f"encoder.encoder.{block_index}.mconv.{base + 2}")
|
|
else:
|
|
x = self.conv1d(x, f"encoder.encoder.{block_index}.mconv.{base}.conv", stride, dilation, padding, 1, False)
|
|
x = self.bn(x, f"encoder.encoder.{block_index}.mconv.{base + 1}")
|
|
if r + 1 != repeat:
|
|
x = F.relu(x)
|
|
if cfg.get("se", False):
|
|
x = self.squeeze_excite(x, block_index, repeat)
|
|
if cfg["residual"]:
|
|
residual_stride = int(cfg["stride"]) if cfg.get("residual_mode") == "stride_add" else 1
|
|
res = self.conv1d(residual_x, f"encoder.encoder.{block_index}.res.0.0.conv", residual_stride, 1, 0, 1, False)
|
|
res = self.bn(res, f"encoder.encoder.{block_index}.res.0.1")
|
|
x = x + res
|
|
return F.relu(x)
|
|
|
|
def compute_features(self, waveform: np.ndarray) -> tuple[torch.Tensor, int]:
|
|
xt = torch.from_numpy(waveform.astype(np.float32, copy=False)).to(self.device).unsqueeze(0)
|
|
stft = torch.stft(xt, n_fft=self.n_fft, hop_length=self.hop, win_length=self.win, window=self.window, center=True, pad_mode="constant", normalized=False, onesided=True, return_complex=True)
|
|
mel = torch.matmul(self.fb, stft.abs().pow(2.0)[0]).transpose(0, 1)
|
|
feats = torch.log(mel + self.LOG_EPS)
|
|
raw_frames = int(feats.shape[0])
|
|
mean = feats.mean(dim=0, keepdim=True)
|
|
var = ((feats - mean) ** 2).sum(dim=0, keepdim=True) / max(raw_frames - 1, 1)
|
|
feats = (feats - mean) / (torch.sqrt(var) + self.NORM_EPS)
|
|
if self.pad_to > 1:
|
|
padded = ((raw_frames + self.pad_to - 1) // self.pad_to) * self.pad_to
|
|
if padded != raw_frames:
|
|
feats = F.pad(feats, (0, 0, 0, padded - raw_frames))
|
|
return feats, raw_frames
|
|
|
|
def output_frames(self, input_frames: int) -> int:
|
|
frames = input_frames
|
|
for cfg in self.jasper:
|
|
stride = int(cfg["stride"])
|
|
if stride > 1:
|
|
frames = (frames + stride - 1) // stride
|
|
return frames
|
|
|
|
def infer_features(self, features: torch.Tensor) -> torch.Tensor:
|
|
x = features.float().t().unsqueeze(0)
|
|
for block_index, cfg in enumerate(self.jasper):
|
|
x = self.block(x, block_index, cfg)
|
|
logits = self.conv1d(x, "decoder.decoder_layers.0", 1, 1, 0, 1, True)
|
|
return logits.squeeze(0).transpose(0, 1).contiguous()
|
|
|
|
def greedy_decode(self, logits: torch.Tensor) -> str:
|
|
rows = logits.detach().cpu().numpy()
|
|
ids: list[int] = []
|
|
prev = -1
|
|
for row in rows:
|
|
best = int(np.argmax(row))
|
|
if best == prev:
|
|
continue
|
|
prev = best
|
|
if best != self.blank_id:
|
|
ids.append(best)
|
|
text = ""
|
|
for idx in ids:
|
|
piece = self.vocab[idx]
|
|
if piece.startswith("##"):
|
|
text += piece[2:]
|
|
elif not text or text[-1] in {"'", "-"} or piece in {".", ",", "!", "?", ":", ";", "'", "\"", ")", "]", "}", "-", "/", "\\"}:
|
|
text += piece
|
|
else:
|
|
text += " " + piece
|
|
return text
|
|
|
|
def transcribe(self, audio_path: Path) -> str:
|
|
waveform, sample_rate = read_wav_mono_f32(audio_path)
|
|
waveform = resample_linear_mono(waveform, sample_rate, self.sample_rate)
|
|
features, raw_frames = self.compute_features(waveform)
|
|
logits = self.infer_features(features)
|
|
return self.greedy_decode(logits[: self.output_frames(raw_frames)])
|
|
|
|
|
|
class RefMarbleNet:
|
|
FEATURE_DIM = 80
|
|
TARGET_SR = 16000
|
|
N_FFT = 512
|
|
HOP = 160
|
|
WIN = 400
|
|
LOG_EPS = float(np.ldexp(1.0, -24))
|
|
BN_EPS = 1.0e-3
|
|
OUTPUT_STRIDE = 2
|
|
|
|
def __init__(self, model_dir: Path, device: torch.device):
|
|
self.device = device
|
|
self.tensors: dict[str, torch.Tensor] = {}
|
|
with safe_open(str(model_dir / "marblenet_vad.safetensors"), framework="pt", device="cpu") as f:
|
|
for key in f.keys():
|
|
self.tensors[key] = f.get_tensor(key).float().contiguous().to(device)
|
|
self.window = self.tensors["preprocessor.featurizer.window"]
|
|
self.fb = self.tensors["preprocessor.featurizer.fb"].view(self.FEATURE_DIM, self.N_FFT // 2 + 1)
|
|
|
|
def bn(self, x: torch.Tensor, prefix: str) -> torch.Tensor:
|
|
w = self.tensors[f"{prefix}.weight"].view(1, -1, 1)
|
|
b = self.tensors[f"{prefix}.bias"].view(1, -1, 1)
|
|
mean = self.tensors[f"{prefix}.running_mean"].view(1, -1, 1)
|
|
var = self.tensors[f"{prefix}.running_var"].view(1, -1, 1)
|
|
return (x - mean) / torch.sqrt(var + self.BN_EPS) * w + b
|
|
|
|
def conv1d(self, x: torch.Tensor, prefix: str, stride: int, dilation: int, padding: int, groups: int, bias: bool) -> torch.Tensor:
|
|
weight = self.tensors[f"{prefix}.weight"]
|
|
if weight.ndim == 2:
|
|
weight = weight.unsqueeze(-1)
|
|
bias_t = self.tensors[f"{prefix}.bias"] if bias else None
|
|
return F.conv1d(x, weight, bias=bias_t, stride=stride, padding=padding, dilation=dilation, groups=groups)
|
|
|
|
def block(self, x: torch.Tensor, block: int, repeat: int, kernel: int, stride: int, dilation: int, residual: bool, separable: bool) -> torch.Tensor:
|
|
residual_x = x
|
|
for r in range(repeat):
|
|
if separable:
|
|
base = r * 5
|
|
x = self.conv1d(x, f"encoder.encoder.{block}.mconv.{base}.conv", stride, dilation, dilation * (kernel - 1) // 2, x.shape[1], False)
|
|
x = self.conv1d(x, f"encoder.encoder.{block}.mconv.{base + 1}.conv", 1, 1, 0, 1, False)
|
|
x = self.bn(x, f"encoder.encoder.{block}.mconv.{base + 2}")
|
|
else:
|
|
base = r * 3
|
|
x = self.conv1d(x, f"encoder.encoder.{block}.mconv.{base}.conv", stride, dilation, dilation * (kernel - 1) // 2, 1, False)
|
|
x = self.bn(x, f"encoder.encoder.{block}.mconv.{base + 1}")
|
|
if r + 1 != repeat:
|
|
x = F.relu(x)
|
|
if residual:
|
|
res = self.conv1d(residual_x, f"encoder.encoder.{block}.res.0.0.conv", 1, 1, 0, 1, False)
|
|
res = self.bn(res, f"encoder.encoder.{block}.res.0.1")
|
|
x = x + res
|
|
return F.relu(x)
|
|
|
|
def detect(self, audio_path: Path, threshold: float = 0.5) -> list[dict[str, float | int]]:
|
|
waveform, sample_rate = read_wav_mono_f32(audio_path)
|
|
waveform = resample_linear_mono(waveform, sample_rate, self.TARGET_SR)
|
|
xt = torch.from_numpy(waveform.astype(np.float32, copy=False)).to(self.device).unsqueeze(0)
|
|
stft = torch.stft(xt, n_fft=self.N_FFT, hop_length=self.HOP, win_length=self.WIN, window=self.window, center=True, pad_mode="constant", normalized=False, onesided=True, return_complex=True)
|
|
mel = torch.matmul(self.fb, stft.abs().pow(2.0)[0]).transpose(0, 1)
|
|
feats = torch.log(mel + self.LOG_EPS)
|
|
if feats.shape[0] % 2 != 0:
|
|
feats = F.pad(feats, (0, 0, 0, 1))
|
|
x = feats.float().t().unsqueeze(0)
|
|
x = self.block(x, 0, repeat=1, kernel=11, stride=2, dilation=1, residual=False, separable=True)
|
|
x = self.block(x, 1, repeat=2, kernel=13, stride=1, dilation=1, residual=True, separable=True)
|
|
x = self.block(x, 2, repeat=2, kernel=15, stride=1, dilation=1, residual=True, separable=True)
|
|
x = self.block(x, 3, repeat=2, kernel=17, stride=1, dilation=1, residual=True, separable=True)
|
|
x = self.block(x, 4, repeat=1, kernel=29, stride=1, dilation=2, residual=False, separable=True)
|
|
x = self.block(x, 5, repeat=1, kernel=1, stride=1, dilation=1, residual=False, separable=False)
|
|
logits = self.conv1d(x, "decoder.layer0", 1, 1, 0, 1, True).squeeze(0).transpose(0, 1).contiguous().detach().cpu().numpy()
|
|
if logits.shape[1] == 1:
|
|
probabilities = 1.0 / (1.0 + np.exp(-logits[:, 0]))
|
|
else:
|
|
shifted = logits[:, :2] - logits[:, :2].max(axis=1, keepdims=True)
|
|
exp = np.exp(shifted)
|
|
probabilities = exp[:, 1] / exp.sum(axis=1)
|
|
segments: list[dict[str, float | int]] = []
|
|
active = False
|
|
start_frame = 0
|
|
confidence_sum = 0.0
|
|
confidence_count = 0
|
|
for frame, probability in enumerate(probabilities):
|
|
if float(probability) >= threshold:
|
|
if not active:
|
|
active = True
|
|
start_frame = frame
|
|
confidence_sum = 0.0
|
|
confidence_count = 0
|
|
confidence_sum += float(probability)
|
|
confidence_count += 1
|
|
continue
|
|
if active:
|
|
segments.append({
|
|
"start_sample": int(start_frame * self.HOP * self.OUTPUT_STRIDE * sample_rate / self.TARGET_SR),
|
|
"end_sample": int(frame * self.HOP * self.OUTPUT_STRIDE * sample_rate / self.TARGET_SR),
|
|
"confidence": confidence_sum / confidence_count,
|
|
})
|
|
active = False
|
|
if active:
|
|
segments.append({
|
|
"start_sample": int(start_frame * self.HOP * self.OUTPUT_STRIDE * sample_rate / self.TARGET_SR),
|
|
"end_sample": int(len(probabilities) * self.HOP * self.OUTPUT_STRIDE * sample_rate / self.TARGET_SR),
|
|
"confidence": confidence_sum / confidence_count,
|
|
})
|
|
return segments
|
|
|
|
|
|
RunnerFn = Callable[[dict[str, Any], argparse.Namespace], tuple[float, float]]
|
|
LoaderFn = Callable[[dict[str, Any], argparse.Namespace], tuple[Any, RunnerFn]]
|
|
|
|
|
|
def qwen3_asr_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import torch
|
|
|
|
set_threads(args.threads)
|
|
use_transformers457()
|
|
sys.path.insert(0, str(REPO_ROOT / "reference" / "Qwen3-ASR"))
|
|
from qwen_asr import Qwen3ASRModel
|
|
|
|
device_map = f"cuda:{args.device}" if args.backend == "cuda" else "cpu"
|
|
model = Qwen3ASRModel.from_pretrained(
|
|
str(ensure_absolute_path(case["model"])),
|
|
device_map=device_map,
|
|
)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
started = time.perf_counter()
|
|
model.transcribe(
|
|
audio=request["audio"],
|
|
context=request.get("context", ""),
|
|
language=request.get("language") or None,
|
|
return_time_stamps=False,
|
|
)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, input_audio_duration_sec(request)
|
|
|
|
return model, run
|
|
|
|
|
|
def ace_step_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
set_threads(args.threads)
|
|
ref_root = REPO_ROOT / "reference" / "ACE-Step-1.5"
|
|
if str(REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
if str(ref_root) not in sys.path:
|
|
sys.path.insert(0, str(ref_root))
|
|
from acestep.handler import AceStepHandler
|
|
from acestep.inference import GenerationConfig, GenerationParams, generate_music
|
|
from acestep.llm_inference import LLMHandler
|
|
|
|
backend = "cuda" if args.backend == "cuda" else "cpu"
|
|
checkpoint_dir = ensure_absolute_path(case["model"])
|
|
os.environ["ACESTEP_CHECKPOINTS_DIR"] = str(checkpoint_dir)
|
|
dit_model_path = case.get("load_options", {}).get("ace_step.dit_model_path", "acestep-v15-turbo")
|
|
lm_model_path = case.get("load_options", {}).get("ace_step.lm_model_path", "acestep-5Hz-lm-1.7B")
|
|
|
|
dit_handler = AceStepHandler()
|
|
init_message, ok = dit_handler.initialize_service(
|
|
project_root=str(checkpoint_dir),
|
|
config_path=dit_model_path,
|
|
device=backend,
|
|
force_dtype=None,
|
|
use_flash_attention=False,
|
|
compile_model=False,
|
|
offload_to_cpu=False,
|
|
offload_dit_to_cpu=False,
|
|
quantization=None,
|
|
)
|
|
if not ok:
|
|
raise RuntimeError(f"ACE-Step DiT init failed: {init_message}")
|
|
|
|
llm_handler = LLMHandler()
|
|
init_message, ok = llm_handler.initialize(
|
|
checkpoint_dir=str(checkpoint_dir),
|
|
lm_model_path=lm_model_path,
|
|
backend="pt",
|
|
device=backend,
|
|
offload_to_cpu=False,
|
|
dtype=None,
|
|
)
|
|
if not ok:
|
|
raise RuntimeError(f"ACE-Step LM init failed: {init_message}")
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
options = dict(request.get("options", {}))
|
|
params_dict = {
|
|
"task_type": request.get("task_route", "text2music"),
|
|
"caption": request.get("text", ""),
|
|
"lyrics": request.get("lyrics", ""),
|
|
"vocal_language": request.get("language", "unknown"),
|
|
"duration": request.get("duration_seconds"),
|
|
"inference_steps": request.get("num_inference_steps"),
|
|
"guidance_scale": request.get("guidance_scale"),
|
|
"seed": request.get("seed"),
|
|
"src_audio": request.get("audio"),
|
|
"reference_audio": request.get("reference_audio"),
|
|
"instruction": request.get("instruction"),
|
|
"audio_codes": request.get("audio_codes", ""),
|
|
"repainting_start": request.get("repaint_start"),
|
|
"repainting_end": request.get("repaint_end"),
|
|
"lm_cfg_scale": options.get("lm_cfg_scale"),
|
|
"lm_temperature": options.get("lm_temperature"),
|
|
"lm_top_k": options.get("lm_top_k"),
|
|
"lm_top_p": options.get("lm_top_p"),
|
|
}
|
|
negative_prompt = request.get("negative_prompt")
|
|
if negative_prompt is not None:
|
|
params_dict["lm_negative_prompt"] = negative_prompt
|
|
params_dict = {key: value for key, value in params_dict.items() if value is not None}
|
|
params = GenerationParams(**params_dict)
|
|
config = GenerationConfig(
|
|
batch_size=int(request.get("batch_size", 1)),
|
|
allow_lm_batch=False,
|
|
use_random_seed=not ("seed" in request),
|
|
seeds=[int(request.get("seed", 0))],
|
|
noise_file=None,
|
|
lm_batch_chunk_size=1,
|
|
constrained_decoding_debug=False,
|
|
audio_format="wav",
|
|
)
|
|
started = time.perf_counter()
|
|
results = generate_music(
|
|
dit_handler=dit_handler,
|
|
llm_handler=llm_handler,
|
|
params=params,
|
|
config=config,
|
|
)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
if not results.success or not results.audios:
|
|
raise RuntimeError(results.error or "ACE-Step returned no output")
|
|
first_audio = results.audios[0]
|
|
audio = np.asarray(first_audio["tensor"], dtype=np.float32).reshape(-1)
|
|
sample_rate = int(first_audio["sample_rate"])
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, audio_duration_sec(audio, sample_rate)
|
|
|
|
return (dit_handler, llm_handler), run
|
|
|
|
|
|
def qwen3_forced_aligner_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import torch
|
|
|
|
set_threads(args.threads)
|
|
use_transformers457()
|
|
sys.path.insert(0, str(REPO_ROOT / "reference" / "Qwen3-ASR"))
|
|
from qwen_asr.inference.qwen3_forced_aligner import Qwen3ForcedAligner
|
|
|
|
device_map = f"cuda:{args.device}" if args.backend == "cuda" else "cpu"
|
|
model = Qwen3ForcedAligner.from_pretrained(
|
|
str(ensure_absolute_path(case["model"])),
|
|
device_map=device_map,
|
|
)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
started = time.perf_counter()
|
|
model.align(
|
|
audio=request["audio"],
|
|
text=request["text"],
|
|
language=request.get("language") or "English",
|
|
)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, input_audio_duration_sec(request)
|
|
|
|
return model, run
|
|
|
|
|
|
def qwen3_tts_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import torch
|
|
|
|
set_threads(args.threads)
|
|
use_transformers457()
|
|
sys.path.insert(0, str(REPO_ROOT / "reference" / "Qwen3-TTS"))
|
|
from qwen_tts import Qwen3TTSModel
|
|
|
|
device_map = f"cuda:{args.device}" if args.backend == "cuda" else "cpu"
|
|
model = Qwen3TTSModel.from_pretrained(
|
|
str(ensure_absolute_path(case["model"])),
|
|
device_map=device_map,
|
|
)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
started = time.perf_counter()
|
|
if case["task"] == "vdes":
|
|
audio_list, sample_rate = model.generate_voice_design(
|
|
text=request["text"],
|
|
instruct=request["instruct"],
|
|
language=request.get("language", "English"),
|
|
max_new_tokens=int(request.get("max_tokens", 512)),
|
|
do_sample=bool(request.get("do_sample", True)),
|
|
)
|
|
elif request.get("speaker"):
|
|
audio_list, sample_rate = model.generate_custom_voice(
|
|
text=request["text"],
|
|
speaker=request["speaker"],
|
|
instruct=request.get("instruct", ""),
|
|
language=request.get("language", "English"),
|
|
max_new_tokens=int(request.get("max_tokens", 512)),
|
|
do_sample=bool(request.get("do_sample", True)),
|
|
)
|
|
else:
|
|
prompt = model.create_voice_clone_prompt(
|
|
ref_audio=request["voice_ref"],
|
|
ref_text=request["reference_text"],
|
|
)
|
|
audio_list, sample_rate = model.generate_voice_clone(
|
|
text=request["text"],
|
|
voice_clone_prompt=prompt,
|
|
language=request.get("language", "English"),
|
|
max_new_tokens=int(request.get("max_tokens", 512)),
|
|
do_sample=bool(request.get("do_sample", True)),
|
|
)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, audio_duration_sec(np.asarray(audio_list[0], dtype=np.float32), int(sample_rate))
|
|
|
|
return model, run
|
|
|
|
|
|
def chatterbox_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import torch
|
|
|
|
set_threads(args.threads)
|
|
torch.backends.cudnn.enabled = False
|
|
sys.path.insert(0, str(REPO_ROOT / "reference" / "chatterbox" / "src"))
|
|
language = "en"
|
|
request_language = ""
|
|
if case["requests"]:
|
|
request_language = case["requests"][0].get("language", "en")
|
|
if request_language.lower() == "en":
|
|
from chatterbox.tts import ChatterboxTTS
|
|
else:
|
|
from chatterbox.mtl_tts import ChatterboxMultilingualTTS as ChatterboxTTS
|
|
language = request_language
|
|
device = f"cuda:{args.device}" if args.backend == "cuda" else "cpu"
|
|
model = ChatterboxTTS.from_local(ensure_absolute_path(case["model"]), device=device)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
started = time.perf_counter()
|
|
if language == "en":
|
|
wav = model.generate(request["text"], audio_prompt_path=request["voice_ref"])
|
|
else:
|
|
wav = model.generate(request["text"], language_id=language, audio_prompt_path=request["voice_ref"])
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
audio = wav.detach().cpu().numpy().astype(np.float32)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, audio_duration_sec(audio, int(model.sr))
|
|
|
|
return model, run
|
|
|
|
|
|
def miocodec_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import torch
|
|
|
|
set_threads(args.threads)
|
|
sys.path.insert(0, str(REPO_ROOT / "reference" / "MioCodec" / "src"))
|
|
from miocodec import MioCodecModel, load_audio
|
|
|
|
model_dir = ensure_absolute_path(case["model"])
|
|
model = MioCodecModel.from_pretrained(
|
|
config_path=str(model_dir / "config.yaml"),
|
|
weights_path=str(model_dir / "model.safetensors"),
|
|
).eval()
|
|
device = torch.device(f"cuda:{args.device}" if args.backend == "cuda" else "cpu")
|
|
model = model.to(device)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
source = load_audio(request["audio"], sample_rate=model.config.sample_rate).to(device)
|
|
target = load_audio(request["voice_ref"], sample_rate=model.config.sample_rate).to(device)
|
|
started = time.perf_counter()
|
|
with torch.inference_mode():
|
|
decoded = model.voice_conversion(source, target)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
audio = decoded.detach().cpu().float().numpy()
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, audio_duration_sec(audio, int(model.config.sample_rate))
|
|
|
|
return model, run
|
|
|
|
|
|
def miotts_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
set_threads(args.threads)
|
|
ref_root = REPO_ROOT / "reference" / "MioTTS-Inference"
|
|
codec_root = REPO_ROOT / "reference" / "MioCodec" / "src"
|
|
if str(codec_root) not in sys.path:
|
|
sys.path.insert(0, str(codec_root))
|
|
if str(ref_root) not in sys.path:
|
|
sys.path.insert(0, str(ref_root))
|
|
from miocodec import MioCodecModel
|
|
from miotts_server.audio import load_reference_audio_path
|
|
from miotts_server.codec import MioCodecService
|
|
from miotts_server.text import normalize_text
|
|
from miotts_server.token_parser import parse_speech_tokens
|
|
|
|
device = f"cuda:{args.device}" if args.backend == "cuda" else "cpu"
|
|
model_path = ensure_absolute_path(case["model"])
|
|
tokenizer = AutoTokenizer.from_pretrained(str(model_path), local_files_only=True)
|
|
model = AutoModelForCausalLM.from_pretrained(str(model_path), local_files_only=True).to(device).eval()
|
|
codec = MioCodecService(model_id=str(REPO_ROOT / "models" / "MioCodec-25Hz-44.1kHz-v2"), device=device, presets_dir=REPO_ROOT / "presets")
|
|
codec_model_dir = REPO_ROOT / "models" / "MioCodec-25Hz-44.1kHz-v2"
|
|
codec._codec = MioCodecModel.from_pretrained(
|
|
config_path=str(codec_model_dir / "config.yaml"),
|
|
weights_path=str(codec_model_dir / "model.safetensors"),
|
|
).eval().to(device)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
text = normalize_text(request["text"])
|
|
messages = [{"role": "user", "content": text}]
|
|
inputs = tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize=True,
|
|
add_generation_prompt=True,
|
|
return_dict=True,
|
|
return_tensors="pt",
|
|
).to(device)
|
|
reference_waveform = load_reference_audio_path(request["voice_ref"], codec.sample_rate).to(device)
|
|
started = time.perf_counter()
|
|
with torch.inference_mode():
|
|
output_ids = model.generate(
|
|
**inputs,
|
|
do_sample=bool(request.get("do_sample", True)),
|
|
max_new_tokens=int(request.get("max_tokens", 700)),
|
|
temperature=0.8,
|
|
top_p=1.0,
|
|
repetition_penalty=1.0,
|
|
pad_token_id=tokenizer.eos_token_id,
|
|
eos_token_id=tokenizer.eos_token_id,
|
|
)
|
|
prompt_length = inputs["input_ids"].shape[1]
|
|
generated_text = tokenizer.decode(output_ids[0, prompt_length:], skip_special_tokens=False)
|
|
speech_tokens = parse_speech_tokens(generated_text)
|
|
audio = codec.synthesize(speech_tokens, reference_waveform=reference_waveform)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
audio_tensor = audio if isinstance(audio, torch.Tensor) else torch.tensor(audio)
|
|
audio_np = audio_tensor.detach().cpu().numpy().astype(np.float32, copy=False).reshape(-1)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, audio_duration_sec(audio_np, codec.sample_rate)
|
|
|
|
return (model, tokenizer, codec), run
|
|
|
|
|
|
def voxcpm2_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import torch
|
|
|
|
set_threads(args.threads)
|
|
sys.path.insert(0, str(REPO_ROOT / "reference" / "VoxCPM" / "src"))
|
|
from voxcpm import VoxCPM
|
|
|
|
model = VoxCPM.from_pretrained(str(ensure_absolute_path(case["model"])), load_denoiser=False)
|
|
if args.backend == "cuda":
|
|
torch.cuda.set_device(args.device)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
kwargs = {
|
|
"text": request["text"],
|
|
"cfg_value": float(request.get("guidance_scale", 2.0)),
|
|
"inference_timesteps": int(request.get("num_inference_steps", 10)),
|
|
}
|
|
if request.get("voice_ref"):
|
|
kwargs["reference_wav_path"] = request["voice_ref"]
|
|
started = time.perf_counter()
|
|
wav = model.generate(**kwargs)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
audio = np.asarray(wav, dtype=np.float32)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, audio_duration_sec(audio, int(model.tts_model.sample_rate))
|
|
|
|
return model, run
|
|
|
|
|
|
def kokoro_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import torch
|
|
from safetensors.torch import load_file
|
|
|
|
set_threads(args.threads)
|
|
ref_root = REPO_ROOT / "reference" / "kokoro"
|
|
if str(ref_root) not in sys.path:
|
|
sys.path.insert(0, str(ref_root))
|
|
from kokoro.pipeline import KPipeline
|
|
from kokoro.model import KModel
|
|
|
|
model_dir = ensure_absolute_path(case["model"])
|
|
device = f"cuda:{args.device}" if args.backend == "cuda" else "cpu"
|
|
checkpoint_path = model_dir / "kokoro-v1_0.safetensors"
|
|
original_torch_load = torch.load
|
|
|
|
def patched_torch_load(path, *load_args, **load_kwargs):
|
|
if Path(path) == checkpoint_path:
|
|
flat_state = load_file(str(checkpoint_path), device="cpu")
|
|
grouped: dict[str, dict[str, object]] = {}
|
|
for name, tensor in flat_state.items():
|
|
prefix, rest = name.split(".", 1)
|
|
grouped.setdefault(prefix, {})[rest] = tensor
|
|
return grouped
|
|
return original_torch_load(path, *load_args, **load_kwargs)
|
|
|
|
torch.load = patched_torch_load
|
|
try:
|
|
model = KModel(
|
|
repo_id="hexgrad/Kokoro-82M",
|
|
config=str(model_dir / "config.json"),
|
|
model=str(checkpoint_path),
|
|
).to(device).eval()
|
|
finally:
|
|
torch.load = original_torch_load
|
|
voice_id = case.get("session_options", {}).get("voice", "af_heart")
|
|
lang_code = voice_id[0]
|
|
pipeline = KPipeline(lang_code=lang_code, model=model, repo_id="hexgrad/Kokoro-82M")
|
|
voices = json.loads((model_dir / "voices.json").read_text(encoding="utf-8"))
|
|
voice_info = voices[voice_id]
|
|
voice_pack = np.fromfile(model_dir / "voices" / voice_info["path"], dtype=np.float32).reshape(
|
|
int(voice_info["rows"]),
|
|
int(voice_info["cols"]),
|
|
)
|
|
voice_tensor = torch.from_numpy(voice_pack).unsqueeze(1)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
started = time.perf_counter()
|
|
outputs = [
|
|
item
|
|
for item in pipeline(
|
|
request["text"],
|
|
voice=voice_tensor,
|
|
speed=1.0,
|
|
split_pattern=r"\n+",
|
|
)
|
|
if getattr(item, "audio", None) is not None
|
|
]
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
if not outputs:
|
|
raise RuntimeError("Kokoro produced no audio")
|
|
audio = np.concatenate([np.asarray(item.audio, dtype=np.float32) for item in outputs])
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, audio_duration_sec(audio, 24000)
|
|
|
|
return (pipeline, voice_tensor), run
|
|
|
|
|
|
def pocket_tts_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import yaml
|
|
import torch
|
|
|
|
set_threads(args.threads)
|
|
ref_root = REPO_ROOT / "reference" / "pocket-tts"
|
|
if str(ref_root) not in sys.path:
|
|
sys.path.insert(0, str(ref_root))
|
|
from pocket_tts import TTSModel
|
|
|
|
model_dir = ensure_absolute_path(case["model"])
|
|
language = case.get("session_options", {}).get("language", "english")
|
|
config_src = ref_root / "pocket_tts" / "config" / f"{language}.yaml"
|
|
config = yaml.safe_load(config_src.read_text(encoding="utf-8"))
|
|
bundle_dir = model_dir / "languages" / language if (model_dir / "languages" / language).exists() else model_dir
|
|
config["weights_path"] = str(bundle_dir / "model.safetensors")
|
|
config["weights_path_without_voice_cloning"] = str(bundle_dir / "model.safetensors")
|
|
config["flow_lm"]["lookup_table"]["tokenizer_path"] = str(bundle_dir / "tokenizer.model")
|
|
config_dir = REPO_ROOT / "build" / "perf" / "pocket_tts"
|
|
config_dir.mkdir(parents=True, exist_ok=True)
|
|
config_path = config_dir / f"{language}.yaml"
|
|
config_path.write_text(yaml.safe_dump(config, sort_keys=False), encoding="utf-8")
|
|
device = torch.device(f"cuda:{args.device}" if args.backend == "cuda" else "cpu")
|
|
model = TTSModel.load_model(config=str(config_path)).to(device).eval()
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
voice_state = model.get_state_for_audio_prompt(request["voice_ref"])
|
|
started = time.perf_counter()
|
|
audio = model.generate_audio(voice_state, request["text"])
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
audio_np = audio.detach().cpu().numpy().astype(np.float32, copy=False)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, audio_duration_sec(audio_np, int(model.sample_rate))
|
|
|
|
return model, run
|
|
|
|
|
|
def omnivoice_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import torch
|
|
|
|
set_threads(args.threads)
|
|
ref_root = REPO_ROOT / "reference" / "OmniVoice"
|
|
if str(ref_root) not in sys.path:
|
|
sys.path.insert(0, str(ref_root))
|
|
from omnivoice.models.omnivoice import OmniVoice
|
|
|
|
device_map = f"cuda:{args.device}" if args.backend == "cuda" else "cpu"
|
|
model = OmniVoice.from_pretrained(str(ensure_absolute_path(case["model"])), device_map=device_map)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
started = time.perf_counter()
|
|
audio = model.generate(
|
|
text=request["text"],
|
|
language=request.get("language", "en"),
|
|
ref_audio=request.get("voice_ref"),
|
|
ref_text=request.get("reference_text"),
|
|
instruct=request.get("instruct"),
|
|
num_step=int(request.get("num_inference_steps", 16)),
|
|
guidance_scale=float(request.get("guidance_scale", 2.0)),
|
|
audio_chunk_duration=float(request.get("options", {}).get("audio_chunk_duration", 15.0)),
|
|
)[0]
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
audio_np = np.asarray(audio, dtype=np.float32)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, audio_duration_sec(audio_np, int(model.sampling_rate))
|
|
|
|
return model, run
|
|
|
|
|
|
def moss_tts_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import torch
|
|
use_transformers457()
|
|
from transformers import AutoConfig, AutoModelForCausalLM
|
|
|
|
set_threads(args.threads)
|
|
device = f"cuda:{args.device}" if args.backend == "cuda" else "cpu"
|
|
model_path = ensure_absolute_path(case["model"])
|
|
config = AutoConfig.from_pretrained(str(model_path), trust_remote_code=True, local_files_only=True)
|
|
attention_impl = "sdpa" if args.backend == "cuda" else "eager"
|
|
config.attn_implementation = attention_impl
|
|
config.local_transformer_attn_implementation = attention_impl
|
|
if getattr(config, "gpt2_config", None) is not None:
|
|
config.gpt2_config._attn_implementation = attention_impl
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
str(model_path),
|
|
config=config,
|
|
trust_remote_code=True,
|
|
local_files_only=True,
|
|
).to(device).eval()
|
|
force_attention_implementation(model, attention_impl)
|
|
audio_tokenizer_path = str(REPO_ROOT / "models" / "MOSS-Audio-Tokenizer-Nano")
|
|
text_tokenizer_path = str(model_path / "tokenizer.model") if (model_path / "tokenizer.model").is_file() else str(model_path)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
out_path = REPO_ROOT / "build" / "perf" / "moss_tts_request.wav"
|
|
started = time.perf_counter()
|
|
model.inference(
|
|
text=request["text"],
|
|
output_audio_path=str(out_path),
|
|
mode="voice_clone",
|
|
prompt_audio_path=request["voice_ref"],
|
|
text_tokenizer_path=text_tokenizer_path,
|
|
audio_tokenizer_pretrained_name_or_path=audio_tokenizer_path,
|
|
device=device,
|
|
nq=16,
|
|
max_new_frames=int(request.get("max_tokens", 120)),
|
|
do_sample=bool(request.get("do_sample", False)),
|
|
use_kv_cache=True,
|
|
)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
audio, sample_rate = read_wav_mono_f32(out_path)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, audio_duration_sec(audio, sample_rate)
|
|
|
|
return model, run
|
|
|
|
|
|
def demucs_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import torch
|
|
|
|
set_threads(args.threads)
|
|
ref_root = REPO_ROOT / "reference" / "demucs"
|
|
if str(ref_root) not in sys.path:
|
|
sys.path.insert(0, str(ref_root))
|
|
from demucs.api import Separator
|
|
|
|
device = f"cuda:{args.device}" if args.backend == "cuda" else "cpu"
|
|
model_path = ensure_absolute_path(case["model"])
|
|
model_name = "htdemucs"
|
|
repo_arg: Path | None = model_path
|
|
manifest_path = model_path / "manifest.json"
|
|
if manifest_path.is_file():
|
|
repo_arg = None
|
|
separator = Separator(
|
|
model=model_name,
|
|
repo=repo_arg,
|
|
device=device,
|
|
shifts=0,
|
|
overlap=0.25,
|
|
split=True,
|
|
segment=None,
|
|
jobs=0,
|
|
progress=False,
|
|
)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
started = time.perf_counter()
|
|
separator.separate_audio_file(Path(request["audio"]))
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, input_audio_duration_sec(request)
|
|
|
|
return separator, run
|
|
|
|
|
|
def roformer_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import types
|
|
import yaml
|
|
import soundfile as sf
|
|
import torch
|
|
|
|
set_threads(args.threads)
|
|
ref_root = REPO_ROOT / "reference" / "Mel-Band-Roformer-Vocal-Model"
|
|
if str(ref_root) not in sys.path:
|
|
sys.path.insert(0, str(ref_root))
|
|
if "ml_collections" not in sys.modules:
|
|
class ConfigDict(dict):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__()
|
|
payload = dict(*args, **kwargs)
|
|
for key, value in payload.items():
|
|
self[key] = self._convert(value)
|
|
|
|
@classmethod
|
|
def _convert(cls, value):
|
|
if isinstance(value, dict):
|
|
return cls(value)
|
|
if isinstance(value, list):
|
|
return [cls._convert(item) for item in value]
|
|
if isinstance(value, tuple):
|
|
return tuple(cls._convert(item) for item in value)
|
|
return value
|
|
|
|
def __getattr__(self, name):
|
|
return self[name]
|
|
|
|
def __setattr__(self, name, value):
|
|
self[name] = self._convert(value)
|
|
module = types.ModuleType("ml_collections")
|
|
module.ConfigDict = ConfigDict
|
|
sys.modules["ml_collections"] = module
|
|
from ml_collections import ConfigDict
|
|
from inference import run_folder
|
|
from utils import get_model_from_config
|
|
|
|
ckpt_path = ensure_absolute_path("models/melbandroformer/MelBandRoformer.ckpt")
|
|
config_path = REPO_ROOT / "reference" / "Mel-Band-Roformer-Vocal-Model" / "configs" / "config_vocals_mel_band_roformer.yaml"
|
|
config = ConfigDict(yaml.load(config_path.read_text(encoding="utf-8"), Loader=yaml.FullLoader))
|
|
model = get_model_from_config("mel_band_roformer", config)
|
|
state = torch.load(ckpt_path, map_location="cpu")
|
|
model.load_state_dict(state, strict=False)
|
|
device = torch.device(f"cuda:{args.device}" if args.backend == "cuda" else "cpu")
|
|
model = model.to(device=device).eval()
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
work_dir = REPO_ROOT / "build" / "perf" / "roformer"
|
|
input_dir = work_dir / "input"
|
|
store_dir = work_dir / "store"
|
|
input_dir.mkdir(parents=True, exist_ok=True)
|
|
store_dir.mkdir(parents=True, exist_ok=True)
|
|
staged = input_dir / Path(request["audio"]).name
|
|
shutil.copyfile(request["audio"], staged)
|
|
started = time.perf_counter()
|
|
run_folder(model, argparse.Namespace(input_folder=str(input_dir), store_dir=str(store_dir)), config, device, verbose=False)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, input_audio_duration_sec(request)
|
|
|
|
return model, run
|
|
|
|
|
|
def seed_vc_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import os
|
|
import importlib
|
|
import torch
|
|
from types import SimpleNamespace
|
|
|
|
set_threads(args.threads)
|
|
ref_root = REPO_ROOT / "reference" / "SeedVC"
|
|
sys.path.insert(0, str(ref_root))
|
|
bigvgan_module = importlib.import_module("modules.bigvgan.bigvgan")
|
|
bigvgan_from_pretrained = bigvgan_module.BigVGAN._from_pretrained.__func__
|
|
|
|
def patched_bigvgan_from_pretrained(
|
|
cls,
|
|
*,
|
|
model_id,
|
|
revision,
|
|
cache_dir,
|
|
force_download,
|
|
proxies=None,
|
|
resume_download=False,
|
|
local_files_only=False,
|
|
token=None,
|
|
map_location="cpu",
|
|
strict=False,
|
|
use_cuda_kernel=False,
|
|
**model_kwargs,
|
|
):
|
|
return bigvgan_from_pretrained(
|
|
cls,
|
|
model_id=model_id,
|
|
revision=revision,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
resume_download=resume_download,
|
|
local_files_only=local_files_only,
|
|
token=token,
|
|
map_location=map_location,
|
|
strict=strict,
|
|
use_cuda_kernel=use_cuda_kernel,
|
|
**model_kwargs,
|
|
)
|
|
|
|
bigvgan_module.BigVGAN._from_pretrained = classmethod(patched_bigvgan_from_pretrained)
|
|
cwd = os.getcwd()
|
|
os.chdir(ref_root)
|
|
try:
|
|
module = load_module("seedvc_reference_inference_v2_perf", ref_root / "inference_v2.py")
|
|
finally:
|
|
os.chdir(cwd)
|
|
model_root = ensure_absolute_path(case["model"])
|
|
if (model_root / "v2" / "ar.safetensors").is_file():
|
|
official_root = REPO_ROOT / "models" / "SeedVC"
|
|
if official_root.is_dir():
|
|
model_root = official_root
|
|
device = f"cuda:{args.device}" if args.backend == "cuda" else "cpu"
|
|
load_args = SimpleNamespace(
|
|
ar_checkpoint_path=str(model_root / "seed-vc/v2/ar_base.pth"),
|
|
cfm_checkpoint_path=str(model_root / "seed-vc/v2/cfm_small.pth"),
|
|
compile=False,
|
|
output=str(REPO_ROOT / "build" / "perf" / "seedvc"),
|
|
)
|
|
cwd = os.getcwd()
|
|
os.chdir(ref_root)
|
|
try:
|
|
wrapper = module.load_v2_models(load_args)
|
|
module.vc_wrapper_v2 = wrapper
|
|
finally:
|
|
os.chdir(cwd)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
run_args = SimpleNamespace(
|
|
diffusion_steps=int(request.get("num_inference_steps", 30)),
|
|
length_adjust=float(request.get("options", {}).get("length_adjust", 1.0)),
|
|
intelligibility_cfg_rate=float(request.get("options", {}).get("intelligibility_cfg_rate", 0.7)),
|
|
similarity_cfg_rate=float(request.get("options", {}).get("similarity_cfg_rate", 0.7)),
|
|
top_p=0.9,
|
|
temperature=1.0,
|
|
repetition_penalty=1.0,
|
|
convert_style=False,
|
|
anonymization_only=False,
|
|
)
|
|
started = time.perf_counter()
|
|
cwd_inner = os.getcwd()
|
|
os.chdir(ref_root)
|
|
try:
|
|
audio = module.convert_voice_v2(request["audio"], request["voice_ref"], run_args)
|
|
finally:
|
|
os.chdir(cwd_inner)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
sample_rate, waveform = audio
|
|
audio_np = np.asarray(waveform, dtype=np.float32).reshape(-1)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, audio_duration_sec(audio_np, int(sample_rate))
|
|
|
|
return wrapper, run
|
|
|
|
|
|
def vevo2_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import torch
|
|
|
|
set_threads(args.threads)
|
|
use_transformers457()
|
|
enable_legacy_llama_config_positional_args()
|
|
ref_root = REPO_ROOT / "reference" / "Amphion"
|
|
if str(ref_root) not in sys.path:
|
|
sys.path.insert(0, str(ref_root))
|
|
from models.svc.vevo2.vevo2_utils import Vevo2InferencePipeline
|
|
|
|
model_root = ensure_absolute_path(case["model"])
|
|
device = torch.device(f"cuda:{args.device}" if args.backend == "cuda" else "cpu")
|
|
cwd = os.getcwd()
|
|
os.chdir(ref_root)
|
|
try:
|
|
pipeline = Vevo2InferencePipeline(
|
|
prosody_tokenizer_ckpt_path=str(model_root / "tokenizer" / "prosody_fvq512_6.25hz"),
|
|
content_style_tokenizer_ckpt_path=str(model_root / "tokenizer" / "contentstyle_fvq16384_12.5hz"),
|
|
ar_cfg_path=str(model_root / "contentstyle_modeling" / "posttrained" / "amphion_config.json"),
|
|
ar_ckpt_path=str(model_root / "contentstyle_modeling" / "posttrained"),
|
|
fmt_cfg_path=str(model_root / "acoustic_modeling" / "fm_emilia101k_singnet7k_repa" / "config.json"),
|
|
fmt_ckpt_path=str(model_root / "acoustic_modeling" / "fm_emilia101k_singnet7k_repa"),
|
|
vocoder_cfg_path=str(model_root / "vocoder" / "config.json"),
|
|
vocoder_ckpt_path=str(model_root / "vocoder"),
|
|
device=device,
|
|
)
|
|
finally:
|
|
os.chdir(cwd)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
started = time.perf_counter()
|
|
cwd_inner = os.getcwd()
|
|
os.chdir(ref_root)
|
|
try:
|
|
audio = pipeline.inference_ar_and_fm(
|
|
target_text=request["target_text"],
|
|
prosody_wav_path=request.get("style_ref"),
|
|
style_ref_wav_path=request.get("style_ref"),
|
|
style_ref_wav_text=request.get("style_ref_text", ""),
|
|
timbre_ref_wav_path=request.get("target_voice"),
|
|
top_k=int(request.get("top_k", 25)),
|
|
top_p=float(request.get("top_p", 0.8)),
|
|
temperature=float(request.get("temperature", 1.0)),
|
|
flow_matching_steps=int(request.get("num_inference_steps", 32)),
|
|
display_audio=False,
|
|
)
|
|
finally:
|
|
os.chdir(cwd_inner)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
audio_np = np.asarray(audio, dtype=np.float32).reshape(-1)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, audio_duration_sec(audio_np, 24000)
|
|
|
|
return pipeline, run
|
|
|
|
|
|
def parakeet_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
set_threads(args.threads)
|
|
driver = load_module("parakeet_test_driver_perf", REPO_ROOT / "reference" / "parakeet-tdt" / "tools" / "test_driver.py")
|
|
driver_args = argparse.Namespace(
|
|
model=str(ensure_absolute_path(case["model"])),
|
|
audio="",
|
|
backend=args.backend,
|
|
device=args.device,
|
|
mode="perf",
|
|
run_mode="offline",
|
|
trace_log=None,
|
|
timing_log=None,
|
|
batch_size=1,
|
|
timestamps=False,
|
|
timestamp_levels="char,word,segment",
|
|
self_attention_model="rel_pos_local_attn",
|
|
att_context_left=256,
|
|
att_context_right=256,
|
|
chunk_secs=2.0,
|
|
left_context_secs=10.0,
|
|
right_context_secs=2.0,
|
|
clean_groundtruth_text=False,
|
|
langid="en",
|
|
)
|
|
driver.configure_runtime(driver_args)
|
|
nemo_asr, open_dict, torch, _timing_log_scalar, _trace_log_scalar = driver.import_runtime_modules()
|
|
torch.set_num_threads(args.threads)
|
|
if hasattr(torch, "set_num_interop_threads"):
|
|
try:
|
|
torch.set_num_interop_threads(1)
|
|
except RuntimeError:
|
|
pass
|
|
model = driver.resolve_model(nemo_asr, driver_args.model)
|
|
driver.configure_device(torch, model, driver_args.backend)
|
|
driver.configure_timestamps(model, open_dict)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
started = time.perf_counter()
|
|
model.transcribe(
|
|
[request["audio"]],
|
|
batch_size=1,
|
|
return_hypotheses=False,
|
|
timestamps=False,
|
|
)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, input_audio_duration_sec(request)
|
|
|
|
return model, run
|
|
|
|
|
|
def sortformer_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import librosa
|
|
import torch
|
|
|
|
set_threads(args.threads)
|
|
ref_root = REPO_ROOT / "reference" / "nemo"
|
|
if str(ref_root) not in sys.path:
|
|
sys.path.insert(0, str(ref_root))
|
|
from nemo.collections.asr.models import SortformerEncLabelModel
|
|
from nemo.collections.asr.parts.mixins.diarization import DiarizeConfig
|
|
|
|
map_location = torch.device(f"cuda:{args.device}" if args.backend == "cuda" else "cpu")
|
|
model_path = ensure_absolute_path(case["model"])
|
|
model_name = str(model_path)
|
|
if model_name.endswith(".nemo"):
|
|
model = SortformerEncLabelModel.restore_from(model_name, map_location=map_location).to(map_location)
|
|
elif model_path.is_dir():
|
|
model = SortformerEncLabelModel.from_pretrained(model_name=f"nvidia/{model_path.name}", map_location=map_location).to(map_location)
|
|
else:
|
|
model = SortformerEncLabelModel.from_pretrained(model_name=model_name, map_location=map_location).to(map_location)
|
|
model.eval()
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
audio_np, sample_rate = librosa.load(request["audio"], sr=None, mono=False)
|
|
started = time.perf_counter()
|
|
diarize_cfg = DiarizeConfig(
|
|
session_len_sec=20.0,
|
|
batch_size=1,
|
|
num_workers=0,
|
|
sample_rate=sample_rate,
|
|
verbose=False,
|
|
include_tensor_outputs=False,
|
|
)
|
|
model.diarize(
|
|
audio=audio_np,
|
|
sample_rate=sample_rate,
|
|
override_config=diarize_cfg,
|
|
)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, input_audio_duration_sec(request)
|
|
|
|
return model, run
|
|
|
|
|
|
def citrinet_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
set_threads(args.threads)
|
|
device = torch.device(f"cuda:{args.device}" if args.backend == "cuda" else "cpu")
|
|
model = RefCitrinet(ensure_absolute_path(case["model"]), device)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
started = time.perf_counter()
|
|
model.transcribe(Path(request["audio"]))
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, input_audio_duration_sec(request)
|
|
|
|
return model, run
|
|
|
|
|
|
def marblenet_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
set_threads(args.threads)
|
|
device = torch.device(f"cuda:{args.device}" if args.backend == "cuda" else "cpu")
|
|
model = RefMarbleNet(ensure_absolute_path(case["model"]), device)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
started = time.perf_counter()
|
|
model.detect(Path(request["audio"]))
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, input_audio_duration_sec(request)
|
|
|
|
return model, run
|
|
|
|
|
|
def silero_vad_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]:
|
|
import torch
|
|
|
|
set_threads(args.threads)
|
|
sys.path.insert(0, str(REPO_ROOT / "reference" / "silero-vad" / "src"))
|
|
from silero_vad import get_speech_timestamps, load_silero_vad
|
|
|
|
model = load_silero_vad(onnx=False)
|
|
device = torch.device(f"cuda:{args.device}" if args.backend == "cuda" else "cpu")
|
|
if args.backend == "cuda":
|
|
model = model.to(device)
|
|
|
|
def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]:
|
|
waveform, sample_rate = read_wav_mono_f32(Path(request["audio"]))
|
|
audio = torch.from_numpy(waveform.astype(np.float32, copy=False)).to(device)
|
|
started = time.perf_counter()
|
|
get_speech_timestamps(audio, model, sampling_rate=sample_rate, return_seconds=False, visualize_probs=False)
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize(args.device)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return wall_ms, audio_duration_sec(waveform, sample_rate)
|
|
|
|
return model, run
|
|
|
|
|
|
LOADERS: dict[str, LoaderFn] = {
|
|
"ace_step": ace_step_loader,
|
|
"qwen3_asr": qwen3_asr_loader,
|
|
"qwen3_forced_aligner": qwen3_forced_aligner_loader,
|
|
"qwen3_tts": qwen3_tts_loader,
|
|
"chatterbox": chatterbox_loader,
|
|
"miocodec": miocodec_loader,
|
|
"miotts": miotts_loader,
|
|
"voxcpm2": voxcpm2_loader,
|
|
"kokoro_tts": kokoro_loader,
|
|
"pocket_tts": pocket_tts_loader,
|
|
"omnivoice": omnivoice_loader,
|
|
"moss_tts_nano": moss_tts_loader,
|
|
"seed_vc": seed_vc_loader,
|
|
"vevo2": vevo2_loader,
|
|
"parakeet_tdt": parakeet_loader,
|
|
"sortformer_diar": sortformer_loader,
|
|
"citrinet_asr": citrinet_loader,
|
|
"marblenet_vad": marblenet_loader,
|
|
"htdemucs": demucs_loader,
|
|
"mel_band_roformer": roformer_loader,
|
|
"silero_vad": silero_vad_loader,
|
|
}
|
|
|
|
|
|
def measure_cold(loader: LoaderFn, case: dict[str, Any], request: dict[str, Any], args: argparse.Namespace) -> dict[str, Any]:
|
|
wall_ms: list[float] = []
|
|
duration_sec = 0.0
|
|
for _ in range(3):
|
|
started = time.perf_counter()
|
|
_model, run = loader(case, args)
|
|
run_wall_ms, duration_sec = run(request, args)
|
|
total_ms = (time.perf_counter() - started) * 1000.0
|
|
if total_ms < run_wall_ms:
|
|
total_ms = run_wall_ms
|
|
wall_ms.append(total_ms)
|
|
average_wall_ms = float(sum(wall_ms) / len(wall_ms))
|
|
return {
|
|
"wall_ms": wall_ms,
|
|
"average_wall_ms": average_wall_ms,
|
|
"duration_sec": duration_sec,
|
|
"rtf": (average_wall_ms / 1000.0 / duration_sec) if duration_sec > 0.0 else 0.0,
|
|
}
|
|
|
|
|
|
def measure_warm(loader: LoaderFn, case: dict[str, Any], warmup_request: dict[str, Any], request: dict[str, Any], args: argparse.Namespace) -> dict[str, Any]:
|
|
_model, run = loader(case, args)
|
|
run(warmup_request, args)
|
|
wall_ms: list[float] = []
|
|
duration_sec = 0.0
|
|
for _ in range(3):
|
|
run_wall_ms, duration_sec = run(request, args)
|
|
wall_ms.append(run_wall_ms)
|
|
average_wall_ms = float(sum(wall_ms) / len(wall_ms))
|
|
return {
|
|
"wall_ms": wall_ms,
|
|
"average_wall_ms": average_wall_ms,
|
|
"duration_sec": duration_sec,
|
|
"rtf": (average_wall_ms / 1000.0 / duration_sec) if duration_sec > 0.0 else 0.0,
|
|
}
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
perf_cases = load_perf_cases(args.perf_cases)
|
|
request_cases = read_json(args.request_cases)
|
|
timestamp = time.strftime("%Y%m%d-%H%M%S")
|
|
summary_out = args.summary_out or (REPO_ROOT / "build" / "logs" / "perf" / f"python_{args.backend}_{timestamp}" / "summary.json")
|
|
summary_out.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
summaries: list[dict[str, Any]] = []
|
|
for perf_case in perf_cases:
|
|
if args.id and args.id != perf_case.id:
|
|
continue
|
|
if perf_case.family not in LOADERS:
|
|
raise RuntimeError(f"python perf harness does not support family yet: {perf_case.family}")
|
|
case = find_case(request_cases, perf_case.case_id)
|
|
request = resolve_request_paths(materialize_perf_request(
|
|
find_request(case, perf_case.request_id),
|
|
perf_case.request_overrides,
|
|
perf_case.request_remove_keys,
|
|
))
|
|
if perf_case.warmup_request_id:
|
|
warmup_request = resolve_request_paths(materialize_perf_request(
|
|
find_request(case, perf_case.warmup_request_id),
|
|
perf_case.warmup_request_overrides,
|
|
perf_case.warmup_request_remove_keys,
|
|
))
|
|
else:
|
|
temp_dir = (summary_out.parent / "_warmup_temp" / perf_case.id).resolve()
|
|
temp_dir.mkdir(parents=True, exist_ok=True)
|
|
warmup_request = derive_warmup_request(request, temp_dir)
|
|
cold = measure_cold(LOADERS[perf_case.family], case, request, args)
|
|
warm = measure_warm(LOADERS[perf_case.family], case, warmup_request, request, args)
|
|
summary = {
|
|
"id": perf_case.id,
|
|
"family": perf_case.family,
|
|
"case_id": perf_case.case_id,
|
|
"request_id": perf_case.request_id,
|
|
"task": case["task"],
|
|
"mode": case["mode"],
|
|
"cold": cold,
|
|
"warm": warm,
|
|
}
|
|
summaries.append(summary)
|
|
print(
|
|
f"{perf_case.id} cold_avg_ms={cold['average_wall_ms']} cold_rtf={cold['rtf']} "
|
|
f"warm_avg_ms={warm['average_wall_ms']} warm_rtf={warm['rtf']}",
|
|
flush=True,
|
|
)
|
|
|
|
payload = {
|
|
"backend": args.backend,
|
|
"device": args.device,
|
|
"threads": args.threads,
|
|
"cases": summaries,
|
|
}
|
|
summary_out.write_text(json.dumps(payload, ensure_ascii=False) + "\n", encoding="utf-8")
|
|
print(f"summary_json={summary_out}", flush=True)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|