asgan_quickstart_colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZRZ67uyEWiMG"
      },
      "outputs": [],
      "source": [
        "!pip install fairseq librosa fastprogress"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install --upgrade omegaconf"
      ],
      "metadata": {
        "id": "S5YpRt4YcPdM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# remember to restart runtime after installing dependencies.\n",
        "import torch\n",
        "from torch import Tensor\n",
        "import IPython.display as ipd\n",
        "import librosa.display\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "model = torch.hub.load('RF5/simple-asgan', 'asgan_hubert_sc09_6', trust_repo=True)\n",
        "model = model.eval()"
      ],
      "metadata": {
        "id": "9tPwIey7Wy8-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# The below returns a batch of (4, 16000) one second waveforms \n",
        "# that you can directly save as .wav files.\n",
        "with torch.inference_mode():\n",
        "    audio = model.unconditional_generate(4)"
      ],
      "metadata": {
        "id": "BK2k46gAXEqW"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%matplotlib inline\n",
        "fig, axs = plt.subplots(nrows=2, ncols=2, figsize=(16, 12))\n",
        "axs = axs.flatten()\n",
        "for i in range(4):\n",
        "    print(f\"Sample {i}\")\n",
        "    ipd.display(ipd.Audio(data=audio[i].cpu().numpy(), rate=model.hifigan.h.sampling_rate))\n",
        "    librosa.display.waveplot(audio[i].cpu().numpy(), ax=axs[i])\n",
        "    axs[i].set_title(f\"Sample {i}\")\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "BdGKYRPWXWsT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "rXG7ksAQZqnt"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}