{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# OmniVoice Quick Start\n\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/k2-fsa/OmniVoice/blob/master/docs/OmniVoice.ipynb)\n\nThis notebook demonstrates the basic usage of [OmniVoice](https://github.com/k2-fsa/OmniVoice), a massively multilingual zero-shot TTS model supporting 600+ languages.\n\n**Contents:**\n1. Installation\n2. Option A — Gradio Demo (interactive web UI, no code needed)\n3. Option B — Python API\n   - 3.1 Load Model\n   - 3.2 Voice Cloning\n   - 3.3 Voice Design\n   - 3.4 Auto Voice"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installation\n",
    "\n",
    "Colab already provides a compatible PyTorch + CUDA environment, so we only need to install OmniVoice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install omnivoice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 2. Option A — Gradio Demo\n\nLaunch an interactive web UI with a public Gradio link. The `--share` flag creates a temporary public URL so you can access the demo from any browser.\n\n> **If you prefer to use the Python API directly, skip to Option B below.**"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!omnivoice-demo --share"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 3. Option B — Python API\n\n### 3.1 Load Model"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "from omnivoice import OmniVoice\nimport soundfile as sf\nimport torch\nfrom IPython.display import Audio, display\n\nmodel = OmniVoice.from_pretrained(\n    \"k2-fsa/OmniVoice\",\n    device_map=\"cuda:0\",\n    dtype=torch.float16,\n    load_asr=True,\n)"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "### 3.2 Voice Cloning\n\nClone a voice from a short (3-10s) reference audio clip. Upload your own `ref.wav` or use any audio file.\n\n`ref_text` is optional — if omitted, the model uses Whisper ASR to auto-transcribe it."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import files\n",
    "\n",
    "print(\"Upload a reference audio file (wav/mp3/flac):\")\n",
    "uploaded = files.upload()\n",
    "ref_audio_path = list(uploaded.keys())[0]\n",
    "print(f\"Uploaded: {ref_audio_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = model.generate(\n",
    "    text=\"Hello, this is a test of zero-shot voice cloning.\",\n",
    "    ref_audio=ref_audio_path,\n",
    "    # ref_text=\"Transcription of the reference audio.\",  # optional\n",
    ")\n",
    "\n",
    "sf.write(\"clone_out.wav\", audio[0], 24000)\n",
    "display(Audio(audio[0], rate=24000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "### 3.3 Voice Design\n\nDescribe the desired voice with speaker attributes — no reference audio needed.\n\nSupported attributes: gender, age, pitch, style (whisper), English accent, Chinese dialect. See [docs/voice-design.md](https://github.com/k2-fsa/OmniVoice/blob/master/docs/voice-design.md) for the full list."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = model.generate(\n",
    "    text=\"Hello, this is a test of zero-shot voice design.\",\n",
    "    instruct=\"female, low pitch, british accent\",\n",
    ")\n",
    "\n",
    "sf.write(\"design_out.wav\", audio[0], 24000)\n",
    "display(Audio(audio[0], rate=24000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "### 3.4 Auto Voice\n\nLet the model choose a voice automatically — no reference audio or instruct needed."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = model.generate(\n",
    "    text=\"This is a sentence generated with automatic voice selection.\",\n",
    ")\n",
    "\n",
    "sf.write(\"auto_out.wav\", audio[0], 24000)\n",
    "display(Audio(audio[0], rate=24000))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}