{ "cells": [ { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\"\n", "\n", "import json \n", "import glob\n", "import torch\n", "import funcy\n", "import IPython\n", "import random\n", "import tempfile\n", "import torchaudio\n", "import numpy as np\n", "import pyloudnorm as pyln \n", "\n", "from tqdm import tqdm\n", "from dac.model.dac4 import DAC\n", "from suno_utils.audio import Audio\n", "from suno_utils.utils.s3 import read_from_s3\n", "from suno_utils.utils.text import write_jsonl, read_jsonl\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 18, "metadata": {}, "outputs": [], "source": [ "# load metas from genius \n", "# iterate over metas and get audio from s3\n", "# download and open audio with torchaudio\n", "# then apply highpass filter to the audio\n", "# then codec cycle both the original and the corrupted audio\n", "# save the vae latents as npz files (locally)\n", "# separete process to create a train and val memmap for these latents with new metas\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "VAE_RATE_HZ = 25\n", "\n", "# load the vae model\n", "device = \"cuda:0\"\n", "#device = \"cpu\"\n", "if VAE_RATE_HZ == 100:\n", " checkpoint_filepath = \"s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth\"\n", " #checkpoint_filepath = \"/app/suno/christian/checkpoints/dac/100hz_vae_peaq_kl_0.005/best/dac/weights.pth\"\n", " NUM_VAE_TOKENS = 3000\n", "elif VAE_RATE_HZ == 25:\n", " checkpoint_filepath = \"s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth\"\n", " NUM_VAE_TOKENS = 750\n", "else:\n", " raise ValueError(f\"VAE_RATE_HZ must be 100 or 25, got {VAE_RATE_HZ}\")\n", "load_f = funcy.partial(torch.load, map_location=\"cpu\")\n", "\n", "if checkpoint_filepath.startswith(\"s3://\"):\n", " sd = read_from_s3(checkpoint_filepath, read_f=load_f)\n", "else:\n", " sd = load_f(checkpoint_filepath)\n", "\n", "sd[\"metadata\"][\"kwargs\"] = {\n", " k: v\n", " for k, v in sd[\"metadata\"][\"kwargs\"].items()\n", " if k in DAC.__init__.__code__.co_varnames\n", "}\n", "model_100hz = DAC(**sd[\"metadata\"][\"kwargs\"])\n", "model_100hz.load_state_dict(sd[\"state_dict\"])\n", "model_100hz.eval()\n", "model_100hz.to(device)\n", "\n", "# load the semantic model\n", "from suno_utils.tasks.mert_25 import (\n", " preload_models as preload_semantic_models,\n", " encode as encode_semantic,\n", " EMBEDDING_RATE as SEMANTIC_HZ,\n", ")\n", "\n", "print(\"loading semantic model...\")\n", "semantic_model_filepath=\"s3://suno-data/georg/models/semantic/mert_25.pt\"\n", "semantic_clusters_filepath=\"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\"\n", "_ = preload_semantic_models(\n", " checkpoint_filepath=semantic_model_filepath,\n", " centroids_filepath=semantic_clusters_filepath,\n", " device=\"cuda\",\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "base_metas = \"/home/christian/code/christian/metadata/genius_hq_metas_filtered.jsonl\"\n", "metas = read_jsonl(base_metas)\n", "print(f\"loaded {len(metas):,} metas\")\n", "\n", "# also load the lyric alignments \n", "genius_alignments_filepath = (\n", " \"/home/tony/Work/tony/hoot/tmp/genius_hq_alignments_t30_v1.jsonl\"\n", ")\n", "if os.path.exists(genius_alignments_filepath):\n", " print(\"loading genius alignments\")\n", " genius_alignments = read_jsonl(genius_alignments_filepath, progress=False)\n", " genius_alignments_map = {a[0]: a[1] for a in genius_alignments}\n", "else:\n", " genius_alignments_map = {}\n", "# merge alignments into single map\n", "alignments_map = {**genius_alignments_map}\n", "print(f\"loaded {len(alignments_map):,} alignments\")\n", "\n", "# create metas map\n", "metas_map = {meta[\"id\"]: meta for meta in metas}\n", "print(f\"loaded {len(metas_map):,} metas\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# create metas map\n", "metas_map = {meta[\"id\"]: meta for meta in metas}\n", "print(f\"loaded {len(metas_map):,} metas\")\n" ] }, { "cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [], "source": [ "def corrupt_audio(audio, sr):\n", " audio_corrupt = audio.clone()\n", " # 10% chance to add highpass\n", " if np.random.uniform() < 0.25:\n", " hpf_cutoff = np.random.uniform(20, 1000)\n", " audio_corrupt = torchaudio.functional.highpass_biquad(\n", " audio_corrupt, sr, hpf_cutoff\n", " )\n", " # 10% chance to add lowpass\n", " if np.random.uniform() < 0.25: \n", " lpf_cutoff = np.random.uniform(1000, 20000)\n", " audio_corrupt = torchaudio.functional.lowpass_biquad(\n", " audio_corrupt, sr, lpf_cutoff\n", " ) \n", "\n", " # 10% chance to make mono \n", " if np.random.uniform() < 0.7: \n", " audio_corrupt = audio_corrupt.mean(dim=0, keepdim=True)\n", " audio_corrupt = audio_corrupt.repeat(2, 1)\n", "\n", " # 10% chance to add noise\n", " if np.random.uniform() < 0.2:\n", " noise_level_db = -np.random.uniform(42, 96)\n", " noise_signal = torch.randn_like(audio_corrupt) * 10**(noise_level_db / 20)\n", " # apply filter to noise signal\n", " lpf_cutoff = np.random.uniform(1000, 20000)\n", " noise_signal = torchaudio.functional.lowpass_biquad(noise_signal, sr, lpf_cutoff)\n", " audio_corrupt = audio_corrupt + noise_signal\n", "\n", "\n", " # 5% chance to add contrast\n", " if np.random.uniform() < 0.05:\n", " contrast = np.random.uniform(0, 100.0)\n", " audio_corrupt = torchaudio.functional.contrast(audio_corrupt, enhancement_amount=contrast)\n", "\n", " # 10% add codec artifacts\n", " if np.random.uniform() < 1.0:\n", " compression = np.random.choice([32, 64, 128])\n", " audio_corrupt = torchaudio.functional.apply_codec(audio_corrupt, sr, \"mp3\", compression)\n", " \n", " # 10% chance to add distortion\n", " if np.random.uniform() < 0.3:\n", " drive_db = np.random.uniform(6, 24)\n", " else:\n", " drive_db = 0\n", "\n", " drive_lin = 10**(drive_db / 20)\n", " audio_corrupt *= drive_lin\n", "\n", " if np.random.uniform() < 0.5:\n", " audio_corrupt = torch.tanh(audio_corrupt)\n", " else:\n", " audio_corrupt = audio_corrupt.clamp(-1, 1)\n", "\n", " return audio_corrupt" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "meter = pyln.Meter(48000)\n", "\n", "def normalize_and_encode(audio):\n", " target_db = -16.0\n", " # peak normalize\n", " audio = audio / audio.abs().max().clamp(min=1e-5)\n", " audio = audio * 10 ** (target_db / 20.0)\n", " audio = audio.to(\"cuda:0\")\n", " with torch.no_grad():\n", " vae_latents = model_100hz.encode(audio.unsqueeze(0))[\"z\"].cpu()\n", " return vae_latents\n", "\n", "\n", "out_dir = \"/home/christian/data/dpo/original_highpass_25hz_30s_v1\"\n", "os.makedirs(out_dir, exist_ok=True)\n", "\n", "num_samples = 1000\n", "# shuffle metas\n", "random.seed(42)\n", "random.shuffle(metas)\n", "\n", "def process_meta(meta):\n", " meta_id = meta[\"id\"]\n", " npz_filepath = os.path.join(out_dir, f\"{meta_id}.npz\")\n", "\n", " if os.path.exists(npz_filepath):\n", " return meta_id\n", "\n", " if meta_id not in alignments_map:\n", " return None\n", "\n", " # apply highpass filter to the audio\n", " audio, sr = read_from_s3(meta[\"audio_filepath\"], read_f=torchaudio.load)\n", "\n", " # ensure stereo\n", " if audio.shape[0] == 1:\n", " audio = audio.repeat(2, 1)\n", "\n", " # resample to 48k\n", " audio_48k = torchaudio.functional.resample(audio, sr, 48000)\n", "\n", " # crop to first 30 seconds\n", " audio_48k_chunk = audio_48k[:, : int(48000 * 30.01)]\n", "\n", " # apply highpass filter\n", " audio_48k_chunk_corrupted = corrupt_audio(audio_48k_chunk, 48000)\n", "\n", " # normalize both audios\n", " vae_latents_original = normalize_and_encode(audio_48k_chunk)\n", " vae_latents_corrupted = normalize_and_encode(audio_48k_chunk_corrupted)\n", "\n", " if vae_latents_original.shape[-1] > NUM_VAE_TOKENS:\n", " vae_latents_original = vae_latents_original[0, :, :NUM_VAE_TOKENS]\n", " if vae_latents_corrupted.shape[-1] > NUM_VAE_TOKENS:\n", " vae_latents_corrupted = vae_latents_corrupted[0, :, :NUM_VAE_TOKENS]\n", "\n", " # resample to 24k\n", " audio_24k_chunk = torchaudio.functional.resample(audio_48k_chunk, 48000, 24000)\n", " audio_24k_chunk_corrupted = torchaudio.functional.resample(audio_48k_chunk_corrupted, 48000, 24000)\n", "\n", " audio_24k_chunk_mono = audio_24k_chunk.mean(dim=0).unsqueeze(0)\n", " audio_24k_chunk_corrupted_mono = audio_24k_chunk_corrupted.mean(dim=0).unsqueeze(0)\n", "\n", " # peak normalize to -16 db\n", " target_db = -16.0\n", " audio_24k_chunk_mono = audio_24k_chunk_mono / audio_24k_chunk_mono.abs().max().clamp(min=1e-5)\n", " audio_24k_chunk_hpf_mono = audio_24k_chunk_hpf_mono / audio_24k_chunk_hpf_mono.abs().max().clamp(min=1e-5)\n", " audio_24k_chunk_mono = audio_24k_chunk_mono * 10 ** (target_db / 20.0)\n", " audio_24k_chunk_hpf_mono = audio_24k_chunk_hpf_mono * 10 ** (target_db / 20.0)\n", "\n", " # encode semantic\n", " semantic_codes = encode_semantic(\n", " [audio_24k_chunk_mono]\n", " )[0].astype(np.uint16)[:, 0]\n", "\n", " semantic_codes_corrupted = encode_semantic(\n", " [audio_24k_chunk_hpf_mono]\n", " )[0].astype(np.uint16)[:, 0]\n", "\n", " # save the vae latents as npz files (locally)\n", " np.savez(npz_filepath, \n", " vae_latents_original=vae_latents_original, \n", " vae_latents_corrupted=vae_latents_corrupted, \n", " semantic_codes=semantic_codes, \n", " semantic_codes_corrupted=semantic_codes_corrupted)\n", " \n", " return meta_id\n", "\n", "from concurrent.futures import ThreadPoolExecutor, as_completed\n", "\n", "with ThreadPoolExecutor(max_workers=32) as executor:\n", " futures = [executor.submit(process_meta, meta) for meta in metas[:z]]\n", " for future in tqdm(as_completed(futures), total=num_samples):\n", " meta_id = future.result()\n", "\n" ] }, { "cell_type": "code", "execution_count": 21, "metadata": {}, "outputs": [], "source": [ "# paralleized \n", "import gc\n", "import multiprocessing as mp\n", "from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor\n", "\n", "BATCH_SIZE = 8\n", "NUM_THREADS = 32 # Adjust based on your system\n", "\n", "def load_audio(meta):\n", " \"\"\"Just load the audio file - I/O bound operation\"\"\"\n", " try:\n", " meta_id = meta[\"id\"]\n", " if meta_id not in alignments_map:\n", " return None\n", " \n", " # skip if npz file exists\n", " npz_filepath = os.path.join(out_dir, f\"{meta_id}.npz\")\n", " if os.path.exists(npz_filepath):\n", " return None\n", " \n", " audio, sr = read_from_s3(meta[\"audio_filepath\"], read_f=torchaudio.load)\n", "\n", " # ensure stereo\n", " if audio.shape[0] == 1:\n", " audio = audio.repeat(2, 1)\n", " elif audio.shape[0] > 2:\n", " audio = audio[:2, :]\n", "\n", " # resample to 48k\n", " audio = torchaudio.functional.resample(audio, sr, 48000)\n", "\n", " # crop to first 30 seconds\n", " audio_48k = audio[:, :int(48000 * 30.01)]\n", "\n", " # Apply highpass filter\n", " audio_48k_corrupted = corrupt_audio(audio_48k, 48000)\n", " \n", " # Normalize for VAE\n", " audio_48k_normalized = audio_48k / audio_48k.abs().max().clamp(min=1e-5)\n", " audio_48k_corrupted_normalized = audio_48k_corrupted / audio_48k_corrupted.abs().max().clamp(min=1e-5)\n", " \n", " # Create 24k versions for semantic\n", " audio_24k = torchaudio.functional.resample(audio_48k, 48000, 24000)\n", " audio_24k_corrupted = torchaudio.functional.resample(audio_48k_corrupted, 48000, 24000)\n", " \n", " audio_24k_mono = audio_24k.mean(dim=0).unsqueeze(0)\n", " audio_24k_corrupted_mono = audio_24k_corrupted.mean(dim=0).unsqueeze(0)\n", "\n", " return {\n", " \"id\": meta[\"id\"],\n", " \"audio_48k\": audio_48k_normalized,\n", " \"audio_48k_corrupted\": audio_48k_corrupted_normalized,\n", " \"audio_24k_mono\": audio_24k_mono,\n", " \"audio_24k_corrupted_mono\": audio_24k_corrupted_mono\n", " }\n", "\n", " except Exception as e:\n", " print(f\"Error loading {meta['id']}: {str(e)}\")\n", " return None\n", " \n", " \n", "\n", "def process_batch(batch_data, model_100hz, out_dir):\n", " \"\"\"Process a batch through models and save results\"\"\"\n", " if not batch_data:\n", " return\n", " \n", " # Stack for VAE processing\n", " original_batch = torch.stack([data[\"audio_48k\"] for data in batch_data]).cuda()\n", " corrupted_batch = torch.stack([data[\"audio_48k_corrupted\"] for data in batch_data]).cuda()\n", " \n", " # Process through VAE\n", " with torch.no_grad():\n", " vae_original = model_100hz.encode(original_batch)[\"z\"].cpu()\n", " vae_corrupted = model_100hz.encode(corrupted_batch)[\"z\"].cpu()\n", " \n", " # Clear GPU memory\n", " del original_batch, corrupted_batch\n", " torch.cuda.empty_cache()\n", " \n", " # Process semantic codes\n", " semantic_original = encode_semantic([d[\"audio_24k_mono\"] for d in batch_data])[0]\n", " semantic_corrupted = encode_semantic([d[\"audio_24k_corrupted_mono\"] for d in batch_data])[0]\n", " \n", " # Save results\n", " for idx, data in enumerate(batch_data):\n", " vae_orig = vae_original[idx, :, :NUM_VAE_TOKENS]\n", " vae_corrupt = vae_corrupted[idx, :, :NUM_VAE_TOKENS]\n", " \n", " np.savez(\n", " os.path.join(out_dir, f\"{data['id']}.npz\"),\n", " vae_latents_original=vae_orig,\n", " vae_latents_corrupted=vae_corrupt,\n", " semantic_codes=semantic_original.astype(np.uint16)[:, 0],\n", " semantic_codes_corrupted=semantic_corrupted.astype(np.uint16)[:, 0]\n", " )\n", " #print(f\"saved {len(batch_data)} files\")\n", " \n", " # Clear memory\n", " del vae_original, vae_corrupted, semantic_original, semantic_corrupted\n", " torch.cuda.empty_cache()\n", "\n", "\n", "def clear_memory():\n", " \"\"\"Aggressive memory cleanup\"\"\"\n", " gc.collect()\n", " torch.cuda.empty_cache()\n", " if torch.cuda.is_available():\n", " torch.cuda.reset_peak_memory_stats()\n", "\n", "def main(metas, out_dir, num_samples):\n", " # Setup\n", " os.makedirs(out_dir, exist_ok=True)\n", " random.seed(10)\n", " random.shuffle(metas)\n", " subset_metas = metas[:num_samples]\n", " \n", " # Move model to GPU\n", " model_100hz.to(\"cuda:0\")\n", " \n", " # Process in batches\n", " with ThreadPoolExecutor(max_workers=NUM_THREADS) as executor:\n", " for i in tqdm(range(0, len(subset_metas), BATCH_SIZE)):\n", " batch_metas = subset_metas[i:i + BATCH_SIZE]\n", " clear_memory() \n", "\n", " # Parallel loading of audio files\n", " loaded_batch = list(executor.map(load_audio, batch_metas))\n", "\n", " batch_data = [data for data in loaded_batch if data is not None]\n", " \n", " # Process through models and save\n", " process_batch(batch_data, model_100hz, out_dir)" ] }, { "cell_type": "code", "execution_count": 25, "metadata": {}, "outputs": [], "source": [ "import gc\n", "import multiprocessing as mp\n", "import traceback\n", "\n", "BATCH_SIZE = 4\n", "NUM_PROCESSES = min(mp.cpu_count(), 16)\n", "\n", "def load_audio(meta):\n", " \"\"\"Load and process audio data with error handling.\"\"\"\n", " try:\n", " meta_id = meta[\"id\"]\n", " if meta_id not in alignments_map:\n", " return None\n", "\n", " #npz_filepath = os.path.join(out_dir, f\"{meta_id}.npz\")\n", " #if os.path.exists(npz_filepath):\n", " # return None\n", "\n", " audio, sr = read_from_s3(meta[\"audio_filepath\"], read_f=torchaudio.load)\n", "\n", " if audio.shape[0] == 1:\n", " audio = audio.repeat(2, 1)\n", " elif audio.shape[0] > 2:\n", " audio = audio[:2, :]\n", "\n", " audio = torchaudio.functional.resample(audio, sr, 48000)\n", " audio_48k = audio[:, :int(48000 * 30.01)]\n", "\n", " audio_48k_corrupted = corrupt_audio(audio_48k, 48000)\n", " \n", " audio_48k_normalized = audio_48k / audio_48k.abs().max().clamp(min=1e-5)\n", " audio_48k_corrupted_normalized = audio_48k_corrupted / audio_48k_corrupted.abs().max().clamp(min=1e-5)\n", "\n", " audio_24k = torchaudio.functional.resample(audio_48k, 48000, 24000)\n", " audio_24k_corrupted = torchaudio.functional.resample(audio_48k_corrupted, 48000, 24000)\n", "\n", " audio_24k_mono = audio_24k.mean(dim=0).unsqueeze(0)\n", " audio_24k_corrupted_mono = audio_24k_corrupted.mean(dim=0).unsqueeze(0)\n", "\n", " return {\n", " \"id\": meta[\"id\"],\n", " \"audio_48k\": audio_48k_normalized,\n", " \"audio_48k_corrupted\": audio_48k_corrupted_normalized,\n", " \"audio_24k_mono\": audio_24k_mono,\n", " \"audio_24k_corrupted_mono\": audio_24k_corrupted_mono\n", " }\n", "\n", " except Exception as e:\n", " print(f\"Error loading {meta['id']}: {str(e)}\\n{traceback.format_exc()}\")\n", " return None\n", "\n", "def main(metas, out_dir, num_samples):\n", " os.makedirs(out_dir, exist_ok=True)\n", " random.seed(10)\n", " random.shuffle(metas)\n", " subset_metas = metas[:num_samples]\n", " \n", " model_100hz.to(\"cuda:0\")\n", " \n", " with mp.Pool(NUM_PROCESSES) as pool:\n", " for i in tqdm(range(0, len(subset_metas), BATCH_SIZE)):\n", " batch_metas = subset_metas[i:i + BATCH_SIZE]\n", " \n", " # Use map_async with error tracking\n", " result = pool.map_async(load_audio, batch_metas)\n", " loaded_batch = result.get() # This will raise exceptions if any\n", " \n", " batch_data = [data for data in loaded_batch if data is not None]\n", " \n", " # Process through models and save\n", " process_batch(batch_data, model_100hz, out_dir)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "num_samples = 4000\n", "out_dir = \"/home/christian/data/dpo/genius_hq_corrupt_dpo_25hz_30_v2_npz\"\n", "main(metas, out_dir, num_samples)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "data = np.load(os.path.join(out_dir, f\"{meta_id}.npz\"))\n", "print(data[\"vae_latents_original\"].dtype, data[\"vae_latents_corrupted\"].shape, data[\"semantic_codes\"].shape, data[\"semantic_codes_corrupted\"].dtype)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "npz_dir = \"/home/christian/data/dpo/genius_hq_corrupt_dpo_25hz_30_v1_npz\"\n", "out_dir = \"/home/christian/data/dpo/genius_hq_corrupt_dpo_25hz_30_v3\"\n", "os.makedirs(out_dir, exist_ok=True)\n", "# now we create a train and val memmap for these latents with new metas\n", "\n", "# find all npz files in the out_dir\n", "npz_files = glob.glob(os.path.join(npz_dir, \"*.npz\"))\n", "print(len(npz_files))\n", "\n", "SEMANTIC_MEMMAP_SIZE = 750\n", "VAE_MEMMAP_SIZE = 750\n", "VAE_DIM = 128\n", "\n", "# split into train and val\n", "npz_files_train = npz_files[:int(len(npz_files) * 0.9)]\n", "npz_files_val = npz_files[int(len(npz_files) * 0.9):] \n", "print(f\"found {len(npz_files):,} npz files, {len(npz_files_train):,} train, {len(npz_files_val):,} val\")\n", "\n", "for dset, npz_files in [(\"tr\", npz_files_train), (\"val\", npz_files_val)]:\n", " out_mm_vae_filepath = os.path.join(out_dir, f\"data_vae_{dset}.bin\")\n", " out_mm_semantic_filepath = os.path.join(out_dir, f\"data_semantic_{dset}.bin\")\n", " out_metas_filepath = os.path.join(out_dir, f\"metas_{dset}.jsonl\")\n", "\n", " new_metas = []\n", "\n", " # initial write\n", " out_mm_semantic = np.memmap(\n", " out_mm_semantic_filepath,\n", " dtype=np.uint16,\n", " mode=\"w+\",\n", " shape=(1),\n", " )\n", " out_mm_vae = np.memmap(\n", " out_mm_vae_filepath,\n", " dtype=np.float16,\n", " mode=\"w+\",\n", " shape=(1),\n", " )\n", "\n", " n_offs_s = 0\n", " n_offs_v = 0\n", "\n", " to_write_len_s = SEMANTIC_MEMMAP_SIZE * len(npz_files) * 2\n", " to_write_len_v = VAE_MEMMAP_SIZE * VAE_DIM * len(npz_files) * 2\n", "\n", " print(to_write_len_s, to_write_len_v)\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", "\n", "\n", " for npz_file in tqdm(npz_files):\n", " data = np.load(npz_file)\n", " meta_id = os.path.basename(npz_file).replace(\".npz\", \"\")\n", " meta = metas_map[meta_id]\n", "\n", " # construct text aligned\n", " if meta_id in alignments_map:\n", " alignments = alignments_map[meta_id]\n", " else:\n", " alignments = []\n", "\n", " text_aligned = \"\"\n", " for text_segment in alignments:\n", " start_s = text_segment[\"start_s\"]\n", " end_s = text_segment[\"end_s\"]\n", " text_aligned = text_segment[\"text\"]\n", "\n", " # construct new meta\n", " new_meta = meta.copy()\n", " new_meta[\"text\"] = meta.get(\"lyrics\", \"\")\n", " new_meta[\"text_aligned\"] = text_aligned\n", " new_meta[\"tags\"] = meta.get(\"tags_text\", [])\n", " new_meta[\"n_vae_tokens\"] = VAE_MEMMAP_SIZE\n", " new_meta.pop(\"lyrics\", None)\n", "\n", " # append it twince, once for original and once for corrupted\n", " new_metas.append(new_meta)\n", " new_metas.append(new_meta)\n", "\n", " # write to memmap\n", " arr_s = data[\"semantic_codes\"]\n", " arr_v = data[\"vae_latents_original\"].astype(np.float16)\n", " arr_v = np.swapaxes(arr_v, 0, 1)\n", " out_mm_semantic[n_offs_s : n_offs_s + arr_s.size] = arr_s.reshape(-1)\n", " out_mm_vae[n_offs_v : n_offs_v + arr_v.size] = arr_v.reshape(-1)\n", " n_offs_s += arr_s.size\n", " n_offs_v += arr_v.size\n", "\n", " arr_s = data[\"semantic_codes_corrupted\"]\n", " arr_v = data[\"vae_latents_corrupted\"].astype(np.float16)\n", " arr_v = np.swapaxes(arr_v, 0, 1)\n", " out_mm_semantic[n_offs_s : n_offs_s + arr_s.size] = arr_s.reshape(-1)\n", " out_mm_vae[n_offs_v : n_offs_v + arr_v.size] = arr_v.reshape(-1)\n", " n_offs_s += arr_s.size\n", " n_offs_v += arr_v.size\n", "\n", " # write it once\n", " out_mm_semantic.flush()\n", " out_mm_vae.flush()\n", " del out_mm_semantic, out_mm_vae\n", " write_jsonl(new_metas, out_metas_filepath) \n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load the memmaps and check the shapes\n", "print(out_dir)\n", "metas_tr = read_jsonl(os.path.join(out_dir, f\"metas_tr.jsonl\"))\n", "print(len(metas_tr)/ 2)\n", "mm_semantic_tr = np.memmap(os.path.join(out_dir, f\"data_semantic_tr.bin\"), dtype=np.uint16, mode=\"r\")\n", "mm_vae_tr = np.memmap(os.path.join(out_dir, f\"data_vae_tr.bin\"), dtype=np.float16, mode=\"r\")\n", "\n", "\n", "mm_vae_tr = mm_vae_tr.reshape(-1, VAE_MEMMAP_SIZE, VAE_DIM)\n", "print(mm_vae_tr.shape)\n", "\n", "mm_semantic_tr = mm_semantic_tr.reshape(-1, SEMANTIC_MEMMAP_SIZE)\n", "print(mm_semantic_tr.shape)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "# select a random idx on 0, 2, 4, 6, ...\n", "rand_idx = np.arange(0, len(metas_tr), 2)\n", "rand_idx = np.random.choice(rand_idx, replace=False)\n", "\n", "\n", "print(rand_idx)\n", "print(metas_tr[rand_idx])\n", "vae_seq = mm_vae_tr[rand_idx]\n", "print(vae_seq.shape)\n", "vae_seq = torch.from_numpy(vae_seq.copy()).unsqueeze(0).float().cuda()\n", "\n", "with torch.no_grad():\n", " audio = model_100hz.decode(vae_seq.permute(0, 2, 1))[0].detach().cpu() \n", "audio /= audio.abs().max().clamp(1e-8)\n", "print(audio.mean())\n", "\n", "IPython.display.display(IPython.display.Audio(audio.numpy(), rate=48000))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print(metas_tr[rand_idx + 1])\n", "vae_seq = mm_vae_tr[rand_idx + 1]\n", "vae_seq = torch.from_numpy(vae_seq.copy()).unsqueeze(0).float().cuda()\n", "\n", "audio = model_100hz.decode(vae_seq.permute(0, 2, 1))[0].detach().cpu() \n", "audio /= audio.abs().max().clamp(1e-8)\n", "print(audio.mean())\n", "\n", "IPython.display.display(IPython.display.Audio(audio.numpy(), rate=48000))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# read npz file and decode\n", "npz_filepath = \"/home/christian/data/dpo/original_highpass_v2/6d0a1d6a-4d02-4cc8-bcfa-65e0e591e6db.npz\"\n", "data = np.load(npz_filepath)\n", "print(data[\"vae_latents_original\"].shape, data[\"vae_latents_corrupted\"].shape, data[\"semantic_codes\"].shape, data[\"semantic_codes_corrupted\"].shape)\n", "\n", "vae_seq = data[\"vae_latents_original\"].astype(np.float16)\n", "vae_seq = torch.from_numpy(vae_seq.copy()).unsqueeze(0).float().cuda()\n", "\n", "with torch.no_grad():\n", " audio = model_100hz.decode(vae_seq)[0].detach().cpu() \n", "audio /= audio.abs().max().clamp(1e-8)\n", "print(audio.mean())\n", "\n", "IPython.display.display(IPython.display.Audio(audio.numpy(), rate=48000))\n", "\n", "\n", "vae_seq = data[\"vae_latents_corrupted\"].astype(np.float16)\n", "vae_seq = torch.from_numpy(vae_seq.copy()).unsqueeze(0).float().cuda()\n", "\n", "with torch.no_grad():\n", " audio = model_100hz.decode(vae_seq)[0].detach().cpu() \n", "audio /= audio.abs().max().clamp(1e-8)\n", "print(audio.mean())\n", "\n", "IPython.display.display(IPython.display.Audio(audio.numpy(), rate=48000))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "suno_env2", "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.15" } }, "nbformat": 4, "nbformat_minor": 2 }