Dance_Diffusion.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/Harmonai-org/sample-generator/blob/main/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 v0.10\n",
        "\n",
        "Welcome to the Dance Diffusion beta!\n",
        "\n",
        "Dance Diffusion is the first in a suite of generative audio tools for producers and musicians to be released by Harmonai. For more info or to get involved in the development of these tools, please visit https://harmonai.org and fill out the form on the front page.\n",
        "\n",
        "[Click here to ensure you are using the latest version](https://colab.research.google.com/github/Harmonai-org/sample-generator/blob/main/Dance_Diffusion.ipynb)\n",
        "\n",
        "**Audio diffusion tools in this notebook**:\n",
        "\n",
        "- Unconditional random audio sample generation\n",
        "- Audio sample regeneration/style transfer using a single audio file\n",
        "- Audio interpolation between two audio files"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1iZwJ9ong-pH"
      },
      "source": [
        "# Instructions\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-ITvsXU6hCAx"
      },
      "source": [
        "## Before anything else\n",
        "- Run the \"Setup\" section\n",
        "- Sign in to the Google Drive account you want to save your models in\n",
        "- Select the model you want to sample from in the \"Model settings\" section, this determines the length and sound of your samples\n",
        "- The `save_to_wandb` option futher down adds the ability to log your audio generations to [Weights & Biases](https://www.wandb.ai/site), an experiment tracking and model and data versioning tool.\n",
        "\n",
        "## For random sample generation\n",
        "- Choose the number of random samples you would like Dance Diffusion to generate for you \n",
        "- Choose the number of diffusion steps you would like Dance Diffusion to execute\n",
        "- Make sure the \"skip_for_run_all\" checkbox is unchecked\n",
        "- Run the cell under the \"Generate new sounds\" header\n",
        "\n",
        "## To regenerate your own sounds\n",
        "- Enter the path to an audio file you want to regenerate, or upload when prompted\n",
        "- Make sure the \"skip_for_run_all\" checkbox is unchecked\n",
        "- Run the cell under the \"Regenerate your own sounds\" header\n",
        "\n",
        "## To interpolate between two different sounds\n",
        "- Enter the paths to two audio files you want to interpolate between, or upload them when prompted\n",
        "- Make sure the \"skip_for_run_all\" checkbox is unchecked\n",
        "- Run the cell under the \"Interpolate between sounds\" header"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pJkAc1j4pfAt"
      },
      "source": [
        "### Credits & License"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cjHsmepZqNMs"
      },
      "source": [
        "\n",
        "\n",
        "Original notebook by Zach Evans (https://github.com/zqevans, https://twitter.com/zqevans). \n",
        "\n",
        "Overall structure and setup code taken from Disco Diffusion (https://www.discodiffusion.com)\n",
        "\n",
        "Interpolation and audio display code from CRASH inference notebook (https://github.com/simonrouard/CRASH)\n",
        "\n",
        "Spruced up by Chris the Wizard (https://twitter.com/chris_wizard)\n",
        "\n"
      ]
    },
    {
      "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": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "aBF6kI89LJ5O"
      },
      "outputs": [],
      "source": [
        "#@title <- View Changelog\n",
        "skip_for_run_all = True #@param {type: 'boolean'}\n",
        "\n",
        "if skip_for_run_all == False:\n",
        "  print(\n",
        "      '''\n",
        "\n",
        "  v0.1 Update: Jul 30, 2022 - zqevans\n",
        "      \n",
        "      - Added Dance Diffusion notebook\n",
        "\n",
        "  v0.2 Update: Aug 8, 2022 - zqevans\n",
        "\n",
        "      - Moved to models trained on \n",
        "\n",
        "  v0.3 Update: Aug 11, 2022 - zqevans\n",
        "\n",
        "      - Reverted to old model architecture\n",
        "      - Fixed CRASH sampling code\n",
        "\n",
        "  v0.4 Update: Aug 16, 2022 - zqevans\n",
        "\n",
        "      - Added jmann-small-190k model\n",
        "\n",
        "  v0.5 Update: Aug 17, 2022 - zqevans\n",
        "\n",
        "      - Added interpolations\n",
        "    \n",
        "  v0.6 Update: Aug 18, 2022 - zqevans\n",
        "\n",
        "      - Fixed bug in interpolations\n",
        "\n",
        "  v0.7 Update: Aug 20, 2022 - zqevans\n",
        "      - Added maestro-150k model\n",
        "      - Added unlocked-250k model\n",
        "      - Improved documentation\n",
        "\n",
        "  v0.7.1 Update: Aug 21, 2022 - chris the wizard\n",
        "      - Added introduction\n",
        "      - Added instructions\n",
        "      - Added skips for sections\n",
        "      - Added upload prompts for audio files\n",
        "      - Removed stale demos\n",
        "\n",
        "  v0.8 Update: Aug 24, 2022 - zqevans\n",
        "      - Added Honk model\n",
        "      - Removed Rave Archive model\n",
        "      - Added sample length multipliers and batch sizes to regeneration and interpolation\n",
        "\n",
        "  v0.9 Update: Aug 24, 2022 - zqevans\n",
        "      - Added glitch.cool model\n",
        "      - Added jmann-large model\n",
        "      - Regenerated sounds are now output individually\n",
        "      - Added custom model sample_size and sample_rate options\n",
        "\n",
        "  v0.10 Update: Sep 26, 2022 - morganmcg1\n",
        "      - Added optional, off by default, Weights & Biases logging of the generated audio samples\n",
        "    '''\n",
        "  )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lU97ZiP7nSKS"
      },
      "source": [
        "# Setup\n",
        "Run everything in this section before any generation"
      ]
    },
    {
      "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",
        "    # Using the !wget command instead of the subprocess to get the loading bar\n",
        "    !wget $url -O $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",
        "    #@markdown Check to connect your Google Drive\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",
        "        ai_root = '/content/drive/MyDrive/AI'\n",
        "        root_path = f'{ai_root}/Dance_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'{ai_root}/models'\n",
        "        createPath(model_path)\n",
        "else:\n",
        "    model_path = f'{root_path}/models'\n",
        "    createPath(model_path)\n",
        "\n",
        "# libraries = f'{root_path}/libraries'\n",
        "# createPath(libraries)\n",
        "\n",
        "#@markdown Check the box below to save your generated audio to [Weights & Biases](https://wandb.ai/site)\n",
        "save_to_wandb = False #@param {type: \"boolean\"}\n",
        "\n",
        "if save_to_wandb:\n",
        "    print(\"\\nInstalling wandb...\")\n",
        "    os.system(\"pip install -qqq wandb --upgrade\")\n",
        "    import wandb\n",
        "    # Check if logged in to wandb\n",
        "    try:\n",
        "      import netrc\n",
        "      netrc.netrc().hosts['api.wandb.ai']\n",
        "      wandb.login()\n",
        "    except:\n",
        "      print(\"\\nPlease log in to Weights & Biases...\")\n",
        "      print(\"1. Sign up for a free wandb account here: https://www.wandb.ai/site\")\n",
        "      print(\"2. Enter your wandb API key, from https://wandb.ai/authorize, in the field below to log in: \\n\")\n",
        "      wandb.login()"
      ]
    },
    {
      "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",
        "!git clone --recursive https://github.com/crowsonkb/v-diffusion-pytorch\n",
        "!pip install /content/sample-generator\n",
        "!pip install /content/v-diffusion-pytorch\n",
        "!pip install ipywidgets==7.7.1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "cellView": "form",
        "id": "haxvUGZ0VpzA"
      },
      "outputs": [],
      "source": [
        "#@title Imports and definitions\n",
        "from prefigure.prefigure import get_all_args\n",
        "from contextlib import contextmanager\n",
        "from copy import deepcopy\n",
        "import math\n",
        "from pathlib import Path\n",
        "from google.colab import files\n",
        "\n",
        "import sys\n",
        "import gc\n",
        "\n",
        "from diffusion import sampling\n",
        "import torch\n",
        "from torch import optim, nn\n",
        "from torch.nn import functional as F\n",
        "from torch.utils import data\n",
        "from tqdm import trange\n",
        "from einops import rearrange\n",
        "\n",
        "import torchaudio\n",
        "from audio_diffusion.models import DiffusionAttnUnet1D\n",
        "import numpy as np\n",
        "\n",
        "import random\n",
        "import matplotlib.pyplot as plt\n",
        "import IPython.display as ipd\n",
        "from audio_diffusion.utils import Stereo, PadCrop\n",
        "from glob import glob\n",
        "\n",
        "#@title Model code\n",
        "class DiffusionUncond(nn.Module):\n",
        "    def __init__(self, global_args):\n",
        "        super().__init__()\n",
        "\n",
        "        self.diffusion = DiffusionAttnUnet1D(global_args, n_attn_layers = 4)\n",
        "        self.diffusion_ema = deepcopy(self.diffusion)\n",
        "        self.rng = torch.quasirandom.SobolEngine(1, scramble=True)\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import IPython.display as ipd\n",
        "\n",
        "def plot_and_hear(audio, sr):\n",
        "    display(ipd.Audio(audio.cpu().clamp(-1, 1), rate=sr))\n",
        "    plt.plot(audio.cpu().t().numpy())\n",
        "  \n",
        "def load_to_device(path, sr):\n",
        "    audio, file_sr = torchaudio.load(path)\n",
        "    if sr != file_sr:\n",
        "      audio = torchaudio.transforms.Resample(file_sr, sr)(audio)\n",
        "    audio = audio.to(device)\n",
        "    return audio\n",
        "\n",
        "def get_alphas_sigmas(t):\n",
        "    \"\"\"Returns the scaling factors for the clean image (alpha) and for the\n",
        "    noise (sigma), given a timestep.\"\"\"\n",
        "    return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2)\n",
        "\n",
        "def get_crash_schedule(t):\n",
        "    sigma = torch.sin(t * math.pi / 2) ** 2\n",
        "    alpha = (1 - sigma ** 2) ** 0.5\n",
        "    return alpha_sigma_to_t(alpha, sigma)\n",
        "\n",
        "def t_to_alpha_sigma(t):\n",
        "    \"\"\"Returns the scaling factors for the clean image and for the noise, given\n",
        "    a timestep.\"\"\"\n",
        "    return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2)\n",
        "\n",
        "def alpha_sigma_to_t(alpha, sigma):\n",
        "    \"\"\"Returns a timestep, given the scaling factors for the clean image and for\n",
        "    the noise.\"\"\"\n",
        "    return torch.atan2(sigma, alpha) / math.pi * 2\n",
        "\n",
        "#@title Args\n",
        "sample_size = 65536 \n",
        "sample_rate = 48000   \n",
        "latent_dim = 0             \n",
        "\n",
        "class Object(object):\n",
        "    pass\n",
        "\n",
        "args = Object()\n",
        "args.sample_size = sample_size\n",
        "args.sample_rate = sample_rate\n",
        "args.latent_dim = latent_dim\n",
        "\n",
        "#@title Logging\n",
        "def get_one_channel(audio_data, channel):\n",
        "  '''\n",
        "  Takes a numpy audio array and returns 1 channel\n",
        "  '''\n",
        "  # Check if the audio has more than 1 channel \n",
        "  if len(audio_data.shape) > 1:\n",
        "    is_stereo = True      \n",
        "    if np.argmax(audio_data.shape)==0:\n",
        "        audio_data = audio_data[:,channel] \n",
        "    else:\n",
        "        audio_data = audio_data[channel,:]\n",
        "  else:\n",
        "    is_stereo = False\n",
        "\n",
        "  return audio_data\n",
        "\n",
        "def log_audio_to_wandb(\n",
        "    generated, model_name, custom_ckpt_path, steps, batch_size, sample_rate, sample_size, \n",
        "    generated_all=None, channel=0, original_sample=None, gen_type='new_sounds', noise_level=None, sample_length_mult=None, file_path=None\n",
        "    ):\n",
        "\n",
        "    print('\\nSaving your audio generations to Weights & Biases...')\n",
        "\n",
        "    # Get model name\n",
        "    if model_name == \"custom\":\n",
        "      wandb_model_name = custom_ckpt_path\n",
        "    else:\n",
        "      wandb_model_name = model_name\n",
        "    \n",
        "    # Create config to log to wandb\n",
        "    wandb_config = {\n",
        "      \"model\":model_name,\n",
        "      \"steps\":steps,\n",
        "      \"batch_size\":batch_size,\n",
        "      \"sample_rate\":sample_rate,\n",
        "      \"sample_size\":sample_size,\n",
        "      \"channel\":channel,\n",
        "      \"gen_type\":gen_type,\n",
        "      \"noise_level\":noise_level,\n",
        "      \"sample_length_mult\":sample_length_mult,\n",
        "      \"file_path\":file_path\n",
        "    }\n",
        "\n",
        "    # Create a new wandb run\n",
        "    wandb.init(project='harmonai-audio-gen', config=wandb_config)\n",
        "    wandb_run_url = wandb.run.get_url()\n",
        "\n",
        "    # Create a Weights & Biases Table\n",
        "    audio_generations_table = wandb.Table(columns=['audio', 'steps', 'model', 'batch_size', \n",
        "      'sample_rate', 'sample_size', 'duration'])\n",
        "\n",
        "    # Add each individual generated sample to a wandb Table\n",
        "    for idx, g in enumerate(generated.cpu().numpy()):\n",
        "    \n",
        "      # Check if the audio has more than 1 channel \n",
        "      if idx==0:  \n",
        "        if len(g.shape) > 1:\n",
        "          stereo = True      \n",
        "        else:\n",
        "          stereo = False\n",
        "\n",
        "      if stereo:\n",
        "        g = g[channel]\n",
        "\n",
        "      duration = np.max(g.shape) / sample_rate \n",
        "      wandb_audio =  wandb.Audio(g, sample_rate=sample_rate, caption=wandb_model_name)\n",
        "      audio_generations_table.add_data(wandb_audio, steps, wandb_model_name, batch_size, \n",
        "        sample_rate, sample_size, duration)\n",
        "\n",
        "    # Log the samples Tables and finish the wandb run\n",
        "    wandb.log({f'{gen_type}/harmonai_generations' : audio_generations_table})\n",
        "    \n",
        "    # Log the combined samples in another wandb Table\n",
        "    if generated_all is not None:\n",
        "      g_all = get_one_channel(generated_all, channel)\n",
        "      duration_all = np.max(g_all.shape) / sample_rate \n",
        "      audio_all_generations_table = wandb.Table(columns=['audio', 'steps', 'model', 'batch_size', \n",
        "        'sample_rate', 'sample_size', 'duration'])\n",
        "      wandb_all_audio = wandb.Audio(g_all.cpu().numpy(), sample_rate=sample_rate, caption=wandb_model_name)\n",
        "      audio_all_generations_table.add_data(wandb_all_audio, steps, wandb_model_name, batch_size, \n",
        "        sample_rate, sample_size, duration_all)\n",
        "      wandb.log({f'{gen_type}/all_harmonai_generations' : audio_all_generations_table})\n",
        "\n",
        "    if original_sample is not None:\n",
        "      original_sample = get_one_channel(original_sample, channel)\n",
        "      audio_original_sample_table = wandb.Table(columns=['audio', 'file_path'])\n",
        "      wandb_original_audio = wandb.Audio(original_sample, sample_rate=sample_rate)\n",
        "      audio_original_sample_table.add_data(wandb_original_audio, file_path)\n",
        "      wandb.log({f'{gen_type}/original_sample' : audio_original_sample_table})\n",
        "    \n",
        "    wandb.finish()\n",
        "\n",
        "    print(f'Your audio generations are saved in Weights & Biases here: {wandb_run_url}\\n')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SMQ8vYNQO22Y"
      },
      "source": [
        "# Model settings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LWxBqHH_Yjvt"
      },
      "source": [
        "Select the model you want to sample from:\n",
        "---\n",
        "Model name | Description | Sample rate | Output samples\n",
        "--- | --- | --- | ---\n",
        "glitch-440k |Trained on clips from samples provided by [glitch.cool](https://glitch.cool) | 48000 | 65536\n",
        "jmann-small-190k |Trained on clips from Jonathan Mann's [Song-A-Day](https://songaday.world/) project | 48000 | 65536\n",
        "jmann-large-580k |Trained on clips from Jonathan Mann's [Song-A-Day](https://songaday.world/) project | 48000 | 131072\n",
        "maestro-150k |Trained on piano clips from the [MAESTRO](https://magenta.tensorflow.org/datasets/maestro) dataset | 16000 | 65536\n",
        "unlocked-250k |Trained on clips from the [Unlocked Recordings](https://archive.org/details/unlockedrecordings) dataset | 16000 | 65536\n",
        "honk-140k |Trained on recordings of the Canada Goose from [xeno-canto](https://xeno-canto.org/) | 16000 | 65536\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "JHsHQcc6rHu7"
      },
      "outputs": [],
      "source": [
        "from urllib.parse import urlparse\n",
        "import hashlib\n",
        "#@title Create the model\n",
        "model_name = \"glitch-440k\" #@param [\"glitch-440k\", \"jmann-small-190k\", \"jmann-large-580k\", \"maestro-150k\", \"unlocked-250k\", \"honk-140k\", \"custom\"]\n",
        "\n",
        "#@markdown ###Custom options\n",
        "\n",
        "#@markdown If you have a custom fine-tuned model, choose \"custom\" above and enter a path to the model checkpoint here\n",
        "\n",
        "#@markdown These options will not affect non-custom models\n",
        "custom_ckpt_path = ''#@param {type: 'string'}\n",
        "\n",
        "custom_sample_rate = 16000 #@param {type: 'number'}\n",
        "custom_sample_size = 65536 #@param {type: 'number'}\n",
        "\n",
        "models_map = {\n",
        "\n",
        "    \"glitch-440k\": {'downloaded': False,\n",
        "                         'sha': \"48caefdcbb7b15e1a0b3d08587446936302535de74b0e05e0d61beba865ba00a\", \n",
        "                         'uri_list': [\"https://model-server.zqevans2.workers.dev/gwf-440k.ckpt\"],\n",
        "                         'sample_rate': 48000,\n",
        "                         'sample_size': 65536\n",
        "                         },\n",
        "    \"jmann-small-190k\": {'downloaded': False,\n",
        "                         'sha': \"1e2a23a54e960b80227303d0495247a744fa1296652148da18a4da17c3784e9b\", \n",
        "                         'uri_list': [\"https://model-server.zqevans2.workers.dev/jmann-small-190k.ckpt\"],\n",
        "                         'sample_rate': 48000,\n",
        "                         'sample_size': 65536\n",
        "                         },\n",
        "    \"jmann-large-580k\": {'downloaded': False,\n",
        "                         'sha': \"6b32b5ff1c666c4719da96a12fd15188fa875d6f79f8dd8e07b4d54676afa096\", \n",
        "                         'uri_list': [\"https://model-server.zqevans2.workers.dev/jmann-large-580k.ckpt\"],\n",
        "                         'sample_rate': 48000,\n",
        "                         'sample_size': 131072\n",
        "                         },\n",
        "    \"maestro-150k\": {'downloaded': False,\n",
        "                         'sha': \"49d9abcae642e47c2082cec0b2dce95a45dc6e961805b6500204e27122d09485\", \n",
        "                         'uri_list': [\"https://model-server.zqevans2.workers.dev/maestro-uncond-150k.ckpt\"],\n",
        "                         'sample_rate': 16000,\n",
        "                         'sample_size': 65536\n",
        "                         },\n",
        "    \"unlocked-250k\": {'downloaded': False,\n",
        "                         'sha': \"af337c8416732216eeb52db31dcc0d49a8d48e2b3ecaa524cb854c36b5a3503a\", \n",
        "                         'uri_list': [\"https://model-server.zqevans2.workers.dev/unlocked-uncond-250k.ckpt\"],\n",
        "                         'sample_rate': 16000,\n",
        "                         'sample_size': 65536\n",
        "                         },\n",
        "    \"honk-140k\": {'downloaded': False,\n",
        "                         'sha': \"a66847844659d287f55b7adbe090224d55aeafdd4c2b3e1e1c6a02992cb6e792\", \n",
        "                         'uri_list': [\"https://model-server.zqevans2.workers.dev/honk-140k.ckpt\"],\n",
        "                         'sample_rate': 16000,\n",
        "                         'sample_size': 65536\n",
        "                         },\n",
        "}\n",
        "\n",
        "#@markdown If you're having issues with model downloads, check this to compare the SHA:\n",
        "check_model_SHA = True #@param{type:\"boolean\"}\n",
        "\n",
        "def get_model_filename(diffusion_model_name):\n",
        "    model_uri = models_map[diffusion_model_name]['uri_list'][0]\n",
        "    model_filename = os.path.basename(urlparse(model_uri).path)\n",
        "    return model_filename\n",
        "\n",
        "def download_model(diffusion_model_name, uri_index=0):\n",
        "    if diffusion_model_name != 'custom':\n",
        "        model_filename = get_model_filename(diffusion_model_name)\n",
        "        model_local_path = os.path.join(model_path, model_filename)\n",
        "        if os.path.exists(model_local_path) and check_model_SHA:\n",
        "            print(f'Checking {diffusion_model_name} File')\n",
        "            with open(model_local_path, \"rb\") as f:\n",
        "                bytes = f.read() \n",
        "                hash = hashlib.sha256(bytes).hexdigest()\n",
        "                print(f'SHA: {hash}')\n",
        "            if hash == models_map[diffusion_model_name]['sha']:\n",
        "                print(f'{diffusion_model_name} SHA matches')\n",
        "                models_map[diffusion_model_name]['downloaded'] = True\n",
        "            else:\n",
        "                print(f\"{diffusion_model_name} SHA doesn't match. Will redownload it.\")\n",
        "        elif os.path.exists(model_local_path) and not check_model_SHA or models_map[diffusion_model_name]['downloaded']:\n",
        "            print(f'{diffusion_model_name} already downloaded. If the file is corrupt, enable check_model_SHA.')\n",
        "            models_map[diffusion_model_name]['downloaded'] = True\n",
        "\n",
        "        if not models_map[diffusion_model_name]['downloaded']:\n",
        "            for model_uri in models_map[diffusion_model_name]['uri_list']:\n",
        "                wget(model_uri, model_local_path)\n",
        "                with open(model_local_path, \"rb\") as f:\n",
        "                  bytes = f.read() \n",
        "                  hash = hashlib.sha256(bytes).hexdigest()\n",
        "                  print(f'SHA: {hash}')\n",
        "                if os.path.exists(model_local_path):\n",
        "                    models_map[diffusion_model_name]['downloaded'] = True\n",
        "                    return\n",
        "                else:\n",
        "                    print(f'{diffusion_model_name} model download from {model_uri} failed. Will try any fallback uri.')\n",
        "            print(f'{diffusion_model_name} download failed.')\n",
        "\n",
        "if model_name == \"custom\":\n",
        "  ckpt_path = custom_ckpt_path\n",
        "  args.sample_size = custom_sample_size\n",
        "  args.sample_rate = custom_sample_rate\n",
        "else:\n",
        "  model_info = models_map[model_name]\n",
        "  download_model(model_name)\n",
        "  ckpt_path = f'{model_path}/{get_model_filename(model_name)}'\n",
        "  args.sample_size = model_info[\"sample_size\"]\n",
        "  args.sample_rate = model_info[\"sample_rate\"]\n",
        "\n",
        "print(\"Creating the model...\")\n",
        "model = DiffusionUncond(args)\n",
        "model.load_state_dict(torch.load(ckpt_path)[\"state_dict\"])\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "model = model.requires_grad_(False).to(device)\n",
        "print(\"Model created\")\n",
        "\n",
        "# # Remove non-EMA\n",
        "del model.diffusion\n",
        "\n",
        "model_fn = model.diffusion_ema\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_GQK9yZHTr_z"
      },
      "source": [
        "# Generate new sounds\n",
        "\n",
        "Feeding white noise into the model to be denoised creates novel sounds in the \"space\" of the training data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "zntGqLTJq6xU"
      },
      "outputs": [],
      "source": [
        "#@markdown How many audio clips to create\n",
        "batch_size =  16#@param {type:\"number\"}\n",
        "\n",
        "#@markdown Number of steps (100 is a good start, more steps trades off speed for quality)\n",
        "steps = 100 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown Check the box below to save your generated audio to [Weights & Biases](https://www.wandb.ai/site)\n",
        "save_new_generations_to_wandb = False #@param {type: \"boolean\"}\n",
        "\n",
        "#@markdown Check the box below to skip this section when running all cells\n",
        "skip_for_run_all = False #@param {type: \"boolean\"}\n",
        "\n",
        "if not skip_for_run_all:\n",
        "  torch.cuda.empty_cache()\n",
        "  gc.collect()\n",
        "\n",
        "  # Generate random noise to sample from\n",
        "  noise = torch.randn([batch_size, 2, args.sample_size]).to(device)\n",
        "\n",
        "  t = torch.linspace(1, 0, steps + 1, device=device)[:-1]\n",
        "  step_list = get_crash_schedule(t)\n",
        "\n",
        "  # Generate the samples from the noise\n",
        "  generated = sampling.iplms_sample(model_fn, noise, step_list, {})\n",
        "\n",
        "  # Hard-clip the generated audio\n",
        "  generated = generated.clamp(-1, 1)\n",
        "\n",
        "  # Put the demos together\n",
        "  generated_all = rearrange(generated, 'b d n -> d (b n)')\n",
        "\n",
        "  print(\"All samples\")\n",
        "  plot_and_hear(generated_all, args.sample_rate)\n",
        "  for ix, gen_sample in enumerate(generated):\n",
        "    print(f'sample #{ix + 1}')\n",
        "    display(ipd.Audio(gen_sample.cpu(), rate=args.sample_rate))\n",
        "\n",
        "  # If Weights & Biases logging enabled, save generations\n",
        "  if save_new_generations_to_wandb:\n",
        "    # Check if logged in to wandb\n",
        "    try:\n",
        "      import netrc\n",
        "      netrc.netrc().hosts['api.wandb.ai']\n",
        "\n",
        "      log_audio_to_wandb(generated, model_name, custom_ckpt_path, steps, batch_size, \n",
        "      args.sample_rate, args.sample_size, generated_all=generated_all)\n",
        "    except:\n",
        "      print(\"Not logged in to Weights & Biases, please tick the `save_to_wandb` box at the top of this notebook and run that cell again to log in to Weights & Biases first\")\n",
        "\n",
        "else:\n",
        "  print(\"Skipping section, uncheck 'skip_for_run_all' to enable\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v0WKP7ku67vn"
      },
      "source": [
        "# Regenerate your own sounds\n",
        "By adding noise to an audio file and running it through the model to be denoised, new details will be created, pulling the audio closer to the \"sonic space\" of the model. The more noise you add, the more the sound will change.\n",
        "\n",
        "The effect of this is a kind of \"style transfer\" on the audio. For those familiar with image generation models, this is analogous to an \"init image\"."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "bKgS7vZc4lN9"
      },
      "outputs": [],
      "source": [
        "\n",
        "#@markdown Enter a path to an audio file you want to alter, or leave blank to upload a file (.wav or .flac)\n",
        "file_path = \"\" #@param{type:\"string\"}\n",
        "\n",
        "#@markdown Total number of steps (100 is a good start, more steps trades off speed for quality)\n",
        "steps = 100#@param {type:\"number\"}\n",
        "\n",
        "#@markdown How much (0-1) to re-noise the original sample. Adding more noise (a higher number) means a bigger change to the input audio\n",
        "noise_level = 0.6#@param {type:\"number\"}\n",
        "\n",
        "#@markdown Multiplier on the default sample length from the model, allows for longer audio clips at the expense of VRAM\n",
        "sample_length_mult = 4#@param {type:\"number\"}\n",
        "\n",
        "#@markdown How many variations to create\n",
        "batch_size = 4 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown Check the box below to save your generated audio to [Weights & Biases](https://www.wandb.ai/site)\n",
        "save_own_generations_to_wandb = False #@param {type: \"boolean\"}\n",
        "\n",
        "#@markdown Check the box below to skip this section when running all cells\n",
        "skip_for_run_all = False #@param {type: \"boolean\"}\n",
        "\n",
        "effective_length = args.sample_size * sample_length_mult\n",
        "\n",
        "if not skip_for_run_all:\n",
        "  torch.cuda.empty_cache()\n",
        "  gc.collect()\n",
        "\n",
        "  if file_path == \"\":\n",
        "    print(\"No file path provided, please upload a file\")\n",
        "    uploaded = files.upload()\n",
        "    file_path = list(uploaded.keys())[0]\n",
        "\n",
        "  augs = torch.nn.Sequential(\n",
        "    PadCrop(effective_length, randomize=True),\n",
        "    Stereo()\n",
        "  )\n",
        "\n",
        "  audio_sample = load_to_device(file_path, args.sample_rate)\n",
        "\n",
        "  audio_sample = augs(audio_sample).unsqueeze(0).repeat([batch_size, 1, 1])\n",
        "\n",
        "  print(\"Initial audio sample\")\n",
        "  plot_and_hear(audio_sample[0], args.sample_rate)\n",
        "\n",
        "  t = torch.linspace(0, 1, steps + 1, device=device)\n",
        "  step_list = get_crash_schedule(t)\n",
        "  step_list = step_list[step_list < noise_level]\n",
        "\n",
        "  alpha, sigma = t_to_alpha_sigma(step_list[-1])\n",
        "  noised = torch.randn([batch_size, 2, effective_length], device='cuda')\n",
        "  noised = audio_sample * alpha + noised * sigma\n",
        "\n",
        "  generated = sampling.iplms_sample(model_fn, noised, step_list.flip(0)[:-1], {})\n",
        "\n",
        "  print(\"Regenerated audio samples\")\n",
        "  plot_and_hear(rearrange(generated, 'b d n -> d (b n)'), args.sample_rate)\n",
        "\n",
        "  for ix, gen_sample in enumerate(generated):\n",
        "    print(f'sample #{ix + 1}')\n",
        "    display(ipd.Audio(gen_sample.cpu(), rate=args.sample_rate))\n",
        "\n",
        "  # If Weights & Biases logging enabled, save generations\n",
        "  if save_own_generations_to_wandb:\n",
        "    # Check if logged in to wandb\n",
        "    try:\n",
        "      import netrc\n",
        "      netrc.netrc().hosts['api.wandb.ai']\n",
        "\n",
        "      log_audio_to_wandb(generated, model_name, custom_ckpt_path, steps, batch_size, \n",
        "        args.sample_rate, args.sample_size, file_path=file_path, original_sample=audio_sample[0].cpu().numpy(),\n",
        "        noise_level=noise_level, gen_type='own_file')\n",
        "    except:\n",
        "      print(\"Not logged in to Weights & Biases, please tick the `save_to_wandb` box at the top of this notebook and run that cell again to log in to Weights & Biases first\")\n",
        "\n",
        "else:\n",
        "  print(\"Skipping section, uncheck 'skip_for_run_all' to enable\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vW8N8GCCM-yT"
      },
      "source": [
        "# Interpolate between sounds\n",
        "Diffusion models allow for interpolation between inputs through a process of deterministic noising and denoising. \n",
        "\n",
        "By deterministically noising two audio files, interpolating between the results, and deterministically denoising them, we can can create new sounds \"between\" the audio files provided."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "l3Al3thgO5rb"
      },
      "outputs": [],
      "source": [
        "# Interpolation code taken and modified from CRASH\n",
        "def compute_interpolation_in_latent(latent1, latent2, lambd):\n",
        "    '''\n",
        "    Implementation of Spherical Linear Interpolation: https://en.wikipedia.org/wiki/Slerp\n",
        "    latent1: tensor of shape (2, n)\n",
        "    latent2: tensor of shape (2, n)\n",
        "    lambd: list of floats between 0 and 1 representing the parameter t of the Slerp\n",
        "    '''\n",
        "    device = latent1.device\n",
        "    lambd = torch.tensor(lambd)\n",
        "\n",
        "    assert(latent1.shape[0] == latent2.shape[0])\n",
        "\n",
        "    # get the number of channels\n",
        "    nc = latent1.shape[0]\n",
        "    interps = []\n",
        "    for channel in range(nc):\n",
        "    \n",
        "      cos_omega = latent1[channel]@latent2[channel] / \\\n",
        "          (torch.linalg.norm(latent1[channel])*torch.linalg.norm(latent2[channel]))\n",
        "      omega = torch.arccos(cos_omega).item()\n",
        "\n",
        "      a = torch.sin((1-lambd)*omega) / np.sin(omega)\n",
        "      b = torch.sin(lambd*omega) / np.sin(omega)\n",
        "      a = a.unsqueeze(1).to(device)\n",
        "      b = b.unsqueeze(1).to(device)\n",
        "      interps.append(a * latent1[channel] + b * latent2[channel])\n",
        "    return rearrange(torch.cat(interps), \"(c b) n -> b c n\", c=nc) \n",
        "\n",
        "#@markdown Enter the paths to two audio files to interpolate between (.wav or .flac)\n",
        "source_audio_path = \"\" #@param{type:\"string\"}\n",
        "target_audio_path = \"\" #@param{type:\"string\"}\n",
        "\n",
        "#@markdown Total number of steps (100 is a good start, can go lower for more speed/less quality)\n",
        "steps = 100#@param {type:\"number\"}\n",
        "\n",
        "#@markdown Number of interpolated samples\n",
        "n_interps = 12 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown Multiplier on the default sample length from the model, allows for longer audio clips at the expense of VRAM\n",
        "sample_length_mult = 1#@param {type:\"number\"}\n",
        "\n",
        "#@markdown Check the box below to skip this section when running all cells\n",
        "skip_for_run_all = False #@param {type: \"boolean\"}\n",
        "\n",
        "effective_length = args.sample_size * sample_length_mult\n",
        "\n",
        "if not skip_for_run_all:\n",
        "\n",
        "  augs = torch.nn.Sequential(\n",
        "    PadCrop(effective_length, randomize=True),\n",
        "    Stereo()\n",
        "  )\n",
        "\n",
        "  if source_audio_path == \"\":\n",
        "    print(\"No file path provided for the source audio, please upload a file\")\n",
        "    uploaded = files.upload()\n",
        "    source_audio_path = list(uploaded.keys())[0]\n",
        "\n",
        "  audio_sample_1 = load_to_device(source_audio_path, args.sample_rate)\n",
        "\n",
        "  print(\"Source audio sample loaded\")\n",
        "\n",
        "  if target_audio_path == \"\":\n",
        "    print(\"No file path provided for the target audio, please upload a file\")\n",
        "    uploaded = files.upload()\n",
        "    target_audio_path = list(uploaded.keys())[0]\n",
        "\n",
        "  audio_sample_2 = load_to_device(target_audio_path, args.sample_rate)\n",
        "\n",
        "  print(\"Target audio sample loaded\")\n",
        "\n",
        "  audio_samples = augs(audio_sample_1).unsqueeze(0).repeat([2, 1, 1])\n",
        "  audio_samples[1] = augs(audio_sample_2)\n",
        "\n",
        "  print(\"Initial audio samples\")\n",
        "  plot_and_hear(audio_samples[0], args.sample_rate)\n",
        "  plot_and_hear(audio_samples[1], args.sample_rate)\n",
        "\n",
        "  t = torch.linspace(0, 1, steps + 1, device=device)\n",
        "  step_list = get_crash_schedule(t)\n",
        "\n",
        "  reversed = sampling.iplms_sample(model_fn, audio_samples, step_list, {}, is_reverse=True)\n",
        "\n",
        "  latent_series = compute_interpolation_in_latent(reversed[0], reversed[1], [k/n_interps for k in range(n_interps + 2)])\n",
        "\n",
        "  generated = sampling.iplms_sample(model_fn, latent_series, step_list.flip(0)[:-1], {})\n",
        "\n",
        "  # Put the demos together\n",
        "  generated_all = rearrange(generated, 'b d n -> d (b n)')\n",
        "\n",
        "  print(\"Full interpolation\")\n",
        "  plot_and_hear(generated_all, args.sample_rate)\n",
        "  for ix, gen_sample in enumerate(generated):\n",
        "    print(f'sample #{ix + 1}')\n",
        "    display(ipd.Audio(gen_sample.cpu(), rate=args.sample_rate))\n",
        "else:\n",
        "  print(\"Skipping section, uncheck 'skip_for_run_all' to enable\") "
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}