{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "from suno_utils.utils.text import read_jsonl, write_jsonl\n", "from suno_utils.utils.s3 import read_from_s3\n", "import numpy as np\n", "import json\n", "import polars as pl\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load the base metas\n", "metas_s3_filepath = \"s3://suno-data/datasets/bundles/v4/genius/metas_v0.jsonl\"\n", "metas = read_from_s3(metas_s3_filepath, read_f=pl.read_ndjson)\n", "\n", "# load the data cut\n", "data_cut = \"/home/christian/code/christian/metadata/ear/genius_t6.json\"\n", "with open(data_cut, \"r\") as f:\n", " info = json.load(f)\n", "print(len(info['genius']))\n", "\n", "# filter the metas based on the data cut\n", "metas = metas.filter(pl.col(\"id\").is_in(info['genius']))\n", "print(len(metas))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# set the seed\n", "import random\n", "random.seed(42)\n", "\n", "# select some random tracks\n", "rand_metas = metas.sample(10000)\n", "print(len(rand_metas))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# save the sampled metas\n", "sampled_metas_filepath = \"/home/christian/code/christian/metadata/ear/genius_t6_sampled_10k.jsonl\"\n", "write_jsonl(rand_metas.to_dicts(), sampled_metas_filepath)\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# save the sampled metas\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# read the sampled metas\n", "sampled_metas_filepath = \"/home/christian/code/christian/metadata/ear/genius_t6_sampled_10k.jsonl\"\n", "\n", "sampled_metas = read_jsonl(sampled_metas_filepath)\n", "print(len(sampled_metas))\n", "\n", "# split into train and val\n", "metas_tr = sampled_metas[:9000]\n", "metas_val = sampled_metas[9000:]\n", "\n", "print(len(metas_tr))\n", "print(len(metas_val))\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "base_s3_filepath = \"s3://suno-data/christian/outputs/v3-bootstrap-data-t0/\"\n", "local_base_filepath = \"/app2/suno/data/christian/outputs/v3-bootstrap-data-t0/\"\n", "\n", "# get all the directories in local_base_filepath\n", "base_dirs = [f for f in os.listdir(local_base_filepath) if os.path.isdir(os.path.join(local_base_filepath, f))]\n", "\n", "# put metadata into a pandas dataframe\n", "metadata_list = []\n", "\n", "for original_id in original_ids[:1000]:\n", " local_filepath = os.path.join(local_base_filepath, original_id)\n", " if os.path.exists(local_filepath):\n", " clips = []\n", " # get all the metadata files in the directory\n", " metadata_files = [f for f in os.listdir(local_filepath) if f.endswith(\"metadata.npz\")]\n", " for metadata_file in metadata_files:\n", " metadata = analyze_audio(os.path.join(local_filepath, metadata_file))\n", " if \"filename\" not in metadata:\n", " continue\n", " clips.append(metadata)\n", " # only take two clips\n", " if len(clips) > 2:\n", " clips = clips[:2]\n", " if len(clips) < 2:\n", " continue\n", " for clip in clips:\n", " metadata_list.append(clip)\n", "\n", "metadata_df = pd.DataFrame(metadata_list)\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", "# aws s3 sync s3://suno-data/christian/outputs/corrupt/genius_t6_sampled_10k/ /app/suno/data/diff_syn_dpo/genius_t6_sampled_10k_corrupted/npz\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 = \"/app2/suno/data/christian/outputs/v3-bootstrap-data-t0/memmaps\"\n", "\n", "#BASE_S3_DIR = \"s3://suno-data/christian/outputs/v2-infill-data-v1\"\n", "BASE_LOCAL_DIR = \"/app2/suno/data/christian/outputs/v3-bootstrap-data-t0/\"\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 [\"tr\", \"val\"]:\n", "\n", " if dset_type == \"tr\":\n", " metas = metas_tr\n", " else:\n", " metas = metas_val\n", "\n", " new_metas = []\n", "\n", " out_mm_semantic_filepath = os.path.join(OUT_DATA_DIR, f\"data_semantic_{dset_type}.bin\")\n", " out_mm_vae_filepath = os.path.join(OUT_DATA_DIR, f\"data_vae_{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_semantic = np.memmap(\n", " out_mm_semantic_filepath, dtype=np.uint16, mode=\"w+\", shape=(1,)\n", " )\n", "\n", " out_mm_vae = np.memmap(\n", " out_mm_vae_filepath, dtype=np.float16, 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(metas)\n", " pbar.set_description(\"Hours: 0.00\")\n", "\n", " for idx, (meta) in enumerate(pbar):\n", "\n", " # load semantic codes from disk\n", " semantic_codes_filepath = f\"{BASE_LOCAL_DIR}/{meta['id']}_semantic.npz\"\n", " vae_latents_filepath = f\"{BASE_LOCAL_DIR}/{meta['id']}_vae.npz\"\n", " vae_upsampled_filepath = f\"{BASE_LOCAL_DIR}/{meta['id']}_corrupted_vae.npz\"\n", "\n", " try:\n", " semantic_data = np.load(semantic_codes_filepath)[\"codes\"][:, 0]\n", " except Exception as e:\n", " #print(f\"Error loading {semantic_codes_filepath}: {e}\")\n", " continue\n", "\n", " try:\n", " vae_data_pos = np.load(vae_latents_filepath)[\"vae_latents\"].astype(np.float16)\n", " except Exception as e:\n", " #print(f\"Error loading {vae_latents_filepath}: {e}\")\n", " continue\n", "\n", " try:\n", " vae_data_neg = np.load(vae_upsampled_filepath)[\"vae_latents\"].astype(np.float16)\n", " except Exception as e:\n", " #print(f\"Error loading {vae_upsampled_filepath}: {e}\")\n", " continue\n", "\n", " num_chunks = semantic_data.shape[0] // CHUNK_SIZE\n", "\n", " to_write_len_s = semantic_data[:750].size * num_chunks * 2\n", " to_write_len_v = vae_data_pos[:750, :].size * num_chunks * 2\n", " \n", " if to_write_len_s == 0:\n", " continue\n", "\n", " if to_write_len_v == 0:\n", " continue\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", " 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", " # 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", " # we will write the positive and negative as interleaved chunks\n", " for i in range(num_chunks):\n", " for vae_data in [vae_data_neg, vae_data_pos]:\n", " # create a new meta\n", " new_meta = {\n", " \"id\": meta[\"id\"],\n", " \"start_s\": i*CHUNK_SIZE_S,\n", " \"end_s\": (i+1)*CHUNK_SIZE_S,\n", " \"original_duration_s\": semantic_data.shape[0] / SEMANTIC_RATE_HZ,\n", " \"n_vae_tokens\": CHUNK_SIZE,\n", " \"n_semantic_tokens\": CHUNK_SIZE,\n", " \"text\" : meta[\"text\"],\n", " \"tags\" : meta[\"tags\"],\n", " }\n", " new_metas.append(new_meta)\n", "\n", " semantic_chunk = semantic_data[i*CHUNK_SIZE:(i+1)*CHUNK_SIZE]\n", "\n", " out_mm_semantic[n_offs_s : n_offs_s + semantic_chunk.size] = semantic_chunk.reshape(\n", " -1,\n", " )\n", " n_offs_s += semantic_chunk.size\n", "\n", " vae_chunk = vae_data[i*CHUNK_SIZE:(i+1)*CHUNK_SIZE, :]\n", " out_mm_vae[n_offs_v : n_offs_v + vae_chunk.size] = vae_chunk.reshape(\n", " -1,\n", " )\n", " n_offs_v += vae_chunk.size\n", "\n", " print(f\"Total hours of audio added: {total_hours:.4f} for {dset_type} set\")\n", " \n", " write_jsonl(\n", " new_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": [ "# test loading the memmaps \n", "\n", "VAE_DIM = 128\n", "VAE_N_MEMMAP_TOKENS = 750\n", "\n", "SEMANTIC_N_CODEBOOKS = 1\n", "SEMANTIC_N_MEMMAP_TOKENS = 750\n", "\n", "#base_dir = \"/app/suno/data/diffusion_mix/dac_vae_fixed_25hz\"\n", "#base_dir = \"/app/suno/data/diff_syn_dpo/genius_t6_sampled_10k_corrupted\"\n", "#base_dir = \"/app/suno/data/diffusion_v5/v0\"\n", "#base_dir = \"/mnt/localdisk/cjs_shards/\"\n", "base_dir = \"/app/suno/data/diff_syn_dpo/interesting_clips_ahi_d3_20250504\"\n", "\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" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "idx = 2000\n", "\n", "meta = metas[idx]\n", "print(meta[\"tags\"], meta[\"text\"])\n", "vae_chunk_neg = vae_data[idx]\n", "vae_chunk_pos = vae_data[idx+1]\n", "semantic_chunk = semantic_data[idx]\n", "\n", "audio_neg = codec_decode(vae_chunk_neg).play()\n", "audio_pos = codec_decode(vae_chunk_pos).play()\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# test some of the local data\n", "CODEC_FILEPATH = \"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\"\n", "\n", "from suno_utils.tasks.dac_vae_fixed_25hz import (\n", " preload_models as preload_codec_models,\n", " decode as codec_decode,\n", " encode as codec_encode,\n", " decode_stream_to_full_audio,\n", ")\n", "\n", "_ = preload_codec_models(CODEC_FILEPATH)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# create a map of the metas from meta['id'] to the meta dict\n", "meta_id_to_meta = {meta['id']: meta for meta in metas}\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "meta_id = \"Rs6hUrJOjUI\"\n", "\n", "meta = meta_id_to_meta[meta_id]\n", "print(meta)\n", "\n", "# load semantic codes from disk\n", "semantic_codes_filepath = f\"{BASE_LOCAL_DIR}/{meta_id}/{meta_id}_semantic.npz\"\n", "vae_latents_filepath = f\"{BASE_LOCAL_DIR}/{meta_id}/{meta_id}_vae.npz\"\n", "vae_upsampled_filepath = f\"{BASE_LOCAL_DIR}/{meta_id}/{meta_id}_upsampled_vae.npz\"\n", "\n", "semantic_data = np.load(semantic_codes_filepath)[\"codes\"][:, 0]\n", "vae_data_pos = np.load(vae_latents_filepath)[\"vae_latents\"]\n", "vae_data_neg = np.load(vae_upsampled_filepath)[\"vae_latents\"]\n", "\n", "print(semantic_data.shape, vae_data_pos.shape, vae_data_neg.shape)\n", "\n", "audio_pos = codec_decode(vae_data_pos[:750, :])\n", "audio_neg = codec_decode(vae_data_neg[:750, :])\n", "\n", "print(audio_pos.loudness)\n", "print(audio_neg.loudness)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import random\n", "import torchaudio\n", "\n", "def apply_tanh_distortion(\n", " audio: torch.Tensor, sample_rate: float, gain_db: float = 0.0\n", "):\n", " gain_lin = 10 ** (gain_db / 20.0)\n", " return torch.tanh(audio * gain_lin)\n", "\n", "\n", "def apply_clipping_distortion(\n", " audio: torch.Tensor, sample_rate: float, gain_db: float = 0.0\n", "):\n", " gain_lin = 10 ** (gain_db / 20.0)\n", " return (audio * gain_lin).clamp(-1, 1)\n", "\n", "\n", "def apply_audio_codec(\n", " audio: torch.Tensor,\n", " sample_rate: float,\n", " bit_rate: int = 16000,\n", " n_passes: int = 1,\n", " format_str: str = \"mp3\",\n", "):\n", " for _ in range(n_passes):\n", " effector = torchaudio.io.AudioEffector(\n", " format=format_str,\n", " codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate),\n", " )\n", " audio = effector.apply(audio.T, sample_rate).T\n", " return audio\n", "\n", "def apply_noise(\n", " audio: torch.Tensor,\n", " sample_rate: float,\n", " gain_db: float = 0.0,\n", " noise_type: str = \"white\",\n", "):\n", " gain_lin = 10 ** (gain_db / 20.0)\n", " noise = torch.randn_like(audio)\n", "\n", " if noise_type == \"white\":\n", " return audio + gain_lin * noise\n", " elif noise_type == \"pink\":\n", " b = torch.tensor([0.049922035, -0.095993537, 0.050612699, -0.004408786])\n", " a = torch.tensor([1, -2.494956002, 2.017265875, -0.522189400])\n", " noise = torchaudio.functional.filtfilt(noise, a, b)\n", " noise /= noise.abs().max()\n", " return audio + gain_lin * noise\n", " else:\n", " raise ValueError(f\"Invalid noise type: {noise_type}\")\n", "\n", "def apply_highpass(audio: torch.Tensor, sample_rate: float, cutoff_hz: float):\n", " return torchaudio.functional.highpass_biquad(audio, sample_rate, cutoff_hz)\n", "\n", "def apply_lowpass(audio: torch.Tensor, sample_rate: float, cutoff_hz: float):\n", " return torchaudio.functional.lowpass_biquad(audio, sample_rate, cutoff_hz)\n", "\n", "def apply_random_tanh_distortion(audio: torch.Tensor, sample_rate: float):\n", " gain_db = random.uniform(0, 12)\n", " return apply_tanh_distortion(audio, sample_rate, gain_db)\n", "\n", "\n", "def apply_random_clipping_distortion(audio: torch.Tensor, sample_rate: float):\n", " gain_db = random.uniform(0, 12)\n", " return apply_clipping_distortion(audio, sample_rate, gain_db)\n", "\n", "\n", "def apply_random_noise(audio: torch.Tensor, sample_rate: float):\n", " noise_type = random.choice([\"white\", \"pink\"])\n", " noise_gain = random.uniform(-48, -6)\n", " return apply_noise(audio, sample_rate, noise_gain, noise_type)\n", "\n", "\n", "def apply_random_audio_codec(audio: torch.Tensor, sample_rate: float): \n", " format_str = \"mp3\"\n", " bit_rate = random.choice([8000, 16000, 32000, 64000, 128000])\n", " n_passes = random.choice([1, 2])\n", " return apply_audio_codec(audio, sample_rate, bit_rate, n_passes, format_str)\n", "\n", "def apply_random_highpass(audio: torch.Tensor, sample_rate: float):\n", " cutoff_hz = random.uniform(20, 6000)\n", " return apply_highpass(audio, sample_rate, cutoff_hz)\n", "\n", "def apply_random_lowpass(audio: torch.Tensor, sample_rate: float):\n", " cutoff_hz = random.uniform(100, 12000)\n", " return apply_lowpass(audio, sample_rate, cutoff_hz)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# next step is to corrupt the audio \n", "import glob\n", "from joblib import Parallel, delayed\n", "from tqdm import tqdm\n", "\n", "output_dir = \"/home/christian/audio/genius_ear_1k\"\n", "source_filepaths = glob.glob(os.path.join(output_dir, \"*.webm\"))\n", "print(len(source_filepaths))\n", "\n", "def corrupt_audio_and_save(filepath: str, p=0.33):\n", " meta_id = os.path.basename(filepath).split(\".\")[0]\n", " out_filepath = os.path.join(output_dir, f\"{meta_id}__corrupted.mp3\")\n", "\n", " #if os.path.exists(out_filepath):\n", " # return\n", "\n", " audio, sample_rate = torchaudio.load(filepath)\n", "\n", " if random.random() < p:\n", " audio = apply_random_highpass(audio, sample_rate)\n", " if random.random() < p:\n", " audio = apply_random_lowpass(audio, sample_rate)\n", " if random.random() < p:\n", " audio = apply_random_noise(audio, sample_rate)\n", " if random.random() < p:\n", " audio = apply_random_tanh_distortion(audio, sample_rate)\n", " if random.random() < p:\n", " audio = apply_random_clipping_distortion(audio, sample_rate)\n", " if random.random() < p:\n", " audio = apply_random_audio_codec(audio, sample_rate)\n", " # save the audio\n", " torchaudio.save(out_filepath, audio, sample_rate)\n", "\n", "# normal for loop\n", "#for filepath in tqdm(source_filepaths):\n", "# corrupt_audio_and_save(filepath)\n", " \n", "\n", "# process the audio in parallel\n", "_ = Parallel(n_jobs=-1)(\n", " delayed(corrupt_audio_and_save)(filepath) for filepath in tqdm(source_filepaths)\n", ")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# now we need to semantic encode the corrupted audio\n", "from suno_utils.tasks.mert_25 import (\n", " preload_models as preload_semantic_models,\n", " load_model as load_semantic_model,\n", " encode as encode_semantic,\n", " EMBEDDING_RATE as SEMANTIC_HZ,\n", ")\n", "import numpy as np\n", "\n", "\n", "semantic_model_filepath=\"/home/georg/notebooks/gpu_nb/tmp/mert_25.pt\"\n", "semantic_clusters_filepath=\"/home/georg/notebooks/gpu_nb/tmp/mert_25_2x4k.npy\"\n", "_ = preload_semantic_models(semantic_model_filepath, semantic_clusters_filepath)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from suno_utils.audio import Audio\n", "\n", "for filepath in tqdm(source_filepaths):\n", " meta_id = os.path.basename(filepath).split(\".\")[0]\n", " out_filepath = os.path.join(output_dir, f\"{meta_id}__corrupted.mp3\")\n", " audio = Audio.from_file(out_filepath)\n", "\n", " # also encode semantic codes\n", " semantic_codes = encode_semantic(\n", " audio.convert(sample_rate=24_000, byte_width=2, n_channels=1)\n", " ).astype(np.int64)[:, 0]\n", " #semantic_codes = torch.from_numpy(semantic_codes)\n", " \n", " # save the semantic codes\n", " semantic_codes_filepath = os.path.join(output_dir, f\"{meta_id}__corrupted.npz\")\n", " np.savez(semantic_codes_filepath, semantic_codes=semantic_codes)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# here we need to iterate over the pass tracks and read the audio from s3\n", "# once we get the audio we need to apply lowpass filter, and then ideally extract semantic codes and save them out \n", "import polars as pl\n", "import os\n", "import json\n", "base_dir = \"/app/suno/data/diffusion_mix/dac_vae_tuned_25hz/\"\n", "\n", "metas_tr_filepath = os.path.join(base_dir, \"metas_tr.jsonl\")\n", "metas_tr = pl.read_ndjson(metas_tr_filepath)\n", "\n", "metas_val_filepath = os.path.join(base_dir, \"metas_val.jsonl\")\n", "metas_val = pl.read_ndjson(metas_val_filepath)\n", "\n", "print(metas_tr[0][\"id\"])\n", "\n", "info_tr_filepath = os.path.join(base_dir, \"info_tr_t4.json\")\n", "with open(info_tr_filepath, \"r\") as f:\n", " info_tr = json.load(f)\n", "\n", "print(len(info_tr[\"diffusion_mix_fix\"]))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Access DataFrame rows properly\n", "for meta in metas_val.iter_rows(named=True):\n", " # meta is already a dict when using iter_rows(named=True)\n", " print(meta[\"id\"])\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 = base_dir\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 [\"tr\"]:\n", "\n", " if dset_type == \"tr\":\n", " metas = metas_tr\n", " else:\n", " metas = metas_val\n", "\n", " new_metas = []\n", "\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_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(metas.iter_rows(named=True))\n", " pbar.set_description(\"Hours: 0.00\")\n", "\n", " for idx, (meta) in enumerate(pbar):\n", "\n", " # load semantic codes from disk\n", " semantic_codes_filepath = os.path.join(output_dir, f\"{meta['id']}__corrupted.npz\")\n", "\n", " if not os.path.exists(semantic_codes_filepath):\n", " continue\n", "\n", " semantic_data = np.load(semantic_codes_filepath)[\"semantic_codes\"]\n", "\n", " num_chunks = semantic_data.shape[0] // CHUNK_SIZE\n", "\n", " to_write_len_s = semantic_data[:750].size * num_chunks\n", " \n", " if to_write_len_s == 0:\n", " continue\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", " \"id\": meta[\"id\"],\n", " \"start_s\": i*CHUNK_SIZE_S,\n", " \"end_s\": (i+1)*CHUNK_SIZE_S,\n", " \"original_duration_s\": semantic_data.shape[0] / SEMANTIC_RATE_HZ,\n", " \"n_vae_tokens\": CHUNK_SIZE,\n", " \"n_semantic_tokens\": CHUNK_SIZE,\n", " }\n", " new_metas.append(new_meta)\n", "\n", " semantic_chunk = semantic_data[i*CHUNK_SIZE:(i+1)*CHUNK_SIZE]\n", "\n", " out_mm_semantic[n_offs_s : n_offs_s + semantic_chunk.size] = semantic_chunk.reshape(\n", " -1,\n", " )\n", " n_offs_s += semantic_chunk.size\n", "\n", " print(f\"Total hours of audio added: {total_hours:.2f} for {dset_type} set\")\n", " \n", " write_jsonl(\n", " new_metas,\n", " os.path.join(out_metas_filepath),\n", " do_append=True\n", " )\n", "\n", " out_mm_semantic.flush()\n", " del out_mm_semantic, f\n", " gc.collect()\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load metas and semantic codes\n", "test_metas = read_jsonl(os.path.join(OUT_DATA_DIR, f\"metas_corrupted_tr.jsonl\"))\n", "test_semantic_codes = np.memmap(os.path.join(OUT_DATA_DIR, f\"data_corrupted_semantic_tr.bin\"), dtype=np.uint16, mode=\"r\")\n", "\n", "# reshape the semantic codes to a 2d array\n", "test_semantic_codes = test_semantic_codes.reshape(-1, 750)\n", "\n", "# create a dataframe from the metas\n", "test_df = pl.DataFrame(test_metas)\n", "\n", "# add the semantic codes to the dataframe\n", "test_df = test_df.with_columns(pl.Series(name=\"semantic_codes\", values=test_semantic_codes))\n", "\n", "print(len(test_metas))\n", "print(test_semantic_codes.shape)\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "# a, b, a__metadata.npz, a_upsampled_vae.npz\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from suno_utils.audio import Audio\n", "import numpy as np\n", "import os\n", "base_dir = \"/app2/suno/data/christian/outputs/v3-bootstrap-data-t4/\"\n", "\n", "# get all directories in base_dir\n", "dirs = os.listdir(base_dir)\n", "print(len(dirs))\n", "# get all files in each directory\n", "\n", "model_name = \"16n_25hz_v45_infill_shared_flow_resume_1_75m\"\n", "\n", "def process_dir(dirpath):\n", " a_metadata_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_a__metadata.npz\")\n", " a_upsampled_vae_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_a_upsampled_vae.npz\")\n", " b_metadata_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_b__metadata.npz\")\n", " b_upsampled_vae_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_b_upsampled_vae.npz\")\n", "\n", " a_metadata = np.load(a_metadata_filepath, allow_pickle=True)\n", " b_metadata = np.load(b_metadata_filepath, allow_pickle=True)\n", "\n", " # put metadata into dict\n", " a_metadata_dict = {}\n", " b_metadata_dict = {}\n", " for key in a_metadata.keys():\n", " a_metadata_dict[key] = a_metadata[key].tolist()\n", " for key in b_metadata.keys():\n", " b_metadata_dict[key] = b_metadata[key].tolist()\n", "\n", " a_upsampled_vae = np.load(a_upsampled_vae_filepath)\n", " b_upsampled_vae = np.load(b_upsampled_vae_filepath)\n", "\n", " # load the mp3s \n", " a_mp3_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_a.mp3\")\n", " b_mp3_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_b.mp3\")\n", " print(f\"a: ear {a_metadata_dict['ear_score']:0.1f}, shimmer {a_metadata_dict['shimmer_score']:0.1f}, stereo {a_metadata_dict['stereo_width']:0.1f}, hoot_cer {a_metadata_dict['hoot_cer']}\")\n", " a_mp3 = Audio.from_file(a_mp3_filepath, n_channels=2).play()\n", "\n", " print(f\"b: ear {b_metadata_dict['ear_score']:0.1f}, shimmer {b_metadata_dict['shimmer_score']:0.1f}, stereo {b_metadata_dict['stereo_width']:0.1f}, hoot_cer {b_metadata_dict['hoot_cer']}\")\n", " b_mp3 = Audio.from_file(b_mp3_filepath, n_channels=2).play()\n", "\n", " return a_metadata, a_upsampled_vae, b_metadata, b_upsampled_vae" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# we need to construct the metas for " ] }, { "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 }