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
480 lines
18 KiB
Python
480 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import contextlib
|
|
import json
|
|
import random
|
|
import tempfile
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Iterator
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
DEFAULT_MODEL = REPO_ROOT / "models" / "VoxCPM2"
|
|
DEFAULT_REFERENCE_SRC = REPO_ROOT / "reference" / "VoxCPM" / "src"
|
|
DEFAULT_CASE_CATALOG = REPO_ROOT / "tools" / "voxcpm2" / "voxcpm2_warm_bench_cases.json"
|
|
DEFAULT_CASE_NAME = "realistic_mixed_longform"
|
|
DEFAULT_TEXT = "VoxCPM2 generates a clear reference sentence for a benchmark request."
|
|
DENOISER_DISABLED_REASON = "VoxCPM2 denoising is disabled; models/VoxCPM2 must be self-contained."
|
|
|
|
GENERATE_KEYS = {
|
|
"text",
|
|
"prompt_wav_path",
|
|
"prompt_text",
|
|
"reference_wav_path",
|
|
"cfg_value",
|
|
"inference_timesteps",
|
|
"min_len",
|
|
"max_len",
|
|
"normalize",
|
|
"denoise",
|
|
"retry_badcase",
|
|
"retry_badcase_max_times",
|
|
"retry_badcase_ratio_threshold",
|
|
}
|
|
REQUEST_META_KEYS = {
|
|
"seed",
|
|
"id",
|
|
}
|
|
PATH_KEYS = {
|
|
"prompt_wav_path",
|
|
"reference_wav_path",
|
|
}
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Python reference VoxCPM2 warmbench.")
|
|
parser.add_argument("--family", default="voxcpm2")
|
|
parser.add_argument("--model", type=Path, default=DEFAULT_MODEL)
|
|
parser.add_argument("--reference-src", type=Path, default=DEFAULT_REFERENCE_SRC)
|
|
parser.add_argument("--backend", choices=("cpu", "cuda"), default="cuda")
|
|
parser.add_argument("--device", type=int, default=0)
|
|
parser.add_argument("--threads", type=int, default=4)
|
|
parser.add_argument("--warmup", type=int, default=0)
|
|
parser.add_argument("--iterations", type=int, default=1)
|
|
parser.add_argument("--run-mode", choices=("offline", "streaming"), default="offline")
|
|
parser.add_argument("--text", action="append", dest="texts", default=[])
|
|
parser.add_argument("--warmup-text", default=DEFAULT_TEXT)
|
|
parser.add_argument("--case-catalog", type=Path, default=DEFAULT_CASE_CATALOG)
|
|
parser.add_argument("--case-name", default=DEFAULT_CASE_NAME)
|
|
parser.add_argument("--request-json", default="")
|
|
parser.add_argument("--request-sequence-json", default="")
|
|
parser.add_argument("--warmup-request-json", default="")
|
|
parser.add_argument("--cfg-value", type=float, default=2.0)
|
|
parser.add_argument("--inference-timesteps", type=int, default=10)
|
|
parser.add_argument("--min-len", type=int, default=2)
|
|
parser.add_argument("--max-len", type=int, default=4096)
|
|
parser.add_argument("--retry-badcase", choices=("true", "false"), default="true")
|
|
parser.add_argument("--normalize", action="store_true")
|
|
parser.add_argument("--optimize", action="store_true")
|
|
parser.add_argument("--dtype", choices=("float32",), default="float32")
|
|
parser.add_argument("--seed", type=int, default=1234)
|
|
parser.add_argument("--noise-file", default="")
|
|
parser.add_argument("--audio-out", type=Path, default=Path("voxcpm2_python_audio.wav"))
|
|
parser.add_argument("--audio-out-dir", type=Path, default=None)
|
|
parser.add_argument("--output-dir", type=Path, default=None)
|
|
parser.add_argument("--timing-file", type=Path, default=Path("voxcpm2_python_timing.log"))
|
|
return parser.parse_args()
|
|
|
|
|
|
def resolve_repo_path(path: Path) -> Path:
|
|
return path if path.is_absolute() else REPO_ROOT / path
|
|
|
|
|
|
def parse_bool(value: str) -> bool:
|
|
return value in {"1", "true", "yes"}
|
|
|
|
|
|
def request_bool(request: dict[str, Any], key: str, default: bool) -> bool:
|
|
if key not in request:
|
|
return default
|
|
value = request[key]
|
|
if isinstance(value, bool):
|
|
return value
|
|
raise RuntimeError(f"VoxCPM2 request field {key} must be a JSON boolean")
|
|
|
|
|
|
def load_json_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 [require_request_object(item) for item in payload]
|
|
if args.request_json:
|
|
payload = json.loads(args.request_json)
|
|
return [require_request_object(payload)]
|
|
if args.texts:
|
|
return [{"text": text} for text in args.texts]
|
|
return load_case(args.case_catalog, args.case_name)["requests"]
|
|
|
|
|
|
def load_warmup_request(args: argparse.Namespace) -> dict[str, Any]:
|
|
if args.warmup_request_json:
|
|
return require_request_object(json.loads(args.warmup_request_json))
|
|
if not args.request_sequence_json and not args.request_json and not args.texts:
|
|
return load_case(args.case_catalog, args.case_name)["warmup"]
|
|
return {"text": args.warmup_text}
|
|
|
|
|
|
def load_case(catalog_path: Path, case_name: str) -> dict[str, Any]:
|
|
path = resolve_repo_path(catalog_path)
|
|
payload = json.loads(path.read_text(encoding="utf-8"))
|
|
if not isinstance(payload, dict):
|
|
raise RuntimeError(f"invalid VoxCPM2 case catalog: {path}")
|
|
case = payload.get(case_name)
|
|
if not isinstance(case, dict):
|
|
available = ", ".join(sorted(str(key) for key in payload))
|
|
raise RuntimeError(f"unknown VoxCPM2 case name {case_name!r}; available: {available}")
|
|
warmup = case.get("warmup", {"text": DEFAULT_TEXT})
|
|
requests = case.get("requests")
|
|
if not isinstance(warmup, dict):
|
|
raise RuntimeError(f"VoxCPM2 case {case_name!r} warmup must be an object")
|
|
if not isinstance(requests, list) or not requests:
|
|
raise RuntimeError(f"VoxCPM2 case {case_name!r} must contain non-empty requests")
|
|
return {
|
|
"warmup": require_request_object(warmup),
|
|
"requests": [require_request_object(request) for request in requests],
|
|
}
|
|
|
|
|
|
def require_request_object(payload: Any) -> dict[str, Any]:
|
|
if not isinstance(payload, dict):
|
|
raise RuntimeError("VoxCPM2 request JSON must decode to an object")
|
|
return dict(payload)
|
|
|
|
|
|
def normalize_request(request: dict[str, Any], args: argparse.Namespace) -> dict[str, Any]:
|
|
unknown = set(request) - GENERATE_KEYS - REQUEST_META_KEYS
|
|
if unknown:
|
|
raise RuntimeError(f"unknown VoxCPM2 request field(s): {', '.join(sorted(unknown))}")
|
|
out: dict[str, Any] = {
|
|
"text": request.get("text", DEFAULT_TEXT),
|
|
"cfg_value": float(request.get("cfg_value", args.cfg_value)),
|
|
"inference_timesteps": int(request.get("inference_timesteps", args.inference_timesteps)),
|
|
"min_len": int(request.get("min_len", args.min_len)),
|
|
"max_len": int(request.get("max_len", args.max_len)),
|
|
"normalize": request_bool(request, "normalize", args.normalize),
|
|
"denoise": request_bool(request, "denoise", False),
|
|
"retry_badcase": request_bool(request, "retry_badcase", parse_bool(args.retry_badcase)),
|
|
}
|
|
for key in (
|
|
"prompt_text",
|
|
"retry_badcase_max_times",
|
|
"retry_badcase_ratio_threshold",
|
|
):
|
|
if key in request:
|
|
out[key] = request[key]
|
|
for key in PATH_KEYS:
|
|
if key in request and request[key] is not None:
|
|
out[key] = str(resolve_repo_path(Path(str(request[key]))))
|
|
if not isinstance(out["text"], str) or not out["text"].strip():
|
|
raise RuntimeError("VoxCPM2 request requires non-empty text")
|
|
if out["denoise"]:
|
|
raise RuntimeError(DENOISER_DISABLED_REASON)
|
|
return out
|
|
|
|
|
|
def add_reference_source(reference_src: Path) -> None:
|
|
source = resolve_repo_path(reference_src)
|
|
package = source / "voxcpm" / "__init__.py"
|
|
if not package.is_file():
|
|
raise RuntimeError(f"missing VoxCPM reference package: {package}")
|
|
sys.path.insert(0, str(source))
|
|
|
|
|
|
def load_voxcpm_class(reference_src: Path):
|
|
add_reference_source(reference_src)
|
|
from voxcpm import VoxCPM
|
|
import voxcpm
|
|
|
|
module_path = Path(voxcpm.__file__).resolve()
|
|
expected_root = resolve_repo_path(reference_src).resolve()
|
|
try:
|
|
module_path.relative_to(expected_root)
|
|
except ValueError as exc:
|
|
raise RuntimeError(f"VoxCPM imported from {module_path}, expected under {expected_root}") from exc
|
|
return VoxCPM, module_path
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def model_root_with_dtype(model_root: Path, dtype: str) -> Iterator[Path]:
|
|
config_path = model_root / "config.json"
|
|
config = json.loads(config_path.read_text(encoding="utf-8"))
|
|
if config.get("dtype") == dtype:
|
|
yield model_root
|
|
return
|
|
with tempfile.TemporaryDirectory(prefix="voxcpm2-fp32-model-") as tmp:
|
|
tmp_root = Path(tmp)
|
|
for child in model_root.iterdir():
|
|
if child.name == "config.json":
|
|
continue
|
|
target = tmp_root / child.name
|
|
target.symlink_to(child.resolve(), target_is_directory=child.is_dir())
|
|
config["dtype"] = dtype
|
|
(tmp_root / "config.json").write_text(json.dumps(config, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
yield tmp_root
|
|
|
|
|
|
def set_seed(seed: int, backend: str, device: int) -> None:
|
|
import torch
|
|
|
|
torch.manual_seed(seed)
|
|
if backend == "cuda":
|
|
torch.cuda.manual_seed_all(seed)
|
|
np.random.seed(seed & 0xFFFFFFFF)
|
|
random.seed(seed)
|
|
|
|
|
|
def sync_device(backend: str, device: int) -> None:
|
|
if backend == "cuda":
|
|
import torch
|
|
|
|
torch.cuda.synchronize(device)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def controlled_cfm_noise(noise_file: str, in_channels: int) -> Iterator[None]:
|
|
if not noise_file:
|
|
yield
|
|
return
|
|
import torch
|
|
|
|
noise_values = np.fromfile(noise_file, dtype=np.float32)
|
|
if noise_values.size == 0:
|
|
raise RuntimeError(f"VoxCPM2 CFM controlled noise file is empty: {noise_file}")
|
|
original_randn = torch.randn
|
|
offset = 0
|
|
used = False
|
|
|
|
def randn(*args: Any, **kwargs: Any) -> torch.Tensor:
|
|
nonlocal offset, used
|
|
shape = args[0] if args else kwargs.get("size")
|
|
if (
|
|
isinstance(shape, (tuple, list, torch.Size))
|
|
and len(shape) == 3
|
|
and int(shape[0]) == 1
|
|
and int(shape[1]) == in_channels
|
|
):
|
|
count = int(np.prod([int(dim) for dim in shape]))
|
|
if noise_values.size < offset + count:
|
|
raise RuntimeError(
|
|
f"VoxCPM2 CFM controlled noise file is too short: expected at least {offset + count} floats, "
|
|
f"got {noise_values.size}"
|
|
)
|
|
array = noise_values[offset : offset + count].reshape(tuple(int(dim) for dim in shape)).copy()
|
|
offset += count
|
|
used = True
|
|
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("VoxCPM2 CFM controlled noise was not consumed")
|
|
|
|
|
|
def configure_torch(args: argparse.Namespace) -> str:
|
|
import torch
|
|
|
|
torch.set_num_threads(max(1, args.threads))
|
|
if args.backend == "cuda":
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA backend requested but torch.cuda.is_available() is false")
|
|
torch.cuda.set_device(args.device)
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
torch.backends.cudnn.allow_tf32 = False
|
|
return f"cuda:{args.device}"
|
|
return "cpu"
|
|
|
|
|
|
def summarize_audio(audio: np.ndarray, sample_rate: int, wall_ms: float, chunk_count: int) -> dict[str, Any]:
|
|
flat = np.asarray(audio, dtype=np.float32).reshape(-1)
|
|
if flat.size == 0:
|
|
raise RuntimeError("VoxCPM2 warmbench received empty audio")
|
|
return {
|
|
"family": "voxcpm2",
|
|
"sample_rate": int(sample_rate),
|
|
"channels": 1,
|
|
"samples": int(flat.size),
|
|
"duration_sec": float(flat.size / sample_rate) if sample_rate else 0.0,
|
|
"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)),
|
|
"synthesize_wall_ms": float(wall_ms),
|
|
"chunk_count": int(chunk_count),
|
|
}
|
|
|
|
|
|
def write_audio(path: Path, audio: np.ndarray, sample_rate: int) -> np.ndarray:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
sf.write(str(path), np.asarray(audio, dtype=np.float32).reshape(-1), int(sample_rate), subtype="PCM_16")
|
|
written, _ = sf.read(str(path), always_2d=False, dtype="float32")
|
|
return np.asarray(written, dtype=np.float32).reshape(-1)
|
|
|
|
|
|
def run_once(model: Any, request: dict[str, Any], args: argparse.Namespace) -> tuple[np.ndarray, int, float, int]:
|
|
seed = int(request.get("seed", args.seed))
|
|
if seed < 0:
|
|
raise RuntimeError("VoxCPM2 seed must be non-negative")
|
|
set_seed(seed, args.backend, args.device)
|
|
kwargs = normalize_request(request, args)
|
|
sync_device(args.backend, args.device)
|
|
started = time.perf_counter()
|
|
with controlled_cfm_noise(args.noise_file, int(model.tts_model.feat_decoder.in_channels)):
|
|
if args.run_mode == "streaming":
|
|
chunks = [
|
|
np.asarray(chunk, dtype=np.float32).reshape(-1)
|
|
for chunk in model.generate_streaming(**kwargs)
|
|
]
|
|
wav = np.concatenate(chunks) if chunks else np.empty(0, dtype=np.float32)
|
|
chunk_count = len(chunks)
|
|
else:
|
|
wav = model.generate(**kwargs)
|
|
chunk_count = 0
|
|
sync_device(args.backend, args.device)
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
return (
|
|
np.asarray(wav, dtype=np.float32).reshape(-1),
|
|
int(model.tts_model.sample_rate),
|
|
wall_ms,
|
|
chunk_count,
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
if args.family != "voxcpm2":
|
|
raise RuntimeError(f"unsupported VoxCPM2 warmbench family: {args.family}")
|
|
|
|
model_root = resolve_repo_path(args.model)
|
|
if not model_root.is_dir():
|
|
raise RuntimeError(f"VoxCPM2 model root not found: {model_root}")
|
|
|
|
VoxCPM, module_path = load_voxcpm_class(args.reference_src)
|
|
device = configure_torch(args)
|
|
set_seed(args.seed, args.backend, args.device)
|
|
|
|
load_start = time.perf_counter()
|
|
with model_root_with_dtype(model_root, args.dtype) as runtime_model_root:
|
|
model = VoxCPM.from_pretrained(
|
|
str(runtime_model_root),
|
|
load_denoiser=False,
|
|
optimize=args.optimize,
|
|
device=device,
|
|
)
|
|
sync_device(args.backend, args.device)
|
|
load_ms = (time.perf_counter() - load_start) * 1000.0
|
|
|
|
requests = load_json_requests(args)
|
|
warmup_request = load_warmup_request(args)
|
|
output_dir = args.audio_out_dir or args.output_dir
|
|
|
|
timing_lines = [
|
|
f"voxcpm2.reference_module {module_path}",
|
|
f"voxcpm2.backend {args.backend}",
|
|
f"voxcpm2.run_mode {args.run_mode}",
|
|
f"voxcpm2.model_load_ms {load_ms:.6f}",
|
|
"voxcpm2.load_denoiser 0",
|
|
f"voxcpm2.optimize {1 if args.optimize else 0}",
|
|
f"voxcpm2.dtype {args.dtype}",
|
|
f"voxcpm2.case_name {args.case_name}",
|
|
]
|
|
|
|
warmup_outputs: list[tuple[np.ndarray, int, float, int]] = []
|
|
for _ in range(max(0, args.warmup)):
|
|
result = run_once(model, warmup_request, args)
|
|
warmup_outputs.append(result)
|
|
timing_lines.append(f"voxcpm2.synthesize_wall_ms {result[2]:.6f}")
|
|
|
|
last_outputs: list[tuple[np.ndarray, int, float, int]] = []
|
|
for request in requests:
|
|
current: tuple[np.ndarray, int, float, int] | None = None
|
|
for _ in range(max(1, args.iterations)):
|
|
current = run_once(model, request, args)
|
|
timing_lines.append(f"voxcpm2.synthesize_wall_ms {current[2]:.6f}")
|
|
assert current is not None
|
|
last_outputs.append(current)
|
|
|
|
args.timing_file.parent.mkdir(parents=True, exist_ok=True)
|
|
args.timing_file.write_text("\n".join(timing_lines) + "\n", encoding="utf-8")
|
|
|
|
for index, (audio, sample_rate, wall_ms, chunk_count) in enumerate(warmup_outputs):
|
|
print(f"warmup_text[{index}]={warmup_request.get('text', DEFAULT_TEXT)}")
|
|
print(
|
|
f"warmup_summary_json[{index}]="
|
|
+ json.dumps(summarize_audio(audio, sample_rate, wall_ms, chunk_count), ensure_ascii=False, separators=(",", ":"))
|
|
)
|
|
|
|
sequence_steps: list[dict[str, Any]] = []
|
|
for index, (request, output) in enumerate(zip(requests, last_outputs)):
|
|
audio, sample_rate, wall_ms, chunk_count = output
|
|
summary = summarize_audio(audio, sample_rate, wall_ms, chunk_count)
|
|
print(f"text[{index}]={request.get('text', DEFAULT_TEXT)}")
|
|
print(
|
|
f"summary_json[{index}]="
|
|
+ json.dumps(summary, ensure_ascii=False, separators=(",", ":"))
|
|
)
|
|
step_audio_path = ""
|
|
if output_dir is not None:
|
|
out_path = output_dir / f"request_{index}.wav"
|
|
write_audio(out_path, audio, sample_rate)
|
|
step_audio_path = str(out_path)
|
|
print(f"audio_out[{index}]={out_path}")
|
|
stem: dict[str, Any] = {
|
|
"name": "audio",
|
|
"summary": summary,
|
|
}
|
|
if step_audio_path:
|
|
stem["audio"] = step_audio_path
|
|
sequence_steps.append(
|
|
{
|
|
"request_index": index,
|
|
"stems": [stem],
|
|
"metrics": {"wall_ms": wall_ms, "chunk_count": chunk_count},
|
|
}
|
|
)
|
|
|
|
print(
|
|
"summary_json="
|
|
+ json.dumps(
|
|
{
|
|
"family": "voxcpm2",
|
|
"backend": args.backend,
|
|
"mode": args.run_mode,
|
|
"sequence_steps": sequence_steps,
|
|
},
|
|
ensure_ascii=False,
|
|
separators=(",", ":"),
|
|
)
|
|
)
|
|
|
|
final_audio, final_sample_rate, _, _ = last_outputs[-1]
|
|
write_audio(args.audio_out, final_audio, final_sample_rate)
|
|
print(f"audio_out={args.audio_out}")
|
|
print(f"timing_out={args.timing_file}")
|
|
for index, (_audio, _sample_rate, wall_ms, _chunk_count) in enumerate(last_outputs):
|
|
print(f"average[{index}]")
|
|
print(f"voxcpm2.synthesize_wall_ms={wall_ms}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|