{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "ffdbda4b",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-13T03:31:44.652173Z",
"start_time": "2023-10-13T03:31:42.215188Z"
}
},
"outputs": [],
"source": [
"import os\n",
"\n",
"os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
"\n",
"import random\n",
"import json\n",
"import numpy as np\n",
"import tqdm\n",
"import torch\n",
"import funcy\n",
"import time\n",
"import gc\n",
"from scipy.io import wavfile\n",
"import tempfile\n",
"import collections\n",
"from collections import defaultdict\n",
"from joblib import Parallel, delayed\n",
"\n",
"from suno_utils.utils.s3 import _apply_mp\n",
"from suno_utils.audio import Audio\n",
"from suno_utils.tasks.data_loader import load_audio_mp\n",
"from suno_utils.utils.text import write_jsonl, read_jsonl, write_json, read_json\n",
"from suno_utils.utils.s3 import read_from_s3, check_s3_file_exists, open_from_s3\n",
"from suno_utils.audio.conversion import convert_audio_files\n",
"from suno_utils.tasks import demucs\n",
"from suno_utils.audio import Audio\n",
"from suno_utils.tasks import dac\n",
"\n",
"SAMPLE_RATE = 24_000\n",
"EMBEDDING_RATE = 25\n",
"N_CODEBOOKS = 8\n",
"\n",
"OUT_DATA_DIR = \"/app/suno/data/mert_25hz_short\"\n",
"OUT_AUDIO_DIR = os.path.join(OUT_DATA_DIR, \"audio\")\n",
"OUT_AUDIO_DEMUC_DIR = os.path.join(OUT_DATA_DIR, \"audio_demuc\")\n",
"OUT_TSV_DIR = os.path.join(OUT_DATA_DIR, \"audio_tsv\")\n",
"OUT_LABEL_DIR = os.path.join(OUT_DATA_DIR, \"label\")\n",
"OUT_TEMP_DIR = os.path.join(OUT_DATA_DIR, \"temp\")\n",
"\n",
"os.makedirs(OUT_DATA_DIR, exist_ok=True)\n",
"os.makedirs(OUT_AUDIO_DIR, exist_ok=True)\n",
"os.makedirs(OUT_TSV_DIR, exist_ok=True)\n",
"os.makedirs(OUT_LABEL_DIR, exist_ok=True)\n",
"os.makedirs(OUT_TEMP_DIR, exist_ok=True)\n",
"\n",
"MIN_SIZE = SAMPLE_RATE * 5\n",
"MAX_SIZE = SAMPLE_RATE * 12"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "4dbee7d7",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-05T01:18:33.546301Z",
"start_time": "2023-10-05T01:18:33.530810Z"
}
},
"outputs": [],
"source": [
"a = Audio.from_file(\"/app/suno/data/mert_25hz_short/audio_demuc/0_0_vocals.wav\")"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "580aa05c",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-05T01:33:55.388599Z",
"start_time": "2023-10-05T01:33:55.386161Z"
}
},
"outputs": [
{
"data": {
"text/plain": [
"44100"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"a.sample_rate"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "cac5555c",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-05T01:18:37.766120Z",
"start_time": "2023-10-05T01:18:37.667135Z"
}
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"a.play()"
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "72ea3da7",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-05T01:07:04.766732Z",
"start_time": "2023-10-05T01:07:04.764252Z"
}
},
"outputs": [],
"source": [
"# KNOW WTF you are doing\n",
"# with open(os.path.join(os.path.join(OUT_DATA_DIR, \"audio_resample_tsv\"), f\"valid.tsv\"), \"w\") as f:\n",
"# f.write('/app/suno/data/mert_25hz_short/audio_demuc' + \"\\n\")\n",
"# f.write(\n",
"# \"\\n\".join(\n",
"# [\n",
"# fn.replace(\".wav\", \"_vocals.wav\") + \"\\t\" + dur\n",
"# for fn, dur in tsv_data\n",
"# ]\n",
"# )\n",
"# )"
]
},
{
"cell_type": "markdown",
"id": "c729d86d",
"metadata": {},
"source": [
"# File loading"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "66631638",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-05T01:06:40.800444Z",
"start_time": "2023-10-05T01:06:40.370480Z"
}
},
"outputs": [],
"source": [
"# load current data\n",
"tsv_info = []\n",
"with open(os.path.join(OUT_TSV_DIR, \"valid.tsv\"), \"r\") as f:\n",
" for line in f.read().strip().split(\"\\n\"):\n",
" if len(line.strip()) == 0:\n",
" continue\n",
" tsv_info.append(line.strip().split(\"\\t\"))"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "cad52faa",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-05T01:06:41.340085Z",
"start_time": "2023-10-05T01:06:41.337900Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"374220\n"
]
}
],
"source": [
"print(len(tsv_info))"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "741d5c23",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-05T01:06:42.257164Z",
"start_time": "2023-10-05T01:06:41.809830Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"374219\n"
]
}
],
"source": [
"tsv_data_path = tsv_info[0]\n",
"tsv_data = tsv_info[1:]\n",
"print(len(tsv_data))"
]
},
{
"cell_type": "markdown",
"id": "6bd58f17",
"metadata": {},
"source": [
"## Generate the split audios"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "a4131de4",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-02T15:51:23.357114Z",
"start_time": "2023-10-02T15:51:23.352626Z"
}
},
"outputs": [
{
"data": {
"text/plain": [
"[['0.wav', '288000'],\n",
" ['1.wav', '288000'],\n",
" ['2.wav', '288000'],\n",
" ['3.wav', '288000'],\n",
" ['4.wav', '288000'],\n",
" ['5.wav', '288000'],\n",
" ['6.wav', '288000'],\n",
" ['7.wav', '288000'],\n",
" ['8.wav', '288000'],\n",
" ['9.wav', '288000']]"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"tsv_data[:10]"
]
},
{
"cell_type": "markdown",
"id": "f589ac16",
"metadata": {},
"source": [
"# Setup processing "
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "073644f5",
"metadata": {
"ExecuteTime": {
"end_time": "2023-09-27T13:48:29.533961Z",
"start_time": "2023-09-27T13:48:29.531818Z"
}
},
"outputs": [],
"source": [
"# last job: python encode_audio.py --start_index 2361972 --end_index 3149296 "
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "b631efe3",
"metadata": {
"ExecuteTime": {
"end_time": "2023-09-27T13:48:29.909995Z",
"start_time": "2023-09-27T13:48:29.907776Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"python encode_audio.py --start_index 0 --end_index 885739\n",
"python encode_audio.py --start_index 885739 --end_index 1771478\n",
"python encode_audio.py --start_index 1771478 --end_index 2657217\n",
"python encode_audio.py --start_index 2657217 --end_index 3542956\n",
"python encode_audio.py --start_index 3542956 --end_index 4428695\n",
"python encode_audio.py --start_index 4428695 --end_index 5314434\n",
"python encode_audio.py --start_index 5314434 --end_index 6200173\n",
"python encode_audio.py --start_index 6200173 --end_index 7085912\n",
"python encode_audio.py --start_index 7085912 --end_index 7085918\n"
]
}
],
"source": [
"x = 7085918\n",
"# x = 374220\n",
"n_steps = 8\n",
"for i in range(0, x, x // n_steps):\n",
" print(\n",
" f\"python encode_audio.py --start_index {i} --end_index {min(i + x//n_steps, x)}\"\n",
" )"
]
},
{
"cell_type": "markdown",
"id": "02f85bb1",
"metadata": {},
"source": [
"# Demucs"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "1468489d",
"metadata": {
"ExecuteTime": {
"end_time": "2023-09-27T13:48:54.088148Z",
"start_time": "2023-09-27T13:48:54.086504Z"
}
},
"outputs": [],
"source": [
"# input_audio_file = \"252017_23.wav\"\n",
"# sr, audio_arr = wavfile.read(os.path.join(OUT_AUDIO_DIR, input_audio_file))\n",
"# audio = Audio.from_array(audio_arr, sr)\n",
"# vocals, non_vocals = demucs.split_vocals(audio)\n",
"# vocals.to_wav(os.path.join(OUT_AUDIO_DEMUC_DIR, input_audio_file.replace(\".wav\", \"_vocals.wav\")))\n",
"# non_vocals.to_wav(os.path.join(OUT_AUDIO_DEMUC_DIR, input_audio_file.replace(\".wav\", \"_bkg.wav\")))"
]
},
{
"cell_type": "markdown",
"id": "cf3c4b06",
"metadata": {},
"source": [
"# Codec"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "84976f6e",
"metadata": {
"ExecuteTime": {
"end_time": "2023-09-27T13:49:28.666239Z",
"start_time": "2023-09-27T13:49:28.663953Z"
}
},
"outputs": [],
"source": [
"input_audio_file = \"252017_23.wav\""
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "037e2e0d",
"metadata": {
"ExecuteTime": {
"end_time": "2023-09-27T13:49:28.942663Z",
"start_time": "2023-09-27T13:49:28.872582Z"
}
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"test_vocal = Audio.from_file(os.path.join(OUT_AUDIO_DIR, input_audio_file))\n",
"test_vocal.play()"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "f303880d",
"metadata": {
"ExecuteTime": {
"end_time": "2023-09-27T13:49:29.325072Z",
"start_time": "2023-09-27T13:49:29.191912Z"
}
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"test_vocal = Audio.from_file(os.path.join(OUT_AUDIO_DEMUC_DIR, input_audio_file.replace(\".wav\", \"_vocals.wav\")))\n",
"test_vocal.play()"
]
},
{
"cell_type": "markdown",
"id": "5d82acc9",
"metadata": {},
"source": [
"# Validate Demuc output"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "90892bb4",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-13T03:32:13.273382Z",
"start_time": "2023-10-13T03:32:07.774946Z"
}
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/home/tony/anaconda3/envs/suno_env/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:30: UserWarning: torch.nn.utils.weight_norm is deprecated in favor of torch.nn.utils.parametrizations.weight_norm.\n",
" warnings.warn(\"torch.nn.utils.weight_norm is deprecated in favor of torch.nn.utils.parametrizations.weight_norm.\")\n"
]
}
],
"source": [
"dac.preload_models(checkpoint_filepath=\"/home/tony/Data/MERT/codec.pt\", device=\"cuda\")"
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "215b246b",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-13T03:44:38.725653Z",
"start_time": "2023-10-13T03:44:35.414918Z"
}
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"test_original = Audio.from_file(os.path.join(\"/app/suno/data/mert_25hz_speech/audio\", \"11.wav\"))\n",
"test_original.play()"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "ff8c5169",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-13T03:40:55.689648Z",
"start_time": "2023-10-13T03:40:50.199274Z"
}
},
"outputs": [],
"source": [
"x = dac.encode(test_original)"
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "3980b7f4",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-13T03:44:38.729317Z",
"start_time": "2023-10-13T03:44:38.727072Z"
}
},
"outputs": [],
"source": [
"test_output= np.load(f\"/app/suno/data/mert_25hz_speech/audio_label/train_11.npy\")"
]
},
{
"cell_type": "code",
"execution_count": 17,
"id": "185c40bd",
"metadata": {
"ExecuteTime": {
"end_time": "2023-10-13T03:44:51.230451Z",
"start_time": "2023-10-13T03:44:42.555295Z"
}
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"