scripts/benchmark_tts.py

"""
TTS Benchmarking Script

Benchmarks Text-to-Speech (TTS) handlers to compare performance.
Measures: warmup time, inference time, time-to-first-chunk, audio duration, and RTF.

Usage:
    python benchmark_tts.py --text "Hello world" --iterations 3
    python benchmark_tts.py --handlers kokoro qwen3 pocket_tts
"""

import argparse
import json
import logging
import time
from queue import Queue
from threading import Event
from typing import Any, Dict, List, Optional

import numpy as np

from speech_to_speech.pipeline.messages import TTSInput

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

DEFAULT_SAMPLE_RATE = 16000
VALID_QWEN3_MLX_QUANTIZATIONS = ("bf16", "4bit", "6bit", "8bit")


class BenchmarkResult:
    """Stores benchmark results for a single TTS handler."""

    def __init__(self, handler_name: str):
        self.handler_name = handler_name
        self.warmup_time = 0.0
        self.inference_times: list[float] = []
        self.time_to_first_chunk: list[float] = []
        self.audio_durations: list[float] = []
        self.errors: list[str] = []

    def add_inference(self, time_taken: float, audio_duration: float, ttfc: Optional[float] = None):
        self.inference_times.append(time_taken)
        self.audio_durations.append(audio_duration)
        if ttfc is not None:
            self.time_to_first_chunk.append(ttfc)

    def add_error(self, error: str):
        self.errors.append(error)

    def get_stats(self) -> Dict[str, Any]:
        if not self.inference_times:
            return {
                "handler": self.handler_name,
                "status": "failed",
                "errors": self.errors,
            }

        avg_time = float(np.mean(self.inference_times))
        avg_audio = float(np.mean(self.audio_durations))
        avg_rtf = avg_audio / avg_time if avg_time > 0 else 0.0

        stats = {
            "handler": self.handler_name,
            "warmup_time": self.warmup_time,
            "avg_inference_time": avg_time,
            "min_inference_time": float(np.min(self.inference_times)),
            "max_inference_time": float(np.max(self.inference_times)),
            "std_inference_time": float(np.std(self.inference_times)),
            "avg_audio_duration": avg_audio,
            "min_audio_duration": float(np.min(self.audio_durations)),
            "max_audio_duration": float(np.max(self.audio_durations)),
            "std_audio_duration": float(np.std(self.audio_durations)),
            "avg_rtf": avg_rtf,
            "total_iterations": len(self.inference_times),
            "errors": self.errors,
        }

        if self.time_to_first_chunk:
            stats["avg_time_to_first_chunk"] = float(np.mean(self.time_to_first_chunk))
            stats["min_time_to_first_chunk"] = float(np.min(self.time_to_first_chunk))
            stats["max_time_to_first_chunk"] = float(np.max(self.time_to_first_chunk))
            stats["std_time_to_first_chunk"] = float(np.std(self.time_to_first_chunk))

        return stats


