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
446 lines
16 KiB
Python
446 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import contextlib
|
|
import json
|
|
import os
|
|
import random
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Iterator
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
REFERENCE_ROOT = REPO_ROOT / "reference" / "Amphion"
|
|
if str(REFERENCE_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REFERENCE_ROOT))
|
|
if str(REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
from models.svc.vevo2.vevo2_utils import Vevo2InferencePipeline # noqa: E402
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Python reference Vevo2 warmbench.")
|
|
parser.add_argument("--family", default="vevo2")
|
|
parser.add_argument("--model", default=str(REPO_ROOT / "models" / "Vevo2"))
|
|
parser.add_argument("--backend", choices=["cpu", "cuda"], default="cuda")
|
|
parser.add_argument("--device", type=int, default=0)
|
|
parser.add_argument("--threads", type=int, default=8)
|
|
parser.add_argument("--warmup", type=int, default=0)
|
|
parser.add_argument("--iterations", type=int, default=1)
|
|
parser.add_argument("--timing-file", default="")
|
|
parser.add_argument("--output-dir", default="")
|
|
parser.add_argument("--noise-file", default="")
|
|
parser.add_argument("--request-json", default="")
|
|
parser.add_argument("--request-sequence-json", default="")
|
|
return parser.parse_args()
|
|
|
|
|
|
def resolve_path(path: str) -> str:
|
|
value = Path(path)
|
|
return str(value if value.is_absolute() else REPO_ROOT / value)
|
|
|
|
|
|
def normalize_device(backend: str, device_index: int) -> torch.device:
|
|
if backend == "cuda":
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA backend requested but torch.cuda.is_available() is false")
|
|
torch.cuda.set_device(device_index)
|
|
return torch.device(f"cuda:{device_index}")
|
|
return torch.device("cpu")
|
|
|
|
|
|
def sync_device(device: torch.device) -> None:
|
|
if device.type == "cuda":
|
|
torch.cuda.synchronize(device)
|
|
|
|
|
|
def set_request_seed(seed: int, device: torch.device) -> None:
|
|
torch.manual_seed(seed)
|
|
if device.type == "cuda":
|
|
torch.cuda.manual_seed_all(seed)
|
|
np.random.seed(seed & 0xFFFFFFFF)
|
|
random.seed(seed)
|
|
|
|
|
|
def load_requests(args: argparse.Namespace) -> list[dict[str, Any]]:
|
|
if args.request_sequence_json:
|
|
payload = json.loads(args.request_sequence_json)
|
|
if not isinstance(payload, list):
|
|
raise RuntimeError("--request-sequence-json must decode to a list")
|
|
return payload
|
|
if args.request_json:
|
|
payload = json.loads(args.request_json)
|
|
if not isinstance(payload, dict):
|
|
raise RuntimeError("--request-json must decode to an object")
|
|
return [payload]
|
|
raise RuntimeError("Vevo2 warmbench requires --request-json or --request-sequence-json")
|
|
|
|
|
|
def force_module_f32(module: Any, device: torch.device) -> None:
|
|
if module is not None and hasattr(module, "to"):
|
|
module.to(device=device, dtype=torch.float32)
|
|
|
|
|
|
def force_pipeline_f32(pipeline: Vevo2InferencePipeline, device: torch.device) -> None:
|
|
for name in (
|
|
"ar_model",
|
|
"style_tokenizer",
|
|
"content_style_tokenizer",
|
|
"fmt_model",
|
|
"mel_model",
|
|
"vocoder_model",
|
|
"whisper_model",
|
|
):
|
|
force_module_f32(getattr(pipeline, name, None), device)
|
|
if getattr(pipeline, "use_normed_whisper", False):
|
|
pipeline.whisper_mean = pipeline.whisper_mean.to(device=device, dtype=torch.float32)
|
|
pipeline.whisper_std = pipeline.whisper_std.to(device=device, dtype=torch.float32)
|
|
|
|
|
|
def load_pipeline(model_root: Path, device: torch.device) -> Vevo2InferencePipeline:
|
|
content_style_tokenizer_ckpt_path = model_root / "tokenizer" / "contentstyle_fvq16384_12.5hz"
|
|
prosody_tokenizer_ckpt_path = model_root / "tokenizer" / "prosody_fvq512_6.25hz"
|
|
ar_root = model_root / "contentstyle_modeling" / "posttrained"
|
|
fm_root = model_root / "acoustic_modeling" / "fm_emilia101k_singnet7k_repa"
|
|
vocoder_root = model_root / "vocoder"
|
|
pipeline = Vevo2InferencePipeline(
|
|
prosody_tokenizer_ckpt_path=str(prosody_tokenizer_ckpt_path),
|
|
content_style_tokenizer_ckpt_path=str(content_style_tokenizer_ckpt_path),
|
|
ar_cfg_path=str(ar_root / "amphion_config.json"),
|
|
ar_ckpt_path=str(ar_root),
|
|
fmt_cfg_path=str(fm_root / "config.json"),
|
|
fmt_ckpt_path=str(fm_root),
|
|
vocoder_cfg_path=str(vocoder_root / "config.json"),
|
|
vocoder_ckpt_path=str(vocoder_root),
|
|
device=device,
|
|
)
|
|
force_pipeline_f32(pipeline, device)
|
|
return pipeline
|
|
|
|
|
|
FM_ROUTES = {
|
|
"fm",
|
|
"inference_fm",
|
|
"source_audio_to_target_voice",
|
|
"style_preserved_vc",
|
|
"style_preserved_svc",
|
|
"vc",
|
|
"svc",
|
|
}
|
|
|
|
AR_ROUTES = {
|
|
"ar_and_fm",
|
|
"inference_ar_and_fm",
|
|
"text_prosody_to_target_voice",
|
|
"zero_shot_tts",
|
|
"tts",
|
|
"text_to_speech",
|
|
"text_to_singing",
|
|
"svs",
|
|
"singing_voice_synthesis",
|
|
"style_converted_vc",
|
|
"style_converted_svc",
|
|
"style_conversion",
|
|
"editing",
|
|
"speech_editing",
|
|
"singing_editing",
|
|
"singing_style_conversion",
|
|
"melody_control",
|
|
"humming_to_singing",
|
|
"instrument_to_singing",
|
|
}
|
|
|
|
|
|
def route_default_use_prosody(path: str) -> bool:
|
|
return path in {
|
|
"style_converted_vc",
|
|
"style_converted_svc",
|
|
"style_conversion",
|
|
"editing",
|
|
"speech_editing",
|
|
"singing_editing",
|
|
"singing_style_conversion",
|
|
"melody_control",
|
|
"humming_to_singing",
|
|
"instrument_to_singing",
|
|
}
|
|
|
|
|
|
def route_default_use_pitch_shift(path: str) -> bool:
|
|
return path in {
|
|
"style_preserved_vc",
|
|
"style_preserved_svc",
|
|
"vc",
|
|
"svc",
|
|
"style_converted_svc",
|
|
"singing_style_conversion",
|
|
"melody_control",
|
|
"humming_to_singing",
|
|
"instrument_to_singing",
|
|
}
|
|
|
|
|
|
def route_uses_source_as_reference(path: str) -> bool:
|
|
return path in {
|
|
"style_converted_vc",
|
|
"style_converted_svc",
|
|
"style_conversion",
|
|
"editing",
|
|
"speech_editing",
|
|
"singing_editing",
|
|
"singing_style_conversion",
|
|
}
|
|
|
|
|
|
def optional_audio_path(request: dict[str, Any], key: str) -> str | None:
|
|
value = request.get(key)
|
|
return resolve_path(str(value)) if value else None
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def controlled_fm_noise(noise_file: str, mel_dim: int) -> Iterator[None]:
|
|
if not noise_file:
|
|
yield
|
|
return
|
|
noise_values = np.fromfile(noise_file, dtype=np.float32)
|
|
if noise_values.size == 0:
|
|
raise RuntimeError(f"Vevo2 FM controlled noise file is empty: {noise_file}")
|
|
original_randn = torch.randn
|
|
used = False
|
|
|
|
def randn(*args: Any, **kwargs: Any) -> torch.Tensor:
|
|
nonlocal used
|
|
shape = args[0] if args else kwargs.get("size")
|
|
if (
|
|
not used
|
|
and isinstance(shape, (tuple, list, torch.Size))
|
|
and len(shape) == 3
|
|
and int(shape[-1]) == mel_dim
|
|
):
|
|
count = int(np.prod([int(dim) for dim in shape]))
|
|
if noise_values.size < count:
|
|
raise RuntimeError(
|
|
f"Vevo2 FM controlled noise file is too short: expected at least {count} floats, "
|
|
f"got {noise_values.size}"
|
|
)
|
|
used = True
|
|
array = noise_values[:count].reshape(tuple(int(dim) for dim in shape)).copy()
|
|
tensor = torch.from_numpy(array)
|
|
target_device = kwargs.get("device")
|
|
target_dtype = kwargs.get("dtype", torch.float32)
|
|
if target_device is not None:
|
|
tensor = tensor.to(device=target_device)
|
|
if target_dtype is not None:
|
|
tensor = tensor.to(dtype=target_dtype)
|
|
return tensor
|
|
return original_randn(*args, **kwargs)
|
|
|
|
torch.randn = randn
|
|
try:
|
|
yield
|
|
finally:
|
|
torch.randn = original_randn
|
|
if not used:
|
|
raise RuntimeError("Vevo2 FM controlled noise was not consumed")
|
|
|
|
|
|
def summarize_audio(audio: np.ndarray, sample_rate: int) -> dict[str, Any]:
|
|
if audio.ndim == 1:
|
|
frames = int(audio.shape[0])
|
|
channels = 1
|
|
flat = audio.astype(np.float32, copy=False)
|
|
else:
|
|
frames = int(audio.shape[0])
|
|
channels = int(audio.shape[1])
|
|
flat = audio.astype(np.float32, copy=False).reshape(-1)
|
|
if flat.size == 0:
|
|
raise RuntimeError("Vevo2 warmbench summary received empty audio")
|
|
return {
|
|
"sample_rate": int(sample_rate),
|
|
"channels": channels,
|
|
"samples": int(flat.size),
|
|
"frames": frames,
|
|
"sum": float(np.sum(flat, dtype=np.float64)),
|
|
"mean_abs": float(np.mean(np.abs(flat), dtype=np.float64)),
|
|
"rms": float(np.sqrt(np.mean(np.square(flat), dtype=np.float64))),
|
|
"min": float(np.min(flat)),
|
|
"max": float(np.max(flat)),
|
|
}
|
|
|
|
|
|
def write_audio(path: Path, waveform: torch.Tensor | np.ndarray, sample_rate: int = 24000) -> np.ndarray:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
if isinstance(waveform, torch.Tensor):
|
|
audio = waveform.detach().cpu().numpy()
|
|
else:
|
|
audio = np.asarray(waveform)
|
|
audio = np.asarray(audio, dtype=np.float32)
|
|
if audio.ndim == 2 and audio.shape[0] == 1:
|
|
audio = audio[0]
|
|
sf.write(str(path), audio, sample_rate, subtype="PCM_16")
|
|
normalized, _ = sf.read(str(path), always_2d=True)
|
|
return np.asarray(normalized, dtype=np.float32)
|
|
|
|
|
|
def run_request(
|
|
pipeline: Vevo2InferencePipeline,
|
|
request: dict[str, Any],
|
|
output_dir: Path,
|
|
device: torch.device,
|
|
noise_file: str,
|
|
) -> tuple[dict[str, Any], list[str]]:
|
|
seed = int(request.get("seed", 1234))
|
|
if seed < 0:
|
|
raise RuntimeError("Vevo2 warmbench requires a non-negative seed")
|
|
set_request_seed(seed, device)
|
|
path = str(request.get("path", "ar_and_fm"))
|
|
started = time.perf_counter()
|
|
with controlled_fm_noise(noise_file, int(pipeline.fmt_model.mel_dim)):
|
|
use_pitch_shift = bool(request.get("use_pitch_shift", route_default_use_pitch_shift(path)))
|
|
if path in FM_ROUTES:
|
|
audio = pipeline.inference_fm(
|
|
src_wav_path=resolve_path(str(request["source_wav_path"])),
|
|
timbre_ref_wav_path=resolve_path(str(request["timbre_ref_wav_path"])),
|
|
src_wav_text=str(request.get("source_wav_text", "")),
|
|
timbre_ref_wav_text=str(request.get("timbre_ref_wav_text", "")),
|
|
use_pitch_shift=use_pitch_shift,
|
|
used_duration_of_timbre_ref_wav_path=request.get("used_duration_of_timbre_ref_seconds"),
|
|
flow_matching_steps=int(request.get("flow_matching_steps", 32)),
|
|
display_audio=False,
|
|
)
|
|
elif path in AR_ROUTES:
|
|
source_wav_path = optional_audio_path(request, "source_wav_path")
|
|
prosody_wav_path = optional_audio_path(request, "prosody_wav_path")
|
|
style_ref_wav_path = optional_audio_path(request, "style_ref_wav_path")
|
|
timbre_ref_wav_path = optional_audio_path(request, "timbre_ref_wav_path")
|
|
if route_uses_source_as_reference(path):
|
|
prosody_wav_path = prosody_wav_path or source_wav_path
|
|
style_ref_wav_path = style_ref_wav_path or source_wav_path
|
|
timbre_ref_wav_path = timbre_ref_wav_path or source_wav_path
|
|
audio = pipeline.inference_ar_and_fm(
|
|
target_text=str(request["target_text"]),
|
|
prosody_wav_path=prosody_wav_path,
|
|
style_ref_wav_path=style_ref_wav_path,
|
|
style_ref_wav_text=str(request.get("style_ref_text", "")),
|
|
timbre_ref_wav_path=timbre_ref_wav_path,
|
|
use_prosody_code=bool(request.get("use_prosody_code", route_default_use_prosody(path))),
|
|
predict_target_prosody=bool(request.get("predict_target_prosody", False)),
|
|
top_k=int(request.get("top_k", 25)),
|
|
top_p=float(request.get("top_p", 0.8)),
|
|
temperature=float(request.get("temperature", 1.0)),
|
|
use_pitch_shift=use_pitch_shift,
|
|
prosody_wav_shifted_steps=int(request.get("prosody_wav_shifted_steps", 0)),
|
|
style_ref_wav_shifted_steps=int(request.get("style_ref_wav_shifted_steps", 0)),
|
|
target_duration=request.get("target_duration_seconds"),
|
|
used_duration_of_timbre_ref_wav_path=request.get("used_duration_of_timbre_ref_seconds"),
|
|
flow_matching_steps=int(request.get("flow_matching_steps", 32)),
|
|
display_audio=False,
|
|
)
|
|
else:
|
|
raise RuntimeError(f"unsupported Vevo2 warmbench path: {path}")
|
|
sync_device(device)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
|
|
audio_path = output_dir / "audio.wav"
|
|
written = write_audio(audio_path, audio)
|
|
step = {
|
|
"request": request,
|
|
"stems": [
|
|
{
|
|
"name": "audio",
|
|
"audio": str(audio_path),
|
|
"summary": summarize_audio(written, 24000),
|
|
}
|
|
],
|
|
"metrics": {"wall_ms": wall_ms},
|
|
}
|
|
return step, [f"vevo2.wall_ms {wall_ms:.6f}"]
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
if args.family != "vevo2":
|
|
raise RuntimeError(f"unsupported Vevo2 warmbench family: {args.family}")
|
|
|
|
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
|
torch.set_num_threads(max(1, args.threads))
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
torch.backends.cudnn.allow_tf32 = False
|
|
torch.backends.cudnn.benchmark = False
|
|
device = normalize_device(args.backend, args.device)
|
|
set_request_seed(0, device)
|
|
|
|
model_root = Path(args.model)
|
|
if not model_root.is_absolute():
|
|
model_root = REPO_ROOT / model_root
|
|
if not model_root.is_dir():
|
|
raise RuntimeError(f"Vevo2 model root not found: {model_root}")
|
|
noise_file = ""
|
|
if args.noise_file:
|
|
noise_path = Path(args.noise_file)
|
|
noise_file = str(noise_path if noise_path.is_absolute() else REPO_ROOT / noise_path)
|
|
os.chdir(REFERENCE_ROOT)
|
|
pipeline = load_pipeline(model_root, device)
|
|
|
|
requests = load_requests(args)
|
|
output_root = Path(args.output_dir) if args.output_dir else REPO_ROOT / "build" / "logs" / "parity" / "vevo2_python_outputs"
|
|
timing_lines: list[str] = [
|
|
f"vevo2.backend {args.backend}",
|
|
"vevo2.python_tf32_disabled 1",
|
|
"vevo2.python_force_f32 1",
|
|
]
|
|
for warmup_index in range(max(0, args.warmup)):
|
|
_, warmup_timing = run_request(
|
|
pipeline,
|
|
requests[0],
|
|
output_root / "warmup" / f"{warmup_index:02d}",
|
|
device,
|
|
noise_file,
|
|
)
|
|
timing_lines.extend(warmup_timing)
|
|
|
|
steps: list[dict[str, Any]] = []
|
|
for request_index, request in enumerate(requests):
|
|
total_ms = 0.0
|
|
last_step: dict[str, Any] | None = None
|
|
for iteration in range(max(1, args.iterations)):
|
|
step, run_timing = run_request(
|
|
pipeline,
|
|
request,
|
|
output_root / f"request_{request_index:02d}" / f"iter_{iteration:02d}",
|
|
device,
|
|
noise_file,
|
|
)
|
|
total_ms += float(step["metrics"]["wall_ms"])
|
|
last_step = step
|
|
timing_lines.extend(run_timing)
|
|
assert last_step is not None
|
|
last_step = dict(last_step)
|
|
last_step["request_index"] = request_index
|
|
last_step["metrics"] = {"wall_ms": total_ms / float(max(1, args.iterations))}
|
|
print(f"vevo2.wall_ms={last_step['metrics']['wall_ms']}")
|
|
steps.append(last_step)
|
|
|
|
if args.timing_file:
|
|
timing_path = Path(args.timing_file)
|
|
timing_path.parent.mkdir(parents=True, exist_ok=True)
|
|
timing_path.write_text("\n".join(timing_lines) + "\n", encoding="utf-8")
|
|
|
|
summary = {
|
|
"family": "vevo2",
|
|
"backend": args.backend,
|
|
"sequence_steps": steps,
|
|
}
|
|
print(f"summary_json={json.dumps(summary, ensure_ascii=False)}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|