first try

app.py

@@ -0,0 +1,93 @@
+# A simple gradio app that converts music tokens to and from audio using JukeboxVQVAE as the model and Gradio as the UI
+
+from transformers import JukeboxVQVAE
+
+import gradio as gr
+import torch as t
+
+model_id = 'openai/jukebox-5b-lyrics' #@param ['openai/jukebox-1b-lyrics', 'openai/jukebox-5b-lyrics']
+
+if 'google.colab' in sys.modules:
+
+  cache_path = '/content/drive/My Drive/jukebox-webui/_data/' #@param {type:"string"}
+  # Connect to your Google Drive
+  from google.colab import drive
+  drive.mount('/content/drive')
+
+else:
+
+  cache_path = '~/.cache/'
+
+class Convert:
+
+  class TokenList:
+
+    def to_tokens_file(tokens_list):
+      # temporary random file name
+      filename = f"tmp/{t.randint(0, 1000000)}.jt"
+      t.save(validate_tokens_list(tokens_list), filename)
+      return filename
+    
+    def to_audio(tokens_list):
+      return model.decode(validate_tokens_list(tokens_list)[2:], start_level=2).squeeze(-1)
+      # TODO: Implement converting other levels besides 2
+
+  class TokensFile:
+
+    def to_tokens_list(file):
+      return validate_tokens_list(t.load(file))
+    
+    def to_audio(file):
+      return Convert.TokenList.to_audio(Convert.TokensFile.to_tokens_list(file))
+  
+  class Audio:
+
+    def to_tokens_list(audio):
+      return model.encode(audio.unsqueeze(0), start_level=2)
+      # (TODO: Generated by copilot, check if it works)
+
+    def to_tokens_file(audio):
+      return Convert.TokenList.to_tokens_file(Convert.Audio.to_tokens_list(audio))
+
+def init():
+  global model
+
+  model = JukeboxVQVAE.from_pretrained(
+    model_id,
+    device_map = "auto",
+    torch_dtype = t.float16,
+    cache_dir = f"{cache_path}/jukebox/models"
+  )
+
+def validate_tokens_list(tokens_list):
+  # Make sure that:
+  # - tokens_list is a list of exactly 3 torch tensors
+  assert len(tokens_list) == 3, "Invalid file format: expecting a list of 3 tensors"
+
+  # - each has the same number of dimensions
+  assert len(tokens_list[0].shape) == len(tokens_list[1].shape) == len(tokens_list[2].shape), "Invalid file format: each tensor in the list should have the same number of dimensions"
+
+  # - the shape along dimension 0 is the same
+  assert tokens_list[0].shape[0] == tokens_list[1].shape[0] == tokens_list[2].shape[0], "Invalid file format: the shape along dimension 0 should be the same for all tensors in the list"
+
+  # - the shape along dimension 1 increases (or stays the same) as we go from 0 to 2
+  assert tokens_list[0].shape[1] >= tokens_list[1].shape[1] >= tokens_list[2].shape[1], "Invalid file format: the shape along dimension 1 should decrease (or stay the same) as we go from 0 to 2"
+  
+  return tokens_list
+
+
+with gr.Blocks() as ui:
+
+  # File input to upload or download the music tokens file
+  tokens = gr.File(label='music_tokens_file')
+
+  # Audio output to play or upload the generated audio
+  audio = gr.Audio(label='audio')
+  
+  # Buttons to convert from music tokens to audio (primary) and vice versa (secondary)
+  gr.Button(label="Convert tokens to audio", primary=True).click(Convert.TokensFile.to_audio, tokens, audio)
+  gr.Button(label="Convert audio to tokens", primary=False).click(Convert.Audio.to_tokens_file, audio, tokens)
+  
+if __name__ == '__main__':
+  init()
+  ui.launch()

app_to_colab.py