def benchmark_handler(
    handler_name: str,
    text: str,
    iterations: int,
    handler_kwargs: Optional[Dict[str, Any]] = None,
    language_code: Optional[str] = "en",
) -> BenchmarkResult:
    logger.info(f"Benchmarking {handler_name}...")
    result = BenchmarkResult(handler_name)

    try:
        stop_event = Event()
        should_listen = Event()
        queue_in: Queue[Any] = Queue()
        queue_out: Queue[Any] = Queue()

        handler: Any = None
        setup_kwargs = handler_kwargs or {}

        start_setup = time.perf_counter()

        if handler_name == "kokoro":
            from speech_to_speech.TTS.kokoro_handler import KokoroTTSHandler
            setup_kwargs = {"device": "auto", **setup_kwargs}
            handler = KokoroTTSHandler(
                stop_event,
                queue_in=queue_in,
                queue_out=queue_out,
                setup_args=(should_listen,),
                setup_kwargs=setup_kwargs,
            )
        elif handler_name == "pocket_tts":
            from speech_to_speech.TTS.pocket_tts_handler import PocketTTSHandler
            setup_kwargs = {"device": "cpu", **setup_kwargs}
            handler = PocketTTSHandler(
                stop_event,
                queue_in=queue_in,
                queue_out=queue_out,
                setup_args=(should_listen,),
                setup_kwargs=setup_kwargs,
            )
        elif handler_name == "qwen3":
            from speech_to_speech.TTS.qwen3_tts_handler import Qwen3TTSHandler
            setup_kwargs = {
                "device": "cuda",
                "model_name": "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
                "ref_audio": "TTS/ref_audio.wav",
                **setup_kwargs,
            }
            handler = Qwen3TTSHandler(
                stop_event,
                queue_in=queue_in,
                queue_out=queue_out,
                setup_args=(should_listen,),
                setup_kwargs=setup_kwargs,
            )
        elif handler_name == "chatTTS":
            from speech_to_speech.TTS.chatTTS_handler import ChatTTSHandler
            setup_kwargs = {"device": "cuda", **setup_kwargs}
            handler = ChatTTSHandler(
                stop_event,
                queue_in=queue_in,
                queue_out=queue_out,
                setup_args=(should_listen,),
                setup_kwargs=setup_kwargs,
            )
        elif handler_name == "facebookMMS":
            from speech_to_speech.TTS.facebookmms_handler import FacebookMMSTTSHandler
            setup_kwargs = {"device": "cuda", "language": "en", **setup_kwargs}
            handler = FacebookMMSTTSHandler(
                stop_event,
                queue_in=queue_in,
                queue_out=queue_out,
                setup_args=(should_listen,),
                setup_kwargs=setup_kwargs,
            )
        else:
            raise ValueError(f"Unknown handler: {handler_name}")

        result.warmup_time = time.perf_counter() - start_setup
        logger.info(f"Handler {handler_name} initialized and warmed up in {result.warmup_time:.3f}s")

        for i in range(iterations):
            logger.info(f"Iteration {i+1}/{iterations} for {handler_name}")
            start_time = time.perf_counter()
            time_to_first_chunk = None
            first_output = True
            total_samples = 0

            tts_input = TTSInput(text=text, language_code=language_code)
            for chunk in handler.process(tts_input):
                if first_output:
                    time_to_first_chunk = time.perf_counter() - start_time
                    first_output = False
                if chunk is None:
                    continue
                try:
                    total_samples += len(chunk)
                except Exception:
                    pass

            end_time = time.perf_counter()
            time_taken = end_time - start_time
            audio_duration = total_samples / DEFAULT_SAMPLE_RATE if total_samples > 0 else 0.0

            result.add_inference(time_taken, audio_duration, time_to_first_chunk)
            ttfc_str = f", TTFC: {time_to_first_chunk:.4f}s" if time_to_first_chunk else ""
            logger.info(
                f"  Time: {time_taken:.4f}s{ttfc_str}, Audio: {audio_duration:.2f}s, RTF: {audio_duration / time_taken if time_taken > 0 else 0:.2f}"
            )

        handler.cleanup()
        stop_event.set()

    except Exception as e:
        logger.error(f"Error benchmarking {handler_name}: {e}", exc_info=True)
        result.add_error(str(e))

    return result


def normalize_qwen3_mlx_quantizations(values: List[str] | None) -> List[str]:
    if not values:
        return []

    normalized = []
    seen = set()
    for value in values:
        quantization = str(value).strip().lower()
        if quantization in ("default", "none", ""):
            quantization = "bf16"
        if quantization not in VALID_QWEN3_MLX_QUANTIZATIONS:
            raise ValueError(
                "Unsupported qwen3 MLX quantization "
                f"{value!r}. Supported values: {', '.join(VALID_QWEN3_MLX_QUANTIZATIONS)}"
            )
        if quantization in seen:
            continue
        seen.add(quantization)
        normalized.append(quantization)
    return normalized


def build_benchmark_targets(args) -> List[tuple[str, str, Dict[str, Any]]]:
    targets = []
    qwen3_quantizations = normalize_qwen3_mlx_quantizations(args.qwen3_mlx_quantizations)

    for handler_name in args.handlers:
        if handler_name == "qwen3" and qwen3_quantizations:
            for quantization in qwen3_quantizations:
                targets.append(
                    (
                        f"qwen3[{quantization}]",
                        "qwen3",
                        {"mlx_quantization": quantization},
                    )
                )
            continue

        targets.append((handler_name, handler_name, {}))

    return targets


