musiclm-pytorch-demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "m-4e_Cd81EdS"
      ],
      "authorship_tag": "ABX9TyM0D3pu8DotBk/QLdvnr6UN",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "34bafa6d41a44d138c65db3661a72de6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_a7f79a97c0cf4b12aecb80aa8b0d9666",
              "IPY_MODEL_98bbf60ef7e447ef85b7844182f43a3b",
              "IPY_MODEL_0d674cd96b63451d852f9a80352a2808"
            ],
            "layout": "IPY_MODEL_92ea596f89ac4918ace1a4bf45e8512d"
          }
        },
        "a7f79a97c0cf4b12aecb80aa8b0d9666": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_0be21d3ffc5342cc93fbae9004812924",
            "placeholder": "​",
            "style": "IPY_MODEL_217916dcf0bf4e9496560231d1740b1d",
            "value": "#0: 100%"
          }
        },
        "98bbf60ef7e447ef85b7844182f43a3b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_6812f8012fcf439a971042f5596a3500",
            "max": 16,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_82ee072886344c89a41bd72a2683eb1f",
            "value": 16
          }
        },
        "0d674cd96b63451d852f9a80352a2808": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_4c4354c7330f459a8434e70a441d7c2e",
            "placeholder": "​",
            "style": "IPY_MODEL_d17cfc58d2954cac9e413c954f5f70f0",
            "value": " 16/16 [00:00&lt;00:00, 79.58ex/s]"
          }
        },
        "92ea596f89ac4918ace1a4bf45e8512d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0be21d3ffc5342cc93fbae9004812924": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "217916dcf0bf4e9496560231d1740b1d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "6812f8012fcf439a971042f5596a3500": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "82ee072886344c89a41bd72a2683eb1f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "4c4354c7330f459a8434e70a441d7c2e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d17cfc58d2954cac9e413c954f5f70f0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "354218fb52b44f89b14194919236ec66": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_12643f650f244a76b8378f60f547c3e0",
              "IPY_MODEL_d5f6758354584a75b68ef034b7c6422b",
              "IPY_MODEL_b5fcc1bf84bd468f84867b71d6fc4ca5"
            ],
            "layout": "IPY_MODEL_0e7ae6ec5b3843c08a28732e6518eafb"
          }
        },
        "12643f650f244a76b8378f60f547c3e0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_dc27bd504abd4c0c9df4396723b53210",
            "placeholder": "​",
            "style": "IPY_MODEL_32186312a34744b08f0fef24e3c2a5f6",
            "value": "#1: 100%"
          }
        },
        "d5f6758354584a75b68ef034b7c6422b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7bf5b3a589974b0692d3282b25c8e7ac",
            "max": 16,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_72c7c28727964ad1a22951e573bf87a2",
            "value": 16
          }
        },
        "b5fcc1bf84bd468f84867b71d6fc4ca5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_26eb1e51a0ab42fcb9906f8e817bbeda",
            "placeholder": "​",
            "style": "IPY_MODEL_41484d952ba245c08e68b42e4b5375a5",
            "value": " 16/16 [00:00&lt;00:00, 39.17ex/s]"
          }
        },
        "0e7ae6ec5b3843c08a28732e6518eafb": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "dc27bd504abd4c0c9df4396723b53210": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "32186312a34744b08f0fef24e3c2a5f6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "7bf5b3a589974b0692d3282b25c8e7ac": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "72c7c28727964ad1a22951e573bf87a2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "26eb1e51a0ab42fcb9906f8e817bbeda": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "41484d952ba245c08e68b42e4b5375a5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/aaronannecchiarico/musiclm-pytorch/blob/main/musiclm-pytorch-demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Install Requirements"
      ],
      "metadata": {
        "id": "06KyYKyj0fdM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!nvidia-smi\n",
        "\n",
        "# If this doesn't work, there's no GPU available or detected"
      ],
      "metadata": {
        "id": "RuvNPhDSzqdy",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6c927d91-3382-40a1-d378-be33be542d45"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Thu Feb  9 02:50:50 2023       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 510.47.03    Driver Version: 510.47.03    CUDA Version: 11.6     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   51C    P0    25W /  70W |      0MiB / 15360MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2eCvkLB7uN4v",
        "outputId": "a6967c93-057e-452f-bac8-b49c0755c9f9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting musiclm-pytorch\n",
            "  Downloading musiclm_pytorch-0.0.17-py3-none-any.whl (10 kB)\n",
            "Collecting vector-quantize-pytorch>=1.0.0\n",
            "  Downloading vector_quantize_pytorch-1.0.0-py3-none-any.whl (8.7 kB)\n",
            "Requirement already satisfied: torch>=1.6 in /usr/local/lib/python3.8/dist-packages (from musiclm-pytorch) (1.13.1+cu116)\n",
            "Collecting x-clip\n",
            "  Downloading x_clip-0.12.0-py3-none-any.whl (1.4 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.4/1.4 MB\u001b[0m \u001b[31m19.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting audiolm-pytorch>=0.10.4\n",
            "  Downloading audiolm_pytorch-0.11.15-py3-none-any.whl (28 kB)\n",
            "Collecting einops>=0.6\n",
            "  Downloading einops-0.6.0-py3-none-any.whl (41 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m41.6/41.6 KB\u001b[0m \u001b[31m5.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting beartype\n",
            "  Downloading beartype-0.12.0-py3-none-any.whl (754 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m754.5/754.5 KB\u001b[0m \u001b[31m65.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: torchaudio in /usr/local/lib/python3.8/dist-packages (from musiclm-pytorch) (0.13.1+cu116)\n",
            "Collecting accelerate\n",
            "  Downloading accelerate-0.16.0-py3-none-any.whl (199 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m199.7/199.7 KB\u001b[0m \u001b[31m27.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: scikit-learn in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch>=0.10.4->musiclm-pytorch) (1.0.2)\n",
            "Collecting fairseq\n",
            "  Downloading fairseq-0.12.2-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.whl (11.0 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m11.0/11.0 MB\u001b[0m \u001b[31m101.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch>=0.10.4->musiclm-pytorch) (4.64.1)\n",
            "Collecting Mega-pytorch\n",
            "  Downloading Mega_pytorch-0.0.12-py3-none-any.whl (6.4 kB)\n",
            "Collecting local-attention>=1.5.7\n",
            "  Downloading local_attention-1.5.8-py3-none-any.whl (6.9 kB)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch>=0.10.4->musiclm-pytorch) (1.2.0)\n",
            "Collecting transformers\n",
            "  Downloading transformers-4.26.0-py3-none-any.whl (6.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.3/6.3 MB\u001b[0m \u001b[31m16.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting ema-pytorch\n",
            "  Downloading ema_pytorch-0.1.4-py3-none-any.whl (4.2 kB)\n",
            "Collecting sentencepiece\n",
            "  Downloading sentencepiece-0.1.97-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m70.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: typing-extensions in /usr/local/lib/python3.8/dist-packages (from torch>=1.6->musiclm-pytorch) (4.4.0)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.8/dist-packages (from accelerate->musiclm-pytorch) (5.4.8)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.8/dist-packages (from accelerate->musiclm-pytorch) (1.21.6)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.8/dist-packages (from accelerate->musiclm-pytorch) (23.0)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.8/dist-packages (from accelerate->musiclm-pytorch) (6.0)\n",
            "Collecting ftfy\n",
            "  Downloading ftfy-6.1.1-py3-none-any.whl (53 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m53.1/53.1 KB\u001b[0m \u001b[31m7.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: regex in /usr/local/lib/python3.8/dist-packages (from x-clip->musiclm-pytorch) (2022.6.2)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.8/dist-packages (from x-clip->musiclm-pytorch) (0.14.1+cu116)\n",
            "Collecting sacrebleu>=1.4.12\n",
            "  Downloading sacrebleu-2.3.1-py3-none-any.whl (118 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m118.9/118.9 KB\u001b[0m \u001b[31m16.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting bitarray\n",
            "  Downloading bitarray-2.7.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (264 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m264.8/264.8 KB\u001b[0m \u001b[31m28.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting omegaconf<2.1\n",
            "  Downloading omegaconf-2.0.6-py3-none-any.whl (36 kB)\n",
            "Requirement already satisfied: cython in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch>=0.10.4->musiclm-pytorch) (0.29.33)\n",
            "Collecting hydra-core<1.1,>=1.0.7\n",
            "  Downloading hydra_core-1.0.7-py3-none-any.whl (123 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m123.8/123.8 KB\u001b[0m \u001b[31m17.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: cffi in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch>=0.10.4->musiclm-pytorch) (1.15.1)\n",
            "Requirement already satisfied: wcwidth>=0.2.5 in /usr/local/lib/python3.8/dist-packages (from ftfy->x-clip->musiclm-pytorch) (0.2.6)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.8/dist-packages (from Mega-pytorch->audiolm-pytorch>=0.10.4->musiclm-pytorch) (1.7.3)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from scikit-learn->audiolm-pytorch>=0.10.4->musiclm-pytorch) (3.1.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.8/dist-packages (from torchvision->x-clip->musiclm-pytorch) (2.25.1)\n",
            "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /usr/local/lib/python3.8/dist-packages (from torchvision->x-clip->musiclm-pytorch) (7.1.2)\n",
            "Collecting huggingface-hub<1.0,>=0.11.0\n",
            "  Downloading huggingface_hub-0.12.0-py3-none-any.whl (190 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m190.3/190.3 KB\u001b[0m \u001b[31m25.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting tokenizers!=0.11.3,<0.14,>=0.11.1\n",
            "  Downloading tokenizers-0.13.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (7.6 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.6/7.6 MB\u001b[0m \u001b[31m106.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from transformers->audiolm-pytorch>=0.10.4->musiclm-pytorch) (3.9.0)\n",
            "Requirement already satisfied: importlib-resources in /usr/local/lib/python3.8/dist-packages (from hydra-core<1.1,>=1.0.7->fairseq->audiolm-pytorch>=0.10.4->musiclm-pytorch) (5.10.2)\n",
            "Collecting antlr4-python3-runtime==4.8\n",
            "  Downloading antlr4-python3-runtime-4.8.tar.gz (112 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m112.4/112.4 KB\u001b[0m \u001b[31m15.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting portalocker\n",
            "  Downloading portalocker-2.7.0-py2.py3-none-any.whl (15 kB)\n",
            "Collecting colorama\n",
            "  Downloading colorama-0.4.6-py2.py3-none-any.whl (25 kB)\n",
            "Requirement already satisfied: tabulate>=0.8.9 in /usr/local/lib/python3.8/dist-packages (from sacrebleu>=1.4.12->fairseq->audiolm-pytorch>=0.10.4->musiclm-pytorch) (0.8.10)\n",
            "Requirement already satisfied: lxml in /usr/local/lib/python3.8/dist-packages (from sacrebleu>=1.4.12->fairseq->audiolm-pytorch>=0.10.4->musiclm-pytorch) (4.9.2)\n",
            "Requirement already satisfied: pycparser in /usr/local/lib/python3.8/dist-packages (from cffi->fairseq->audiolm-pytorch>=0.10.4->musiclm-pytorch) (2.21)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests->torchvision->x-clip->musiclm-pytorch) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests->torchvision->x-clip->musiclm-pytorch) (2.10)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests->torchvision->x-clip->musiclm-pytorch) (4.0.0)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests->torchvision->x-clip->musiclm-pytorch) (2022.12.7)\n",
            "Requirement already satisfied: zipp>=3.1.0 in /usr/local/lib/python3.8/dist-packages (from importlib-resources->hydra-core<1.1,>=1.0.7->fairseq->audiolm-pytorch>=0.10.4->musiclm-pytorch) (3.12.0)\n",
            "Building wheels for collected packages: antlr4-python3-runtime\n",
            "  Building wheel for antlr4-python3-runtime (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for antlr4-python3-runtime: filename=antlr4_python3_runtime-4.8-py3-none-any.whl size=141231 sha256=b2c2bf947a667137db98fbe6f1f3f70f1f0b908895223c35d04e175aa00b7237\n",
            "  Stored in directory: /root/.cache/pip/wheels/c8/d0/ab/d43c02eaddc5b9004db86950802442ad9a26f279c619e28da0\n",
            "Successfully built antlr4-python3-runtime\n",
            "Installing collected packages: tokenizers, sentencepiece, bitarray, antlr4-python3-runtime, portalocker, omegaconf, ftfy, einops, colorama, beartype, vector-quantize-pytorch, sacrebleu, Mega-pytorch, local-attention, hydra-core, huggingface-hub, ema-pytorch, accelerate, x-clip, transformers, fairseq, audiolm-pytorch, musiclm-pytorch\n",
            "Successfully installed Mega-pytorch-0.0.12 accelerate-0.16.0 antlr4-python3-runtime-4.8 audiolm-pytorch-0.11.15 beartype-0.12.0 bitarray-2.7.0 colorama-0.4.6 einops-0.6.0 ema-pytorch-0.1.4 fairseq-0.12.2 ftfy-6.1.1 huggingface-hub-0.12.0 hydra-core-1.0.7 local-attention-1.5.8 musiclm-pytorch-0.0.17 omegaconf-2.0.6 portalocker-2.7.0 sacrebleu-2.3.1 sentencepiece-0.1.97 tokenizers-0.13.2 transformers-4.26.0 vector-quantize-pytorch-1.0.0 x-clip-0.12.0\n"
          ]
        }
      ],
      "source": [
        "pip install musiclm-pytorch"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "! pip install datasets[audio] yt-dlp"
      ],
      "metadata": {
        "id": "CsHn6W4T094K"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Imports"
      ],
      "metadata": {
        "id": "qKw_h3ohz0Yr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import subprocess\n",
        "import os\n",
        "import requests\n",
        "from pathlib import Path\n",
        "\n",
        "from datasets import load_dataset, Audio\n",
        "\n",
        "import torch\n",
        "import torchaudio\n",
        "from torch.utils.data import Dataset\n",
        "\n",
        "from musiclm_pytorch import MuLaN, MuLaNTrainer, AudioSpectrogramTransformer, TextTransformer, MuLaNEmbedQuantizer, MusicLM\n",
        "from audiolm_pytorch import SoundStream, SoundStreamTrainer, HubertWithKmeans, SemanticTransformer, SemanticTransformerTrainer, HubertWithKmeans, CoarseTransformer, CoarseTransformerWrapper, CoarseTransformerTrainer, FineTransformer, FineTransformerWrapper, FineTransformerTrainer, AudioLM\n",
        "\n"
      ],
      "metadata": {
        "id": "4NVQRFOWzxe4"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Data"
      ],
      "metadata": {
        "id": "4iIM6vrp00GO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def download_clip(\n",
        "    video_identifier,\n",
        "    output_filename,\n",
        "    start_time,\n",
        "    end_time,\n",
        "    tmp_dir='/tmp/musiccaps',\n",
        "    num_attempts=5,\n",
        "    url_base='https://www.youtube.com/watch?v='\n",
        "):\n",
        "    status = False\n",
        "\n",
        "    command = f\"\"\"\n",
        "        yt-dlp --quiet --no-warnings -x --audio-format wav -f bestaudio -o \"{output_filename}\" --download-sections \"*{start_time}-{end_time}\" {url_base}{video_identifier}\n",
        "    \"\"\".strip()\n",
        "\n",
        "    attempts = 0\n",
        "    while True:\n",
        "        try:\n",
        "            output = subprocess.check_output(command, shell=True,\n",
        "                                                stderr=subprocess.STDOUT)\n",
        "        except subprocess.CalledProcessError as err:\n",
        "            attempts += 1\n",
        "            if attempts == num_attempts:\n",
        "                return status, err.output\n",
        "        else:\n",
        "            break\n",
        "\n",
        "    # Check if the video was successfully saved.\n",
        "    status = os.path.exists(output_filename)\n",
        "    return status, 'Downloaded'\n",
        "\n",
        "\n",
        "def main(\n",
        "    data_dir: str,\n",
        "    sampling_rate: int = 44100,\n",
        "    limit: int = None,\n",
        "    num_proc: int = 1,\n",
        "    writer_batch_size: int = 1000,\n",
        "):\n",
        "    \"\"\"\n",
        "    Download the clips within the MusicCaps dataset from YouTube.\n",
        "    Args:\n",
        "        data_dir: Directory to save the clips to.\n",
        "        sampling_rate: Sampling rate of the audio clips.\n",
        "        limit: Limit the number of examples to download.\n",
        "        num_proc: Number of processes to use for downloading.\n",
        "        writer_batch_size: Batch size for writing the dataset. This is per process.\n",
        "    \"\"\"\n",
        "\n",
        "    ds = load_dataset('google/MusicCaps', split='train')\n",
        "    if limit is not None:\n",
        "        print(f\"Limiting to {limit} examples\")\n",
        "        ds = ds.select(range(limit))\n",
        "\n",
        "    data_dir = Path(data_dir)\n",
        "    data_dir.mkdir(exist_ok=True, parents=True)\n",
        "\n",
        "    def process(example):\n",
        "        outfile_path = str(data_dir / f\"{example['ytid']}.wav\")\n",
        "        status = True\n",
        "        if not os.path.exists(outfile_path):\n",
        "            status = False\n",
        "            status, log = download_clip(\n",
        "                example['ytid'],\n",
        "                outfile_path,\n",
        "                example['start_s'],\n",
        "                example['end_s'],\n",
        "            )\n",
        "\n",
        "        example['audio'] = outfile_path\n",
        "        example['download_status'] = status\n",
        "        return example\n",
        "\n",
        "    return ds.map(\n",
        "        process,\n",
        "        num_proc=num_proc,\n",
        "        writer_batch_size=writer_batch_size,\n",
        "        keep_in_memory=False\n",
        "    ).cast_column('audio', Audio(sampling_rate=sampling_rate))"
      ],
      "metadata": {
        "id": "X_Of87mq0-uZ"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "ds = main('./music_data', num_proc=2, limit=32, writer_batch_size=4) # change limit for larger dataset"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 172,
          "referenced_widgets": [
            "34bafa6d41a44d138c65db3661a72de6",
            "a7f79a97c0cf4b12aecb80aa8b0d9666",
            "98bbf60ef7e447ef85b7844182f43a3b",
            "0d674cd96b63451d852f9a80352a2808",
            "92ea596f89ac4918ace1a4bf45e8512d",
            "0be21d3ffc5342cc93fbae9004812924",
            "217916dcf0bf4e9496560231d1740b1d",
            "6812f8012fcf439a971042f5596a3500",
            "82ee072886344c89a41bd72a2683eb1f",
            "4c4354c7330f459a8434e70a441d7c2e",
            "d17cfc58d2954cac9e413c954f5f70f0",
            "354218fb52b44f89b14194919236ec66",
            "12643f650f244a76b8378f60f547c3e0",
            "d5f6758354584a75b68ef034b7c6422b",
            "b5fcc1bf84bd468f84867b71d6fc4ca5",
            "0e7ae6ec5b3843c08a28732e6518eafb",
            "dc27bd504abd4c0c9df4396723b53210",
            "32186312a34744b08f0fef24e3c2a5f6",
            "7bf5b3a589974b0692d3282b25c8e7ac",
            "72c7c28727964ad1a22951e573bf87a2",
            "26eb1e51a0ab42fcb9906f8e817bbeda",
            "41484d952ba245c08e68b42e4b5375a5"
          ]
        },
        "id": "mB_MBq-J1fvJ",
        "outputId": "8a32d767-0a8e-4b11-9d71-5f5ef8b44cee"
      },
      "execution_count": 73,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:datasets.builder:Using custom data configuration google--MusicCaps-7925612b943f961b\n",
            "WARNING:datasets.builder:Found cached dataset csv (/root/.cache/huggingface/datasets/google___csv/google--MusicCaps-7925612b943f961b/0.0.0/6b34fb8fcf56f7c8ba51dc895bfa2bfbe43546f190a60fcf74bb5e8afdcc2317)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Limiting to 32 examples\n",
            "    "
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "#0:   0%|          | 0/16 [00:00<?, ?ex/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "34bafa6d41a44d138c65db3661a72de6"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "#1:   0%|          | 0/16 [00:00<?, ?ex/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "354218fb52b44f89b14194919236ec66"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class TextAudioDataset(Dataset):\n",
        "    def __init__(self, dset):\n",
        "        super().__init__()\n",
        "        self.dset = dset\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.dset)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        ex = self.dset[idx]\n",
        "        caption = ex['caption']\n",
        "        wav, samplerate = torchaudio.load(ex['audio']['path'])\n",
        "        return caption, wav"
      ],
      "metadata": {
        "id": "MGKj-_ztGJiz"
      },
      "execution_count": 81,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Train MuLaN"
      ],
      "metadata": {
        "id": "XOTZ2LWd0X6f"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "audio_transformer = AudioSpectrogramTransformer(\n",
        "    dim = 512,\n",
        "    depth = 6,\n",
        "    heads = 8,\n",
        "    dim_head = 64,\n",
        "    spec_n_fft = 128,\n",
        "    spec_win_length = 24,\n",
        "    spec_aug_stretch_factor = 0.8\n",
        ")\n",
        "\n",
        "text_transformer = TextTransformer(\n",
        "    dim = 512,\n",
        "    depth = 6,\n",
        "    heads = 8,\n",
        "    dim_head = 64\n",
        ")\n",
        "\n",
        "mulan = MuLaN(\n",
        "    audio_transformer = audio_transformer,\n",
        "    text_transformer = text_transformer\n",
        ")\n",
        "\n",
        "trainer = MuLaNTrainer(\n",
        "    mulan = mulan,\n",
        "    dataset = TextAudioDataset(ds),\n",
        "    batch_size = 4\n",
        ")\n",
        "\n",
        "trainer.train()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 425
        },
        "id": "z4TvhrZCuUTA",
        "outputId": "d6d0f708-b47b-47bf-b1a5-32724dd17f02"
      },
      "execution_count": 82,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "training with dataset of 30 samples and validating with randomly splitted 2 samples\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "RuntimeError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-82-829f3e1b0a39>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     27\u001b[0m )\n\u001b[1;32m     28\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 29\u001b[0;31m \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m<@beartype(musiclm_pytorch.trainer.MuLaNTrainer.train) at 0x7fa205ec8940>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(__beartype_func, __beartype_conf, __beartype_get_violation, __beartype_object_14464304, *args, **kwargs)\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/musiclm_pytorch/trainer.py\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(self, log_fn)\u001b[0m\n\u001b[1;32m    345\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    346\u001b[0m         \u001b[0;32mwhile\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msteps\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_train_steps\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 347\u001b[0;31m             \u001b[0mlogs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    348\u001b[0m             \u001b[0mlog_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlogs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    349\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/musiclm_pytorch/trainer.py\u001b[0m in \u001b[0;36mtrain_step\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    311\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    312\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0m_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrad_accum_every\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 313\u001b[0;31m             \u001b[0mdata_kwargs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata_tuple_to_kwargs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdl_iter\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    314\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    315\u001b[0m             \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmulan\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0mdata_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/musiclm_pytorch/trainer.py\u001b[0m in \u001b[0;36mcycle\u001b[0;34m(dl)\u001b[0m\n\u001b[1;32m     55\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mcycle\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdl\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     56\u001b[0m     \u001b[0;32mwhile\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 57\u001b[0;31m         \u001b[0;32mfor\u001b[0m \u001b[0mdata\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mdl\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     58\u001b[0m             \u001b[0;32myield\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     59\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/accelerate/data_loader.py\u001b[0m in \u001b[0;36m__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    386\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    387\u001b[0m                     \u001b[0mcurrent_batch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msend_to_device\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcurrent_batch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 388\u001b[0;31m                 \u001b[0mnext_batch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdataloader_iter\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    389\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mbatch_index\u001b[0m \u001b[0;34m>=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mskip_batches\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    390\u001b[0m                     \u001b[0;32myield\u001b[0m \u001b[0mcurrent_batch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    626\u001b[0m                 \u001b[0;31m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    627\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_reset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[call-arg]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 628\u001b[0;31m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_next_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    629\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_num_yielded\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    630\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dataset_kind\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0m_DatasetKind\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mIterable\u001b[0m \u001b[0;32mand\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m_next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    669\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_next_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    670\u001b[0m         \u001b[0mindex\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_next_index\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 671\u001b[0;31m         \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dataset_fetcher\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfetch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    672\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_pin_memory\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    673\u001b[0m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_pin_memory_device\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/utils/data/_utils/fetch.py\u001b[0m in \u001b[0;36mfetch\u001b[0;34m(self, possibly_batched_index)\u001b[0m\n\u001b[1;32m     59\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     60\u001b[0m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mpossibly_batched_index\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 61\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcollate_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/musiclm_pytorch/trainer.py\u001b[0m in \u001b[0;36minner\u001b[0;34m(data)\u001b[0m\n\u001b[1;32m    118\u001b[0m                 \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdatum\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    119\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 120\u001b[0;31m                 \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdatum\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    121\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    122\u001b[0m             \u001b[0moutputs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/musiclm_pytorch/trainer.py\u001b[0m in \u001b[0;36mcurtail_to_shortest_collate\u001b[0;34m(data)\u001b[0m\n\u001b[1;32m    130\u001b[0m     \u001b[0mmin_len\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mdatum\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mdatum\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    131\u001b[0m     \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mdatum\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mmin_len\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mdatum\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 132\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    133\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    134\u001b[0m \u001b[0;34m@\u001b[0m\u001b[0mcollate_one_or_multiple_tensors\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mRuntimeError\u001b[0m: stack expects each tensor to be equal size, but got [2, 958728] at entry 0 and [2, 441352] at entry 3"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# setup the quantizer with the namespaced conditioning embeddings, unique per quantizer as well as namespace (per transformer)\n",
        "\n",
        "quantizer = MuLaNEmbedQuantizer(\n",
        "    mulan = mulan,                          # pass in trained mulan from above\n",
        "    conditioning_dims = (1024, 1024, 1024), # say all three transformers have model dimensions of 1024\n",
        "    namespaces = ('semantic', 'coarse', 'fine')\n",
        ")"
      ],
      "metadata": {
        "id": "bSO7Q4mwu4lt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Download Hubert"
      ],
      "metadata": {
        "id": "m-4e_Cd81EdS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Create a script that downloads the following files: https://dl.fbaipublicfiles.com/hubert/hubert_base_ls960.pt, https://dl.fbaipublicfiles.com/hubert/hubert_base_ls960_L9_km500.bin\n",
        "# into a new folder in the current directory called hubert.\n",
        "\n",
        "def download_file(url, file_name):\n",
        "    with open(file_name, \"wb\") as file:\n",
        "        response = requests.get(url)\n",
        "        file.write(response.content)\n",
        "\n",
        "def get_hubert():\n",
        "    # Create a folder called hubert\n",
        "    os.mkdir(\"hubert\")\n",
        "\n",
        "    # Download the files\n",
        "    download_file(\"https://dl.fbaipublicfiles.com/hubert/hubert_base_ls960.pt\", \"hubert/hubert_base_ls960.pt\")\n",
        "    download_file(\"https://dl.fbaipublicfiles.com/hubert/hubert_base_ls960_L9_km500.bin\", \"hubert/hubert_base_ls960_L9_km500.bin\")\n"
      ],
      "metadata": {
        "id": "VqZ7DudV8OYo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "get_hubert()"
      ],
      "metadata": {
        "id": "s_tv61QI8VLS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Semantic Transformer Training"
      ],
      "metadata": {
        "id": "B_ANyfo11J3Y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "wav2vec = HubertWithKmeans(\n",
        "    checkpoint_path = './hubert/hubert_base_ls960.pt',\n",
        "    kmeans_path = './hubert/hubert_base_ls960_L9_km500.bin'\n",
        ")\n",
        "\n",
        "semantic_transformer = SemanticTransformer(\n",
        "    num_semantic_tokens = wav2vec.codebook_size,\n",
        "    dim = 1024,\n",
        "    depth = 6,\n",
        "    audio_text_condition = True      # this must be set to True (same for CoarseTransformer and FineTransformers)\n",
        ").cuda()\n",
        "\n",
        "trainer = SemanticTransformerTrainer(\n",
        "    transformer = semantic_transformer,\n",
        "    wav2vec = wav2vec,\n",
        "    audio_conditioner = quantizer,   # pass in the MulanEmbedQuantizer instance above\n",
        "    folder ='./music_data',\n",
        "    batch_size = 1,\n",
        "    data_max_length = 320 * 32,\n",
        "    num_train_steps = 1\n",
        ")\n",
        "\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "vOmxhZ6Au5ZA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Soundstream Training"
      ],
      "metadata": {
        "id": "V_2qdlTD1QqR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "soundstream = SoundStream(\n",
        "    codebook_size = 1024,\n",
        "    rq_num_quantizers = 8,\n",
        ")\n",
        "\n",
        "trainer = SoundStreamTrainer(\n",
        "    soundstream,\n",
        "    folder = './music_data',\n",
        "    batch_size = 4,\n",
        "    grad_accum_every = 8,         # effective batch size of 32\n",
        "    data_max_length = 320 * 32,\n",
        "    save_results_every = 2,\n",
        "    save_model_every = 4,\n",
        "    num_train_steps = 9\n",
        ").cuda()\n",
        "\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "UK7EN7OAsz97"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Coarse Transformer Training"
      ],
      "metadata": {
        "id": "vcfNWiRq1Xue"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "wav2vec = HubertWithKmeans(\n",
        "    checkpoint_path = './hubert/hubert_base_ls960.pt',\n",
        "    kmeans_path = './hubert/hubert_base_ls960_L9_km500.bin'\n",
        ")\n",
        "\n",
        "soundstream = SoundStream(\n",
        "    codebook_size = 1024,\n",
        "    rq_num_quantizers = 8,\n",
        ")\n",
        "\n",
        "soundstream.load(\"./results/soundstream.8.pt\")\n",
        "\n",
        "coarse_transformer = CoarseTransformer(\n",
        "    num_semantic_tokens = wav2vec.codebook_size,\n",
        "    codebook_size = 1024,\n",
        "    num_coarse_quantizers = 3,\n",
        "    dim = 1024,\n",
        "    depth = 6,\n",
        "    audio_text_condition = True\n",
        ")\n",
        "\n",
        "trainer = CoarseTransformerTrainer(\n",
        "    transformer = coarse_transformer,\n",
        "    soundstream = soundstream,\n",
        "    audio_conditioner = quantizer,\n",
        "    wav2vec = wav2vec,\n",
        "    folder = './music_data',\n",
        "    batch_size = 1,\n",
        "    data_max_length = 320 * 32,\n",
        "    save_results_every = 2,\n",
        "    save_model_every = 4,\n",
        "    num_train_steps = 9\n",
        ")\n",
        "# NOTE: I changed num_train_steps to 9 (aka 8 + 1) from 10000 to make things go faster for demo purposes\n",
        "# adjusting save_*_every variables for the same reason\n",
        "\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "4w1ZJ_ipr9qV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Fine Transformer Training"
      ],
      "metadata": {
        "id": "yZMJqr5-1dyG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "soundstream = SoundStream(\n",
        "    codebook_size = 1024,\n",
        "    rq_num_quantizers = 8,\n",
        ")\n",
        "\n",
        "soundstream.load(\"./results/soundstream.8.pt\")\n",
        "\n",
        "fine_transformer = FineTransformer(\n",
        "    num_coarse_quantizers = 3,\n",
        "    num_fine_quantizers = 5,\n",
        "    codebook_size = 1024,\n",
        "    dim = 1024,\n",
        "    depth = 6,\n",
        "    audio_text_condition = True\n",
        ")\n",
        "\n",
        "trainer = FineTransformerTrainer(\n",
        "    transformer = fine_transformer,\n",
        "    soundstream = soundstream,\n",
        "    audio_conditioner = quantizer,\n",
        "    folder = './music_data',\n",
        "    batch_size = 1,\n",
        "    data_max_length = 320 * 32,\n",
        "    num_train_steps = 9\n",
        ")\n",
        "# NOTE: I changed num_train_steps to 9 (aka 8 + 1) from 10000 to make things go faster for demo purposes\n",
        "# adjusting save_*_every variables for the same reason\n",
        "\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "BbSJFGfVty4c"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Generate"
      ],
      "metadata": {
        "id": "fspjCAyr1jjO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from musiclm_pytorch import MusicLM\n",
        "\n",
        "audiolm = AudioLM(\n",
        "    wav2vec = wav2vec,\n",
        "    soundstream = soundstream,\n",
        "    semantic_transformer = semantic_transformer,\n",
        "    coarse_transformer = coarse_transformer,\n",
        "    fine_transformer = fine_transformer\n",
        ")\n",
        "musiclm = MusicLM(\n",
        "    audio_lm = audiolm,\n",
        "    mulan_embed_quantizer = quantizer\n",
        ")\n",
        "\n",
        "music = musiclm(['the crystalline sounds of the piano in a ballroom']) # torch.Tensor"
      ],
      "metadata": {
        "id": "O4tklVJsvBDT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "output_path = \"out.wav\"\n",
        "sample_rate = 44100\n",
        "torchaudio.save(output_path, music.cpu(), sample_rate)"
      ],
      "metadata": {
        "id": "PFbnpBUK2XkS"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}