MusicGen_Transformers.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "e70e6dbb-3211-4ef9-93f6-efaba764ac77",
      "metadata": {
        "id": "e70e6dbb-3211-4ef9-93f6-efaba764ac77"
      },
      "source": [
        "## Prepare the Environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "21d38c22-bb79-495c-8aa9-09ceabb2957a",
      "metadata": {
        "id": "21d38c22-bb79-495c-8aa9-09ceabb2957a"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "66af0411-c18e-4d8b-b6d9-318ff4460e48",
      "metadata": {
        "id": "66af0411-c18e-4d8b-b6d9-318ff4460e48"
      },
      "outputs": [],
      "source": [
        "!pip install --upgrade --quiet pip\n",
        "!pip install --quiet git+https://github.com/huggingface/transformers.git datasets[audio]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b0d87424-9f38-4658-ba47-2a465d52ad77",
      "metadata": {
        "id": "b0d87424-9f38-4658-ba47-2a465d52ad77"
      },
      "outputs": [],
      "source": [
        "from transformers import MusicgenForConditionalGeneration\n",
        "\n",
        "model = MusicgenForConditionalGeneration.from_pretrained(\"facebook/musicgen-small\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9508dee8-39df-46fe-82f3-6cc2e9f21a97",
      "metadata": {
        "id": "9508dee8-39df-46fe-82f3-6cc2e9f21a97"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
        "model.to(device);"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fb7708e8-e4f1-4ab8-b04a-19395d78dea2",
      "metadata": {
        "id": "fb7708e8-e4f1-4ab8-b04a-19395d78dea2"
      },
      "outputs": [],
      "source": [
        "unconditional_inputs = model.get_unconditional_inputs(num_samples=1)\n",
        "\n",
        "audio_values = model.generate(**unconditional_inputs, do_sample=True, max_new_tokens=256)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "15f0bc7c-b899-4e7a-943e-594e73f080ea",
      "metadata": {
        "id": "15f0bc7c-b899-4e7a-943e-594e73f080ea"
      },
      "outputs": [],
      "source": [
        "from IPython.display import Audio\n",
        "\n",
        "sampling_rate = model.config.audio_encoder.sampling_rate\n",
        "Audio(audio_values[0].cpu().numpy(), rate=sampling_rate)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "04291f52-0a75-4ddb-9eff-e853d0f17288",
      "metadata": {
        "id": "04291f52-0a75-4ddb-9eff-e853d0f17288"
      },
      "outputs": [],
      "source": [
        "import scipy\n",
        "\n",
        "scipy.io.wavfile.write(\"musicgen_out.wav\", rate=sampling_rate, data=audio_values[0, 0].cpu().numpy())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d75ad107-e19b-47f3-9cf1-5102ab4ae74a",
      "metadata": {
        "id": "d75ad107-e19b-47f3-9cf1-5102ab4ae74a"
      },
      "outputs": [],
      "source": [
        "audio_length_in_s = 256 / model.config.audio_encoder.frame_rate\n",
        "\n",
        "audio_length_in_s"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5fba4154-13f6-403a-958b-101d6eacfb6e",
      "metadata": {
        "id": "5fba4154-13f6-403a-958b-101d6eacfb6e"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoProcessor\n",
        "\n",
        "processor = AutoProcessor.from_pretrained(\"facebook/musicgen-small\")\n",
        "\n",
        "inputs = processor(\n",
        "    text=[\"80s pop track with bassy drums and synth\", \"90s rock song with loud guitars and heavy drums\"],\n",
        "    padding=True,\n",
        "    return_tensors=\"pt\",\n",
        ")\n",
        "\n",
        "audio_values = model.generate(**inputs.to(device), do_sample=True, guidance_scale=3, max_new_tokens=256)\n",
        "\n",
        "Audio(audio_values[0].cpu().numpy(), rate=sampling_rate)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "56a5c28a-f6c1-4ac8-ae08-6776a2b2c5b8",
      "metadata": {
        "id": "56a5c28a-f6c1-4ac8-ae08-6776a2b2c5b8"
      },
      "outputs": [],
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "dataset = load_dataset(\"sanchit-gandhi/gtzan\", split=\"train\", streaming=True)\n",
        "sample = next(iter(dataset))[\"audio\"]\n",
        "\n",
        "# take the first half of the audio sample\n",
        "sample[\"array\"] = sample[\"array\"][: len(sample[\"array\"]) // 2]\n",
        "\n",
        "inputs = processor(\n",
        "    audio=sample[\"array\"],\n",
        "    sampling_rate=sample[\"sampling_rate\"],\n",
        "    text=[\"80s blues track with groovy saxophone\"],\n",
        "    padding=True,\n",
        "    return_tensors=\"pt\",\n",
        ")\n",
        "\n",
        "audio_values = model.generate(**inputs.to(device), do_sample=True, guidance_scale=3, max_new_tokens=256)\n",
        "\n",
        "Audio(audio_values[0].cpu().numpy(), rate=sampling_rate)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5495f568-51ca-439d-b47b-8b52e89b78f1",
      "metadata": {
        "id": "5495f568-51ca-439d-b47b-8b52e89b78f1"
      },
      "outputs": [],
      "source": [
        "sample = next(iter(dataset))[\"audio\"]\n",
        "\n",
        "# take the first quater of the audio sample\n",
        "sample_1 = sample[\"array\"][: len(sample[\"array\"]) // 4]\n",
        "\n",
        "# take the first half of the audio sample\n",
        "sample_2 = sample[\"array\"][: len(sample[\"array\"]) // 2]\n",
        "\n",
        "inputs = processor(\n",
        "    audio=[sample_1, sample_2],\n",
        "    sampling_rate=sample[\"sampling_rate\"],\n",
        "    text=[\"80s blues track with groovy saxophone\", \"90s rock song with loud guitars and heavy drums\"],\n",
        "    padding=True,\n",
        "    return_tensors=\"pt\",\n",
        ")\n",
        "\n",
        "audio_values = model.generate(**inputs.to(device), do_sample=True, guidance_scale=3, max_new_tokens=256)\n",
        "\n",
        "# post-process to remove padding from the batched audio\n",
        "audio_values = processor.batch_decode(audio_values, padding_mask=inputs.padding_mask)\n",
        "\n",
        "Audio(audio_values[0], rate=sampling_rate)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model.generation_config"
      ],
      "metadata": {
        "id": "0zM4notb8Y1g"
      },
      "id": "0zM4notb8Y1g",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# increase the guidance scale to 4.0\n",
        "model.generation_config.guidance_scale = 4.0\n",
        "\n",
        "# set the max new tokens to 256\n",
        "model.generation_config.max_new_tokens = 256\n",
        "\n",
        "# set the softmax sampling temperature to 1.5\n",
        "model.generation_config.temperature = 1.5"
      ],
      "metadata": {
        "id": "ensSj1IB81dA"
      },
      "id": "ensSj1IB81dA",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "audio_values = model.generate(**inputs.to(device))"
      ],
      "metadata": {
        "id": "KAExrhDl9YvS"
      },
      "id": "KAExrhDl9YvS",
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.6"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}