Chat Mode (Turn-based Chat)

{
  "role": "system | user | assistant",
  "content": "plain text string"
}

{
  "role": "user",
  "content": [
    {"type": "text", "text": "Describe this image"},
    {"type": "image", "data": "<base64>"},
    {"type": "audio", "data": "<base64 PCM float32>", "sample_rate": 16000},
    {"type": "video", "data": "<base64 video file>", "stack_frames": 1}
  ]
}

{
  "messages": [
    {"role": "user", "content": "Hello!"}
  ],
  "streaming": true,
  "generation": {"max_new_tokens": 256, "length_penalty": 1.1},
  "tts": {"enabled": true, "ref_audio_data": "<base64>"},
  "image": {"max_slice_nums": null},
  "omni_mode": false,
  "enable_thinking": false
}

{
  "type": "chunk",
  "text_delta": "Hello",
  "audio_data": "<base64, 24kHz>"
}

{
  "type": "done",
  "text": "Hello! How can I help you?",
  "generated_tokens": 15,
  "input_tokens": 42,
  "audio_data": "<base64, 24kHz>",
  "recording_session_id": "chat_abc123"
}

// -- Reference audio for voice cloning (base64 PCM float32, 16kHz) --
// Loaded from an uploaded file or the server default (/api/default_ref_audio).
const refAudioBase64 = getRefAudioBase64();

const ws = new WebSocket(`wss://${location.host}/ws/chat`);
let fullText = '';

ws.onopen = () => {
  const payload = {
    messages: [
      // System message uses content-list format: [text, audio, text].
      // The audio item embeds the reference voice into the LLM context.
      { role: 'system', content: [
        { type: 'text', text: 'Mimic the voice from the audio sample.' },
        { type: 'audio', data: refAudioBase64 },          // reference voice
        { type: 'text', text: 'You are a helpful assistant. Reply naturally.' },
      ]},
      { role: 'user', content: 'Hello!' },
    ],
    streaming: true,
    generation: { max_new_tokens: 256, temperature: 0.7 },
    tts: {
      enabled: true,
      mode: 'audio_assistant',
      ref_audio_data: refAudioBase64,  // same audio initializes TTS vocoder
    },
  };
  ws.send(JSON.stringify(payload));
};

ws.onmessage = (event) => {
  const msg = JSON.parse(event.data);
  switch (msg.type) {
    case 'prefill_done':
      // Server finished prefilling the prompt into KV Cache
      console.log(`Prefill done, ${msg.input_tokens} input tokens`);
      break;
    case 'chunk':
      // Streaming token: incremental text and/or audio segment
      if (msg.text_delta) {
        fullText += msg.text_delta;
        console.log('Streaming:', fullText);
      }
      if (msg.audio_data) playAudio(msg.audio_data);  // PCM float32, 24kHz
      break;
    case 'done':
      // Generation complete — full text and token stats
      console.log('Final text:', msg.text);
      console.log(`Tokens: ${msg.input_tokens} in, ${msg.generated_tokens} out`);
      ws.close();
      break;
    case 'error':
      console.error('Error:', msg.error);
      ws.close();
      break;
  }
};

ws.onerror = () => console.error('WebSocket error');
ws.onclose = () => console.log('Connection closed');

import asyncio, json, base64
import numpy as np
import websockets

def load_ref_audio(path: str) -> str:
    """Load a WAV file and return base64-encoded PCM float32 at 16kHz."""
    import soundfile as sf
    audio, _ = sf.read(path, dtype="float32", samplerate=16000)
    return base64.b64encode(audio.tobytes()).decode()

async def chat_streaming(
    server_url="wss://localhost:8006/ws/chat",
    ref_audio_path: str | None = "ref.wav",
):
    async with websockets.connect(server_url) as ws:
        # Load reference audio for voice cloning
        ref_b64 = load_ref_audio(ref_audio_path) if ref_audio_path else None

        # System message uses content-list format: [text, audio, text].
        # The audio item embeds the reference voice into the LLM context.
        system_content = "You are a helpful assistant."
        if ref_b64:
            system_content = [
                {"type": "text", "text": "Mimic the voice from the audio sample."},
                {"type": "audio", "data": ref_b64},          # reference voice
                {"type": "text", "text": "You are a helpful assistant. Reply naturally."},
            ]

        # TTS config — ref_audio_data initializes the TTS vocoder separately
        tts_config = {"enabled": True, "mode": "audio_assistant"}
        if ref_b64:
            tts_config["ref_audio_data"] = ref_b64

        await ws.send(json.dumps({
            "messages": [
                {"role": "system", "content": system_content},
                {"role": "user", "content": "Hello!"},
            ],
            "streaming": True,
            "generation": {"max_new_tokens": 256, "temperature": 0.7},
            "tts": tts_config,
        }))

        full_text = ""
        audio_chunks = []

        async for raw in ws:
            msg = json.loads(raw)
            if msg["type"] == "prefill_done":
                print(f"Prefill done, {msg['input_tokens']} input tokens")
            elif msg["type"] == "chunk":
                # Incremental text token
                if msg.get("text_delta"):
                    full_text += msg["text_delta"]
                    print(f"Streaming: {full_text}")
                # Incremental audio segment (PCM float32, 24kHz)
                if msg.get("audio_data"):
                    audio_chunks.append(base64.b64decode(msg["audio_data"]))
            elif msg["type"] == "done":
                print(f"Final: {msg['text']}")
                print(f"Tokens: {msg['input_tokens']} in, {msg['generated_tokens']} out")
                break
            elif msg["type"] == "error":
                print(f"Error: {msg['error']}")
                break

        if audio_chunks:
            pcm = np.frombuffer(b"".join(audio_chunks), dtype=np.float32)
            print(f"Received {len(pcm)/24000:.1f}s of audio")

asyncio.run(chat_streaming())