@@ -0,0 +1,83 @@
+# Script to create a notebook out of `requirements.txt` (installing the dependencies) and `app.py`
+
+import json
+from textwrap import dedent
+
+def create_colab(requirements_file='requirements.txt', app_file='app.py'):
+
+  # cells = []
+
+  requirements_txt = open('requirements.txt', 'r').read().replace('\n', '\\n')
+
+  def text_to_cell(text, cell_type='code'):
+
+    lines = dedent(text).splitlines()
+
+    # add a \n to the end of each line except the last one
+    lines = [ f'{line}\n' for line in lines[:-1] ] + [ lines[-1] ]
+    
+    return dict(
+      metadata={},
+      execution_count=None,
+      outputs=[],
+      cell_type=cell_type,
+      source=lines
+    )
+
+  cells = [
+
+    # Cell to mount drive, install the dependencies etc.
+    text_to_cell(f"""\
+      from google.colab import drive
+      mount_drive = True #@param {{type:"boolean"}}
+      if mount_drive:
+        drive.mount('/content/drive')
+
+      requirements_txt = "{requirements_txt}"
+
+      # Save the requirements.txt file
+      with open('requirements.txt', 'w') as f:
+        f.write(requirements_txt)
+
+      # Install the dependencies
+      %pip install -r requirements.txt
+    """),
+
+    # Cell to run the app
+    text_to_cell(open(app_file, 'r').read())
+
+  ]
+
+  # Add notebook metadata
+  metadata = dict(
+    kernelspec = dict(
+      display_name = 'Python 3',
+      language = 'python',
+      name = 'python3'
+    ),
+    language_info = dict(
+      name = 'python',
+      version = '3.7.5',
+    ),
+    orig_nbformat = 4,
+  )
+
+  # Finalize the notebook
+  notebook = dict(
+    cells=cells,
+    metadata=metadata,
+    nbformat=4,
+    nbformat_minor=2,
+  )
+
+  # Name the notebook the same as the parent directory
+  from pathlib import Path
+  notebook_name = f'{Path().absolute().name}.ipynb'
+
+  # Save the notebook in JSON format
+  open(notebook_name, 'w').write(json.dumps(notebook, indent=2))
+
+  return notebook
+
+if __name__ == '__main__':
+  create_colab()

jukwi-vqvae.ipynb

