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
270 lines
9.9 KiB
Python
270 lines
9.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the official OuteTTS 1.0 HF reference for audio.cpp parity.
|
|
|
|
The script intentionally keeps the official sampling configuration. It resets
|
|
the PyTorch CPU and CUDA generators immediately before each model.generate()
|
|
call, after prompt/profile construction, so no unrelated torch operation can
|
|
consume the request's sampling stream.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import copy
|
|
import json
|
|
import re
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[3]
|
|
DEFAULT_REQUESTS = REPO_ROOT / "tests" / "outetts" / "warm_bench_requests.json"
|
|
|
|
|
|
def framework_chunks(text: str, budget: int | None) -> list[str]:
|
|
text = text.strip()
|
|
if budget is None or len(text) <= budget:
|
|
return [text]
|
|
words = list(re.finditer(r"\S+", text))
|
|
chunks: list[str] = []
|
|
start = 0
|
|
while start < len(words):
|
|
hard_end = start
|
|
while hard_end < len(words):
|
|
width = words[hard_end].end() - words[start].start()
|
|
if width > budget:
|
|
break
|
|
hard_end += 1
|
|
if hard_end == start:
|
|
hard_end += 1
|
|
end = hard_end
|
|
if hard_end < len(words) and hard_end > start + 1:
|
|
for index in range(hard_end, start + 1, -1):
|
|
if words[index - 1].group()[-1:] in ".!?":
|
|
end = index
|
|
break
|
|
if end == hard_end:
|
|
for index in range(hard_end, start + 1, -1):
|
|
if words[index - 1].group()[-1:] in ",;:":
|
|
end = index
|
|
break
|
|
chunks.append(text[words[start].start() : words[end - 1].end()].strip())
|
|
start = end
|
|
return chunks
|
|
|
|
|
|
def resolve_path(value: str, base: Path) -> Path:
|
|
path = Path(value)
|
|
return path if path.is_absolute() else (base / path).resolve()
|
|
|
|
|
|
def reset_seed(torch_mod: Any, seed: int, device: str) -> None:
|
|
torch_mod.manual_seed(seed)
|
|
if device == "cuda":
|
|
torch_mod.cuda.manual_seed_all(seed)
|
|
|
|
|
|
def memory_snapshot(torch_mod: Any, device: str) -> dict[str, float]:
|
|
result: dict[str, float] = {}
|
|
try:
|
|
import psutil
|
|
|
|
result["rss_mb"] = psutil.Process().memory_info().rss / (1024.0 * 1024.0)
|
|
except Exception: # psutil is optional for reference generation.
|
|
pass
|
|
if device == "cuda":
|
|
result["cuda_alloc_mb"] = torch_mod.cuda.memory_allocated() / (1024.0 * 1024.0)
|
|
result["cuda_peak_mb"] = torch_mod.cuda.max_memory_allocated() / (1024.0 * 1024.0)
|
|
result["cuda_reserved_mb"] = torch_mod.cuda.memory_reserved() / (1024.0 * 1024.0)
|
|
return result
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--model", type=Path, required=True)
|
|
parser.add_argument("--dac", type=Path, required=True)
|
|
parser.add_argument("--alignment-json", type=Path, required=True)
|
|
parser.add_argument("--request-file", type=Path, default=DEFAULT_REQUESTS)
|
|
parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda")
|
|
parser.add_argument("--dtype", choices=("fp32", "bf16"), default="fp32")
|
|
parser.add_argument(
|
|
"--out-dir", type=Path, default=REPO_ROOT / "build" / "logs" / "outetts_python_fp32"
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
import torch
|
|
import torchaudio
|
|
import outetts
|
|
|
|
model_path = args.model.resolve()
|
|
dac_path = args.dac.resolve()
|
|
request_path = args.request_file.resolve()
|
|
request_base = request_path.parent.parent.parent if request_path.parent.name == "outetts" else REPO_ROOT
|
|
output_dir = args.out_dir.resolve()
|
|
wav_dir = output_dir / "outputs"
|
|
wav_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
dtype = torch.float32 if args.dtype == "fp32" else torch.bfloat16
|
|
config = outetts.ModelConfig(
|
|
model_path=str(model_path),
|
|
tokenizer_path=str(model_path),
|
|
interface_version=outetts.InterfaceVersion.V3,
|
|
backend=outetts.Backend.HF,
|
|
device=args.device,
|
|
dtype=dtype,
|
|
audio_codec_path=str(dac_path),
|
|
max_seq_length=8192,
|
|
)
|
|
interface = outetts.Interface(config)
|
|
alignment = json.loads(args.alignment_json.read_text(encoding="utf-8"))
|
|
requests = json.loads(request_path.read_text(encoding="utf-8"))["requests"]
|
|
|
|
speakers: dict[tuple[str, str], dict[str, Any]] = {}
|
|
boundaries: dict[str, list[dict[str, Any]]] = {}
|
|
memory: list[dict[str, Any]] = []
|
|
stdout_lines = [
|
|
"family=outetts",
|
|
"reference=official_hf",
|
|
f"dtype={args.dtype}",
|
|
f"device={args.device}",
|
|
]
|
|
|
|
for request in requests:
|
|
name = str(request["name"])
|
|
seed = int(request["seed"])
|
|
max_new_tokens = int(request["max_tokens"])
|
|
profile_ms = 0.0
|
|
speaker = None
|
|
|
|
started = time.perf_counter()
|
|
voice_ref = request.get("voice_ref")
|
|
if voice_ref:
|
|
reference_text = str(request["reference_text"])
|
|
reference_path = resolve_path(str(voice_ref), request_base)
|
|
key = (str(reference_path), reference_text)
|
|
if key not in speakers:
|
|
profile_started = time.perf_counter()
|
|
speakers[key] = interface.audio_processor.create_speaker_from_dict(
|
|
{
|
|
"audio": {"bytes": reference_path.read_bytes()},
|
|
"text": reference_text,
|
|
"words": [
|
|
{
|
|
"word": item["word"],
|
|
"start": item["start_sample"] / 16000.0,
|
|
"end": item["end_sample"] / 16000.0,
|
|
}
|
|
for item in alignment
|
|
],
|
|
}
|
|
)
|
|
speakers[key]["interface_version"] = outetts.InterfaceVersion.V3.value
|
|
profile_ms = (time.perf_counter() - profile_started) * 1000.0
|
|
speaker = speakers[key]
|
|
|
|
chunks = framework_chunks(request["text"], request.get("text_chunk_size"))
|
|
audio_chunks = []
|
|
request_boundaries = []
|
|
if args.device == "cuda":
|
|
torch.cuda.reset_peak_memory_stats()
|
|
|
|
for chunk in chunks:
|
|
current_speaker = copy.deepcopy(speaker)
|
|
prompt_text = interface.prompt_processor.get_completion_prompt(chunk, current_speaker)
|
|
prompt = interface._prepare_prompt(prompt_text)
|
|
prompt_tokens = int(prompt.shape[-1])
|
|
generation = outetts.GenerationConfig(
|
|
text=chunk,
|
|
speaker=current_speaker,
|
|
generation_type=outetts.GenerationType.REGULAR,
|
|
max_length=prompt_tokens + max_new_tokens,
|
|
sampler_config=outetts.SamplerConfig(
|
|
temperature=0.4,
|
|
repetition_penalty=1.1,
|
|
repetition_range=64,
|
|
top_k=40,
|
|
top_p=0.9,
|
|
min_p=0.05,
|
|
),
|
|
)
|
|
|
|
# This reset is deliberately adjacent to generate(). Prompt/profile
|
|
# work above cannot consume the model-sampling RNG stream.
|
|
reset_seed(torch, seed, args.device)
|
|
full_tokens = interface.model.generate(prompt, generation)
|
|
generated = full_tokens[prompt_tokens:]
|
|
codebooks = interface.prompt_processor.extract_audio_from_tokens(generated)
|
|
audio = interface.get_audio(generated).detach().cpu().flatten()
|
|
audio_chunks.append(audio)
|
|
request_boundaries.append(
|
|
{
|
|
"prompt_text": prompt_text,
|
|
"prompt_ids": prompt.detach().cpu().flatten().tolist(),
|
|
"generated_ids": generated,
|
|
"c1": codebooks[0],
|
|
"c2": codebooks[1],
|
|
"audio_frames_24k": int(audio.numel()),
|
|
}
|
|
)
|
|
|
|
if args.device == "cuda":
|
|
torch.cuda.synchronize()
|
|
wall_ms = (time.perf_counter() - started) * 1000.0
|
|
audio = torch.cat(audio_chunks)
|
|
wav_path = wav_dir / f"{name}_1.wav"
|
|
torchaudio.save(str(wav_path), audio.unsqueeze(0), 24000, encoding="PCM_S", bits_per_sample=16)
|
|
audio_seconds = audio.numel() / 24000.0
|
|
rtf = wall_ms / 1000.0 / audio_seconds
|
|
|
|
snapshot = memory_snapshot(torch, args.device)
|
|
snapshot.update(
|
|
{
|
|
"request_id": name,
|
|
"wall_ms": wall_ms,
|
|
"profile_ms": profile_ms,
|
|
"audio_seconds": audio_seconds,
|
|
"rtf": rtf,
|
|
}
|
|
)
|
|
memory.append(snapshot)
|
|
boundaries[name] = request_boundaries
|
|
stdout_lines.extend(
|
|
[
|
|
f"request_id={name}",
|
|
f"[TIMING] request.{name}.wall_ms {wall_ms}",
|
|
f"[TIMING] request.{name}.profile_ms {profile_ms}",
|
|
f"audio_seconds={audio_seconds}",
|
|
f"rtf={rtf}",
|
|
f"audio_out={wav_path}",
|
|
]
|
|
)
|
|
print(f"[REF] {name}: {wall_ms:.2f} ms, {audio_seconds:.5f}s, RTF={rtf:.3f}")
|
|
|
|
(output_dir / "command.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"reference": "official OuteTTS HF",
|
|
"model": str(model_path),
|
|
"dac": str(dac_path),
|
|
"requests": str(request_path),
|
|
"dtype": args.dtype,
|
|
"device": args.device,
|
|
},
|
|
indent=2,
|
|
)
|
|
+ "\n",
|
|
encoding="utf-8",
|
|
)
|
|
(output_dir / "stdout.log").write_text("\n".join(stdout_lines) + "\n", encoding="utf-8")
|
|
(output_dir / "memory.json").write_text(
|
|
json.dumps({"requests": memory}, indent=2) + "\n", encoding="utf-8"
|
|
)
|
|
(output_dir / "boundaries.json").write_text(
|
|
json.dumps(boundaries, indent=2) + "\n", encoding="utf-8"
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|