vibevoice-tts

/**
 * @file index.ts
 * @description VibeVoice TTS plugin — LM Studio plugin entry point.
 *
 * Registers two tools:
 *   - `format`: Basic text cleaning + multi-speaker script formatting
 *   - `prepare`: Full podcast script preparation pipeline
 *
 * VibeVoice does NOT support SSML. This plugin prepares plain text for
 * VibeVoice, handling:
 *   - Text cleaning (smart quotes, emoji, HTML, URLs, annotations)
 *   - Multi-speaker script formatting ("Speaker N:" prefix format)
 *   - Voice/emotion annotations per speaker (for voice reference selection)
 *   - JSON array input for structured speaker data
 *   - Abbreviation expansion, pause markers, direction tags, chunking
 *
 * This plugin does NOT generate audio — it prepares formatted text for
 * downstream VibeVoice TTS consumption.
 */

import {
  tool,
  type Tool,
  type ToolsProviderController,
  type PluginContext,
} from "@lmstudio/sdk";
import { z } from "zod";

import { processText } from "./processor";
import { prepareScript } from "./preparer";
import { narrateScript } from "./narrator";
import type { VibeVoiceResult, PrepareResult, NarrateResult } from "./types";

// ---------------------------------------------------------------------------
// Zod schema for the format tool parameters
// ---------------------------------------------------------------------------

const FormatParams = z.object({
  /** The input text to format for VibeVoice TTS. */
  text: z
    .string()
    .min(1)
    .describe(
      "The text to format for VibeVoice TTS. " +
        'Supports plain text, multi-speaker "Speaker N:" prefix format, ' +
        "or a JSON array of {speaker, text} objects.",
    ),

  /** Speaker mode */
  mode: z
    .enum(["single", "multi", "auto"])
    .optional()
    .default("auto")
    .describe(
      "Speaker mode. " +
        '"single": one speaker, ignores "Speaker N:" prefixes. ' +
        '"multi": parse "Speaker N:" prefixes (warns if none found). ' +
        '"auto": detect prefixes automatically; fallback to single speaker.',
    ),

  /** Number of speakers */
  speakers: z
    .number()
    .int()
    .min(1)
    .max(4)
    .optional()
    .default(1)
    .describe(
      "Number of speakers (1-4). Used for voice/emotion array indexing. " +
        "Default: 1.",
    ),

  /** Voice names per speaker */
  voices: z
    .array(z.string())
    .optional()
    .describe(
      'Voice names per speaker, e.g. ["Alice", "Frank"]. ' +
        "Used as documentation — the user selects the matching voice reference " +
        "audio for each VibeVoice speaker slot.",
    ),

  /** Emotion labels per speaker */
  emotions: z
    .array(
      z.enum([
        "neutral",
        "happy",
        "sad",
        "angry",
        "excited",
        "calm",
        "thoughtful",
        "whisper",
        "serious",
        "energetic",
        "warm",
        "cold",
        "mysterious",
        "dramatic",
      ]),
    )
    .optional()
    .describe(
      "Emotion labels per speaker for voice reference selection. " +
        "VibeVoice does NOT support inline emotion tags — emotion comes from " +
        "the reference audio selected for each speaker slot. " +
        "E.g. Speaker 1 = calm, Speaker 2 = excited.",
    ),

  /** Emotion-to-speaker-slot mapping (PsiPi method) */
  emotion_map: z
    .record(
      z.enum([
        "neutral",
        "happy",
        "sad",
        "angry",
        "excited",
        "calm",
        "thoughtful",
        "whisper",
        "serious",
        "energetic",
        "warm",
        "cold",
        "mysterious",
        "dramatic",
      ]),
      z.number().int().min(1).max(4),
    )
    .optional()
    .describe(
      "Map emotions to speaker slots (PsiPi method). Lines annotated with " +
        "[emotion] or (emotion) are stripped of the tag and reassigned to " +
        "the mapped slot. Default: neutral/calm to 1, happy/sad to 2, " +
        "angry/excited to 3. Use the same voice with different emotional " +
        "reference clips per slot to vary emotion.",
    ),

  /** Apply text cleaning */
  clean: z
    .boolean()
    .optional()
    .default(true)
    .describe(
      "Apply the VibeVoice text-cleaning pipeline. " +
        "Strips smart quotes, Chinese quotation marks, emoji, HTML/XML tags, " +
        "URLs (replaced with [URL]), non-speech annotations, " +
        "tabs, and excess whitespace. Default: true.",
    ),

  /** Output format */
  output: z
    .enum(["plain", "script", "json"])
    .optional()
    .default("script")
    .describe(
      "Output format. " +
        '"plain": just the cleaned text without speaker prefixes. ' +
        '"script": "Speaker N:" prefix format (suitable for VibeVoice). ' +
        '"json": structured array of {speaker, text} objects.',
    ),
});