def print_results(results: List[BenchmarkResult]):
    print("\n" + "=" * 80)
    print("TTS BENCHMARK RESULTS")
    print("=" * 80)

    for result in results:
        stats = result.get_stats()
        print(f"\nHandler: {stats['handler']}")
        print("-" * 80)

        if stats.get("status") == "failed":
            print("  Status: FAILED")
            print(f"  Errors: {stats['errors']}")
            continue

        print(f"  Warmup Time:          {stats['warmup_time']:.4f}s")
        print(f"  Avg Inference Time:   {stats['avg_inference_time']:.4f}s")
        print(f"  Min Inference Time:   {stats['min_inference_time']:.4f}s")
        print(f"  Max Inference Time:   {stats['max_inference_time']:.4f}s")
        print(f"  Std Deviation:        {stats['std_inference_time']:.4f}s")

        print(f"  Avg Audio Duration:   {stats['avg_audio_duration']:.2f}s")
        print(f"  Min Audio Duration:   {stats['min_audio_duration']:.2f}s")
        print(f"  Max Audio Duration:   {stats['max_audio_duration']:.2f}s")
        print(f"  Std Audio Duration:   {stats['std_audio_duration']:.4f}s")
        print(f"  Avg RTF:              {stats['avg_rtf']:.2f}")

        if "avg_time_to_first_chunk" in stats:
            print("\n  Time to First Chunk:")
            print(f"    Avg TTFC:           {stats['avg_time_to_first_chunk']:.4f}s")
            print(f"    Min TTFC:           {stats['min_time_to_first_chunk']:.4f}s")
            print(f"    Max TTFC:           {stats['max_time_to_first_chunk']:.4f}s")
            print(f"    Std TTFC:           {stats['std_time_to_first_chunk']:.4f}s")

        print(f"\n  Total Iterations:     {stats['total_iterations']}")

        if stats["errors"]:
            print(f"  Errors: {stats['errors']}")

    print("\n" + "=" * 80)
    print("COMPARISON (Average Inference Time)")
    print("=" * 80)

    successful_results = [r for r in results if r.inference_times]
    if successful_results:
        sorted_results = sorted(successful_results, key=lambda x: np.mean(x.inference_times))
        fastest = sorted_results[0]
        fastest_time = np.mean(fastest.inference_times)

        for result in sorted_results:
            avg_time = np.mean(result.inference_times)
            speedup = avg_time / fastest_time
            print(f"  {result.handler_name:25s}: {avg_time:.4f}s  ({speedup:.2f}x slower than fastest)")


def save_results(results: List[BenchmarkResult], output_file: str):
    data = {
        "results": [r.get_stats() for r in results],
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
    }

    with open(output_file, "w") as f:
        json.dump(data, f, indent=2)

    logger.info(f"Results saved to: {output_file}")


def main():
    parser = argparse.ArgumentParser(description="Benchmark TTS handlers")
    parser.add_argument(
        "--text",
        type=str,
        default="Hello from the speech to speech benchmark. This is a latency test.",
        help="Text to synthesize",
    )
    parser.add_argument(
        "--handlers",
        nargs="+",
        default=["kokoro", "qwen3", "pocket_tts"],
        help="List of handlers to benchmark (kokoro, qwen3, pocket_tts, chatTTS, facebookMMS)",
    )
    parser.add_argument(
        "--iterations",
        type=int,
        default=3,
        help="Number of iterations per handler (default: 3)",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="tts_benchmark_results.json",
        help="Output JSON file for results (default: tts_benchmark_results.json)",
    )
    parser.add_argument(
        "--language_code",
        type=str,
        default="en",
        help="Language code to pass to TTS handlers (default: en)",
    )
    parser.add_argument(
        "--qwen3_mlx_quantizations",
        nargs="+",
        default=None,
        help=(
            "Optional list of Apple Silicon MLX Qwen3-TTS quantizations to benchmark "
            "as separate variants. Supported values: bf16, 4bit, 6bit, 8bit."
        ),
    )

    args = parser.parse_args()

    if not args.handlers:
        logger.error("No handlers provided")
        return

    try:
        targets = build_benchmark_targets(args)
    except ValueError as e:
        logger.error(str(e))
        return

    results = []
    for result_name, handler_name, handler_kwargs in targets:
        result = benchmark_handler(
            handler_name,
            args.text,
            args.iterations,
            handler_kwargs=handler_kwargs,
            language_code=args.language_code,
        )
        result.handler_name = result_name
        results.append(result)

    print_results(results)
    save_results(results, args.output)

    logger.info("TTS benchmarking complete!")


if __name__ == "__main__":
    main()