Dance_Diffusion.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YRwrORLsarkk"
      },
      "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.12\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 or recording\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",
        "- Select the sampler you want to use.\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",
        "- Record a file, enter the path to an audio file you want to regenerate, or upload a file 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\n",
        "\n",
        "## To regenerate sounds using built-in audio recording widget\n",
        "- Enter a path to save your audio recordings\n",
        "- Enter the number of audio recordings you want to combine into one\n",
        "- Run the cell under the \"Regenerate your own sound from the recording\" header\n",
        "- Make sure the \"skip_for_run_all\" checkbox is unchecked"
      ]
    },
    {
      "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",
        "Audio recording widget added by Niels Bantilan (https://twitter.com/cosmicBboy)\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",
        "  v0.11 Update: Oct 27, 2022 - cosmicBboy\n",
        "      - Added a gradio audio recording widget in the \"Regenerate your own sounds\" section\n",
        "      \n",
        "  v0.12 Update: Nov 3, 2022 - zqevans\n",
        "      - Added k-diffusion samplers\n",
        "      - Added option to increase unconditional generation length\n",
        "    '''\n",
        "  )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9ufsXWZk4Nkv"
      },
      "source": [
        "# Install dependencies and restart notebook\n",
        "\n",
        "**The session will crash to restart automatically, don't worry about the error. This is to ensure the proper dependencies are installed**\n",
        "\n",
        "**Once the notebook restarts, continue with the steps below**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "y9BS0ks1oEgP"
      },
      "outputs": [],
      "source": [
        "from IPython.display import clear_output  \n",
        "import os, signal\n",
        "\n",
        "#@title Install and restart\n",
        "\n",
        "#@markdown \n",
        "\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/v-diffusion-pytorch\n",
        "!pip install ipywidgets==7.7.1 gradio\n",
        "!pip install k-diffusion\n",
        "!pip install /content/sample-generator\n",
        "\n",
        "os.kill(os.getpid(), signal.SIGKILL)"
      ]
    },
    {
      "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.\")\n",
        "    is_colab = True\n",
        "    #@markdown Check to connect your Google Drive\n",
        "    google_drive = True #@param {type:\"boolean\"}\n",
        "    if google_drive:\n",
        "      print(\"Using Google Drive.\")\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": "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 os, signal, 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')\n",
        "\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",
        "import k_diffusion as K\n",
        "\n",
        "#@title Create the model\n",
        "model_name = \"maestro-150k\" #@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 = False #@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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Knzr6CYmtaI_"
      },
      "source": [
        "Select the sampler you want to use\n",
        "---\n",
        "Sampler name | Notes\n",
        "--- | ---\n",
        "v-ddim | This is the sampler used in the training notebook. Needs more steps, but more reliable.\n",
        "v-iplms | Similar to above, but sometimes lower quality.\n",
        "k-heun | Needs fewer steps, but ideal sigma_min and sigma_max need to be found. Doesn't work with all models.\n",
        "k-dpmpp_2s_ancestral | Fastest sampler, but you may have to find new sigmas. Recommended min & max sigmas: 0.01, 80\n",
        "k-lms | \"\n",
        "k-dpm-2 | \"\n",
        "k-dpm-fast | \"\n",
        "k-dpm-adaptive | Takes in extra parameters for quality, step count is non-deterministic"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iyR4w86-ei_z"
      },
      "outputs": [],
      "source": [
        "#@title Sampler options\n",
        "sampler_type = \"v-iplms\" #@param [\"v-ddim\", \"v-iplms\", \"k-heun\", \"k-dpmpp_2s_ancestral\", \"k-lms\", \"k-dpm-2\", \"k-dpm-fast\", \"k-dpm-adaptive\"]\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown <h3>Advanced settings</h3>\n",
        "#@markdown \n",
        "#@markdown **V-diffusion settings**\n",
        "eta = 0 #@param {type: \"number\"}\n",
        "#@markdown **K-diffusion settings**\n",
        "sigma_min = 0.0001 #@param {type: \"number\"}\n",
        "sigma_max = 1 #@param {type: \"number\"}\n",
        "rho=7. #@param {type: \"number\"}\n",
        "#@markdown k-dpm-adaptive settings\n",
        "rtol = 0.01 #@param {type: \"number\"}\n",
        "atol = 0.01 #@param {type: \"number\"}\n",
        "\n",
        "def sample(model_fn, noise, steps=100, sampler_type=\"v-iplms\", noise_level = 1.0):\n",
        "  #Check for k-diffusion\n",
        "  if sampler_type.startswith('k-'):\n",
        "    denoiser = K.external.VDenoiser(model_fn)\n",
        "    sigmas = K.sampling.get_sigmas_karras(steps, sigma_min, sigma_max, rho, device=device)\n",
        "\n",
        "  elif sampler_type.startswith(\"v-\"):\n",
        "    t = torch.linspace(1, 0, steps + 1, device=device)[:-1]\n",
        "    step_list = get_crash_schedule(t)\n",
        "\n",
        "  if sampler_type == \"v-ddim\":\n",
        "    return sampling.sample(model_fn, noise, step_list, eta, {})\n",
        "  elif sampler_type == \"v-iplms\":\n",
        "    return sampling.iplms_sample(model_fn, noise, step_list, {})\n",
        "\n",
        "  elif sampler_type == \"k-heun\":\n",
        "    return K.sampling.sample_heun(denoiser, noise, sigmas, disable=False)\n",
        "  elif sampler_type == \"k-lms\":\n",
        "    return K.sampling.sample_lms(denoiser, noise, sigmas, disable=False)\n",
        "  elif sampler_type == \"k-dpmpp_2s_ancestral\":\n",
        "    return K.sampling.sample_dpmpp_2s_ancestral(denoiser, noise, sigmas, disable=False)\n",
        "  elif sampler_type == \"k-dpm-2\":\n",
        "    return K.sampling.sample_dpm_2(denoiser, noise, sigmas, disable=False)\n",
        "  elif sampler_type == \"k-dpm-fast\":\n",
        "    return K.sampling.sample_dpm_fast(denoiser, noise, sigma_min, sigma_max, steps, disable=False)\n",
        "  elif sampler_type == \"k-dpm-adaptive\":\n",
        "    return K.sampling.sample_dpm_adaptive(denoiser, noise, sigma_min, sigma_max, rtol=rtol, atol=atol, disable=False)\n",
        "\n",
        "def resample(model_fn, audio, steps=100, sampler_type=\"v-iplms\", noise_level = 1.0):\n",
        "  #Noise the input\n",
        "\n",
        "  if sampler_type.startswith(\"v-\"):\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 * alpha + noised * sigma\n",
        "    noise = noised\n",
        "\n",
        "  elif sampler_type.startswith(\"k-\"):\n",
        "    denoiser = K.external.VDenoiser(model_fn)\n",
        "    noised = audio + torch.randn_like(audio) * noise_level\n",
        "    sigmas = K.sampling.get_sigmas_karras(steps, sigma_min, noise_level, rho, device=device)\n",
        "\n",
        "  # Denoise\n",
        "  if sampler_type == \"v-iplms\":\n",
        "    return sampling.iplms_sample(model_fn, noised, step_list.flip(0)[:-1], {})\n",
        "\n",
        "  if sampler_type == \"v-ddim\":\n",
        "    return sampling.sample(model_fn, noise, step_list.flip(0)[:-1], eta, {})\n",
        "\n",
        "  elif sampler_type == \"k-heun\":\n",
        "    return K.sampling.sample_heun(denoiser, noised, sigmas, disable=False)\n",
        "\n",
        "  elif sampler_type == \"k-dpmpp_2s_ancestral\":\n",
        "    return K.sampling.sample_dpmpp_2s_ancestral(denoiser, noised, sigmas, disable=False)\n",
        "\n",
        "  elif sampler_type == \"k-lms\":\n",
        "    return K.sampling.sample_lms(denoiser, noised, sigmas, disable=False)\n",
        "\n",
        "  elif sampler_type == \"k-dpm-2\":\n",
        "    return K.sampling.sample_dpm_2(denoiser, noised, sigmas, s_noise=0., disable=False)\n",
        "\n",
        "  elif sampler_type == \"k-dpm-fast\":\n",
        "    return K.sampling.sample_dpm_fast(denoiser, noised, sigma_min, noise_level, steps, disable=False)\n",
        "\n",
        "  elif sampler_type == \"k-dpm-adaptive\":\n",
        "    return K.sampling.sample_dpm_adaptive(denoiser, noised, sigma_min, noise_level, rtol=rtol, atol=atol, disable=False)\n",
        "\n",
        "def reverse_sample(model_fn, audio, steps=100, sampler_type=\"v-iplms\", noise_level = 1.0):\n",
        "  \n",
        "  if sampler_type.startswith(\"v-\"):\n",
        "    t = torch.linspace(0, 1, steps + 1, device=device)\n",
        "    step_list = get_crash_schedule(t)\n",
        "\n",
        "    if sampler_type == \"v-iplms\":\n",
        "      return sampling.iplms_sample(model_fn, audio_samples, step_list, {}, is_reverse=True)\n",
        "\n",
        "    if sampler_type == \"v-ddim\":\n",
        "      return sampling.sample(model_fn, noise, step_list, eta, {}, is_reverse=True)\n",
        "\n",
        "  elif sampler_type.startswith(\"k-\"):\n",
        "    denoiser = K.external.VDenoiser(model_fn)\n",
        "    sigmas = K.sampling.get_sigmas_karras(steps, sigma_min, noise_level, rho, device=device)\n",
        "\n",
        "  # Denoise\n",
        "  if sampler_type == \"k-heun\":\n",
        "    return K.sampling.sample_heun(denoiser, audio, sigmas.flip(0)[:-1], disable=False)\n",
        "  elif sampler_type == \"k-lms\":\n",
        "    return K.sampling.sample_lms(denoiser, audio, sigmas.flip(0)[:-1], disable=False)\n",
        "  elif sampler_type == \"k-dpmpp_2s_ancestral\":\n",
        "    return K.sampling.sample_dpmpp_2s_ancestral(denoiser, audio, sigmas.flip(0)[:-1], disable=False)\n",
        "  elif sampler_type == \"k-dpm-2\":\n",
        "    return K.sampling.sample_dpm_2(denoiser, audio, sigmas.flip(0)[:-1], s_noise=0., disable=False)\n",
        "  elif sampler_type == \"k-dpm-fast\":\n",
        "    return K.sampling.sample_dpm_fast(denoiser, audio, noise_level, sigma_min, steps, disable=False)\n",
        "\n",
        "  elif sampler_type == \"k-dpm-adaptive\":\n",
        "    return K.sampling.sample_dpm_adaptive(denoiser, audio, noise_level, sigma_min, rtol=rtol, atol=atol, disable=False)\n",
        "\n",
        "\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 =  4#@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 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 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",
        "effective_length = sample_length_mult * args.sample_size\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, effective_length]).to(device)\n",
        "\n",
        "  generated = sample(model_fn, noise, steps, sampler_type)\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": "oTK_JxQ0arky"
      },
      "outputs": [],
      "source": [
        "#@title Record audio or enter a filepath to a prerecorded audio file\n",
        "import torch\n",
        "import torchaudio\n",
        "from typing import Iterable, Tuple\n",
        "import gradio as gr\n",
        "\n",
        "Audio = Tuple[int, np.ndarray]\n",
        "\n",
        "#@markdown Check the box below to create an audio recording interface below\n",
        "record_audio = False #@param {type: \"boolean\"}\n",
        "\n",
        "#@markdown If you left \"record_audio\" blank, 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 Number of audio recordings to combine into one clip. Only applies if the \"record_audio\" box is checked.\n",
        "n_audio_recordings = 1 #@param{type:\"number\"}\n",
        "\n",
        "# this is a global variable to be filled in by the generate_from_audio callback\n",
        "recording_file_path = \"\"\n",
        "\n",
        "\n",
        "def combine_audio(*audio_iterable: Iterable[Audio]) -> Audio:\n",
        "    \"\"\"Combines an iterable of audio signals into one.\"\"\"\n",
        "    max_len = max([x.shape for _, x in audio_iterable])\n",
        "    combined_audio = np.zeros(max_len, dtype=np.int32)\n",
        "    for _, a in audio_iterable:\n",
        "        combined_audio[:a.shape[0]] = combined_audio[:a.shape[0]] * .5 + a * .5\n",
        "    return combined_audio\n",
        "\n",
        "\n",
        "def generate_from_audio(file_path: str, *audio_iterable: Iterable[Audio]):\n",
        "    sample_rate = audio_iterable[0][0]\n",
        "    combined_audio = combine_audio(*audio_iterable)\n",
        "    tensor = torch.from_numpy(\n",
        "        np.concatenate(\n",
        "            [\n",
        "                combined_audio.reshape(1, -1),\n",
        "                combined_audio.reshape(1, -1)\n",
        "            ],\n",
        "            axis=0,\n",
        "        )\n",
        "    )\n",
        "    global recording_file_path\n",
        "    recording_file_path = file_path\n",
        "    torchaudio.save(\n",
        "        file_path,\n",
        "        tensor,\n",
        "        sample_rate=sample_rate,\n",
        "        format=\"wav\"\n",
        "    )\n",
        "    return (sample_rate, combined_audio), file_path\n",
        "\n",
        "if record_audio:\n",
        "    recording_interface = gr.Interface(\n",
        "        fn=generate_from_audio,\n",
        "        inputs=[\n",
        "            gr.Textbox(\n",
        "                \"/content/recording.wav\",\n",
        "                label=\"save recording to filepath\",\n",
        "            ),\n",
        "            *[\n",
        "                gr.Audio(source=\"microphone\", label=f\"audio clip {i}\")\n",
        "                for i in range(1, n_audio_recordings + 1)\n",
        "            ]\n",
        "        ],\n",
        "        outputs=[\n",
        "            gr.Audio(label=\"combined output audio\"),\n",
        "            gr.File(label=\"output file\"),\n",
        "        ],\n",
        "        allow_flagging=\"never\",\n",
        "    )\n",
        "\n",
        "    recording_interface.launch();\n",
        "elif 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",
        "if not record_audio:\n",
        "    print(f\"Using file_path: {file_path} to regenerate new sounds.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "bKgS7vZc4lN9"
      },
      "outputs": [],
      "source": [
        "#@title Generate new sounds from recording\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.3#@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 = 2#@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",
        "  augs = torch.nn.Sequential(\n",
        "    PadCrop(effective_length, randomize=True),\n",
        "    Stereo()\n",
        "  )\n",
        "\n",
        "  fp = recording_file_path if record_audio else file_path\n",
        "\n",
        "  audio_sample = load_to_device(fp, 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",
        "  generated = resample(model_fn, audio_sample, steps, sampler_type, noise_level=noise_level)\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=fp, 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",
        "  reversed = reverse_sample(model_fn, audio_samples, steps)\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 = sample(model_fn, latent_series, steps) \n",
        "  \n",
        "  #sampling.iplms_sample(, 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": {
      "provenance": []
    },
    "gpuClass": "standard",
    "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.9.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}