type FormatParamsType = z.infer<typeof FormatParams>;

// ---------------------------------------------------------------------------
// Zod schema for the prepare tool parameters
// ---------------------------------------------------------------------------

const PrepareParams = z.object({
  text: z
    .string()
    .min(1)
    .describe(
      "The raw podcast script text to prepare for VibeVoice TTS. " +
        "Supports plain text, multi-speaker Speaker N: format, " +
        "pause markers like [pause 1s], and direction tags like [emphasis].",
    ),

  chunkSize: z
    .number()
    .int()
    .min(100)
    .max(2000)
    .optional()
    .default(500)
    .describe(
      "Target words per chunk for incremental TTS generation. " +
        "500 words is optimal for VibeVoice quality. Default: 500.",
    ),

  paragraphPauses: z
    .boolean()
    .optional()
    .default(true)
    .describe(
      "Insert empty-line pauses between paragraphs. " +
        "VibeVoice interprets empty lines as natural pauses. Default: true.",
    ),

  sentencePauses: z
    .boolean()
    .optional()
    .default(false)
    .describe(
      "Split long paragraphs into shorter ones (max 3 sentences each). " +
        "Improves TTS quality for dense text. Default: false.",
    ),

  maxSentencesPerParagraph: z
    .number()
    .int()
    .min(1)
    .max(10)
    .optional()
    .default(3)
    .describe(
      "Maximum sentences per paragraph before auto-splitting. " +
        "Only used when sentencePauses=true. Default: 3.",
    ),

  expandAbbreviations: z
    .boolean()
    .optional()
    .default(true)
    .describe(
      "Expand common abbreviations (e.g. becomes for example, AI becomes A I). " +
        "First occurrence only. Improves TTS pronunciation. Default: true.",
    ),

  preserveDirections: z
    .boolean()
    .optional()
    .default(true)
    .describe(
      "Extract direction tags ([emphasis], [slower], etc.) as metadata " +
        "hints instead of stripping them. Guides TTS param adjustments. Default: true.",
    ),

  output: z
    .enum(["plain", "script", "json"])
    .optional()
    .default("script")
    .describe(
      "Output format. " +
        '"plain": just the prepared text. ' +
        '"script": Speaker N: prefix format. ' +
        '"json": structured chunk array. Default: script.',
    ),
});

type PrepareParamsType = z.infer<typeof PrepareParams>;

// ---------------------------------------------------------------------------
// Zod schema for the narrate tool parameters
// ---------------------------------------------------------------------------

const NarrateParams = z.object({
  text: z
    .string()
    .min(1)
    .describe(
      "The raw text to narrate for VibeVoice TTS. " +
        "Supports plain text, multi-speaker format, " +
        "expression tags like <laughs>, and direction tags.",
    ),

  style: z
    .enum([
      "storytelling",
      "children",
      "christmas-tale",
      "news",
      "technical",
      "dramatic",
    ])
    .optional()
    .default("storytelling")
    .describe(
      "Narration style preset. " +
        '"storytelling": balanced pauses, 0.9x speed, 145 WPM. ' +
        '"children": longer pauses, 0.85x speed, 130 WPM. ' +
        '"christmas-tale": warm pauses, 0.88x speed, 135 WPM. ' +
        '"news": short pauses, 1.0x speed, 170 WPM. ' +
        '"technical": moderate pauses, 0.95x speed, 155 WPM. ' +
        '"dramatic": variable pauses, 0.92x speed, 140 WPM.',
    ),

  mood: z
    .enum([
      "warm",
      "mysterious",
      "energetic",
      "somber",
      "playful",
      "serious",
      "neutral",
    ])
    .optional()
    .default("neutral")
    .describe(
      "Mood modifier within the style. Adjusts pause patterns and direction tag density. Default: neutral.",
    ),

  language: z
    .enum(["en", "pl", "de", "fr", "es", "it", "pt", "ja", "ko", "zh"])
    .optional()
    .default("en")
    .describe(
      "Language code for language-specific processing (abbreviations, foreign words). Default: en.",
    ),

  chunkSize: z
    .number()
    .int()
    .min(100)
    .max(2000)
    .optional()
    .default(500)
    .describe(
      "Target words per chunk for incremental TTS generation. Default: 500.",
    ),

  expandAbbreviations: z
    .boolean()
    .optional()
    .default(true)
    .describe(
      "Expand language-specific abbreviations (e.g., np. → na przykład). Default: true.",
    ),

  phoneticHints: z
    .boolean()
    .optional()
    .default(true)
    .describe(
      "Insert phonetic hints for foreign words and acronyms. Default: true.",
    ),

  detectStructure: z
    .boolean()
    .optional()
    .default(true)
    .describe(
      "Detect narrative structure (dialogue, scenes, chapters). Default: true.",
    ),

  autoExpressions: z
    .boolean()
    .optional()
    .default(true)
    .describe(
      "Auto-insert expression tags based on narrative context. Default: true.",
    ),

  output: z
    .enum(["plain", "script", "json"])
    .optional()
    .default("script")
    .describe(
      "Output format. plain: just text. script: Speaker N: format. json: structured chunks. Default: script.",
    ),
});

