Jukebox_Diffusion.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "## Jukebox Diffusion\n",
        "by **Jeff Sontag**\n",
        "\n",
        "https://medium.com/@jeffsontagmusic/jukebox-diffusion-cbe22ff3cd47\n",
        "\n",
        "https://github.com/jmoso13/jukebox-diffusion\n",
        "\n",
        "\n",
        "notebook by [johannezz](https://twitter.com/_johannezz)\n",
        "\n",
        "Learn the basics here:\n",
        "[JBDiff Quick Guide](https://docs.google.com/document/d/1vEc2vCQVpN5Ev2BcWk80D8Bk7O0N3XjjFCDCN1Jarao/edit?usp=sharing)"
      ],
      "metadata": {
        "id": "DTBhNttwzJCM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown (optional) Connect to Google drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "id": "TIFByqhuuTr-",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "fRceO45_mRr7"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Install\n",
        "!git clone -b all-requirements https://github.com/deeplearn-art/jukebox-diffusion\n",
        "!pip install -e jukebox-diffusion\n",
        "%cd jukebox-diffusion\n",
        "!pip install -r requirements.txt\n",
        "import os\n",
        "os.mkdir('/content/tmp_save')\n",
        "!python download_ckpts.py\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "#@markdown ##Settings\n",
        "\n",
        "project_name = 'jbdiff-fun'#@param {type: 'string'}\n",
        "context_audio = '/content/audio.wav'#@param {type: 'string'}\n",
        "#@markdown if you use init audio, the length will be the same as init, otherwise:\n",
        "length_in_seconds = 8 #@param {type: 'number'}\n",
        "save_dir = '/content/drive/MyDrive/AI/Jbdiff'#@param {type: 'string'}\n",
        "upsampling_level = 'middle'  #@param [\"low\", \"middle\", \"high\"]\n",
        "\n",
        "# @markdown ---\n",
        "#@markdown Jukebox base layer diffusion\n",
        "noise_seed = -1 #@param {type: 'number'}\n",
        "\n",
        "noise_style = 'constant'  #@param [\"random\", \"constant\", \"region\", \"walk\"]\n",
        "# @markdown How far to wander around init noise (affects only Region and Walk styles)\n",
        "noise_step_size = 0.7 #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "# @markdown ---\n",
        "# @markdown  Dance Diffusion\n",
        "dd_noise_seed = -1 #@param {type: 'number'}\n",
        "\n",
        "dd_noise_style = 'walk'  #@param [\"random\", \"constant\", \"region\", \"walk\"]\n",
        "\n",
        "dd_noise_step_size = 0.15 # @param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "# @markdown ---\n",
        "token_multiplier = 4 # @param {type:\"slider\", min:0, max:32, step:2}\n",
        "# @markdown ---\n",
        "# @markdown Optional:\n",
        "init_audio = ''#@param {type: 'string'}\n",
        "#@markdown diffusion strength, 1 returns the original\n",
        "init_strength = 0.12 #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "\n",
        "# @markdown ---\n",
        "update_lowest_context = False #@param {type: 'boolean'}"
      ],
      "metadata": {
        "id": "TbrHxrNE5gh-",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ##Generate!\n",
        "if upsampling_level == 'high':\n",
        "  level = \"012\"\n",
        "elif upsampling_level == 'low':\n",
        "  level = \"2\"\n",
        "else:\n",
        "  level = \"12\"\n",
        "args = f\"--seconds-length {length_in_seconds} --context-audio {context_audio} --save-dir {save_dir} --project-name {project_name} --levels {level} \"\n",
        "if init_audio != '':\n",
        "  args += f\"--init-audio {init_audio} --init-strength {init_strength} \"\n",
        "if noise_seed > 0:\n",
        "  args += f\"--noise-seed {noise_seed} \"\n",
        "if dd_noise_seed > 0:\n",
        "  args += f\"--dd-noise-seed {dd_noise_seed} \"\n",
        "args += f\"--noise-style {noise_style} --noise-step-size {noise_step_size} \"\n",
        "args += f\"--dd-noise-style {dd_noise_style} --dd-noise-step-size {dd_noise_step_size} \"\n",
        "if token_multiplier > 0:\n",
        "  args += f\"--token-multiplier {token_multiplier} \"\n",
        "args += f\"--update-lowest-context {update_lowest_context} \"\n",
        "print(f'\\n{args}\\n')\n",
        "\n",
        "!python sample.py {args}\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "4lXG_ha3qzsT",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# shutdown\n",
        "from google.colab import runtime\n",
        "runtime.unassign()"
      ],
      "metadata": {
        "id": "Ifnj36FfNdj0"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}