{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "import glob\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\"\n", "import time\n", "import torch\n", "import torchaudio\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "from suno_utils.tasks.ear import load_model\n" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "base_dir = \"/app/suno/data/diff_sft/v45_2b_step_2_600_000/pos_interesting_clips_up_u_1_20241201_full_v2\"" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# group the filepaths based on the base s3 id \n", "base_s3_ids = [filepath.split(\"_upsample_\")[0] for filepath in filepaths]\n", "\n", "# get the unique base s3 ids\n", "unique_base_s3_ids = list(set(base_s3_ids))\n", "\n", "# group the filepaths based on the base s3 id\n", "grouped_filepaths = {base_s3_id: [] for base_s3_id in unique_base_s3_ids}\n", "for filepath in filepaths:\n", " base_s3_id = filepath.split(\"_upsample_\")[0]\n", " grouped_filepaths[base_s3_id].append(filepath)\n", "\n", "print(len(grouped_filepaths))" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [ "from suno_utils.utils.s3 import list_s3_dir\n", "\n", "base_dir = \"s3://suno-data/christian/outputs/v45_2b_step_2_600_000/pos_interesting_clips_up_u_1_20241201_full_v2\"\n", "filepaths = list_s3_dir(base_dir)\n", "filepaths = [filepath[0] for filepath in filepaths]\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# group the filepaths based on the base s3 id\n", "base_s3_ids = [filepath.split(\"_upsample_\")[0] for filepath in filepaths]\n", "print(len(base_s3_ids))\n", "# get the unique base s3 ids\n", "unique_base_s3_ids = list(set(base_s3_ids))\n", "print(len(unique_base_s3_ids))\n", "# group the filepaths based on the base s3 id\n", "grouped_filepaths = {base_s3_id: [] for base_s3_id in unique_base_s3_ids}\n", "for filepath in filepaths:\n", " if \".mp3\" in filepath:\n", " base_s3_id = filepath.split(\"_upsample_\")[0]\n", " grouped_filepaths[base_s3_id].append(filepath)\n", "\n", "len(grouped_filepaths)" ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [], "source": [ "import json\n", "# save the index of grouped filepaths\n", "with open(\"/home/christian/code/christian/metadata/sft/pos_interesting_clips_up_u_1_20241201_full_v2_grouped_filepaths.json\", \"w\") as f:\n", " json.dump(grouped_filepaths, f)\n" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "from suno_utils.utils.s3 import list_s3_dir, read_from_s3\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "base_dir = \"s3://suno-data/christian/outputs/v45_2b_step_2_600_000/pos_interesting_clips_up_u_1_20241201_full_v2/ear_score\"\n", "filepaths = list_s3_dir(base_dir)\n", "filepaths = [filepath[0] for filepath in filepaths]\n", "print(len(filepaths))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# download the ear_scores.json files to local\n", "# get the ear_scores from the local files\n", "\n", "# aws s3 sync s3://suno-data/christian/outputs/v45_2b_step_2_600_000/pos_interesting_clips_up_u_1_20241201_full_v2/ear_scores /app/suno/data/diff_sft/v45_2b_step_2_600_000/pos_interesting_clips_up_u_1_20241201_full_v2/ear_scores\n", "\n", "ear_score_filepaths = glob.glob(\"/app/suno/data/diff_sft/v45_2b_step_2_600_000/pos_interesting_clips_up_u_1_20241201_full_v2/ear_scores/*.json\")\n", "print(len(ear_score_filepaths))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# get all ear_scores.json files\n", "#ear_scores_filepaths = [filepath for filepath in filepaths if \"ear_scores.json\" in filepath]\n", "#print(len(ear_scores_filepaths))\n", "from tqdm import tqdm\n", "import json\n", "# load the ear_scores.json files\n", "ear_scores = {}\n", "for filepath in tqdm(ear_score_filepaths):\n", " if \"s3://\" in filepath:\n", " s3_filepath = f\"s3://suno-data/{filepath}\"\n", " actual_base_s3_id = filepath.split(\"/\")[-1].replace(\"_ear_scores.json\", \"\")\n", " ear_score = read_from_s3(s3_filepath)\n", " #print(ear_score)\n", " else:\n", " actual_base_s3_id = filepath.split(\"/\")[-1].replace(\"_ear_scores.json\", \"\")\n", " ear_score = json.load(open(filepath))\n", " ear_scores[actual_base_s3_id] = ear_score\n", "\n", "print(len(ear_scores))\n", "\n", "# split into tr and val\n", "# put 500 into val set\n", "# put the rest into tr set\n", "import random\n", "\n", "# Get all base_s3_ids\n", "all_base_s3_ids = list(ear_scores.keys())\n", "print(f\"Total number of base_s3_ids: {len(all_base_s3_ids)}\")\n", "\n", "# Set random seed for reproducibility\n", "random.seed(42)\n", "\n", "# Shuffle the list of base_s3_ids\n", "random.shuffle(all_base_s3_ids)\n", "\n", "# Split into validation and training sets\n", "val_base_s3_ids = all_base_s3_ids[:500]\n", "tr_base_s3_ids = all_base_s3_ids[500:]\n", "\n", "# Create dictionaries for validation and training sets\n", "val_ear_scores = {base_id: ear_scores[base_id] for base_id in val_base_s3_ids}\n", "tr_ear_scores = {base_id: ear_scores[base_id] for base_id in tr_base_s3_ids}\n", "\n", "print(f\"Number of samples in validation set: {len(val_ear_scores)}\")\n", "print(f\"Number of samples in training set: {len(tr_ear_scores)}\")\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# to create a memmap for sft we will select the higheset scoring upsample_id for each base_s3_id\n", "# we also need to grab the correct vae latents and semantic codes and text prompt\n", "from tqdm import tqdm\n", "from suno_utils.utils.text import write_jsonl\n", "import gc\n", "import sys\n", "import shutil\n", "\n", "SEMANTIC_RATE_HZ = 25\n", "CHUNK_SIZE_S = 30\n", "CHUNK_SIZE = int(CHUNK_SIZE_S * SEMANTIC_RATE_HZ)\n", "OUT_DATA_DIR = \"/app/suno/data/diff_sft/syn_sft_v1\"\n", "LOCAL_DATA_DIR = \"/app/suno/data/diff_sft/v45_2b_step_2_600_000/pos_interesting_clips_up_u_1_20241201_full_v2\"\n", "\n", "if not os.path.exists(OUT_DATA_DIR):\n", " os.makedirs(OUT_DATA_DIR, exist_ok=True)\n", "else:\n", " #shutil.rmtree(OUT_DATA_DIR)\n", " os.makedirs(OUT_DATA_DIR, exist_ok=True)\n", "\n", "for dset_type in [\"val\", \"tr\"]:\n", " metas = []\n", "\n", " if dset_type == \"tr\":\n", " dset_ear_scores = tr_ear_scores\n", " else:\n", " dset_ear_scores = val_ear_scores\n", "\n", " out_mm_vae_filepath = os.path.join(OUT_DATA_DIR, f\"data_vae_{dset_type}.bin\")\n", " out_mm_semantic_filepath = os.path.join(OUT_DATA_DIR, f\"data_semantic_{dset_type}.bin\")\n", " out_metas_filepath = os.path.join(OUT_DATA_DIR, f\"metas_{dset_type}.jsonl\")\n", "\n", " n_offs_v = 0\n", " n_offs_s = 0\n", " to_write_len_v = 0\n", " to_write_len_s = 0\n", " total_hours = 0 # Counter for total hours of audio\n", "\n", " out_mm_vae = np.memmap(\n", " out_mm_vae_filepath, dtype=np.float16, mode=\"w+\", shape=(1,)\n", " )\n", " out_mm_semantic = np.memmap(\n", " out_mm_semantic_filepath, dtype=np.uint16, mode=\"w+\", shape=(1,)\n", " )\n", "\n", " # clear the metas file\n", " with open(out_metas_filepath, \"w\") as f:\n", " f.write(\"\")\n", "\n", " # Create a tqdm progress bar with hours counter\n", " pbar = tqdm(dset_ear_scores.items())\n", " pbar.set_description(\"Hours: 0.00\")\n", "\n", " for idx, (base_s3_id, ear_score) in enumerate(pbar):\n", " #for upsample_id, score in ear_score.items():\n", " # print(upsample_id, score[\"mean_score\"])\n", " # get the upsample_id with the highest mean score\n", " upsample_id = max(ear_score.keys(), key=lambda k: ear_score[k][\"mean_score\"])\n", " #print(upsample_id, ear_score[upsample_id][\"mean_score\"])\n", " \n", " # get the filepaths for vae, semantic and metadata\n", " if LOCAL_DATA_DIR is not None:\n", " vae_filepath = os.path.join(LOCAL_DATA_DIR, f\"{base_s3_id}_upsample_{upsample_id}_vae.npz\")\n", " semantic_filepath = os.path.join(LOCAL_DATA_DIR, f\"{base_s3_id}_upsample_{upsample_id}_semantic.npz\")\n", " metadata_filepath = os.path.join(LOCAL_DATA_DIR, f\"{base_s3_id}_upsample_{upsample_id}_metadata.json\")\n", " try:\n", " metadata = json.load(open(metadata_filepath))\n", " semantic_data = np.load(semantic_filepath)[\"arr_0\"]\n", " vae_data = np.load(vae_filepath)[\"arr_0\"]\n", " except Exception as e:\n", " print(e)\n", " sys.exit()\n", " \n", " else:\n", " vae_filepath = f\"s3://suno-data/christian/outputs/v45_2b_step_2_600_000/pos_interesting_clips_up_u_1_20241201_full_v2/{base_s3_id}_upsample_{upsample_id}_vae.npz\"\n", " semantic_filepath = f\"s3://suno-data/christian/outputs/v45_2b_step_2_600_000/pos_interesting_clips_up_u_1_20241201_full_v2/{base_s3_id}_upsample_{upsample_id}_semantic.npz\"\n", " metadata_filepath = f\"s3://suno-data/christian/outputs/v45_2b_step_2_600_000/pos_interesting_clips_up_u_1_20241201_full_v2/{base_s3_id}_upsample_{upsample_id}_metadata.json\"\n", " # read the metadata, semantic and vae data\n", " try:\n", " metadata = json.loads(read_from_s3(metadata_filepath))\n", " semantic_data = read_from_s3(semantic_filepath, read_f=np.load)[\"arr_0\"]\n", " vae_data = read_from_s3(vae_filepath, read_f=np.load)[\"arr_0\"] \n", " except Exception as e:\n", " print(e)\n", " sys.exit()\n", "\n", " # check that vae and semantic data are the same length\n", " assert vae_data.shape[0] == semantic_data.shape[0]\n", "\n", " # convert vae_data to float16\n", " vae_data = vae_data.astype(np.float16)\n", "\n", " num_chunks = vae_data.shape[0] // CHUNK_SIZE\n", "\n", " to_write_len_v = vae_data[:750,:].size * num_chunks\n", " to_write_len_s = semantic_data[:750].size * num_chunks\n", " \n", " if to_write_len_s == 0 or to_write_len_v == 0:\n", " continue\n", "\n", " out_mm_vae = np.memmap(\n", " out_mm_vae_filepath,\n", " dtype=np.float16,\n", " mode=\"r+\",\n", " shape=(n_offs_v + to_write_len_v,),\n", " )\n", " \n", " out_mm_semantic = np.memmap(\n", " out_mm_semantic_filepath,\n", " dtype=np.uint16,\n", " mode=\"r+\",\n", " shape=(n_offs_s + to_write_len_s,),\n", " )\n", "\n", " # Add to total hours counter\n", " audio_duration_hours = (num_chunks * CHUNK_SIZE_S) / 3600\n", " total_hours += audio_duration_hours\n", " \n", " # Update progress bar description with current total hours\n", " pbar.set_description(f\"Hours: {total_hours:.2f}\")\n", "\n", " for i in range(num_chunks):\n", " # create a new meta\n", " new_meta = {\n", " \"base_s3_id\": base_s3_id,\n", " \"upsample_id\": upsample_id,\n", " \"start_s\": i*CHUNK_SIZE_S,\n", " \"end_s\": (i+1)*CHUNK_SIZE_S,\n", " \"original_duration_s\": vae_data.shape[0] / SEMANTIC_RATE_HZ,\n", " \"n_vae_tokens\": CHUNK_SIZE,\n", " \"n_semantic_tokens\": CHUNK_SIZE,\n", " }\n", " for key, value in metadata.items():\n", " new_meta[key] = value\n", " metas.append(new_meta)\n", "\n", " vae_chunk = vae_data[i*CHUNK_SIZE:(i+1)*CHUNK_SIZE, :]\n", " semantic_chunk = semantic_data[i*CHUNK_SIZE:(i+1)*CHUNK_SIZE]\n", "\n", " # convert vae_chunk to float16\n", "\n", " out_mm_vae[n_offs_v : n_offs_v + vae_chunk.size] = vae_chunk.reshape(\n", " -1,\n", " )\n", " out_mm_semantic[n_offs_s : n_offs_s + semantic_chunk.size] = semantic_chunk.reshape(\n", " -1,\n", " )\n", "\n", " n_offs_s += semantic_chunk.size\n", " n_offs_v += vae_chunk.size\n", "\n", " print(f\"Total hours of audio added: {total_hours:.2f} for {dset_type} set\")\n", " \n", " write_jsonl(\n", " metas,\n", " os.path.join(out_metas_filepath),\n", " do_append=True\n", " )\n", "\n", " out_mm_semantic.flush()\n", " out_mm_vae.flush()\n", " del out_mm_semantic, out_mm_vae, f\n", " gc.collect()\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "VAE_DIM = 128\n", "VAE_N_MEMMAP_TOKENS = 750\n", "\n", "SEMANTIC_N_CODEBOOKS = 1\n", "SEMANTIC_N_MEMMAP_TOKENS = 750\n", "\n", "from suno_utils.utils.text import read_jsonl\n", "\n", "base_dir = \"/app/suno/data/diff_sft/syn_sft_v1\"\n", "metas = read_jsonl(f\"{base_dir}/metas_val.jsonl\", progress=True)\n", "vae_memmap_filepath = f\"{base_dir}/data_vae_val.bin\"\n", "semantic_memmap_filepath = f\"{base_dir}/data_semantic_val.bin\"\n", "\n", "# load memmaps\n", "vae_memmap = np.memmap(vae_memmap_filepath, dtype=np.float16, mode=\"r\")\n", "semantic_memmap = np.memmap(semantic_memmap_filepath, dtype=np.uint16, mode=\"r\")\n", "\n", "# reshape memmaps\n", "vae_data = vae_memmap.reshape(-1, VAE_N_MEMMAP_TOKENS, VAE_DIM)\n", "semantic_data = semantic_memmap.reshape(-1, SEMANTIC_N_MEMMAP_TOKENS, SEMANTIC_N_CODEBOOKS)\n", "\n", "print(vae_data.shape, semantic_data.shape, len(metas))\n", "\n", "#del vae_memmap, semantic_memmap\n", "#gc.collect()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from suno_utils.tasks.dac_vae_fixed_25hz import decode, preload_models\n", "_ = preload_models(checkpoint_filepath=\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "idx = 238\n", "print(metas[idx])\n", "audio = decode(vae_data[idx])\n", "audio.play()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "suno_env", "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.10.9" } }, "nbformat": 4, "nbformat_minor": 2 }