tts_player.py

from __future__ import annotations

import json
from pathlib import Path


_WORKLET_SOURCE = (Path(__file__).parent / "web" / "PCMPlayerWorklet.js").read_text()

TTS_PLAYER_TEMPLATE = """
<div class="audex-player">
  <div class="audex-player-status">Ready</div>
  <button class="audex-player-enable" type="button">Enable playback</button>
  <audio class="audex-player-audio" controls></audio>
  <a class="audex-player-download" hidden download="audex-tts.wav">Download WAV</a>
</div>
"""

TTS_PLAYER_CSS = """
.audex-player {
  display: grid;
  gap: 0.75rem;
  padding: 1rem;
  border: 1px solid var(--border-color-primary);
  border-radius: var(--radius-lg);
}
.audex-player-status {
  color: var(--body-text-color-subdued);
}
.audex-player-audio {
  width: 100%;
}
.audex-player-download {
  color: var(--link-text-color);
  font-weight: 600;
}
"""

TTS_PLAYER_JS = """
const WORKLET_SOURCE = __WORKLET_SOURCE__;
const status = element.querySelector(".audex-player-status");
const enable = element.querySelector(".audex-player-enable");
const audio = element.querySelector(".audex-player-audio");
const download = element.querySelector(".audex-player-download");
let context = null;
let player = null;
let chunks = [];
let wavUrl = null;
let lastSequence = -1;

const ensurePlayer = async () => {
  if (context) {
    await context.resume();
    return;
  }
  context = new AudioContext();
  const blob = new Blob([WORKLET_SOURCE], { type: "text/javascript" });
  const moduleUrl = URL.createObjectURL(blob);
  await context.audioWorklet.addModule(moduleUrl);
  URL.revokeObjectURL(moduleUrl);
  player = new AudioWorkletNode(
    context,
    "audex-pcm-player",
    { outputChannelCount: [1] },
  );
  player.connect(context.destination);
  player.port.onmessage = (event) => {
    if (event.data.type === "underrun") {
      status.textContent = "Buffering generated speech…";
    }
  };
  await context.resume();
};

const decodePcm = (encoded) => {
  const binary = atob(encoded);
  const bytes = new Uint8Array(binary.length);
  for (let index = 0; index < binary.length; index += 1) {
    bytes[index] = binary.charCodeAt(index);
  }
  return new Float32Array(bytes.buffer);
};

const resample = (input, sourceRate, targetRate) => {
  if (sourceRate === targetRate) return input;
  const output = new Float32Array(Math.max(1, Math.round(input.length * targetRate / sourceRate)));
  const ratio = sourceRate / targetRate;
  for (let index = 0; index < output.length; index += 1) {
    const position = index * ratio;
    const left = Math.floor(position);
    const right = Math.min(left + 1, input.length - 1);
    const fraction = position - left;
    output[index] = input[left] * (1 - fraction) + input[right] * fraction;
  }
  return output;
};

const buildWav = (parts, sampleRate) => {
  const length = parts.reduce((total, part) => total + part.length, 0);
  const buffer = new ArrayBuffer(44 + length * 2);
  const view = new DataView(buffer);
  const write = (offset, text) => {
    for (let index = 0; index < text.length; index += 1) {
      view.setUint8(offset + index, text.charCodeAt(index));
    }
  };
  write(0, "RIFF");
  view.setUint32(4, 36 + length * 2, true);
  write(8, "WAVE");
  write(12, "fmt ");
  view.setUint32(16, 16, true);
  view.setUint16(20, 1, true);
  view.setUint16(22, 1, true);
  view.setUint32(24, sampleRate, true);
  view.setUint32(28, sampleRate * 2, true);
  view.setUint16(32, 2, true);
  view.setUint16(34, 16, true);
  write(36, "data");
  view.setUint32(40, length * 2, true);
  let offset = 44;
  for (const part of parts) {
    for (const sample of part) {
      const clipped = Math.max(-1, Math.min(1, sample));
      view.setInt16(offset, clipped < 0 ? clipped * 32768 : clipped * 32767, true);
      offset += 2;
    }
  }
  return new Blob([buffer], { type: "audio/wav" });
};

const reset = async () => {
  chunks = [];
  lastSequence = -1;
  audio.removeAttribute("src");
  audio.load();
  download.hidden = true;
  if (wavUrl) URL.revokeObjectURL(wavUrl);
  wavUrl = null;
  await ensurePlayer();
  player.port.postMessage({ type: "reset" });
  status.textContent = "Waiting for speech tokens…";
};

const consume = async () => {
  const value = props.value || {};
  if (value.sequence === lastSequence) return;
  lastSequence = value.sequence;
  if (value.reset) await reset();
  if (value.pcm) {
    await ensurePlayer();
    const pcm = decodePcm(value.pcm);
    chunks.push(pcm.slice());
    const playback = resample(pcm, value.sample_rate || 16000, context.sampleRate);
    player.port.postMessage({ type: "audio", samples: playback.buffer }, [playback.buffer]);
    status.textContent = `Streaming ${value.token_count || 0} speech tokens…`;
  }
  if (value.done && chunks.length) {
    const wav = buildWav(chunks, value.sample_rate || 16000);
    wavUrl = URL.createObjectURL(wav);
    audio.src = wavUrl;
    download.href = wavUrl;
    download.hidden = false;
    status.textContent = `Complete · ${value.token_count || 0} speech tokens`;
  }
};

const runButton = document.getElementById("audex-run");
if (runButton) runButton.addEventListener("click", () => ensurePlayer());
enable.addEventListener("click", async () => {
  await ensurePlayer();
  status.textContent = "Playback enabled";
});
const stopButton = document.getElementById("audex-stop");
if (stopButton) {
  stopButton.addEventListener("click", () => {
    if (player) player.port.postMessage({ type: "reset" });
    status.textContent = "Stopped";
  });
}
watch("value", consume);
consume();
""".replace("__WORKLET_SOURCE__", json.dumps(_WORKLET_SOURCE))


def player_value(
    sequence: int,
    *,
    pcm: str | None = None,
    token_count: int = 0,
    done: bool = False,
    reset: bool = False,
) -> dict[str, object]:
    return {
        "sequence": sequence,
        "sample_rate": 16000,
        "pcm": pcm,
        "token_count": token_count,
        "done": done,
        "reset": reset,
    }