Hybrid Demucs Music Source Separation.ipynb

hybrid-demucs-music-source-separation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Hybrid Demucs Music Source Separation.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/gandresto/2602f93ddc702b7098a71ce5972947d6/hybrid-demucs-music-source-separation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uLghRcAdqaph"
      },
      "source": [
        "# Hybrid Demucs from Colab\n",
        "\n",
        "This supports the Demucs source separation model (https://github.com/facebookresearch/demucs/)\n",
        "This is only for separation with pre-trained models, not training!\n",
        "\n",
        "You can either upload files manually (slow) or link your Google Drive account."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "79JbZGcAqX3p"
      },
      "source": [
        "!pip install demucs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a-2p6OoOrOZK"
      },
      "source": [
        "# Please BE VERY CAREFUL, this will link your entire drive.\n",
        "# So don't edit code, except the one that says 'Customize the following options',\n",
        "# or you might mess up your files.\n",
        "# IF YOU DO NO WANT TO LINK DRIVE, please see below for an alternative!\n",
        "from google.colab import drive\n",
        "drive.mount('/gdrive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "znCvBifRrO-b"
      },
      "source": [
        "# Customize the following options!\n",
        "mp3 = True\n",
        "mp3_rate = 320\n",
        "model = \"mdx_extra\"\n",
        "extensions = [\"mp3\", \"wav\", \"ogg\", \"flac\"]\n",
        "in_path = '/gdrive/MyDrive/demucs/'\n",
        "out_path = '/gdrive/MyDrive/demucs_separated/'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kMKN9BN4r37E",
        "cellView": "form"
      },
      "source": [
        "#@title Useful functions, don't forget to execute\n",
        "import io\n",
        "from pathlib import Path\n",
        "import select\n",
        "from shutil import rmtree\n",
        "import subprocess as sp\n",
        "import sys\n",
        "from typing import Dict, Tuple, Optional, IO\n",
        "\n",
        "from google.colab import files\n",
        "\n",
        "def find_files(in_path):\n",
        "    out = []\n",
        "    for file in Path(in_path).iterdir():\n",
        "        if file.suffix.lower().lstrip(\".\") in extensions:\n",
        "            out.append(file)\n",
        "    return out\n",
        "\n",
        "def copy_process_streams(process: sp.Popen):\n",
        "    def raw(stream: Optional[IO[bytes]]) -> IO[bytes]:\n",
        "        assert stream is not None\n",
        "        if isinstance(stream, io.BufferedIOBase):\n",
        "            stream = stream.raw\n",
        "        return stream\n",
        "\n",
        "    p_stdout, p_stderr = raw(process.stdout), raw(process.stderr)\n",
        "    stream_by_fd: Dict[int, Tuple[IO[bytes], io.StringIO, IO[str]]] = {\n",
        "        p_stdout.fileno(): (p_stdout, sys.stdout),\n",
        "        p_stderr.fileno(): (p_stderr, sys.stderr),\n",
        "    }\n",
        "    fds = list(stream_by_fd.keys())\n",
        "\n",
        "    while fds:\n",
        "        # `select` syscall will wait until one of the file descriptors has content.\n",
        "        ready, _, _ = select.select(fds, [], [])\n",
        "        for fd in ready:\n",
        "            p_stream, std = stream_by_fd[fd]\n",
        "            raw_buf = p_stream.read(2 ** 16)\n",
        "            if not raw_buf:\n",
        "                fds.remove(fd)\n",
        "                continue\n",
        "            buf = raw_buf.decode()\n",
        "            std.write(buf)\n",
        "            std.flush()\n",
        "\n",
        "def separate(inp=None, outp=None):\n",
        "    inp = inp or in_path\n",
        "    outp = outp or out_path\n",
        "    cmd = [\"python3\", \"-m\", \"demucs.separate\", \"-o\", str(outp), \"-n\", model]\n",
        "    if mp3:\n",
        "        cmd += [\"--mp3\", f\"--mp3-bitrate={mp3_rate}\"]\n",
        "    files = [str(f) for f in find_files(inp)]\n",
        "    if not files:\n",
        "        print(f\"No valid audio files in {in_path}\")\n",
        "        return\n",
        "    print(\"Going to separate the files:\")\n",
        "    print('\\n'.join(files))\n",
        "    print(\"With command: \", \" \".join(cmd))\n",
        "    p = sp.Popen(cmd + files, stdout=sp.PIPE, stderr=sp.PIPE)\n",
        "    copy_process_streams(p)\n",
        "    p.wait()\n",
        "    if p.returncode != 0:\n",
        "        print(\"Command failed, something went wrong.\")\n",
        "\n",
        "\n",
        "def from_upload():\n",
        "    out_path = Path('separated')\n",
        "    in_path = Path('tmp_in')\n",
        "    \n",
        "    if in_path.exists():\n",
        "        rmtree(in_path)\n",
        "    in_path.mkdir()\n",
        "    \n",
        "    if out_path.exists():\n",
        "        rmtree(out_path)\n",
        "    out_path.mkdir()\n",
        "    \n",
        "    uploaded = files.upload()\n",
        "    for name, content in uploaded.items():\n",
        "        (in_path / name).write_bytes(content)\n",
        "    separate(in_path, out_path)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gr9OJvf-tYyt"
      },
      "source": [
        "# This can be quite slow, in particular the loading, and saving from GDrive. Please be patient!\n",
        "# This is from google drive! Also, this will separate all the files inside the MyDrive/demucs folder,\n",
        "# so when you are happy with the results, remove the songs from there.\n",
        "separate()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v__3gMJawTD0"
      },
      "source": [
        "# This is manual upload and download :)\n",
        "from_upload()\n",
        "!zip -r separated.zip separated\n",
        "files.download('./separated.zip')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}