feat: 大屏后端（FastAPI+MQTT+sim_engine）
- REST :10085 / MQTT 控制 / SSE 回退 / AI+语音 - 数据源规划：接核心服务端 + 算力引擎真实数据
This commit is contained in:
+385
@@ -0,0 +1,385 @@
|
||||
"""
|
||||
STT Benchmarking Script
|
||||
|
||||
This script benchmarks different Speech-to-Text (STT) handlers to compare their performance.
|
||||
Measures: inference time, warmup time, memory usage, and transcription quality.
|
||||
|
||||
Usage:
|
||||
python benchmark_stt.py --audio_file path/to/audio.wav --iterations 10
|
||||
python benchmark_stt.py --audio_file path/to/audio.wav --handlers whisper mlx-audio-whisper
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from pathlib import Path
|
||||
from queue import Queue
|
||||
from threading import Event
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
from speech_to_speech.pipeline.messages import VADAudio
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BenchmarkResult:
|
||||
"""Stores benchmark results for a single STT 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_token: list[float] = []
|
||||
self.transcriptions: list[str] = []
|
||||
self.errors: list[str] = []
|
||||
|
||||
def add_inference(self, time_taken: float, transcription: Any, ttft: Optional[float] = None):
|
||||
self.inference_times.append(time_taken)
|
||||
self.transcriptions.append(transcription)
|
||||
if ttft is not None:
|
||||
self.time_to_first_token.append(ttft)
|
||||
|
||||
def add_error(self, error: str):
|
||||
self.errors.append(error)
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""Calculate statistics from benchmark results."""
|
||||
if not self.inference_times:
|
||||
return {
|
||||
"handler": self.handler_name,
|
||||
"status": "failed",
|
||||
"errors": self.errors,
|
||||
}
|
||||
|
||||
stats = {
|
||||
"handler": self.handler_name,
|
||||
"warmup_time": self.warmup_time,
|
||||
"avg_inference_time": np.mean(self.inference_times),
|
||||
"min_inference_time": np.min(self.inference_times),
|
||||
"max_inference_time": np.max(self.inference_times),
|
||||
"std_inference_time": np.std(self.inference_times),
|
||||
"total_iterations": len(self.inference_times),
|
||||
"errors": self.errors,
|
||||
"sample_transcription": self.transcriptions[0] if self.transcriptions else None,
|
||||
}
|
||||
|
||||
# Add time to first token stats if available
|
||||
if self.time_to_first_token:
|
||||
stats["avg_time_to_first_token"] = np.mean(self.time_to_first_token)
|
||||
stats["min_time_to_first_token"] = np.min(self.time_to_first_token)
|
||||
stats["max_time_to_first_token"] = np.max(self.time_to_first_token)
|
||||
stats["std_time_to_first_token"] = np.std(self.time_to_first_token)
|
||||
|
||||
return stats
|
||||
|
||||
|
||||
def load_audio(audio_path: str) -> np.ndarray:
|
||||
"""Load audio file and return as numpy array."""
|
||||
logger.info(f"Loading audio from: {audio_path}")
|
||||
audio, sample_rate = sf.read(audio_path)
|
||||
|
||||
# Convert to mono if stereo
|
||||
if len(audio.shape) > 1:
|
||||
audio = audio.mean(axis=1)
|
||||
|
||||
# Resample to 16kHz if needed (most whisper models expect 16kHz)
|
||||
if sample_rate != 16000:
|
||||
logger.warning(f"Audio sample rate is {sample_rate}Hz, resampling to 16000Hz")
|
||||
try:
|
||||
import librosa
|
||||
audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=16000)
|
||||
except ImportError:
|
||||
logger.error("librosa not installed. Please install it with: pip install librosa")
|
||||
logger.error("Attempting scipy resampling as fallback...")
|
||||
from scipy import signal
|
||||
# Calculate resampling ratio
|
||||
num_samples = int(len(audio) * 16000 / sample_rate)
|
||||
audio = signal.resample(audio, num_samples)
|
||||
|
||||
return audio.astype(np.float32)
|
||||
|
||||
|
||||
def benchmark_handler(
|
||||
handler_name: str,
|
||||
audio: np.ndarray,
|
||||
iterations: int,
|
||||
handler_kwargs: Optional[Dict[str, Any]] = None
|
||||
) -> BenchmarkResult:
|
||||
"""Benchmark a single STT handler."""
|
||||
logger.info(f"Benchmarking {handler_name}...")
|
||||
result = BenchmarkResult(handler_name)
|
||||
|
||||
try:
|
||||
# Create queues and events for handler
|
||||
stop_event = Event()
|
||||
queue_in: Queue[Any] = Queue()
|
||||
queue_out: Queue[Any] = Queue()
|
||||
|
||||
handler: Any = None
|
||||
if handler_name == "whisper":
|
||||
from speech_to_speech.STT.whisper_stt_handler import WhisperSTTHandler
|
||||
setup_kwargs = handler_kwargs or {
|
||||
"model_name": "distil-whisper/distil-large-v3",
|
||||
"device": "cuda",
|
||||
"torch_dtype": "float16",
|
||||
}
|
||||
handler = WhisperSTTHandler(
|
||||
stop_event,
|
||||
queue_in=queue_in,
|
||||
queue_out=queue_out,
|
||||
setup_kwargs=setup_kwargs
|
||||
)
|
||||
|
||||
elif handler_name == "whisper-mlx":
|
||||
from speech_to_speech.STT.lightning_whisper_mlx_handler import LightningWhisperSTTHandler
|
||||
setup_kwargs = handler_kwargs or {
|
||||
"model_name": "large-v3",
|
||||
"device": "mps",
|
||||
}
|
||||
handler = LightningWhisperSTTHandler(
|
||||
stop_event,
|
||||
queue_in=queue_in,
|
||||
queue_out=queue_out,
|
||||
setup_kwargs=setup_kwargs
|
||||
)
|
||||
|
||||
elif handler_name == "mlx-audio-whisper":
|
||||
from speech_to_speech.STT.mlx_audio_whisper_handler import MLXAudioWhisperSTTHandler
|
||||
setup_kwargs = handler_kwargs or {
|
||||
"model_name": "mlx-community/whisper-large-v3-turbo",
|
||||
}
|
||||
handler = MLXAudioWhisperSTTHandler(
|
||||
stop_event,
|
||||
queue_in=queue_in,
|
||||
queue_out=queue_out,
|
||||
setup_kwargs=setup_kwargs
|
||||
)
|
||||
|
||||
elif handler_name == "faster-whisper":
|
||||
from speech_to_speech.STT.faster_whisper_handler import FasterWhisperSTTHandler
|
||||
setup_kwargs = handler_kwargs or {
|
||||
"model_name": "large-v3",
|
||||
"device": "auto",
|
||||
"compute_type": "float16",
|
||||
}
|
||||
handler = FasterWhisperSTTHandler(
|
||||
stop_event,
|
||||
queue_in=queue_in,
|
||||
queue_out=queue_out,
|
||||
setup_kwargs=setup_kwargs
|
||||
)
|
||||
|
||||
elif handler_name == "moonshine":
|
||||
from archive.STT.moonshine_handler import MoonshineSTTHandler
|
||||
handler = MoonshineSTTHandler(
|
||||
stop_event,
|
||||
queue_in=queue_in,
|
||||
queue_out=queue_out,
|
||||
)
|
||||
|
||||
elif handler_name == "parakeet-tdt":
|
||||
from speech_to_speech.STT.parakeet_tdt_handler import ParakeetTDTSTTHandler
|
||||
setup_kwargs = handler_kwargs or {
|
||||
"device": "mps",
|
||||
"enable_live_transcription": False,
|
||||
}
|
||||
handler = ParakeetTDTSTTHandler(
|
||||
stop_event,
|
||||
queue_in=queue_in,
|
||||
queue_out=queue_out,
|
||||
setup_kwargs=setup_kwargs
|
||||
)
|
||||
|
||||
elif handler_name == "parakeet-tdt-progressive":
|
||||
from speech_to_speech.STT.parakeet_tdt_handler import ParakeetTDTSTTHandler
|
||||
setup_kwargs = handler_kwargs or {
|
||||
"device": "mps",
|
||||
"enable_live_transcription": True,
|
||||
"live_transcription_update_interval": 0.25,
|
||||
}
|
||||
handler = ParakeetTDTSTTHandler(
|
||||
stop_event,
|
||||
queue_in=queue_in,
|
||||
queue_out=queue_out,
|
||||
setup_kwargs=setup_kwargs
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown handler: {handler_name}")
|
||||
|
||||
# Warmup is done in handler setup
|
||||
logger.info(f"Handler {handler_name} initialized and warmed up")
|
||||
|
||||
# Additional warmup on the actual audio (excluded from timings)
|
||||
for _ in handler.process(VADAudio(audio=audio)):
|
||||
pass
|
||||
|
||||
# Run benchmark iterations
|
||||
for i in range(iterations):
|
||||
logger.info(f"Iteration {i+1}/{iterations} for {handler_name}")
|
||||
|
||||
start_time = time.perf_counter()
|
||||
|
||||
# Process audio
|
||||
transcription = None
|
||||
time_to_first_token = None
|
||||
first_output = True
|
||||
|
||||
for output in handler.process(VADAudio(audio=audio)):
|
||||
# Measure time to first token
|
||||
if first_output:
|
||||
time_to_first_token = time.perf_counter() - start_time
|
||||
first_output = False
|
||||
|
||||
if isinstance(output, tuple):
|
||||
transcription = output[0] # (text, language)
|
||||
else:
|
||||
transcription = output
|
||||
|
||||
end_time = time.perf_counter()
|
||||
|
||||
time_taken = end_time - start_time
|
||||
result.add_inference(time_taken, transcription, time_to_first_token)
|
||||
|
||||
ttft_str = f", TTFT: {time_to_first_token:.4f}s" if time_to_first_token else ""
|
||||
text_preview = str(transcription)[:50] if transcription is not None else "(none)"
|
||||
logger.info(f" Time: {time_taken:.4f}s{ttft_str}, Text: {text_preview}...")
|
||||
|
||||
# Cleanup
|
||||
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 print_results(results: List[BenchmarkResult]):
|
||||
"""Print benchmark results in a formatted table."""
|
||||
print("\n" + "="*80)
|
||||
print("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 time to first token stats if available
|
||||
if 'avg_time_to_first_token' in stats:
|
||||
print("\n Time to First Token:")
|
||||
print(f" Avg TTFT: {stats['avg_time_to_first_token']:.4f}s")
|
||||
print(f" Min TTFT: {stats['min_time_to_first_token']:.4f}s")
|
||||
print(f" Max TTFT: {stats['max_time_to_first_token']:.4f}s")
|
||||
print(f" Std TTFT: {stats['std_time_to_first_token']:.4f}s")
|
||||
|
||||
print(f"\n Total Iterations: {stats['total_iterations']}")
|
||||
print(f" Sample Transcription: {stats['sample_transcription']}")
|
||||
|
||||
if stats['errors']:
|
||||
print(f" Errors: {stats['errors']}")
|
||||
|
||||
# Comparison table
|
||||
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):
|
||||
"""Save benchmark results to JSON file."""
|
||||
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 STT handlers")
|
||||
parser.add_argument(
|
||||
"--audio_file",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to audio file for benchmarking"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--handlers",
|
||||
nargs="+",
|
||||
default=["whisper", "whisper-mlx", "mlx-audio-whisper", "faster-whisper", "parakeet-tdt", "parakeet-tdt-progressive"],
|
||||
help="List of handlers to benchmark (default: all)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iterations",
|
||||
type=int,
|
||||
default=5,
|
||||
help="Number of iterations per handler (default: 5)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
default="stt_benchmark_results.json",
|
||||
help="Output JSON file for results (default: stt_benchmark_results.json)"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Validate audio file exists
|
||||
if not Path(args.audio_file).exists():
|
||||
logger.error(f"Audio file not found: {args.audio_file}")
|
||||
return
|
||||
|
||||
# Load audio
|
||||
audio = load_audio(args.audio_file)
|
||||
logger.info(f"Audio loaded: {len(audio)} samples, {len(audio)/16000:.2f}s duration")
|
||||
|
||||
# Run benchmarks
|
||||
results = []
|
||||
for handler_name in args.handlers:
|
||||
result = benchmark_handler(handler_name, audio, args.iterations)
|
||||
results.append(result)
|
||||
|
||||
# Print and save results
|
||||
print_results(results)
|
||||
save_results(results, args.output)
|
||||
|
||||
logger.info("Benchmarking complete!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+395
@@ -0,0 +1,395 @@
|
||||
"""
|
||||
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()
|
||||
+136
@@ -0,0 +1,136 @@
|
||||
import socket
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from queue import Queue
|
||||
|
||||
import numpy as np
|
||||
import sounddevice as sd
|
||||
from transformers import HfArgumentParser
|
||||
|
||||
|
||||
@dataclass
|
||||
class ListenAndPlayArguments:
|
||||
send_rate: int = field(default=16000, metadata={"help": "In Hz. Default is 16000."})
|
||||
recv_rate: int = field(default=16000, metadata={"help": "In Hz. Default is 16000."})
|
||||
list_play_chunk_size: int = field(
|
||||
default=1024,
|
||||
metadata={"help": "The size of data chunks (in bytes). Default is 1024."},
|
||||
)
|
||||
host: str = field(
|
||||
default="localhost",
|
||||
metadata={
|
||||
"help": "The hostname or IP address for listening and playing. Default is 'localhost'."
|
||||
},
|
||||
)
|
||||
send_port: int = field(
|
||||
default=12345,
|
||||
metadata={"help": "The network port for sending data. Default is 12345."},
|
||||
)
|
||||
recv_port: int = field(
|
||||
default=12346,
|
||||
metadata={"help": "The network port for receiving data. Default is 12346."},
|
||||
)
|
||||
|
||||
|
||||
def listen_and_play(
|
||||
send_rate=16000,
|
||||
recv_rate=44100,
|
||||
list_play_chunk_size=1024,
|
||||
host="localhost",
|
||||
send_port=12345,
|
||||
recv_port=12346,
|
||||
):
|
||||
send_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
send_socket.connect((host, send_port))
|
||||
|
||||
recv_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
recv_socket.connect((host, recv_port))
|
||||
|
||||
print("Recording and streaming...")
|
||||
|
||||
stop_event = threading.Event()
|
||||
recv_queue = Queue()
|
||||
send_queue = Queue()
|
||||
|
||||
# Pre-generate a static dither buffer (±1 LSB, -96 dB) to keep the audio
|
||||
# sink active without calling numpy inside the real-time audio callback.
|
||||
dither_bytes = np.random.randint(
|
||||
-1, 2, size=list_play_chunk_size, dtype=np.int16
|
||||
).tobytes()
|
||||
|
||||
def callback_recv(outdata, frames, time, status):
|
||||
if not recv_queue.empty():
|
||||
data = recv_queue.get()
|
||||
outdata[: len(data)] = data
|
||||
outdata[len(data) :] = b"\x00" * (len(outdata) - len(data))
|
||||
else:
|
||||
outdata[:] = dither_bytes
|
||||
|
||||
def callback_send(indata, frames, time, status):
|
||||
if recv_queue.empty():
|
||||
data = bytes(indata)
|
||||
send_queue.put(data)
|
||||
|
||||
def send(stop_event, send_queue):
|
||||
while not stop_event.is_set():
|
||||
data = send_queue.get()
|
||||
send_socket.sendall(data)
|
||||
|
||||
def recv(stop_event, recv_queue):
|
||||
def receive_full_chunk(conn, chunk_size):
|
||||
data = b""
|
||||
while len(data) < chunk_size:
|
||||
packet = conn.recv(chunk_size - len(data))
|
||||
if not packet:
|
||||
return None # Connection has been closed
|
||||
data += packet
|
||||
return data
|
||||
|
||||
while not stop_event.is_set():
|
||||
data = receive_full_chunk(recv_socket, list_play_chunk_size * 2)
|
||||
if data:
|
||||
recv_queue.put(data)
|
||||
|
||||
try:
|
||||
send_stream = sd.RawInputStream(
|
||||
samplerate=send_rate,
|
||||
channels=1,
|
||||
dtype="int16",
|
||||
blocksize=list_play_chunk_size,
|
||||
callback=callback_send,
|
||||
)
|
||||
recv_stream = sd.RawOutputStream(
|
||||
samplerate=recv_rate,
|
||||
channels=1,
|
||||
dtype="int16",
|
||||
blocksize=list_play_chunk_size,
|
||||
callback=callback_recv,
|
||||
)
|
||||
threading.Thread(target=send_stream.start).start()
|
||||
threading.Thread(target=recv_stream.start).start()
|
||||
|
||||
send_thread = threading.Thread(target=send, args=(stop_event, send_queue))
|
||||
send_thread.start()
|
||||
recv_thread = threading.Thread(target=recv, args=(stop_event, recv_queue))
|
||||
recv_thread.start()
|
||||
|
||||
input("Press Enter to stop...")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("Finished streaming.")
|
||||
|
||||
finally:
|
||||
stop_event.set()
|
||||
# Given that socket::recv is blocking in receive_data_chunk, shut it down to allow the thread to continue.
|
||||
recv_socket.shutdown(socket.SHUT_RDWR)
|
||||
recv_thread.join()
|
||||
send_thread.join()
|
||||
send_socket.close()
|
||||
recv_socket.close()
|
||||
print("Connection closed.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = HfArgumentParser((ListenAndPlayArguments,)) # type: ignore[arg-type]
|
||||
(listen_and_play_kwargs,) = parser.parse_args_into_dataclasses()
|
||||
listen_and_play(**vars(listen_and_play_kwargs))
|
||||
@@ -0,0 +1,385 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from queue import Empty, Queue
|
||||
from threading import Event, Lock
|
||||
from typing import Any, Optional
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
|
||||
@dataclass
|
||||
class ListenAndPlayRealtimeArguments:
|
||||
host: str = field(
|
||||
default="192.168.1.9",
|
||||
metadata={"help": "Realtime server host. Default is 192.168.1.9."},
|
||||
)
|
||||
port: int = field(
|
||||
default=8765,
|
||||
metadata={"help": "Realtime server port. Default is 8765."},
|
||||
)
|
||||
model: str = field(
|
||||
default="local",
|
||||
metadata={"help": "Model name sent to the OpenAI-compatible realtime client."},
|
||||
)
|
||||
api_key: str = field(
|
||||
default="test-key",
|
||||
metadata={"help": "API key for the OpenAI SDK client. Local server ignores it."},
|
||||
)
|
||||
base_url: Optional[str] = field(
|
||||
default=None,
|
||||
metadata={"help": "Optional HTTP base URL, e.g. http://192.168.1.9:8765/v1"},
|
||||
)
|
||||
websocket_base_url: Optional[str] = field(
|
||||
default=None,
|
||||
metadata={"help": "Optional WS base URL, e.g. ws://192.168.1.9:8765/v1"},
|
||||
)
|
||||
send_rate: int = field(
|
||||
default=16000,
|
||||
metadata={"help": "Microphone sample rate in Hz. Default is 16000."},
|
||||
)
|
||||
recv_rate: int = field(
|
||||
default=16000,
|
||||
metadata={"help": "Speaker sample rate in Hz. Default is 16000."},
|
||||
)
|
||||
chunk_size: int = field(
|
||||
default=1024,
|
||||
metadata={"help": "Audio callback block size in samples. Default is 1024."},
|
||||
)
|
||||
input_device: Optional[int] = field(
|
||||
default=None,
|
||||
metadata={"help": "Optional sounddevice input device index."},
|
||||
)
|
||||
output_device: Optional[int] = field(
|
||||
default=None,
|
||||
metadata={"help": "Optional sounddevice output device index."},
|
||||
)
|
||||
instructions: Optional[str] = field(
|
||||
default=None,
|
||||
metadata={"help": "Optional session instructions to apply on connect."},
|
||||
)
|
||||
voice: Optional[str] = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": (
|
||||
"TTS voice sent as session.audio.output.voice. "
|
||||
"Local Kokoro: e.g. bm_fable, af_heart, am_adam. "
|
||||
"OpenAI Realtime: e.g. marin, cedar, alloy."
|
||||
),
|
||||
},
|
||||
)
|
||||
print_json: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "Print raw event payloads in addition to friendly logs."},
|
||||
)
|
||||
block_mic_during_playback: bool = field(
|
||||
default=False,
|
||||
metadata={
|
||||
"help": "If set, pause microphone capture while speaker audio is playing. Disabled by default so barge-in works."
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _make_client(args: ListenAndPlayRealtimeArguments) -> AsyncOpenAI:
|
||||
base_url = args.base_url or f"http://{args.host}:{args.port}/v1"
|
||||
websocket_base_url = args.websocket_base_url or f"ws://{args.host}:{args.port}/v1"
|
||||
return AsyncOpenAI(
|
||||
api_key=args.api_key,
|
||||
base_url=base_url,
|
||||
websocket_base_url=websocket_base_url,
|
||||
)
|
||||
|
||||
|
||||
def _build_session_update(args: ListenAndPlayRealtimeArguments) -> dict:
|
||||
def maybe_pcm_format(rate: int) -> Optional[dict]:
|
||||
# The OpenAI realtime Pydantic models only validate audio/pcm at 24 kHz.
|
||||
# Our local pipeline defaults to 16 kHz internally when format is omitted,
|
||||
# so omit the field for the common local case instead of sending an
|
||||
# invalid 16 kHz declaration.
|
||||
if rate == 16000:
|
||||
return None
|
||||
if rate == 24000:
|
||||
return {"type": "audio/pcm", "rate": 24000}
|
||||
raise ValueError(
|
||||
f"Unsupported rate {rate}. Use 16000 for the local pipeline default "
|
||||
f"or 24000 to match the OpenAI realtime audio format schema."
|
||||
)
|
||||
|
||||
input_cfg = {
|
||||
"turn_detection": {"type": "server_vad", "interrupt_response": True},
|
||||
}
|
||||
output_cfg: dict[str, Any] = {}
|
||||
|
||||
input_format = maybe_pcm_format(args.send_rate)
|
||||
output_format = maybe_pcm_format(args.recv_rate)
|
||||
if input_format is not None:
|
||||
input_cfg["format"] = input_format
|
||||
if output_format is not None:
|
||||
output_cfg["format"] = output_format
|
||||
if args.voice:
|
||||
output_cfg["voice"] = args.voice
|
||||
|
||||
session = {
|
||||
"type": "realtime",
|
||||
"audio": {
|
||||
"input": input_cfg,
|
||||
"output": output_cfg,
|
||||
},
|
||||
}
|
||||
if args.instructions:
|
||||
session["instructions"] = args.instructions
|
||||
return {"type": "session.update", "session": session}
|
||||
|
||||
|
||||
async def listen_and_play_realtime(args: ListenAndPlayRealtimeArguments) -> None:
|
||||
import sounddevice as sd
|
||||
|
||||
client = _make_client(args)
|
||||
|
||||
mic_queue: Queue[bytes] = Queue(maxsize=128)
|
||||
stop_event = Event()
|
||||
playback_buffer = bytearray()
|
||||
playback_lock = Lock()
|
||||
speaker_active_until = [0.0]
|
||||
partial_user_text = ""
|
||||
live_user_width = 0
|
||||
saw_user_speech = False
|
||||
|
||||
def render_live_user_text(text: str, final: bool = False) -> None:
|
||||
nonlocal live_user_width
|
||||
line = f"USER: {text}"
|
||||
padded = line
|
||||
if live_user_width > len(line):
|
||||
padded += " " * (live_user_width - len(line))
|
||||
|
||||
if final:
|
||||
print(f"\r{padded}", flush=True)
|
||||
live_user_width = 0
|
||||
return
|
||||
|
||||
print(f"\r{padded}", end="", flush=True)
|
||||
live_user_width = len(line)
|
||||
|
||||
def clear_live_user_text() -> None:
|
||||
nonlocal live_user_width
|
||||
if live_user_width == 0:
|
||||
return
|
||||
print("\r" + (" " * live_user_width) + "\r", end="", flush=True)
|
||||
live_user_width = 0
|
||||
|
||||
def clear_playback_buffer() -> None:
|
||||
speaker_active_until[0] = 0.0
|
||||
with playback_lock:
|
||||
playback_buffer.clear()
|
||||
|
||||
def callback_recv(outdata, _frames, _time_info, status):
|
||||
if status:
|
||||
print(f"Speaker status: {status}", flush=True)
|
||||
|
||||
needed = len(outdata)
|
||||
with playback_lock:
|
||||
available = min(needed, len(playback_buffer))
|
||||
if available:
|
||||
outdata[:available] = playback_buffer[:available]
|
||||
del playback_buffer[:available]
|
||||
if available < needed:
|
||||
outdata[available:] = b"\x00" * (needed - available)
|
||||
|
||||
def callback_send(indata, _frames, _time_info, status):
|
||||
if status:
|
||||
print(f"Mic status: {status}", flush=True)
|
||||
|
||||
if args.block_mic_during_playback:
|
||||
with playback_lock:
|
||||
speaker_active = bool(playback_buffer)
|
||||
if speaker_active or time.monotonic() < speaker_active_until[0]:
|
||||
return
|
||||
|
||||
try:
|
||||
mic_queue.put_nowait(bytes(indata))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def send_audio(conn):
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
chunk = await asyncio.to_thread(mic_queue.get, True, 0.1)
|
||||
except Empty:
|
||||
continue
|
||||
|
||||
await conn.send(
|
||||
{
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": base64.b64encode(chunk).decode("ascii"),
|
||||
}
|
||||
)
|
||||
|
||||
async def receive_events(conn):
|
||||
nonlocal partial_user_text, saw_user_speech
|
||||
|
||||
while not stop_event.is_set():
|
||||
event = await conn.recv()
|
||||
|
||||
if args.print_json:
|
||||
try:
|
||||
print(f"EVENT: {event.model_dump_json()}", flush=True)
|
||||
except Exception:
|
||||
print(f"EVENT: {event}", flush=True)
|
||||
|
||||
if event.type == "session.created":
|
||||
print("Connected.", flush=True)
|
||||
elif event.type == "input_audio_buffer.speech_started":
|
||||
clear_playback_buffer()
|
||||
partial_user_text = ""
|
||||
if saw_user_speech:
|
||||
print("", flush=True)
|
||||
saw_user_speech = True
|
||||
elif event.type == "input_audio_buffer.speech_stopped":
|
||||
pass
|
||||
elif event.type == "conversation.item.input_audio_transcription.delta":
|
||||
# This server currently sends the latest partial hypothesis in
|
||||
# each "delta" event rather than a token-level suffix, so render
|
||||
# the newest snapshot instead of concatenating repeated text.
|
||||
partial_user_text = event.delta.strip()
|
||||
if partial_user_text:
|
||||
render_live_user_text(partial_user_text)
|
||||
elif event.type == "conversation.item.input_audio_transcription.completed":
|
||||
partial_user_text = ""
|
||||
render_live_user_text(event.transcript.strip(), final=True)
|
||||
elif event.type == "response.created":
|
||||
clear_live_user_text()
|
||||
print("ASSISTANT: <response started>", flush=True)
|
||||
elif event.type == "response.output_audio.delta":
|
||||
audio = base64.b64decode(event.delta)
|
||||
with playback_lock:
|
||||
playback_buffer.extend(audio)
|
||||
speaker_active_until[0] = time.monotonic() + max(0.15, len(audio) / (2 * args.recv_rate))
|
||||
elif event.type == "response.output_audio.done":
|
||||
print("ASSISTANT: <audio done>", flush=True)
|
||||
elif event.type == "response.output_audio_transcript.done":
|
||||
print(f"ASSISTANT: {event.transcript}", flush=True)
|
||||
elif event.type == "response.function_call_arguments.done":
|
||||
print(
|
||||
f"TOOL: {event.name} call_id={event.call_id} arguments={event.arguments}",
|
||||
flush=True,
|
||||
)
|
||||
elif event.type == "response.done":
|
||||
if event.response.status == "cancelled":
|
||||
clear_playback_buffer()
|
||||
print(f"ASSISTANT: <response {event.response.status}>", flush=True)
|
||||
elif event.type == "error":
|
||||
clear_live_user_text()
|
||||
print(f"ERROR: {event.error.type}: {event.error.message}", flush=True)
|
||||
else:
|
||||
clear_live_user_text()
|
||||
print(f"EVENT: {event.type}", flush=True)
|
||||
|
||||
async def wait_for_stop():
|
||||
await asyncio.to_thread(input, "Press Enter to stop...\n")
|
||||
stop_event.set()
|
||||
|
||||
input_stream = sd.RawInputStream(
|
||||
samplerate=args.send_rate,
|
||||
channels=1,
|
||||
dtype="int16",
|
||||
blocksize=args.chunk_size,
|
||||
callback=callback_send,
|
||||
device=args.input_device,
|
||||
)
|
||||
output_stream = sd.RawOutputStream(
|
||||
samplerate=args.recv_rate,
|
||||
channels=1,
|
||||
dtype="int16",
|
||||
blocksize=args.chunk_size,
|
||||
callback=callback_recv,
|
||||
device=args.output_device,
|
||||
)
|
||||
|
||||
input_stream.start()
|
||||
output_stream.start()
|
||||
|
||||
try:
|
||||
async with client.realtime.connect(model=args.model) as conn:
|
||||
await conn.send(_build_session_update(args)) # type: ignore[arg-type]
|
||||
|
||||
sender_task = asyncio.create_task(send_audio(conn))
|
||||
receiver_task = asyncio.create_task(receive_events(conn))
|
||||
stopper_task = asyncio.create_task(wait_for_stop())
|
||||
|
||||
done, pending = await asyncio.wait(
|
||||
{sender_task, receiver_task, stopper_task},
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
|
||||
stop_event.set()
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
await asyncio.gather(*pending, return_exceptions=True)
|
||||
|
||||
for task in done:
|
||||
exc = task.exception()
|
||||
if exc is not None:
|
||||
raise exc
|
||||
finally:
|
||||
stop_event.set()
|
||||
clear_live_user_text()
|
||||
input_stream.stop()
|
||||
output_stream.stop()
|
||||
input_stream.close()
|
||||
output_stream.close()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Talk to the local OpenAI-compatible realtime speech pipeline.")
|
||||
defaults = ListenAndPlayRealtimeArguments()
|
||||
parser.add_argument("--host", default=defaults.host)
|
||||
parser.add_argument("--port", type=int, default=defaults.port)
|
||||
parser.add_argument("--model", default=defaults.model)
|
||||
parser.add_argument("--api-key", default=defaults.api_key)
|
||||
parser.add_argument("--base-url", default=defaults.base_url)
|
||||
parser.add_argument("--websocket-base-url", default=defaults.websocket_base_url)
|
||||
parser.add_argument("--send-rate", type=int, default=defaults.send_rate)
|
||||
parser.add_argument("--recv-rate", type=int, default=defaults.recv_rate)
|
||||
parser.add_argument("--chunk-size", type=int, default=defaults.chunk_size)
|
||||
parser.add_argument("--input-device", type=int, default=defaults.input_device)
|
||||
parser.add_argument("--output-device", type=int, default=defaults.output_device)
|
||||
parser.add_argument("--instructions", default=defaults.instructions)
|
||||
parser.add_argument(
|
||||
"--voice",
|
||||
default=defaults.voice,
|
||||
help=("session.audio.output.voice (Kokoro id like bm_fable, or OpenAI name like marin)."),
|
||||
)
|
||||
parser.add_argument("--print-json", action="store_true", default=defaults.print_json)
|
||||
parser.add_argument(
|
||||
"--block-mic-during-playback",
|
||||
action="store_true",
|
||||
default=defaults.block_mic_during_playback,
|
||||
)
|
||||
namespace = parser.parse_args()
|
||||
args = ListenAndPlayRealtimeArguments(
|
||||
host=namespace.host,
|
||||
port=namespace.port,
|
||||
model=namespace.model,
|
||||
api_key=namespace.api_key,
|
||||
base_url=namespace.base_url,
|
||||
websocket_base_url=namespace.websocket_base_url,
|
||||
send_rate=namespace.send_rate,
|
||||
recv_rate=namespace.recv_rate,
|
||||
chunk_size=namespace.chunk_size,
|
||||
input_device=namespace.input_device,
|
||||
output_device=namespace.output_device,
|
||||
instructions=namespace.instructions,
|
||||
voice=namespace.voice,
|
||||
print_json=namespace.print_json,
|
||||
block_mic_during_playback=namespace.block_mic_during_playback,
|
||||
)
|
||||
try:
|
||||
asyncio.run(listen_and_play_realtime(args))
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,569 @@
|
||||
"""Synthetic realtime client(s): single or parallel, single-turn or multi-turn.
|
||||
|
||||
Opens one (or N) websocket connection(s) to /v1/realtime and cycles through
|
||||
``--turns`` prompts per client at a fixed cadence (``--interval`` seconds
|
||||
between turn starts). Each turn:
|
||||
|
||||
1. Synthesize the prompt with macOS ``say`` (cached on disk after first run).
|
||||
2. Stream the prompt audio + trailing silence as input_audio_buffer.append.
|
||||
3. Wait for response.done.
|
||||
4. Sleep until --interval has elapsed since the turn started.
|
||||
|
||||
This single script subsumes two earlier ones:
|
||||
|
||||
* Parallel pool / capacity test:
|
||||
python scripts/synthetic_conversation_realtime_client.py --clients 3 --turns 1
|
||||
→ 3 clients connect simultaneously, each sends one prompt. Surplus clients
|
||||
beyond pool size receive ``session_limit_reached`` and exit cleanly.
|
||||
|
||||
* Single-client soak / multi-turn:
|
||||
python scripts/synthetic_conversation_realtime_client.py --turns 60 --interval 10
|
||||
→ 1 client, ~10 minute conversation, 60 sequential turns.
|
||||
|
||||
* Both at once (soak the pool):
|
||||
python scripts/synthetic_conversation_realtime_client.py --clients 2 --turns 60
|
||||
|
||||
* Soak a Hugging Face Inference Endpoint via its load balancer (10 min, 2 clients):
|
||||
export HF_TOKEN=hf_...
|
||||
python scripts/synthetic_conversation_realtime_client.py \
|
||||
--lb-url https://<your-lb>.us-east-1.aws.endpoints.huggingface.cloud \
|
||||
--clients 2 \
|
||||
--turns 60 \
|
||||
--interval 10 \
|
||||
--log-dir /tmp/hf_endpoint_soak
|
||||
|
||||
Each client uses a per-client prompt offset (coprime shift) so concurrent
|
||||
clients send distinct prompts at each turn — making cross-session leaks
|
||||
trivially detectable in the per-client transcript logs.
|
||||
|
||||
Outputs land under --log-dir:
|
||||
* prompts/prompt_NNN.wav — shared cache of say-synthesized prompts
|
||||
* client_NNN/conversation.txt — per-client transcript with timestamps
|
||||
* client_NNN/conversation.wav — per-client assistant audio concatenated
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import wave
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import websockets
|
||||
from scipy.signal import resample_poly
|
||||
|
||||
logger = logging.getLogger("synthetic_client")
|
||||
|
||||
SAMPLE_RATE_HZ = 16000
|
||||
CHUNK_MS = 20
|
||||
BYTES_PER_SAMPLE = 2 # PCM16
|
||||
CHUNK_BYTES = SAMPLE_RATE_HZ * BYTES_PER_SAMPLE * CHUNK_MS // 1000 # 640
|
||||
# Trailing silence after each prompt so server-side VAD detects speech_stopped
|
||||
# and auto-commits. The local realtime service only supports server VAD.
|
||||
TRAILING_SILENCE_MS = 1500
|
||||
# Per-client prompt offset shift. 7 is coprime with len(PROMPTS)=60, so each
|
||||
# client visits a unique permutation of the prompt list.
|
||||
PROMPT_SHIFT_PER_CLIENT = 7
|
||||
|
||||
# 60 varied prompts — short enough to fit a 10-second turn budget, diverse
|
||||
# enough that responses don't pattern-match. Cycled if --turns > len(PROMPTS).
|
||||
PROMPTS: list[str] = [
|
||||
"What is the capital of France?",
|
||||
"Tell me a joke about robots.",
|
||||
"How does photosynthesis work?",
|
||||
"What is two plus two?",
|
||||
"Who painted the Mona Lisa?",
|
||||
"What is the largest ocean on Earth?",
|
||||
"Recommend a simple pasta recipe.",
|
||||
"What is the speed of light?",
|
||||
"Who wrote Romeo and Juliet?",
|
||||
"What is the boiling point of water in Celsius?",
|
||||
"Tell me one fact about the planet Mars.",
|
||||
"What is the difference between weather and climate?",
|
||||
"How many continents are there?",
|
||||
"What is a haiku?",
|
||||
"What is the chemical formula for water?",
|
||||
"Who was the first president of the United States?",
|
||||
"What is gravity?",
|
||||
"Tell me a fun fact about dolphins.",
|
||||
"What is the tallest mountain on Earth?",
|
||||
"How do magnets work?",
|
||||
"What is the meaning of life in one sentence?",
|
||||
"Recommend a short book to read.",
|
||||
"What is the population of Tokyo, roughly?",
|
||||
"How does a refrigerator stay cold?",
|
||||
"Who invented the telephone?",
|
||||
"What is the smallest country in the world?",
|
||||
"Explain the theory of relativity simply.",
|
||||
"What is the difference between a virus and bacteria?",
|
||||
"Who wrote War and Peace?",
|
||||
"What is the deepest part of the ocean called?",
|
||||
"How does the human heart work?",
|
||||
"What is the most spoken language in the world?",
|
||||
"Tell me one fact about black holes.",
|
||||
"What is the longest river on Earth?",
|
||||
"How do you make a paper airplane?",
|
||||
"What is the chemical symbol for gold?",
|
||||
"Who composed the Fifth Symphony?",
|
||||
"What is the difference between an alligator and a crocodile?",
|
||||
"How many bones are in the human body?",
|
||||
"What is the smallest planet in our solar system?",
|
||||
"Suggest a beginner workout routine.",
|
||||
"What is the busiest airport in the world?",
|
||||
"How does Wi-Fi work in one sentence?",
|
||||
"What is the oldest known written language?",
|
||||
"Tell me about the Great Wall of China briefly.",
|
||||
"What is a neutron star?",
|
||||
"How does sound travel through space?",
|
||||
"What is the most popular sport globally?",
|
||||
"Who painted Starry Night?",
|
||||
"What happens when you mix baking soda and vinegar?",
|
||||
"How fast can a cheetah run?",
|
||||
"What is the difference between fiction and non-fiction?",
|
||||
"Why is the sky blue?",
|
||||
"What is the largest desert in the world?",
|
||||
"How many planets are in our solar system?",
|
||||
"Who discovered penicillin?",
|
||||
"What is a leap year?",
|
||||
"Tell me how rain forms.",
|
||||
"What is the longest-living animal?",
|
||||
"Say goodbye for now.",
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Audio helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def synthesize_with_say(text: str, out_path: Path) -> None:
|
||||
"""Render *text* to a 16kHz mono PCM16 WAV using macOS ``say``."""
|
||||
subprocess.run(
|
||||
[
|
||||
"say",
|
||||
text,
|
||||
"-o",
|
||||
str(out_path),
|
||||
"--file-format=WAVE",
|
||||
"--data-format=LEI16@16000",
|
||||
],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
)
|
||||
|
||||
|
||||
def load_pcm16_mono_16k(path: Path) -> bytes:
|
||||
"""Load any WAV / AIFF and return 16kHz mono PCM16 little-endian bytes."""
|
||||
data, src_rate = sf.read(str(path), always_2d=True)
|
||||
mono = data.mean(axis=1) if data.shape[1] > 1 else data[:, 0]
|
||||
if src_rate != SAMPLE_RATE_HZ:
|
||||
mono = resample_poly(mono, SAMPLE_RATE_HZ, src_rate)
|
||||
mono = np.clip(mono, -1.0, 1.0)
|
||||
pcm16 = (mono * 32767.0).astype(np.int16)
|
||||
return pcm16.tobytes()
|
||||
|
||||
|
||||
def write_wav(path: Path, pcm16_bytes: bytes, rate: int = SAMPLE_RATE_HZ) -> None:
|
||||
with wave.open(str(path), "wb") as w:
|
||||
w.setnchannels(1)
|
||||
w.setsampwidth(BYTES_PER_SAMPLE)
|
||||
w.setframerate(rate)
|
||||
w.writeframes(pcm16_bytes)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-turn ws flow
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def stream_prompt(ws, audio_pcm: bytes) -> None:
|
||||
"""Send prompt chunks at real-time pacing, then trailing silence."""
|
||||
for i in range(0, len(audio_pcm), CHUNK_BYTES):
|
||||
chunk = audio_pcm[i : i + CHUNK_BYTES]
|
||||
if len(chunk) < CHUNK_BYTES:
|
||||
chunk = chunk + b"\x00" * (CHUNK_BYTES - len(chunk))
|
||||
await ws.send(
|
||||
json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(chunk).decode("ascii")})
|
||||
)
|
||||
await asyncio.sleep(CHUNK_MS / 1000.0)
|
||||
silence_chunk = b"\x00" * CHUNK_BYTES
|
||||
silence_payload = base64.b64encode(silence_chunk).decode("ascii")
|
||||
for _ in range(TRAILING_SILENCE_MS // CHUNK_MS):
|
||||
await ws.send(json.dumps({"type": "input_audio_buffer.append", "audio": silence_payload}))
|
||||
await asyncio.sleep(CHUNK_MS / 1000.0)
|
||||
|
||||
|
||||
async def consume_until_response_done(
|
||||
ws,
|
||||
response_audio_out: bytearray,
|
||||
response_timeout_s: float,
|
||||
) -> dict:
|
||||
"""Read server events until response.done (or timeout). Returns turn summary."""
|
||||
info: dict = {
|
||||
"transcript_in": "",
|
||||
"transcript_out": "",
|
||||
"error": None,
|
||||
}
|
||||
deadline = time.monotonic() + response_timeout_s
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
raw = await asyncio.wait_for(ws.recv(), timeout=max(0.1, deadline - time.monotonic()))
|
||||
except asyncio.TimeoutError:
|
||||
info["error"] = "response_timeout"
|
||||
return info
|
||||
event = json.loads(raw)
|
||||
t = event.get("type", "")
|
||||
|
||||
if t == "conversation.item.input_audio_transcription.completed":
|
||||
info["transcript_in"] += event.get("transcript", "")
|
||||
elif t in ("response.audio.delta", "response.output_audio.delta"):
|
||||
delta = event.get("delta", "")
|
||||
if delta:
|
||||
response_audio_out.extend(base64.b64decode(delta))
|
||||
elif t in ("response.audio_transcript.delta", "response.output_audio_transcript.delta"):
|
||||
info["transcript_out"] += event.get("delta", "")
|
||||
elif t in ("response.audio_transcript.done", "response.output_audio_transcript.done"):
|
||||
transcript = event.get("transcript")
|
||||
if transcript:
|
||||
info["transcript_out"] = transcript
|
||||
elif t == "response.done":
|
||||
return info
|
||||
elif t == "error":
|
||||
info["error"] = event.get("error", {}).get("message", "error event")
|
||||
return info
|
||||
info["error"] = "response_timeout"
|
||||
return info
|
||||
|
||||
|
||||
def _truncate_transcript(s: str, max_len: int = 20) -> str:
|
||||
"""Quote *s*; if longer than *max_len*, slice + ellipsis + total char count."""
|
||||
if len(s) <= max_len:
|
||||
return repr(s)
|
||||
return f"{(s[:max_len] + '...')!r} ({len(s)} chars)"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Client driver (one per client)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _lb_allocate_session(
|
||||
lb_url: str,
|
||||
auth_headers: dict[str, str],
|
||||
http: httpx.AsyncClient,
|
||||
) -> dict:
|
||||
"""POST {lb_url}/session → returns dict with connect_url, session_id, session_token."""
|
||||
resp = await http.post(f"{lb_url}/session", headers=auth_headers, timeout=30.0)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
async def _lb_send_event(
|
||||
lb_url: str,
|
||||
session_id: str,
|
||||
session_token: str,
|
||||
event: str,
|
||||
http: httpx.AsyncClient,
|
||||
) -> None:
|
||||
"""POST {lb_url}/internal/sessions/{id}/event with {session_token, event}."""
|
||||
url = f"{lb_url}/internal/sessions/{session_id}/event"
|
||||
resp = await http.post(url, json={"session_token": session_token, "event": event}, timeout=30.0)
|
||||
resp.raise_for_status()
|
||||
|
||||
|
||||
async def run_client(
|
||||
client_id: int,
|
||||
args: argparse.Namespace,
|
||||
ws_url: str,
|
||||
extra_headers: list[tuple[str, str]],
|
||||
audio_files: list[tuple[str, Path]],
|
||||
) -> dict:
|
||||
"""Run one client's multi-turn conversation. Returns per-client summary.
|
||||
|
||||
If --lb-url is set, hits the load balancer first to allocate a slot on a
|
||||
compute node, then connects to the returned connect_url. Otherwise connects
|
||||
directly to ws_url. The LB callback events ("connected"/"disconnected") are
|
||||
sent around the websocket session so the LB doesn't reclaim the slot early.
|
||||
"""
|
||||
summary: dict = {
|
||||
"client_id": client_id,
|
||||
"connected": False,
|
||||
"rejected": False,
|
||||
"completed": 0,
|
||||
"errors": 0,
|
||||
"error_msg": None,
|
||||
}
|
||||
|
||||
log_dir = Path(args.log_dir) / f"client_{client_id:03d}"
|
||||
log_dir.mkdir(parents=True, exist_ok=True)
|
||||
transcript_log = log_dir / "conversation.txt"
|
||||
response_audio = bytearray()
|
||||
|
||||
prefix = f"[c{client_id}]"
|
||||
auth_headers = {k: v for k, v in extra_headers}
|
||||
lb_session_id: str | None = None
|
||||
lb_session_token: str | None = None
|
||||
|
||||
async with httpx.AsyncClient() as http:
|
||||
# Step 1: allocate a slot via the LB (if configured).
|
||||
if args.lb_url:
|
||||
try:
|
||||
alloc = await _lb_allocate_session(args.lb_url, auth_headers, http)
|
||||
except Exception as e: # noqa: BLE001 — diagnostic path
|
||||
summary["error_msg"] = f"LB /session failed: {type(e).__name__}: {e}"
|
||||
logger.info(f"{prefix} {summary['error_msg']}")
|
||||
return summary
|
||||
connect_url = alloc["connect_url"]
|
||||
lb_session_id = alloc["session_id"]
|
||||
lb_session_token = alloc["session_token"]
|
||||
logger.info(
|
||||
f"{prefix} LB allocated session_id={lb_session_id} "
|
||||
f"ws_url={alloc.get('websocket_url', '?')}"
|
||||
)
|
||||
else:
|
||||
connect_url = ws_url
|
||||
|
||||
try:
|
||||
async with websockets.connect(
|
||||
connect_url,
|
||||
max_size=2**24,
|
||||
additional_headers=extra_headers or None,
|
||||
) as ws:
|
||||
summary["connected"] = True
|
||||
|
||||
# Step 3: tell the LB the client actually connected (clears the
|
||||
# pending_timeout_s reaper). Best-effort — failure here logs but
|
||||
# doesn't abort the session.
|
||||
if lb_session_id and lb_session_token:
|
||||
try:
|
||||
await _lb_send_event(args.lb_url, lb_session_id, lb_session_token, "connected", http)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning(f"{prefix} LB 'connected' callback failed: {e}")
|
||||
|
||||
first = json.loads(await asyncio.wait_for(ws.recv(), timeout=10.0))
|
||||
if first.get("type") == "error" and "session_limit_reached" in str(first):
|
||||
summary["rejected"] = True
|
||||
summary["error_msg"] = first.get("error", {}).get("message", "limit reached")
|
||||
logger.warning(f"{prefix} REJECTED: {summary['error_msg']}")
|
||||
elif first.get("type") != "session.created":
|
||||
summary["error_msg"] = f"unexpected first event: {first.get('type')}"
|
||||
logger.error(f"{prefix} ERROR: {summary['error_msg']}")
|
||||
else:
|
||||
sess = first.get("session") or {}
|
||||
session_id = sess.get("id") or first.get("event_id") or "?"
|
||||
logger.info(f"{prefix} connected, session={session_id}")
|
||||
|
||||
with transcript_log.open("w") as log_f:
|
||||
log_f.write(f"# Synthetic conversation, client={client_id}, session={session_id}\n")
|
||||
log_f.write(f"# Target turns={args.turns}, interval={args.interval}s\n\n")
|
||||
|
||||
for turn_idx in range(args.turns):
|
||||
# Per-client prompt offset (coprime shift) so concurrent
|
||||
# clients send distinct prompts at every turn.
|
||||
prompt_idx = (turn_idx + client_id * PROMPT_SHIFT_PER_CLIENT) % len(audio_files)
|
||||
text, wav_path = audio_files[prompt_idx]
|
||||
|
||||
turn_start = time.monotonic()
|
||||
turn_audio = load_pcm16_mono_16k(wav_path)
|
||||
|
||||
logger.info(f"{prefix} turn {turn_idx + 1}/{args.turns} USER: {text!r}")
|
||||
log_f.write(f"[turn {turn_idx + 1}/{args.turns}] USER: {text}\n")
|
||||
|
||||
try:
|
||||
await stream_prompt(ws, turn_audio)
|
||||
except websockets.exceptions.ConnectionClosed as e:
|
||||
logger.info(
|
||||
f"{prefix} turn {turn_idx + 1}/{args.turns} "
|
||||
f"connection closed during send: {e}"
|
||||
)
|
||||
log_f.write(f"[turn {turn_idx + 1}/{args.turns}] CONNECTION_CLOSED: {e}\n")
|
||||
summary["error_msg"] = f"connection_closed: {e}"
|
||||
break
|
||||
|
||||
info = await consume_until_response_done(
|
||||
ws, response_audio, args.response_timeout
|
||||
)
|
||||
turn_elapsed = time.monotonic() - turn_start
|
||||
|
||||
if info["error"]:
|
||||
logger.info(
|
||||
f"{prefix} turn {turn_idx + 1}/{args.turns} ERROR: {info['error']}"
|
||||
)
|
||||
log_f.write(f"[turn {turn_idx + 1}/{args.turns}] ERROR: {info['error']}\n\n")
|
||||
summary["errors"] += 1
|
||||
else:
|
||||
summary["completed"] += 1
|
||||
logger.info(
|
||||
f"{prefix} turn {turn_idx + 1}/{args.turns} "
|
||||
f"ASSISTANT: {_truncate_transcript(info['transcript_out'])}"
|
||||
)
|
||||
log_f.write(f"[turn {turn_idx + 1}/{args.turns}] STT: {info['transcript_in']}\n")
|
||||
log_f.write(
|
||||
f"[turn {turn_idx + 1}/{args.turns}] ASSISTANT: {info['transcript_out']}\n\n"
|
||||
)
|
||||
log_f.flush()
|
||||
|
||||
remaining = args.interval - turn_elapsed
|
||||
if remaining > 0:
|
||||
await asyncio.sleep(remaining)
|
||||
|
||||
except Exception as e: # noqa: BLE001 — diagnostic path
|
||||
summary["error_msg"] = f"{type(e).__name__}: {e}"
|
||||
logger.error(f"{prefix} EXCEPTION: {summary['error_msg']}")
|
||||
finally:
|
||||
# Step 5: release the LB slot, regardless of how the ws session ended.
|
||||
if lb_session_id and lb_session_token:
|
||||
try:
|
||||
await _lb_send_event(
|
||||
args.lb_url, lb_session_id, lb_session_token, "disconnected", http
|
||||
)
|
||||
logger.info(f"{prefix} LB notified: disconnected")
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning(f"{prefix} LB 'disconnected' callback failed: {e}")
|
||||
|
||||
if response_audio:
|
||||
wav_out = log_dir / "conversation.wav"
|
||||
write_wav(wav_out, bytes(response_audio))
|
||||
logger.info(f"{prefix} wrote {len(response_audio)} bytes -> {wav_out}")
|
||||
return summary
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Top-level: synthesize prompts, spawn clients, summarize
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def run_all(args: argparse.Namespace) -> None:
|
||||
log_dir = Path(args.log_dir)
|
||||
log_dir.mkdir(parents=True, exist_ok=True)
|
||||
prompts_dir = log_dir / "prompts"
|
||||
prompts_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
logger.info(f"Synthesizing {len(PROMPTS)} prompts with `say` (cached in {prompts_dir})...")
|
||||
audio_files: list[tuple[str, Path]] = []
|
||||
for i, text in enumerate(PROMPTS):
|
||||
wav_path = prompts_dir / f"prompt_{i:03d}.wav"
|
||||
if not wav_path.exists():
|
||||
synthesize_with_say(text, wav_path)
|
||||
audio_files.append((text, wav_path))
|
||||
|
||||
ws_url = args.url or f"ws://{args.host}:{args.port}/v1/realtime"
|
||||
token = os.environ.get("HF_TOKEN")
|
||||
if not token:
|
||||
raise SystemExit(
|
||||
"HF_TOKEN env var is not set. Export it before running this script "
|
||||
"(e.g. `export HF_TOKEN=hf_...`)."
|
||||
)
|
||||
extra_headers: list[tuple[str, str]] = [("Authorization", f"Bearer {token}")]
|
||||
logger.info("Auth: Bearer token attached from HF_TOKEN env")
|
||||
if args.lb_url:
|
||||
target = f"LB {args.lb_url}"
|
||||
else:
|
||||
target = ws_url
|
||||
logger.info(
|
||||
f"Spawning {args.clients} client(s) against {target}, "
|
||||
f"{args.turns} turns each @ {args.interval:.1f}s interval"
|
||||
)
|
||||
|
||||
summaries = await asyncio.gather(
|
||||
*(
|
||||
run_client(
|
||||
client_id=i,
|
||||
args=args,
|
||||
ws_url=ws_url,
|
||||
extra_headers=extra_headers,
|
||||
audio_files=audio_files,
|
||||
)
|
||||
for i in range(args.clients)
|
||||
)
|
||||
)
|
||||
|
||||
logger.info("\n=== summary ===")
|
||||
for s in summaries:
|
||||
status = "rejected" if s["rejected"] else ("error" if s["error_msg"] else "ok")
|
||||
logger.info(
|
||||
f" c{s['client_id']}: {status:8s} completed={s['completed']}/{args.turns} "
|
||||
f"errors={s['errors']} err={s['error_msg']}"
|
||||
)
|
||||
n_ok = sum(1 for s in summaries if s["connected"] and not s["rejected"] and not s["error_msg"])
|
||||
n_rej = sum(1 for s in summaries if s["rejected"])
|
||||
n_err = sum(1 for s in summaries if s["error_msg"] and not s["rejected"])
|
||||
total_turns = sum(s["completed"] for s in summaries)
|
||||
logger.info(f"=> {n_ok} successful clients, {n_rej} rejected, {n_err} errored")
|
||||
logger.info(f"=> {total_turns} total turns completed across pool")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(levelname)s %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
stream=sys.stdout,
|
||||
)
|
||||
|
||||
if not shutil.which("say"):
|
||||
raise SystemExit("This script requires macOS `say` (not found in PATH).")
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lb-url",
|
||||
default=None,
|
||||
help=(
|
||||
"Base URL of the load balancer (e.g. https://lb.example.com). When set, the script "
|
||||
"POSTs /session to allocate a slot on a compute node, connects to the returned "
|
||||
"connect_url, and sends connected/disconnected events around the session. Takes "
|
||||
"precedence over --url / --host / --port."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--url",
|
||||
default=None,
|
||||
help=(
|
||||
"Full ws:// or wss:// URL to a compute-node realtime endpoint, e.g. "
|
||||
"wss://endpoint.example.com/v1/realtime. Used when --lb-url is not set. "
|
||||
"Overrides --host/--port."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--host", default="192.168.1.9", help="Used only when --url and --lb-url are unset.")
|
||||
parser.add_argument("--port", type=int, default=8765, help="Used only when --url and --lb-url are unset.")
|
||||
parser.add_argument("--clients", type=int, default=1, help="Number of parallel clients (default 1).")
|
||||
parser.add_argument("--turns", type=int, default=60, help="Turns per client (default 60).")
|
||||
parser.add_argument(
|
||||
"--interval",
|
||||
type=float,
|
||||
default=10.0,
|
||||
help="Seconds between turn starts per client. Total runtime ≈ turns × interval. Default 10.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--response-timeout",
|
||||
type=float,
|
||||
default=30.0,
|
||||
help="Per-turn wait for response.done before marking the turn errored.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-dir",
|
||||
default="/tmp/synthetic_conversation",
|
||||
help="Top-level directory for prompt cache and per-client logs.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
if args.clients < 1:
|
||||
parser.error("--clients must be >= 1")
|
||||
if args.turns < 1:
|
||||
parser.error("--turns must be >= 1")
|
||||
asyncio.run(run_all(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user