type NarrateParamsType = z.infer<typeof NarrateParams>;
// ---------------------------------------------------------------------------
// Plugin entry point
// ---------------------------------------------------------------------------

export async function main(context: PluginContext): Promise<void> {
  context.withToolsProvider(toolsProvider);
}

// ---------------------------------------------------------------------------
// Tools provider — registers the `format` and `prepare` tools
// ---------------------------------------------------------------------------

/**
 * Register the format and prepare tools with the LM Studio plugin runtime.
 */
export async function toolsProvider(
  _ctl: ToolsProviderController,
): Promise<Tool[]> {
  const formatTool = tool({
    name: "format",
    description:
      "Format text for VibeVoice TTS (Text-To-Speech) engine output.\n\n" +
      "VibeVoice does NOT support SSML — this tool prepares clean, structured " +
      "plain text in the multi-speaker script format that VibeVoice consumes natively.\n\n" +
      "CAPABILITIES:\n" +
      "  - Text cleaning — strips smart quotes, emoji, HTML tags, URLs, " +
      "non-speech annotations, tabs, excess whitespace\n" +
      '  - Multi-speaker detection — parses "Speaker N:" prefix format ' +
      "or accepts JSON array input\n" +
      "  - Voice/emotion annotations — documents which voice reference " +
      "audio to use per speaker slot\n" +
      '  - Output formats — plain text, script with "Speaker N:" prefixes, ' +
      "or structured JSON\n\n" +
      "EMOTION NOTE:\n" +
      "VibeVoice does NOT support inline emotion tags. Emotion is controlled " +
      "by the reference audio sample selected for each speaker. " +
      "The `emotions` parameter is documentation — it tells the user which " +
      "emotion-level voice to assign to each speaker slot.\n\n" +
      "SPEAKER FORMAT:\n" +
      "  Speaker 1: Hello, how are you?\n" +
      "  Speaker 2: I'm doing great!\n" +
      "  Speaker 1: That's wonderful to hear.\n\n" +
      "This tool does NOT generate audio — it prepares formatted text " +
      "input for VibeVoice TTS engines.",

    parameters: {
      text: FormatParams.shape.text,
      mode: FormatParams.shape.mode,
      speakers: FormatParams.shape.speakers,
      voices: FormatParams.shape.voices,
      emotions: FormatParams.shape.emotions,
      clean: FormatParams.shape.clean,
      output: FormatParams.shape.output,
    },

    implementation: async (
      params: FormatParamsType,
    ): Promise<VibeVoiceResult> => {
      const { text, mode, speakers, voices, emotions, clean, output } = params;

      return processText(text, {
        mode: mode ?? "auto",
        speakers: speakers ?? 1,
        voices: voices ?? [],
        emotions: emotions ?? [],
        clean: clean ?? true,
        output: output ?? "script",
        emotion_map: params.emotion_map,
      });
    },
  });

  const prepareTool = tool({
    name: "prepare",
    description:
      "Prepare a podcast script for VibeVoice TTS generation.\n\n" +
      "This is the FULL preparation pipeline — beyond basic text cleaning, " +
      "it handles everything needed for broadcast-quality TTS output:\n\n" +
      "PREPARATION STEPS:\n" +
      "  - Abbreviation expansion — e.g. becomes for example, AI becomes A I (first use)\n" +
      "  - Direction tag extraction — [emphasis], [slower] become metadata hints\n" +
      "  - Pause marker conversion — [pause 1s] becomes empty lines (VibeVoice pauses)\n" +
      "  - Paragraph splitting — long paragraphs become 2-3 sentence chunks\n" +
      "  - Script chunking — ~500 words per chunk for optimal quality\n" +
      "  - TTS recommendations — temperature, speed, WPM metadata\n\n" +
      "PAUSE MARKERS:\n" +
      "  [pause 0.3s] — between sentences\n" +
      "  [pause 0.5s] — between paragraphs\n" +
      "  [pause 1s]   — between sections\n\n" +
      "DIRECTION TAGS (metadata, not sent to TTS):\n" +
      "  [emphasis] [slower] [faster] [whisper] [louder] [softer]\n\n" +
      "TTS RECOMMENDATIONS:\n" +
      "  Narration: temperature 0.6-0.7, speed 0.95x, 155 WPM\n" +
      "  Dialogue:  temperature 0.8-0.9, speed 0.95x, 160 WPM\n\n" +
      "WORKFLOW: prepare, then generate each chunk separately, " +
      "concatenate with 200ms crossfade, add intro/outro, publish.",

    parameters: {
      text: PrepareParams.shape.text,
      chunkSize: PrepareParams.shape.chunkSize,
      paragraphPauses: PrepareParams.shape.paragraphPauses,
      sentencePauses: PrepareParams.shape.sentencePauses,
      maxSentencesPerParagraph: PrepareParams.shape.maxSentencesPerParagraph,
      expandAbbreviations: PrepareParams.shape.expandAbbreviations,
      preserveDirections: PrepareParams.shape.preserveDirections,
      output: PrepareParams.shape.output,
    },

    implementation: async (
      params: PrepareParamsType,
    ): Promise<PrepareResult> => {
      const {
        text,
        chunkSize,
        paragraphPauses,
        sentencePauses,
        maxSentencesPerParagraph,
        expandAbbreviations,
        preserveDirections,
        output,
      } = params;

      return prepareScript(text, {
        chunkSize: chunkSize ?? 500,
        paragraphPauses: paragraphPauses ?? true,
        sentencePauses: sentencePauses ?? false,
        maxSentencesPerParagraph: maxSentencesPerParagraph ?? 3,
        expandAbbreviations: expandAbbreviations ?? true,
        preserveDirections: preserveDirections ?? true,
        output: output ?? "script",
      });
    },
  });

  const narrateTool = tool({
    name: "narrate",
    description:
      "Prepare text for VibeVoice TTS with style-aware narration.\n\n" +
      "This tool adds NARRATION STYLE, MOOD, and LANGUAGE support:\n\n" +
      "STYLE PRESETS:\n" +
      "  - storytelling: balanced pauses, 0.9x speed, 145 WPM\n" +
      "  - children: longer pauses, 0.85x speed, 130 WPM\n" +
      "  - christmas-tale: warm pauses, 0.88x speed, 135 WPM\n" +
      "  - news: short pauses, 1.0x speed, 170 WPM\n" +
      "  - technical: moderate pauses, 0.95x speed, 155 WPM\n" +
      "  - dramatic: variable pauses, 0.92x speed, 140 WPM\n\n" +
      "MOOD MODIFIERS:\n" +
      "  warm, mysterious, energetic, somber, playful, serious, neutral\n\n" +
      "LANGUAGE PACKS:\n" +
      "  - Polish (pl): abbreviations (np., itd., m.in., tzn.), titles, addresses\n" +
      "  - English (en): baseline abbreviations\n" +
      "  - More languages configurable\n\n" +
      "EXPRESSION TAGS:\n" +
      "  VibeVoice supports expression tags like <laughs>, <sighs>, <giggles>\n" +
      "  Any <tag> pattern is preserved — model can invent new ones freely\n" +
      "  Convert [expression] to <expression> automatically\n\n" +
      "NARRATIVE DETECTION:\n" +
      "  - Dialogue detection (quotes, speaker changes)\n" +
      "  - Scene changes, chapter boundaries\n" +
      "  - Transition phrases, emotional intensity\n" +
      "  - Auto-pause insertion based on structure\n\n" +
      "WORKFLOW: narrate with style, then generate each chunk separately,\n" +
      "concatenate with 200ms crossfade, add intro/outro, publish.",

    parameters: {
      text: NarrateParams.shape.text,
      style: NarrateParams.shape.style,
      mood: NarrateParams.shape.mood,
      language: NarrateParams.shape.language,
      chunkSize: NarrateParams.shape.chunkSize,
      expandAbbreviations: NarrateParams.shape.expandAbbreviations,
      phoneticHints: NarrateParams.shape.phoneticHints,
      detectStructure: NarrateParams.shape.detectStructure,
      autoExpressions: NarrateParams.shape.autoExpressions,
      output: NarrateParams.shape.output,
    },

    implementation: async (
      params: NarrateParamsType,
    ): Promise<NarrateResult> => {
      const {
        text,
        style,
        mood,
        language,
        chunkSize,
        expandAbbreviations,
        phoneticHints,
        detectStructure,
        autoExpressions,
        output,
      } = params;

      return narrateScript({
        text,
        style: style ?? "storytelling",
        mood: mood ?? "neutral",
        language: language ?? "en",
        chunkSize: chunkSize ?? 500,
        expandAbbreviations: expandAbbreviations ?? true,
        phoneticHints: phoneticHints ?? true,
        detectStructure: detectStructure ?? true,
        autoExpressions: autoExpressions ?? true,
        output: output ?? "script",
      });
    },
  });

  return [formatTool, prepareTool, narrateTool];
}