WebSocket API

import asyncio
import aiohttp
import json
import base64
import os


async def generate_speech():
    api_key = os.environ.get("TOGETHER_API_KEY")
    url = (
        "wss://api.together.ai/v1/audio/speech"
        "/websocket?model=hexgrad/Kokoro-82M"
        "&voice=af_alloy"
        "&response_format=pcm"
        "&sample_rate=24000"
    )

    headers = {"Authorization": f"Bearer {api_key}"}

    text_chunks = [
        "Hello, this is a test.",
        "This is the second sentence.",
        "And this is the final one.",
    ]

    audio_chunks = []

    async with aiohttp.ClientSession(headers=headers) as session:
        async with session.ws_connect(url) as ws:
            # Wait for session.created
            msg = await ws.receive()
            session_data = json.loads(msg.data)
            print(f"Session created: {session_data['session']['id']}")

            async def send_text():
                for chunk in text_chunks:
                    await ws.send_json(
                        {
                            "type": "input_text_buffer.append",
                            "text": chunk,
                        }
                    )
                    print(f"Sent: {chunk}")
                    await asyncio.sleep(0.5)
                await ws.send_json({"type": "input_text_buffer.commit"})
                print("Committed")

            async def receive_audio():
                async for msg in ws:
                    if msg.type == aiohttp.WSMsgType.TEXT:
                        data = json.loads(msg.data)
                        mtype = data.get("type", "")

                        if mtype == "conversation.item.audio_output.delta":
                            chunk = base64.b64decode(data.get("delta", ""))
                            audio_chunks.append(chunk)

                        elif mtype == "conversation.item.word_timestamps":
                            words = data.get("words", [])
                            starts = data.get("start_seconds", [])
                            stamps = list(
                                zip(words, [f"{s:.2f}s" for s in starts])
                            )
                            print(f"  timestamps: {stamps}")

                        elif mtype in (
                            "error",
                            "conversation.item.tts.failed",
                        ):
                            err = data.get(
                                "error",
                                data.get("message"),
                            )
                            print(f"Error: {err}")
                            return

                    elif msg.type in (
                        aiohttp.WSMsgType.CLOSE,
                        aiohttp.WSMsgType.CLOSED,
                    ):
                        break

            send_task = asyncio.create_task(send_text())
            recv_task = asyncio.create_task(receive_audio())

            await send_task

            # Wait up to 10s for audio to stop arriving
            deadline = asyncio.get_event_loop().time() + 10
            while asyncio.get_event_loop().time() < deadline:
                await asyncio.sleep(0.1)
                n = len(audio_chunks)
                await asyncio.sleep(0.3)
                if len(audio_chunks) == n:
                    break

            recv_task.cancel()
            try:
                await recv_task
            except asyncio.CancelledError:
                pass

    if audio_chunks:
        pcm = b"".join(audio_chunks)
        with open("output.pcm", "wb") as f:
            f.write(pcm)
        print(
            f"\nAudio saved to output.pcm ({len(pcm):,} bytes, "
            f"{len(pcm)/48000:.1f}s at 24kHz)"
        )
        print("Play with: ffplay -f s16le -ar 24000 output.pcm")
    else:
        print("No audio received")


asyncio.run(generate_speech())