benchmark_voicedesign.py

"""
Speed benchmark for Voice Design mode.

Deliberately mirrors benchmark_autovoice.py: same model, same sentences,
same measurement. The only difference is the instruct string, so the two
result sets can be compared directly to see what voice design costs.

Because the instruct stays constant, the voice should stay constant too --
unlike auto voice, which picks a new speaker on every call.

Run from the project root (omnivoice_lab):
    uv run --project OmniVoice python benchmarks/benchmark_voicedesign.py
"""

import os
import time

# Force huggingface_hub offline BEFORE anything imports it. Every file we
# need is already on disk, so a network call would mean something is wrong.
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

import soundfile as sf
import torch
from dotenv import load_dotenv
from omnivoice import OmniVoice

load_dotenv()

MODEL_DIR = os.getenv("OMNIVOICE_MODEL_DIR")
OUTPUT_DIR = os.path.join("data", "output", "benchmark_voicedesign")

LANGUAGE = "de"

# See scripts/voice_design.py for the full attribute list. Accents are left
# out on purpose: they only apply to English text.
INSTRUCT = "female, young adult, high pitch"

# Identical to benchmark_autovoice.py so the timings line up.
SENTENCES = [
    "Guten Tag, das hier ist ein erster Test der Sprachsynthese.",
    "Der Zug fährt heute leider mit zwanzig Minuten Verspätung ab.",
    "Kannst du mir bitte kurz erklären, wie das genau funktioniert?",
    "Am Wochenende soll das Wetter deutlich besser werden als bisher gedacht.",
    "Ich hätte nicht gedacht, dass das so schnell geht, mein lieber Freund.",
    "Bitte legen Sie die Unterlagen auf den Tisch im Nebenzimmer.",
]


def save_markdown_report(report_path, model_name, device_name, instruct, load_time, rows, sentences):
    first = rows[0]
    rest = rows[1:]
    warm_rtf = (sum(r[3] for r in rest) / sum(r[2] for r in rest)) if rest else first[4]
    speedup = (first[4] / warm_rtf) if warm_rtf > 0 else 1.0
    total_compute = load_time + sum(r[3] for r in rows)
    total_audio = sum(r[2] for r in rows)
    overall_rtf = total_compute / total_audio if total_audio > 0 else 0.0

    lines = [
        "# Benchmark Results: Voice Design",
        "",
        f"- **Timestamp:** {time.strftime('%Y-%m-%d %H:%M:%S')}",
        f"- **Model:** `{model_name}`",
        f"- **Device:** `{device_name}`",
        f"- **Instruct:** `{instruct}`",
        f"- **Language:** `de`",
        f"- **Generation Mode:** Voice Design (textual prompt)",
        "",
        "## Per-Sentence Measurements",
        "",
        "| # | Characters | Audio Duration | GPU Compute Time | RTF | Sentence Text |",
        "|---|---|---|---|---|---|",
    ]

    for (idx, char_count, duration, elapsed, rtf), text in zip(rows, sentences):
        lines.append(f"| {idx} | {char_count} | {duration:.2f} s | {elapsed:.2f} s | {rtf:.3f} | *{text}* |")

    lines.extend([
        "",
        "## Summary Metrics",
        "",
        "| Metric | Measured Value | Description |",
        "|---|---|---|",
        f"| **Model Load Time** | `{load_time:.2f} s` | One-off initialization cost per process |",
        f"| **First Call RTF (Cold Start)** | `{first[4]:.3f}` | Includes CUDA context init and kernel compilation |",
        f"| **Warm RTF (Calls 2–{len(rows)})** | `{warm_rtf:.3f}` | Steady-state generation performance |",
        f"| **Warm-up Speedup** | `{speedup:.1f}x` | Performance gain over cold start |",
        f"| **Overall RTF (incl. Loading)** | `{overall_rtf:.3f}` | Total compute divided by total synthesized audio |",
        "",
    ])

    with open(report_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))


def main():
    if not MODEL_DIR or not os.path.isdir(MODEL_DIR):
        raise SystemExit(
            f"Model directory not found. Please set OMNIVOICE_MODEL_DIR in your environment or .env (current value: {MODEL_DIR})"
        )

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    print(f"Loading model from {MODEL_DIR} ...")
    started = time.perf_counter()
    model = OmniVoice.from_pretrained(
        MODEL_DIR,
        device_map="cuda:0",
        dtype=torch.float16,
    )
    load_time = time.perf_counter() - started
    print(f"Model loaded in {load_time:.1f} s")
    print(f"Instruct: {INSTRUCT!r}\n")

    sample_rate = getattr(model, "sampling_rate", None) or 24000
    device_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
    rows = []

    for index, text in enumerate(SENTENCES, start=1):
        started = time.perf_counter()
        audio = model.generate(text=text, language=LANGUAGE, instruct=INSTRUCT)
        # CUDA is asynchronous: without this the timer would stop before the
        # GPU is actually done and every call would look absurdly fast.
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        elapsed = time.perf_counter() - started

        wav = audio[0]
        duration = len(wav) / sample_rate
        path = os.path.join(OUTPUT_DIR, f"bench_{index:02d}.wav")
        sf.write(path, wav, sample_rate)

        rows.append((index, len(text), duration, elapsed, elapsed / duration))
        print(
            f"{index:>2}  {len(text):>3} chars  "
            f"audio {duration:5.2f} s  compute {elapsed:5.2f} s  "
            f"RTF {elapsed / duration:.3f}"
        )

    first = rows[0]
    rest = rows[1:]

    print("\n--- summary ---")
    print(f"model load           : {load_time:.2f} s (once per process)")
    print(f"first call RTF       : {first[4]:.3f}  <- includes warm-up")

    if rest:
        warm_compute = sum(r[3] for r in rest)
        warm_audio = sum(r[2] for r in rest)
        warm_rtf = warm_compute / warm_audio
        print(f"warm RTF (calls 2-{len(rows)}) : {warm_rtf:.3f}")
        print(f"speed-up vs first    : {first[4] / warm_rtf:.1f}x")

    total_compute = load_time + sum(r[3] for r in rows)
    total_audio = sum(r[2] for r in rows)
    print(f"overall incl. loading: {total_compute / total_audio:.3f}")

    report_path = os.path.join(OUTPUT_DIR, "benchmark_results.md")
    save_markdown_report(report_path, MODEL_DIR, device_name, INSTRUCT, load_time, rows, SENTENCES)
    print(f"\nWrote {len(rows)} audio files and report to {report_path}")


if __name__ == "__main__":
    main()