Live 1 Realtime Diarization

1import asyncio
2import json
3import os
4import subprocess
5
6import websockets
7
8MODEL_ID = ""  # Add Model ID from Baseten Dashboard
9
10endpoint_url = (
11    f"wss://model-{MODEL_ID}.api.baseten.co"
12    "/environments/production/websocket"
13)
14headers = {"Authorization": f"Api-Key {os.getenv('BASETEN_API_KEY')}"}
15audio_url = (
16    "https://test-audios-public.s3.us-west-2.amazonaws.com"
17    "/10-sec-01-podcast.m4a"
18)
19
20# 100 ms of 16 kHz mono float32 PCM; every frame must be exactly this size
21CHUNK_BYTES = 6400
22
23# Decode any audio format to raw 16 kHz mono float32 PCM.
24pcm = subprocess.run(
25    ["ffmpeg", "-loglevel", "error", "-i", audio_url,
26     "-f", "f32le", "-ac", "1", "-ar", "16000", "-"],
27    capture_output=True, check=True,
28).stdout
29
30
31async def send_audio(ws):
32    for i in range(0, len(pcm), CHUNK_BYTES):
33        await ws.send(pcm[i : i + CHUNK_BYTES].ljust(CHUNK_BYTES, b"\0"))
34        await asyncio.sleep(0.1)  # stream at real-time pace
35    # Flushes in-flight audio; the server sends the remaining events,
36    # then closes the socket.
37    await ws.send(json.dumps({"type": "end_of_stream"}))
38
39
40async def main():
41    # A scaled-to-zero deployment takes a few minutes to accept its
42    # first connection.
43    async with websockets.connect(
44        endpoint_url, additional_headers=headers, open_timeout=300
45    ) as ws:
46        sender = asyncio.create_task(send_audio(ws))
47        # Ends when the server closes after end_of_stream.
48        async for message in ws:
49            print(message)
50        await sender
51
52
53asyncio.run(main())