Newer
Older
voice / scripts / benchmark_compile.py
"""Quick test: Fish Speech with and without torch.compile on one sentence."""

import asyncio
import os
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT / "src") not in os.sys.path:
    os.sys.path.insert(0, str(ROOT / "src"))

from voice_tts.config import settings
from voice_tts.tts.fish_speech_backend import FishSpeechEngine


TEXT = "The quick brown fox jumps over the lazy dog."
LANG = "en"
REF = Path(settings.default_voice_ref)


async def run(compile: bool) -> float:
    engine = FishSpeechEngine(
        checkpoint_path=settings.tts_model_path,
        source_root=settings.tts_vocab_path,
        device=settings.device,
        compile=compile,
        chunk_length=settings.fish_chunk_length,
    )
    engine.load()

    # First run (warm-up + compile)
    t0 = time.perf_counter()
    await asyncio.to_thread(
        lambda: asyncio.run(
            engine.synthesize(
                text=TEXT,
                ref_audio_path=REF,
                language=LANG,
                speed=1.0,
                emotion="neutral",
                ref_text=settings.default_ref_text,
            )
        )
    )
    first = time.perf_counter() - t0

    # Second run
    t0 = time.perf_counter()
    await asyncio.to_thread(
        lambda: asyncio.run(
            engine.synthesize(
                text=TEXT,
                ref_audio_path=REF,
                language=LANG,
                speed=1.0,
                emotion="neutral",
                ref_text=settings.default_ref_text,
            )
        )
    )
    second = time.perf_counter() - t0
    return first, second


async def main():
    f1, f2 = await run(compile=False)
    print(f"compile=False first={f1:.2f}s second={f2:.2f}s")
    c1, c2 = await run(compile=True)
    print(f"compile=True first={c1:.2f}s second={c2:.2f}s")


if __name__ == "__main__":
    asyncio.run(main())