voice_audio_multi/tests/personaplex/personaplex_python_warm_bench.py
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
Initial import: audio.cpp with Orange Pi documentation
2026-11-04 19:05:22 +01:00

506 lines
20 KiB
Python

#!/usr/bin/env python3
from __future__ import annotations
import argparse
import json
import random
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import numpy as np
import torch
REPO_ROOT = Path(__file__).resolve().parents[2]
REFERENCE_ROOT = REPO_ROOT / "reference" / "PersonaPlex" / "moshi"
DEFAULT_HF_REPO = "nvidia/personaplex-7b-v1"
DEFAULT_INPUT_WAV = REPO_ROOT / "reference" / "PersonaPlex" / "assets" / "test" / "input_assistant.wav"
DEFAULT_ASSISTANT_PROMPT = (
"You are a wise and friendly teacher. Answer questions or provide advice in a clear and engaging way."
)
@dataclass
class Request:
request_id: str
input_wav: Path
text_prompt: str
voice_prompt: str
seed: int
temp_audio: float
temp_text: float
topk_audio: int
topk_text: int
greedy: bool
@dataclass
class Runtime:
mimi: Any
other_mimi: Any
lm_gen: Any
text_tokenizer: Any
sample_rate: int
frame_size: int
device: str
offline: Any
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Python reference PersonaPlex warmbench.")
parser.add_argument("--model", type=Path, default=None)
parser.add_argument("--hf-repo", default=DEFAULT_HF_REPO)
parser.add_argument("--tokenizer", type=Path, default=None)
parser.add_argument("--moshi-weight", type=Path, default=None)
parser.add_argument("--mimi-weight", type=Path, default=None)
parser.add_argument("--voice-prompt-dir", type=Path, default=None)
parser.add_argument("--voice-prompt", default="NATF2.pt")
parser.add_argument("--text-prompt", default=DEFAULT_ASSISTANT_PROMPT)
parser.add_argument("--input-wav", type=Path, default=DEFAULT_INPUT_WAV)
parser.add_argument("--request-json", default="")
parser.add_argument("--request-sequence-json", default="")
parser.add_argument("--backend", choices=("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=1)
parser.add_argument("--iterations", type=int, default=1)
parser.add_argument("--warmup-input-wav", type=Path, default=DEFAULT_INPUT_WAV)
parser.add_argument("--warmup-text-prompt", default=DEFAULT_ASSISTANT_PROMPT)
parser.add_argument("--warmup-voice-prompt", default="NATF2.pt")
parser.add_argument("--seed", type=int, default=42424242)
parser.add_argument("--temp-audio", type=float, default=0.8)
parser.add_argument("--temp-text", type=float, default=0.7)
parser.add_argument("--topk-audio", type=int, default=250)
parser.add_argument("--topk-text", type=int, default=25)
parser.add_argument("--greedy", action="store_true")
parser.add_argument("--cpu-offload", action="store_true")
parser.add_argument("--audio-out", type=Path, default=Path("personaplex_python_audio.wav"))
parser.add_argument("--audio-out-dir", type=Path, default=None)
parser.add_argument("--text-out-dir", type=Path, default=None)
parser.add_argument("--timing-file", type=Path, default=None)
parser.add_argument("--summary-file", type=Path, default=None)
return parser.parse_args()
def resolve_path(path: Path | None) -> Path | None:
if path is None:
return path
return path if path.is_absolute() else REPO_ROOT / path
def resolve_voice_prompt_path(voice_prompt_dir: str, voice_prompt: str) -> Path:
path = Path(voice_prompt)
if path.is_absolute():
return path
repo_path = REPO_ROOT / path
if repo_path.is_file():
return repo_path
return Path(voice_prompt_dir) / path
def import_reference() -> Any:
if not (REFERENCE_ROOT / "moshi" / "offline.py").is_file():
raise RuntimeError(f"missing PersonaPlex reference code: {REFERENCE_ROOT}")
sys.path.insert(0, str(REFERENCE_ROOT))
from moshi import offline
return offline
def seed_all(seed: int, device: str) -> None:
if seed == -1:
return
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if device.startswith("cuda"):
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = False
def synchronize(device: str) -> None:
if device.startswith("cuda"):
torch.cuda.synchronize(torch.device(device))
def optional_file(path: Path | None, model_dir: Path | None, filename: str) -> str | None:
if path is not None:
resolved = resolve_path(path)
if not resolved.is_file():
raise FileNotFoundError(f"missing file: {resolved}")
return str(resolved)
if model_dir is not None:
candidate = resolve_path(model_dir) / filename
if candidate.is_file():
return str(candidate)
return None
def resolve_voice_prompt_dir(offline: Any, args: argparse.Namespace) -> str:
if args.voice_prompt_dir:
resolved = resolve_path(args.voice_prompt_dir)
if not resolved.is_dir():
raise FileNotFoundError(f"missing voice prompt directory: {resolved}")
return str(resolved)
model_dir = resolve_path(args.model) if args.model is not None else None
if model_dir is not None:
for candidate in (model_dir / "voices", model_dir / "voice_prompts"):
if candidate.is_dir():
return str(candidate)
return offline._get_voice_prompt_dir(None, args.hf_repo)
def load_runtime(args: argparse.Namespace, voice_prompt_dir: str, offline: Any) -> Runtime:
loaders = offline.loaders
if args.backend == "cuda" and not torch.cuda.is_available():
raise RuntimeError("PersonaPlex warmbench requested CUDA, but torch.cuda.is_available() is false")
torch.set_num_threads(max(1, args.threads))
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
device = f"cuda:{args.device}"
torch.cuda.set_device(args.device)
model_dir = resolve_path(args.model) if args.model is not None else None
tokenizer_path = optional_file(args.tokenizer, model_dir, loaders.TEXT_TOKENIZER_NAME)
moshi_weight = optional_file(args.moshi_weight, model_dir, loaders.MOSHI_NAME)
mimi_weight = optional_file(args.mimi_weight, model_dir, loaders.MIMI_NAME)
# Download config.json to mirror the reference offline path and verify repo access before timing.
offline.hf_hub_download(args.hf_repo, "config.json")
if mimi_weight is None:
mimi_weight = offline.hf_hub_download(args.hf_repo, loaders.MIMI_NAME)
if tokenizer_path is None:
tokenizer_path = offline.hf_hub_download(args.hf_repo, loaders.TEXT_TOKENIZER_NAME)
if moshi_weight is None:
moshi_weight = offline.hf_hub_download(args.hf_repo, loaders.MOSHI_NAME)
print(f"[TRACE] personaplex.reference_root {REFERENCE_ROOT}")
print(f"[TRACE] personaplex.hf_repo {args.hf_repo}")
print(f"[TRACE] personaplex.device {device}")
print(f"[TRACE] personaplex.voice_prompt_dir {voice_prompt_dir}")
print(f"[TRACE] personaplex.moshi_weight {moshi_weight}")
print(f"[TRACE] personaplex.mimi_weight {mimi_weight}")
print(f"[TRACE] personaplex.tokenizer {tokenizer_path}")
started = time.perf_counter()
mimi = loaders.get_mimi(mimi_weight, device)
other_mimi = loaders.get_mimi(mimi_weight, device)
text_tokenizer = offline.sentencepiece.SentencePieceProcessor(tokenizer_path)
lm = loaders.get_moshi_lm(moshi_weight, device=device, cpu_offload=args.cpu_offload)
lm.eval()
frame_size = int(mimi.sample_rate / mimi.frame_rate)
lm_gen = offline.LMGen(
lm,
audio_silence_frame_cnt=int(0.5 * mimi.frame_rate),
sample_rate=mimi.sample_rate,
device=device,
frame_rate=mimi.frame_rate,
save_voice_prompt_embeddings=False,
use_sampling=not args.greedy,
temp=args.temp_audio,
temp_text=args.temp_text,
top_k=args.topk_audio,
top_k_text=args.topk_text,
)
mimi.streaming_forever(1)
other_mimi.streaming_forever(1)
lm_gen.streaming_forever(1)
synchronize(device)
load_ms = (time.perf_counter() - started) * 1000.0
print(f"[TIMING] personaplex.model_load_ms {load_ms:.6f}")
started = time.perf_counter()
offline.warmup(mimi, other_mimi, lm_gen, device, frame_size)
synchronize(device)
warmup_ms = (time.perf_counter() - started) * 1000.0
print(f"[TIMING] personaplex.graph_warmup_ms {warmup_ms:.6f}")
return Runtime(
mimi=mimi,
other_mimi=other_mimi,
lm_gen=lm_gen,
text_tokenizer=text_tokenizer,
sample_rate=int(mimi.sample_rate),
frame_size=frame_size,
device=device,
offline=offline,
)
def load_user_audio(runtime: Runtime, input_wav: Path) -> np.ndarray:
resolved = resolve_path(input_wav)
if not resolved.is_file():
raise FileNotFoundError(f"missing input WAV: {resolved}")
return runtime.offline.lm_load_audio(str(resolved), runtime.sample_rate)
def request_from_dict(args: argparse.Namespace, value: dict[str, Any], index: int) -> Request:
return Request(
request_id=str(value.get("id") or value.get("request_id") or f"request{index + 1}"),
input_wav=Path(value.get("input_wav") or value.get("audio") or str(args.input_wav)),
text_prompt=str(value.get("text_prompt") or args.text_prompt),
voice_prompt=str(value.get("voice_prompt") or args.voice_prompt),
seed=int(value.get("seed", args.seed)),
temp_audio=float(value.get("temp_audio", args.temp_audio)),
temp_text=float(value.get("temp_text", args.temp_text)),
topk_audio=int(value.get("topk_audio", args.topk_audio)),
topk_text=int(value.get("topk_text", args.topk_text)),
greedy=bool(value.get("greedy", args.greedy)),
)
def load_requests(args: argparse.Namespace) -> list[Request]:
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 [request_from_dict(args, item, index) for index, item in enumerate(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 [request_from_dict(args, payload, 0)]
return [
Request(
request_id="request1",
input_wav=args.input_wav,
text_prompt=args.text_prompt,
voice_prompt=args.voice_prompt,
seed=args.seed,
temp_audio=args.temp_audio,
temp_text=args.temp_text,
topk_audio=args.topk_audio,
topk_text=args.topk_text,
greedy=args.greedy,
)
]
def decode_text_token(runtime: Runtime, token_id: int) -> str:
if token_id in (0, 1, 2, 3):
return ("EPAD", "BOS", "EOS", "PAD")[token_id]
return runtime.text_tokenizer.id_to_piece(token_id).replace("▁", " ")
def run_request(
runtime: Runtime,
request: Request,
voice_prompt_dir: str,
) -> tuple[np.ndarray, list[str], dict[str, float]]:
voice_prompt_path = resolve_voice_prompt_path(voice_prompt_dir, request.voice_prompt)
if not voice_prompt_path.is_file():
raise FileNotFoundError(f"missing voice prompt: {voice_prompt_path}")
user_audio = load_user_audio(runtime, request.input_wav)
total_target_samples = int(user_audio.shape[-1])
runtime.lm_gen.use_sampling = not request.greedy
runtime.lm_gen.temp = request.temp_audio
runtime.lm_gen.temp_text = request.temp_text
runtime.lm_gen.top_k = request.topk_audio
runtime.lm_gen.top_k_text = request.topk_text
seed_all(request.seed, runtime.device)
synchronize(runtime.device)
started = time.perf_counter()
if str(voice_prompt_path).endswith(".pt"):
runtime.lm_gen.load_voice_prompt_embeddings(str(voice_prompt_path))
else:
runtime.lm_gen.load_voice_prompt(str(voice_prompt_path))
runtime.lm_gen.text_prompt_tokens = runtime.text_tokenizer.encode(
runtime.offline.wrap_with_system_tags(request.text_prompt)
)
runtime.mimi.reset_streaming()
runtime.other_mimi.reset_streaming()
runtime.lm_gen.reset_streaming()
runtime.lm_gen.step_system_prompts(runtime.mimi)
runtime.mimi.reset_streaming()
synchronize(runtime.device)
prompt_ms = (time.perf_counter() - started) * 1000.0
generated_frames: list[np.ndarray] = []
generated_text_tokens: list[str] = []
started = time.perf_counter()
for user_encoded in runtime.offline.lm_encode_from_sphn(
runtime.mimi,
runtime.offline.lm_iterate_audio(
user_audio,
sample_interval_size=runtime.frame_size,
pad=True,
),
max_batch=1,
):
steps = int(user_encoded.shape[-1])
for step_index in range(steps):
tokens = runtime.lm_gen.step(user_encoded[:, :, step_index : step_index + 1])
if tokens is None:
continue
pcm = runtime.offline.decode_tokens_to_pcm(
runtime.mimi,
runtime.other_mimi,
runtime.lm_gen,
tokens,
)
generated_frames.append(pcm)
generated_text_tokens.append(decode_text_token(runtime, int(tokens[0, 0, 0].item())))
synchronize(runtime.device)
generate_ms = (time.perf_counter() - started) * 1000.0
if not generated_frames:
raise RuntimeError("PersonaPlex produced no audio frames")
output_pcm = np.concatenate(generated_frames, axis=-1)
if output_pcm.shape[-1] > total_target_samples:
output_pcm = output_pcm[:total_target_samples]
elif output_pcm.shape[-1] < total_target_samples:
pad_len = total_target_samples - output_pcm.shape[-1]
output_pcm = np.concatenate([output_pcm, np.zeros(pad_len, dtype=output_pcm.dtype)], axis=-1)
wall_ms = prompt_ms + generate_ms
duration_ms = 1000.0 * total_target_samples / float(runtime.sample_rate)
return output_pcm.astype(np.float32, copy=False), generated_text_tokens, {
"prompt_ms": prompt_ms,
"generate_ms": generate_ms,
"wall_ms": wall_ms,
"audio_duration_ms": duration_ms,
"rtf": wall_ms / max(duration_ms, 1e-6),
"x_realtime": duration_ms / max(wall_ms, 1e-6),
}
def write_timing(path: Path, rows: list[dict[str, Any]]) -> None:
if not path:
return
resolved = resolve_path(path)
resolved.parent.mkdir(parents=True, exist_ok=True)
with resolved.open("w", encoding="utf-8") as output:
for row in rows:
output.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
def write_summary(path: Path, rows: list[dict[str, Any]]) -> None:
if not path:
return
resolved = resolve_path(path)
resolved.parent.mkdir(parents=True, exist_ok=True)
resolved.write_text(json.dumps(rows, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
def output_path(args: argparse.Namespace, request_id: str, suffix: str) -> Path:
if suffix == ".wav" and args.audio_out_dir:
return resolve_path(args.audio_out_dir) / f"{request_id}.wav"
if suffix == ".json" and args.text_out_dir:
return resolve_path(args.text_out_dir) / f"{request_id}.json"
if suffix == ".wav":
return resolve_path(args.audio_out)
return resolve_path(args.audio_out).with_suffix(".json")
def main() -> int:
args = parse_args()
offline = import_reference()
voice_prompt_dir = resolve_voice_prompt_dir(offline, args)
runtime = load_runtime(args, voice_prompt_dir, offline)
warmup_request = Request(
request_id="warmup",
input_wav=args.warmup_input_wav,
text_prompt=args.warmup_text_prompt,
voice_prompt=args.warmup_voice_prompt,
seed=args.seed,
temp_audio=args.temp_audio,
temp_text=args.temp_text,
topk_audio=args.topk_audio,
topk_text=args.topk_text,
greedy=args.greedy,
)
for index in range(max(0, args.warmup)):
_, _, timing = run_request(runtime, warmup_request, voice_prompt_dir)
print(
f"[TIMING] warmup{index + 1}.wall_ms {timing['wall_ms']:.6f} "
f"rtf {timing['rtf']:.6f} x_realtime {timing['x_realtime']:.6f}"
)
requests = load_requests(args)
timing_rows: list[dict[str, Any]] = []
summary_rows: list[dict[str, Any]] = []
for iteration in range(max(1, args.iterations)):
for request in requests:
request_id = request.request_id if args.iterations == 1 else f"{request.request_id}_iter{iteration + 1}"
audio, text_tokens, timing = run_request(runtime, request, voice_prompt_dir)
wav_path = output_path(args, request_id, ".wav")
text_path = output_path(args, request_id, ".json")
wav_path.parent.mkdir(parents=True, exist_ok=True)
text_path.parent.mkdir(parents=True, exist_ok=True)
runtime.offline.sphn.write_wav(str(wav_path), audio, runtime.sample_rate)
text_path.write_text(json.dumps(text_tokens, ensure_ascii=False) + "\n", encoding="utf-8")
row = {
"request": request_id,
"input_wav": str(resolve_path(request.input_wav)),
"voice_prompt": request.voice_prompt,
"sample_rate": runtime.sample_rate,
**timing,
}
timing_rows.append(row)
summary_rows.append(
{
**row,
"output_wav": str(wav_path),
"output_text": str(text_path),
"samples": int(audio.shape[-1]),
"text_token_count": len(text_tokens),
"mean_abs": float(np.abs(audio).mean()) if audio.size else 0.0,
"rms": float(np.sqrt(np.mean(np.square(audio, dtype=np.float64)))) if audio.size else 0.0,
}
)
print(
f"[TIMING] request.{request_id}.prompt_ms {timing['prompt_ms']:.6f}\n"
f"[TIMING] request.{request_id}.generate_ms {timing['generate_ms']:.6f}\n"
f"[TIMING] request.{request_id}.wall_ms {timing['wall_ms']:.6f}\n"
f"[TIMING] request.{request_id}.audio_duration_ms {timing['audio_duration_ms']:.6f}\n"
f"[TIMING] request.{request_id}.rtf {timing['rtf']:.6f}\n"
f"[TIMING] request.{request_id}.x_realtime {timing['x_realtime']:.6f}"
)
write_timing(args.timing_file, timing_rows)
write_summary(args.summary_file, summary_rows)
sequence_steps: list[dict[str, Any]] = []
for row in summary_rows:
sequence_steps.append({
"id": row["request"],
"stems": [{
"name": "audio",
"audio": row["output_wav"],
"summary": {
"sample_rate": row["sample_rate"],
"channels": 1,
"frames": row["samples"],
"samples": row["samples"],
"mean_abs": row["mean_abs"],
"rms": row["rms"],
},
}],
"metrics": {
"prompt_ms": row["prompt_ms"],
"generate_ms": row["generate_ms"],
"wall_ms": row["wall_ms"],
"audio_duration_ms": row["audio_duration_ms"],
"rtf": row["rtf"],
"x_realtime": row["x_realtime"],
},
})
print("summary_json=" + json.dumps(
{"family": "personaplex", "backend": args.backend, "sequence_steps": sequence_steps},
ensure_ascii=False,
separators=(",", ":"),
))
return 0
if __name__ == "__main__":
raise SystemExit(main())