@@ -0,0 +1,140 @@
+{
+  "cells": [
+    {
+      "metadata": {},
+      "execution_count": null,
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "from google.colab import drive\n",
+        "mount_drive = True #@param {type:\"boolean\"}\n",
+        "if mount_drive:\n",
+        "  drive.mount('/content/drive')\n",
+        "\n",
+        "requirements_txt = \"git+https://github.com/ArthurZucker/transformers.git@jukebox\\naccelerate\\nbitsandbytes==0.31.8\\ngradio\"\n",
+        "\n",
+        "# Save the requirements.txt file\n",
+        "with open('requirements.txt', 'w') as f:\n",
+        "  f.write(requirements_txt)\n",
+        "\n",
+        "# Install the dependencies\n",
+        "%pip install -r requirements.txt"
+      ]
+    },
+    {
+      "metadata": {},
+      "execution_count": null,
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "# A simple gradio app that converts music tokens to and from audio using JukeboxVQVAE as the model and Gradio as the UI\n",
+        "\n",
+        "from transformers import JukeboxVQVAE\n",
+        "\n",
+        "import gradio as gr\n",
+        "import torch as t\n",
+        "\n",
+        "model_id = 'openai/jukebox-5b-lyrics' #@param ['openai/jukebox-1b-lyrics', 'openai/jukebox-5b-lyrics']\n",
+        "\n",
+        "if 'google.colab' in sys.modules:\n",
+        "\n",
+        "  cache_path = '/content/drive/My Drive/jukebox-webui/_data/' #@param {type:\"string\"}\n",
+        "  # Connect to your Google Drive\n",
+        "  from google.colab import drive\n",
+        "  drive.mount('/content/drive')\n",
+        "\n",
+        "else:\n",
+        "\n",
+        "  cache_path = '~/.cache/'\n",
+        "\n",
+        "class Convert:\n",
+        "\n",
+        "  class TokenList:\n",
+        "\n",
+        "    def to_tokens_file(tokens_list):\n",
+        "      # temporary random file name\n",
+        "      filename = f\"tmp/{t.randint(0, 1000000)}.jt\"\n",
+        "      t.save(validate_tokens_list(tokens_list), filename)\n",
+        "      return filename\n",
+        "\n",
+        "    def to_audio(tokens_list):\n",
+        "      return model.decode(validate_tokens_list(tokens_list)[2:], start_level=2).squeeze(-1)\n",
+        "      # TODO: Implement converting other levels besides 2\n",
+        "\n",
+        "  class TokensFile:\n",
+        "\n",
+        "    def to_tokens_list(file):\n",
+        "      return validate_tokens_list(t.load(file))\n",
+        "\n",
+        "    def to_audio(file):\n",
+        "      return Convert.TokenList.to_audio(Convert.TokensFile.to_tokens_list(file))\n",
+        "\n",
+        "  class Audio:\n",
+        "\n",
+        "    def to_tokens_list(audio):\n",
+        "      return model.encode(audio.unsqueeze(0), start_level=2)\n",
+        "      # (TODO: Generated by copilot, check if it works)\n",
+        "\n",
+        "    def to_tokens_file(audio):\n",
+        "      return Convert.TokenList.to_tokens_file(Convert.Audio.to_tokens_list(audio))\n",
+        "\n",
+        "def init():\n",
+        "  global model\n",
+        "\n",
+        "  model = JukeboxVQVAE.from_pretrained(\n",
+        "    model_id,\n",
+        "    device_map = \"auto\",\n",
+        "    torch_dtype = t.float16,\n",
+        "    cache_dir = f\"{cache_path}/jukebox/models\"\n",
+        "  )\n",
+        "\n",
+        "def validate_tokens_list(tokens_list):\n",
+        "  # Make sure that:\n",
+        "  # - tokens_list is a list of exactly 3 torch tensors\n",
+        "  assert len(tokens_list) == 3, \"Invalid file format: expecting a list of 3 tensors\"\n",
+        "\n",
+        "  # - each has the same number of dimensions\n",
+        "  assert len(tokens_list[0].shape) == len(tokens_list[1].shape) == len(tokens_list[2].shape), \"Invalid file format: each tensor in the list should have the same number of dimensions\"\n",
+        "\n",
+        "  # - the shape along dimension 0 is the same\n",
+        "  assert tokens_list[0].shape[0] == tokens_list[1].shape[0] == tokens_list[2].shape[0], \"Invalid file format: the shape along dimension 0 should be the same for all tensors in the list\"\n",
+        "\n",
+        "  # - the shape along dimension 1 increases (or stays the same) as we go from 0 to 2\n",
+        "  assert tokens_list[0].shape[1] >= tokens_list[1].shape[1] >= tokens_list[2].shape[1], \"Invalid file format: the shape along dimension 1 should decrease (or stay the same) as we go from 0 to 2\"\n",
+        "\n",
+        "  return tokens_list\n",
+        "\n",
+        "\n",
+        "with gr.Blocks() as ui:\n",
+        "\n",
+        "  # File input to upload or download the music tokens file\n",
+        "  tokens = gr.File(label='music_tokens_file')\n",
+        "\n",
+        "  # Audio output to play or upload the generated audio\n",
+        "  audio = gr.Audio(label='audio')\n",
+        "\n",
+        "  # Buttons to convert from music tokens to audio (primary) and vice versa (secondary)\n",
+        "  gr.Button(label=\"Convert tokens to audio\", primary=True).click(Convert.TokensFile.to_audio, tokens, audio)\n",
+        "  gr.Button(label=\"Convert audio to tokens\", primary=False).click(Convert.Audio.to_tokens_file, audio, tokens)\n",
+        "\n",
+        "if __name__ == '__main__':\n",
+        "  init()\n",
+        "  ui.launch()"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python",
+      "version": "3.7.5"
+    },
+    "orig_nbformat": 4
+  },
+  "nbformat": 4,
+  "nbformat_minor": 2
+}

requirements.txt

@@ -0,0 +1,4 @@
+git+https://github.com/ArthurZucker/transformers.git@jukebox
+accelerate
+bitsandbytes==0.31.8
+gradio