Finetune_Dance_Diffusion.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f39_5v4kEAMo"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Harmonai-org/sample-generator/blob/main/Finetune_Dance_Diffusion.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HHcTRGvUmoME"
      },
      "source": [
        "# Dance Diffusion finetune"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u97w34BXmust"
      },
      "source": [
        "Licensed under the MIT License\n",
        "\n",
        "Copyright (c) 2022 Zach Evans\n",
        "\n",
        "Permission is hereby granted, free of charge, to any person obtaining a copy\n",
        "of this software and associated documentation files (the \"Software\"), to deal\n",
        "in the Software without restriction, including without limitation the rights\n",
        "to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
        "copies of the Software, and to permit persons to whom the Software is\n",
        "furnished to do so, subject to the following conditions:\n",
        "\n",
        "The above copyright notice and this permission notice shall be included in\n",
        "all copies or substantial portions of the Software.\n",
        "\n",
        "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
        "IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
        "FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
        "AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
        "LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
        "OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\n",
        "THE SOFTWARE.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lU97ZiP7nSKS"
      },
      "source": [
        "# Set Up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "mxb-qgh0nUOf"
      },
      "outputs": [],
      "source": [
        "#@title Check GPU Status\n",
        "import subprocess\n",
        "simple_nvidia_smi_display = True#@param {type:\"boolean\"}\n",
        "if simple_nvidia_smi_display:\n",
        "    #!nvidia-smi\n",
        "    nvidiasmi_output = subprocess.run(['nvidia-smi', '-L'], stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
        "    print(nvidiasmi_output)\n",
        "else:\n",
        "    #!nvidia-smi -i 0 -e 0\n",
        "    nvidiasmi_output = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
        "    print(nvidiasmi_output)\n",
        "    nvidiasmi_ecc_note = subprocess.run(['nvidia-smi', '-i', '0', '-e', '0'], stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
        "    print(nvidiasmi_ecc_note)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "T_mFtzHvnlJL"
      },
      "outputs": [],
      "source": [
        "#@title Prepare folders\n",
        "import subprocess, os, sys, ipykernel\n",
        "\n",
        "def gitclone(url, targetdir=None):\n",
        "    if targetdir:\n",
        "        res = subprocess.run(['git', 'clone', url, targetdir], stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
        "    else:\n",
        "        res = subprocess.run(['git', 'clone', url], stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
        "    print(res)\n",
        "\n",
        "def pipi(modulestr):\n",
        "    res = subprocess.run(['pip', 'install', modulestr], stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
        "    print(res)\n",
        "\n",
        "def pipie(modulestr):\n",
        "    res = subprocess.run(['git', 'install', '-e', modulestr], stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
        "    print(res)\n",
        "\n",
        "def wget(url, outputdir):\n",
        "    res = subprocess.run(['wget', url, '-P', f'{outputdir}'], stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
        "    print(res)\n",
        "\n",
        "try:\n",
        "    from google.colab import drive\n",
        "    print(\"Google Colab detected. Using Google Drive.\")\n",
        "    is_colab = True\n",
        "    google_drive = True #@param {type:\"boolean\"}\n",
        "    #@markdown Click here if you'd like to save the diffusion model checkpoint file to (and/or load from) your Google Drive:\n",
        "    save_models_to_google_drive = True #@param {type:\"boolean\"}\n",
        "except:\n",
        "    is_colab = False\n",
        "    google_drive = False\n",
        "    save_models_to_google_drive = False\n",
        "    print(\"Google Colab not detected.\")\n",
        "\n",
        "if is_colab:\n",
        "    if google_drive is True:\n",
        "        drive.mount('/content/drive')\n",
        "        root_path = '/content/drive/MyDrive/AI/Bass_Diffusion'\n",
        "    else:\n",
        "        root_path = '/content'\n",
        "else:\n",
        "    root_path = os.getcwd()\n",
        "\n",
        "import os\n",
        "def createPath(filepath):\n",
        "    os.makedirs(filepath, exist_ok=True)\n",
        "\n",
        "initDirPath = f'{root_path}/init_audio'\n",
        "createPath(initDirPath)\n",
        "outDirPath = f'{root_path}/audio_out'\n",
        "createPath(outDirPath)\n",
        "\n",
        "if is_colab:\n",
        "    if google_drive and not save_models_to_google_drive or not google_drive:\n",
        "        model_path = '/content/models'\n",
        "        createPath(model_path)\n",
        "    if google_drive and save_models_to_google_drive:\n",
        "        model_path = f'{root_path}/models'\n",
        "        createPath(model_path)\n",
        "else:\n",
        "    model_path = f'{root_path}/models'\n",
        "    createPath(model_path)\n",
        "\n",
        "#@markdown Click here if you'd like to save the diffusion model checkpoint file to your [Weights & Biases](www.wandb.ai/site) account:\n",
        "save_models_to_wandb = False #@param {type:\"boolean\"}\n",
        "save_wandb_str = '--save-wandb all' if save_models_to_wandb else ''\n",
        "if save_models_to_wandb == True:\n",
        "    print('Saving model checkpoints in wandb')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "y9BS0ks1oEgP"
      },
      "outputs": [],
      "source": [
        "#@title Install dependencies\n",
        "!git clone https://github.com/harmonai-org/sample-generator\n",
        "!pip install /content/sample-generator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0xq2TJzIPTcJ"
      },
      "source": [
        "# Train"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "oxJFFEZ8CD8g"
      },
      "outputs": [],
      "source": [
        "#@markdown Log in to [Weights & Biases](https://wandb.ai/) for run tracking\n",
        "!wandb login"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-Q0XrS0HEmch"
      },
      "outputs": [],
      "source": [
        "#@markdown Name for the finetune project, used as the W&B project name, as well as the directory for the saved checkpoints\n",
        "NAME=\"dd-drums-finetune\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Path to the directory of audio data to use for fine-tuning\n",
        "TRAINING_DIR=\"/content/drive/MyDrive/Audio/Drums\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Path to the checkpoint to fine-tune \n",
        "CKPT_PATH=\"/content/drive/MyDrive/AI/models/jmann-small-190k.ckpt\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Directory path for saving the fine-tuned outputs\n",
        "OUTPUT_DIR=\"/content/drive/MyDrive/AI/models/DanceDiffusion/finetune\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Number of training steps between demos\n",
        "DEMO_EVERY=250 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown Number of training steps between saving model checkpoints\n",
        "CHECKPOINT_EVERY=500 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown Sample rate to train at\n",
        "SAMPLE_RATE=48000 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown Number of audio samples per training sample\n",
        "SAMPLE_SIZE=65536 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown If true, the audio samples provided will be randomly cropped to SAMPLE_SIZE samples\n",
        "#@markdown\n",
        "#@markdown Turn off if you want to ensure the training data always starts at the beginning of the audio files (good for things like drum one-shots)\n",
        "RANDOM_CROP=True #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown Batch size to fine-tune (make it as high as it can go for your GPU)\n",
        "BATCH_SIZE=2 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown Accumulate gradients over n batches, useful for training on one GPU. \n",
        "#@markdown\n",
        "#@markdown Effective batch size is BATCH_SIZE * ACCUM_BATCHES.\n",
        "#@markdown\n",
        "#@markdown Also increases the time between demos and saved checkpoints\n",
        "ACCUM_BATCHES=4 #@param {type:\"number\"}\n",
        "\n",
        "random_crop_str = f\"--random-crop True\" if RANDOM_CROP else \"\"\n",
        "\n",
        "# Escape spaces in paths\n",
        "CKPT_PATH = CKPT_PATH.replace(f\" \", f\"\\ \")\n",
        "OUTPUT_DIR = f\"{OUTPUT_DIR}/{NAME}\".replace(f\" \", f\"\\ \")\n",
        "\n",
        "%cd /content/sample-generator/\n",
        "\n",
        "ckpt_path_str = f\"--ckpt-path {CKPT_PATH}\" if not CKPT_PATH ==\"\" else \"\"\n",
        "\n",
        "!python3 /content/sample-generator/train_uncond.py $ckpt_path_str \\\n",
        "                                                          --name $NAME \\\n",
        "                                                          --training-dir $TRAINING_DIR \\\n",
        "                                                          --sample-size $SAMPLE_SIZE \\\n",
        "                                                          --accum-batches $ACCUM_BATCHES \\\n",
        "                                                          --sample-rate $SAMPLE_RATE \\\n",
        "                                                          --batch-size $BATCH_SIZE \\\n",
        "                                                          --demo-every $DEMO_EVERY \\\n",
        "                                                          --checkpoint-every $CHECKPOINT_EVERY \\\n",
        "                                                          --num-workers 2 \\\n",
        "                                                          $save_wandb_str \\\n",
        "                                                          --num-gpus 1 \\\n",
        "                                                          $random_crop_str \\\n",
        "                                                          --save-path $OUTPUT_DIR"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "HHcTRGvUmoME"
      ],
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3.8.8 ('ml')",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.8.8"
    },
    "vscode": {
      "interpreter": {
        "hash": "9ce29cb5aeaaebb09d381145ff1103741c9816dbccbfcb1b857cd83265b8bd84"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}