{ "cells": [ { "cell_type": "markdown", "id": "575ecb86", "metadata": {}, "source": [] }, { "cell_type": "code", "execution_count": null, "id": "fcf338ba", "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\"" ] }, { "cell_type": "code", "execution_count": null, "id": "9947c3ca", "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torchaudio\n", "#import polars as pl\n", "from suno_utils.utils.text import read_jsonl\n", "from suno_utils.audio import Audio\n", "from pathlib import Path\n", "from suno_utils.utils.opusfile import OpusFile" ] }, { "cell_type": "code", "execution_count": null, "id": "b4e53557", "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, "id": "b771829b", "metadata": {}, "outputs": [], "source": [ "#filepath = \"/app2/suno/data/auk_v0/metas_v8_tr_mini.jsonl\"\n", "filepath = \"/app2/suno/data/diffusion/v1/metas_v9_tr_filtered.jsonl\"\n", "\n", "metas = read_jsonl(filepath)\n", "#print(len(metas))\n", "#metas = pl.read_ndjson(filepath)\n", "print(len(metas))" ] }, { "cell_type": "code", "execution_count": null, "id": "14d2cccc", "metadata": {}, "outputs": [], "source": [ "#filepath = \"/app2/suno/data/auk_v0/metas_v8_tr_mini.jsonl\"\n", "filepath = \"/app2/suno/data/diffusion/v1/metas_v9_val_filtered.jsonl\"\n", "\n", "val_metas = read_jsonl(filepath)\n", "#print(len(metas))\n", "#metas = pl.read_ndjson(filepath)\n", "print(len(val_metas))" ] }, { "cell_type": "code", "execution_count": null, "id": "638bed98", "metadata": {}, "outputs": [], "source": [ "def _read_opus_window(path: str, start_s: float, total_len_s: float, sample_rate: int = 48000) -> \"Audio\":\n", " buf_size = int(sample_rate * total_len_s)\n", " float_arr = OpusFile(path=path).read(\n", " buf_size=buf_size,\n", " float_samples=True,\n", " from_position=int(sample_rate * max(0.0, start_s)),\n", " )\n", " return Audio.from_array_float(float_arr, sample_rate=sample_rate, max_allowed_val=12)\n", "\n", "from tqdm import tqdm\n", "\n", "results = []\n", "for meta in tqdm(val_metas):\n", " if \"podcast\" in meta[\"id\"]:\n", " try:\n", " #_read_opus_window(meta[\"local_filepath\"], 0.0, meta[\"duration_s\"])\n", " Audio.from_file(meta[\"local_filepath\"])\n", " results.append(True)\n", " except:\n", " results.append(False)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "2e4c33e7", "metadata": {}, "outputs": [], "source": [ "print(sum(results), len(results))" ] }, { "cell_type": "code", "execution_count": null, "id": "2c1a7e6e", "metadata": {}, "outputs": [], "source": [ "# lets first build a mapping of artist_ids to meta indices\n", "artist_id_to_meta_idx = {}\n", "for idx, meta in enumerate(metas):\n", " if \"artist_ids\" in meta:\n", " if len(meta[\"artist_ids\"]) == 1: # only do if there is a single artist\n", " for artist_id in meta[\"artist_ids\"]:\n", " if artist_id not in artist_id_to_meta_idx:\n", " artist_id_to_meta_idx[artist_id] = []\n", " artist_id_to_meta_idx[artist_id].append(idx)\n", "\n", "print(len(artist_id_to_meta_idx))" ] }, { "cell_type": "code", "execution_count": null, "id": "c74d9d5f", "metadata": {}, "outputs": [], "source": [ "# lets first build a mapping of artist_ids to meta indices\n", "artist_id_to_meta_idx = {}\n", "for idx, meta in enumerate(metas):\n", " if \"artist_ids\" in meta:\n", " if len(meta[\"artist_ids\"]) == 1: # only do if there is a single artist\n", " for artist_id in meta[\"artist_ids\"]:\n", " if artist_id not in artist_id_to_meta_idx:\n", " artist_id_to_meta_idx[artist_id] = []\n", " artist_id_to_meta_idx[artist_id].append(idx)\n", "\n", "# but what we actually want is a mapping of artist_ids to vox stem paths\n", "artist_id_to_vox_paths = {}\n", "for artist_id, meta_indices in artist_id_to_meta_idx.items():\n", " for meta_idx in meta_indices:\n", " meta = metas[meta_idx]\n", " if \"stems\" in meta and \"Vocals\" in meta[\"stems\"]:\n", " if artist_id not in artist_id_to_vox_paths:\n", " artist_id_to_vox_paths[artist_id] = []\n", " artist_id_to_vox_paths[artist_id].append(meta[\"stems\"][\"Vocals\"])\n", "\n", "print(f\"Found {len(artist_id_to_vox_paths)} artist_ids with vox stem paths\")" ] }, { "cell_type": "code", "execution_count": null, "id": "e99ff41b", "metadata": {}, "outputs": [], "source": [ "from typing import List, Dict, Optional, Any, Tuple\n", "\n", "def _build_artist_vox_mappings(\n", " metas: List[Dict],\n", ") -> Tuple[Dict[str, List[int]], Dict[str, Optional[Dict[str, Any]]]]:\n", " \"\"\"\n", " Build mappings from artist IDs to metadata indices and vox stem paths with duration.\n", "\n", " Args:\n", " metas: List of metadata dictionaries\n", "\n", " Returns:\n", " Tuple of (artist_id_to_meta_idx, artist_id_to_vox_paths)\n", " - artist_id_to_meta_idx: Maps artist_id to list of meta indices\n", " - artist_id_to_vox_paths: Maps artist_id to dict with 'path' and 'duration_s' (or None)\n", " \"\"\"\n", " # Build mapping of artist_ids to meta indices\n", " artist_id_to_meta_idx = {}\n", " for idx, meta in enumerate(metas):\n", " if \"artist_ids\" in meta:\n", " if len(meta[\"artist_ids\"]) == 1: # only do if there is a single artist\n", " for artist_id in meta[\"artist_ids\"]:\n", " if artist_id not in artist_id_to_meta_idx:\n", " artist_id_to_meta_idx[artist_id] = []\n", " artist_id_to_meta_idx[artist_id].append(idx)\n", "\n", " # Build mapping of artist_ids to vox stem paths with duration\n", " artist_id_to_vox_paths = {}\n", " for artist_id, meta_indices in artist_id_to_meta_idx.items():\n", " for meta_idx in meta_indices:\n", " meta = metas[meta_idx]\n", " if \"stems\" in meta and \"Vocals\" in meta[\"stems\"]:\n", " # Extract duration_s from the parent metadata\n", " duration_s = meta.get(\"duration_s\", None)\n", " if artist_id not in artist_id_to_vox_paths:\n", " artist_id_to_vox_paths[artist_id] = []\n", " artist_id_to_vox_paths[artist_id].append({\n", " \"path\": meta[\"stems\"][\"Vocals\"],\n", " \"duration_s\": duration_s,\n", " })\n", "\n", " return artist_id_to_meta_idx, artist_id_to_vox_paths" ] }, { "cell_type": "code", "execution_count": null, "id": "61bfd3a1", "metadata": {}, "outputs": [], "source": [ "artist_id_to_meta_idx, artist_id_to_vox_paths = _build_artist_vox_mappings(metas)\n", "print(len(artist_id_to_vox_paths))\n", "print(f\"Found {len([k for k, v in artist_id_to_vox_paths.items() if v is not None])} artist_ids with vox stem paths\")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "16a89417", "metadata": {}, "outputs": [], "source": [ "# filter artist_id_to_meta_idx by the keys of artist_id_to_vox_paths\n", "filtered_artist_id_to_meta_idx = {artist_id: artist_id_to_meta_idx[artist_id] for artist_id in artist_id_to_vox_paths.keys()}\n", "print(len(filtered_artist_id_to_meta_idx))" ] }, { "cell_type": "code", "execution_count": null, "id": "68c7a915", "metadata": {}, "outputs": [], "source": [ "# Count occurrences of each unique language string per artist id\n", "from collections import Counter\n", "\n", "artist_id_to_lang_counter = {}\n", "for artist_id, meta_indices in artist_id_to_meta_idx.items():\n", " if \"podcast\" in artist_id: \n", " continue\n", " if meta_indices is not None and len(meta_indices) > 1:\n", " lang_counter = Counter()\n", " for meta_idx in meta_indices:\n", " meta = metas[meta_idx]\n", " meta_lang = meta.get(\"lang\", None)\n", " if meta_lang is not None:\n", " lang_counter[meta_lang] += 1\n", " if lang_counter:\n", " artist_id_to_lang_counter[artist_id] = dict(lang_counter)\n", "\n", "unique_langs = set()\n", "for lang_counts in artist_id_to_lang_counter.values():\n", " unique_langs.update(lang_counts.keys())\n", "unique_langs = sorted(unique_langs)\n", "print(len(unique_langs))\n" ] }, { "cell_type": "code", "execution_count": null, "id": "d4a9101b", "metadata": {}, "outputs": [], "source": [ "METAS_DIR = \"/app/suno/tmp\"\n", "discogs_subset_metas_map = {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"raw_discogs_subset_metas.jsonl\"))}\n", "print(len(discogs_subset_metas_map))" ] }, { "cell_type": "code", "execution_count": null, "id": "84e5110d", "metadata": {}, "outputs": [], "source": [ "from tqdm import tqdm\n", "sft_subset_metas = []\n", "for meta in tqdm(filtered_sft_subset_metas):\n", " has_stems = \"stems\" in meta\n", " if meta[\"id\"] in discogs_subset_metas_map:\n", " has_views = discogs_subset_metas_map[meta[\"id\"]].get(\"views\", 0) != 0\n", " else:\n", " has_views = False\n", " has_lang = meta.get(\"lang\", None) is not None\n", " # check if we have artist_ids\n", " has_one_artist = len(meta.get(\"artist_ids\", [])) == 1\n", " if has_stems and has_views and has_lang and has_one_artist:\n", " sft_subset_metas.append(meta)\n", "\n", "print(len(sft_subset_metas))\n" ] }, { "cell_type": "code", "execution_count": null, "id": "44f5daa5", "metadata": {}, "outputs": [], "source": [ "# plot the view distribution from sft_subset_metas\n", "view_counts = []\n", "for meta in sft_subset_metas:\n", " if meta[\"id\"] in discogs_subset_metas_map:\n", " views = discogs_subset_metas_map[meta[\"id\"]].get(\"views\", 0)\n", " meta[\"views\"] = views\n", " if views != 0:\n", " view_counts.append(views)" ] }, { "cell_type": "code", "execution_count": null, "id": "42471651", "metadata": {}, "outputs": [], "source": [ "filtered_sft_subset_metas = [meta for meta in sft_subset_metas if meta[\"views\"] > 500_000]\n", "print(len(filtered_sft_subset_metas))" ] }, { "cell_type": "code", "execution_count": null, "id": "2471dd4b", "metadata": {}, "outputs": [], "source": [ "# Sample N = 1000 random metas from sft_subset_metas, trying to balance the proportion of languages\n", "\n", "import random\n", "from collections import defaultdict\n", "\n", "N = 4000\n", "\n", "# Step 1: Organize metas by language\n", "lang_to_metas = defaultdict(list)\n", "for meta in filtered_sft_subset_metas:\n", " lang = meta.get(\"lang\", None)\n", " if lang is not None:\n", " lang_to_metas[lang].append(meta)\n", "\n", "# Step 2: Determine per-language quota (even split, allow roundoff)\n", "num_langs = len(lang_to_metas)\n", "base_per_lang = N // num_langs\n", "extra = N % num_langs\n", "\n", "per_lang_quota = {}\n", "for i, lang in enumerate(sorted(lang_to_metas.keys())):\n", " per_lang_quota[lang] = base_per_lang + (1 if i < extra else 0)\n", "\n", "# Step 3: For each language, randomly sample up to the quota\n", "sampled_metas = []\n", "for lang, quota in per_lang_quota.items():\n", " metas = lang_to_metas[lang]\n", " if len(metas) <= quota:\n", " sampled_metas.extend(metas)\n", " else:\n", " sampled_metas.extend(random.sample(metas, quota))\n", "\n", "print(f\"Sampled {len(sampled_metas)} metas\")\n", "lang_counts = defaultdict(int)\n", "for meta in sampled_metas:\n", " lang_counts[meta['lang']] += 1\n", "print(\"Per-language counts:\", dict(lang_counts))\n", "\n", "import matplotlib.pyplot as plt\n", "\n", "langs = list(lang_counts.keys())\n", "counts = [lang_counts[lang] for lang in langs]\n", "\n", "print(len(sampled_metas))\n", "\n", "# Sort langs and counts together by count (descending)\n", "sorted_langs_counts = sorted(zip(langs, counts), key=lambda x: x[1], reverse=True)\n", "sorted_langs, sorted_counts = zip(*sorted_langs_counts)\n", "\n", "plt.figure(figsize=(20, 6))\n", "plt.bar(sorted_langs, sorted_counts, color='skyblue')\n", "plt.xlabel(\"Language\")\n", "plt.ylabel(\"Number of Metas\")\n", "plt.title(\"Number of Sampled Metas per Language\")\n", "plt.xticks(rotation=45, ha=\"right\")\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "id": "76d3efd5", "metadata": {}, "outputs": [], "source": [ "from suno_utils.utils.text import write_jsonl\n", "# save these metas as a subset\n", "output_filepath = \"/home/christian/code/christian/metadata/metas_with_stems_lang_bal_one_artist_1k.jsonl\"\n", "write_jsonl(sampled_metas[:1000], output_filepath)" ] }, { "cell_type": "code", "execution_count": null, "id": "7c80f824", "metadata": {}, "outputs": [], "source": [ "unique_artist_ids = set()\n", "for meta in sampled_metas:\n", " artist_ids = meta.get(\"artist_ids\", [])\n", " if len(artist_ids) == 1:\n", " unique_artist_ids.update(artist_ids)\n", "print(len(unique_artist_ids))\n" ] }, { "cell_type": "code", "execution_count": null, "id": "fd7bfd83", "metadata": {}, "outputs": [], "source": [ "sampled_metas[0]" ] }, { "cell_type": "code", "execution_count": null, "id": "1103d9e6", "metadata": {}, "outputs": [], "source": [ "discogs_subset_metas_map[list(discogs_subset_metas_map.keys())[0]]" ] }, { "cell_type": "code", "execution_count": null, "id": "9fd54e1b", "metadata": {}, "outputs": [], "source": [ "from suno_utils.audio import Audio\n", "\n", "filepath = \"/app2/suno/data/podcast/audio_chunks/chunk_69_3882db2e-0e06-4bd0-9165-4a64fa9fe8d9.mp3.mp3\"\n", "audio = Audio.from_file(filepath, n_channels=2)\n", "audio.play()\n", "print(audio.duration_s)\n", "audio = audio.pad_to_length(30.02)\n", "print(audio.duration_s)\n", "audio.play()" ] }, { "cell_type": "code", "execution_count": null, "id": "a060885f", "metadata": {}, "outputs": [], "source": [ "# Sample up to num_samples artist ids, trying to **aggressively** balance across languages,\n", "# but this time, **heavily upweight non-English** by giving them a much larger per-language cap\n", "# (e.g. 4x English's cap), and *lowering* the English quota to a minimum.\n", "# Also, report how many samples we get per language.\n", "\n", "import random\n", "from collections import Counter, defaultdict\n", "\n", "num_samples = 1000\n", "\n", "# Specify which language code is English\n", "EN_LANG_CODE = \"en\"\n", "\n", "# 1. Map language to artist ids with at least one song with stems AND lang\n", "lang_to_artist_ids = {lang: [] for lang in unique_langs}\n", "artist_id_has_stems_and_lang = {}\n", "\n", "for artist_id, meta_indices in artist_id_to_meta_idx.items():\n", " has_stems_and_lang = False\n", " if meta_indices:\n", " for idx in meta_indices:\n", " meta = metas[idx]\n", " has_stems = \"stems\" in meta\n", " has_lang = meta.get(\"lang\", None) is not None\n", " if has_stems and has_lang:\n", " has_stems_and_lang = True\n", " break\n", " artist_id_has_stems_and_lang[artist_id] = has_stems_and_lang\n", "\n", "for artist_id, lang_counts in artist_id_to_lang_counter.items():\n", " if artist_id_has_stems_and_lang.get(artist_id, False):\n", " for lang in lang_counts.keys():\n", " lang_to_artist_ids[lang].append(artist_id)\n", "\n", "# 2. Flat list of all eligible artist ids\n", "all_artist_ids = set([aid for aid, ok in artist_id_has_stems_and_lang.items() if ok])\n", "\n", "################################################################################\n", "# 3. Upweight non-English artist selection!\n", "################################################################################\n", "\n", "# Heuristic: set English to a **minimum** quota (maybe 1/6 total), and all others\n", "# get an even split of the rest.\n", "english_percentage = 1/6\n", "max_english_quota = int(num_samples * english_percentage)\n", "n_non_english = max(1, len(unique_langs) - (1 if EN_LANG_CODE in unique_langs else 0))\n", "max_non_english_quota = (num_samples - max_english_quota) // n_non_english\n", "# For robustness:\n", "per_lang_quota = {lang: (max_english_quota if lang == EN_LANG_CODE else max_non_english_quota) for lang in unique_langs}\n", "\n", "# Track selected artist ids strictly per-language, upweighting non-English\n", "selected_per_lang = {lang: set() for lang in unique_langs}\n", "used_artist_ids = set()\n", "\n", "# Shuffle to randomize selection within each lang bucket\n", "for lang in unique_langs:\n", " artist_list = list(set(lang_to_artist_ids[lang])) # unique\n", " random.shuffle(artist_list)\n", " take = 0\n", " for artist_id in artist_list:\n", " # Don't double-count artists who are in multiple lang buckets (pick first lang to add them)\n", " if artist_id in used_artist_ids:\n", " continue\n", " selected_per_lang[lang].add(artist_id)\n", " used_artist_ids.add(artist_id)\n", " take += 1\n", " if take >= per_lang_quota[lang]:\n", " break\n", "\n", "# Gather final artist id list, maintaining language upweighting for non-English\n", "sampled_artist_ids = []\n", "for lang in unique_langs:\n", " sampled_artist_ids.extend(selected_per_lang[lang])\n", "\n", "# If we have more than num_samples after upweighting, trim randomly\n", "if len(sampled_artist_ids) > num_samples:\n", " random.shuffle(sampled_artist_ids)\n", " sampled_artist_ids = sampled_artist_ids[:num_samples]\n", "# If too few, fill from remaining eligible ids (prioritizing non-English, then English last)\n", "elif len(sampled_artist_ids) < num_samples:\n", " remaining_ids = list(all_artist_ids - set(sampled_artist_ids))\n", " # sort: non-English first, English last\n", " def sort_key(aid):\n", " langs = artist_id_to_lang_counter.get(aid, {})\n", " return (EN_LANG_CODE in langs, aid) # False < True\n", " remaining_ids_sorted = sorted(remaining_ids, key=sort_key)\n", " for aid in remaining_ids_sorted:\n", " sampled_artist_ids.append(aid)\n", " if len(sampled_artist_ids) >= num_samples:\n", " break\n", "\n", "print(f\"Sampled {len(sampled_artist_ids)} artist ids with stems AND lang, aggressively upweighted non-English across {len(unique_langs)} languages.\")\n", "print(\"Per-lang quotas (for upweighting):\")\n", "for lang in unique_langs:\n", " print(f\" {lang}: {per_lang_quota[lang]}\")\n", "\n", "################################################################################\n", "# 6. For each artist, randomly select one song id (metadata index) from their list, must have stems AND lang\n", "################################################################################\n", "sampled_song_metas = []\n", "sampled_artist_id_to_langs = {}\n", "for artist_id in sampled_artist_ids:\n", " indices = artist_id_to_meta_idx.get(artist_id, [])\n", " song_meta_with_stems_and_lang = []\n", " if indices:\n", " for idx in indices:\n", " meta = metas[idx]\n", " # Additional check: meta id present in discogs_subset_metas_map and has a valid views count\n", " meta_id = meta.get(\"id\") or meta.get(\"track_id\") or meta.get(\"song_id\")\n", " has_stems = \"stems\" in meta\n", " has_lang = meta.get(\"lang\", None) is not None\n", " has_discogs_view_count = (\n", " meta_id is not None and meta_id in discogs_subset_metas_map and discogs_subset_metas_map[meta_id].get(\"views\", None) is not None\n", " )\n", " if has_discogs_view_count:\n", " views = discogs_subset_metas_map[meta_id].get(\"views\", 0)\n", " else:\n", " views = 0\n", " if has_stems and has_lang and has_discogs_view_count and views > 100_000:\n", " song_meta_with_stems_and_lang.append(meta)\n", " if song_meta_with_stems_and_lang:\n", " meta = random.choice(song_meta_with_stems_and_lang)\n", " meta_id = meta.get(\"id\") or meta.get(\"track_id\") or meta.get(\"song_id\")\n", " if meta_id and meta_id in discogs_subset_metas_map:\n", " meta[\"views\"] = discogs_subset_metas_map[meta_id].get(\"views\", None)\n", " sampled_song_metas.append(meta)\n", " sampled_artist_id_to_langs[artist_id] = list(artist_id_to_lang_counter.get(artist_id, {}).keys())\n", "\n", "print(f\"Selected {len(sampled_song_metas)} song ids with stems AND lang, one per sampled artist.\")\n", "\n", "# Report count of sampled artists per language\n", "lang_sample_counter = Counter()\n", "for artist_id, langs in sampled_artist_id_to_langs.items():\n", " for lang in langs:\n", " lang_sample_counter[lang] += 1\n", "\n", "print(\"Sample count per language in sampled_artist_ids:\")\n", "for lang in unique_langs:\n", " print(f\" {lang}: {lang_sample_counter[lang]}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "4666c74d", "metadata": {}, "outputs": [], "source": [ "from suno_utils.utils.text import write_jsonl\n", "# save these metas as a subset\n", "output_filepath = \"/home/christian/code/christian/metadata/metas_with_stems_lang_bal_1k.jsonl\"\n", "write_jsonl(sampled_metas, output_filepath)" ] }, { "cell_type": "code", "execution_count": null, "id": "2d3e7d7a", "metadata": {}, "outputs": [], "source": [ "sampled_song_metas[500]" ] }, { "cell_type": "code", "execution_count": null, "id": "71ea76f7", "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "\n", "# Count the occurrence of each unique lang in the sampled_song_metas\n", "lang_counter = Counter()\n", "for meta in sampled_song_metas:\n", " lang = meta.get(\"lang\", None)\n", " if lang:\n", " lang_counter[lang] += 1\n", "\n", "# Sort languages by frequency (descending)\n", "sorted_lang_counts = sorted(lang_counter.items(), key=lambda x: x[1], reverse=True)\n", "langs = [item[0] for item in sorted_lang_counts]\n", "counts = [item[1] for item in sorted_lang_counts]\n", "\n", "# Plot the sorted histogram\n", "plt.figure(figsize=(8, 4))\n", "plt.bar(langs, counts)\n", "plt.xlabel(\"Language\")\n", "plt.ylabel(\"Count\")\n", "plt.title(\"Histogram of language occurrence in sampled_song_metas (sorted)\")\n", "plt.xticks(rotation=45)\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "id": "97d79b0b", "metadata": {}, "outputs": [], "source": [ "for artist_id, result in artist_id_to_lang_counter.items():\n", " print(artist_id, result)" ] }, { "cell_type": "code", "execution_count": null, "id": "7e24ca4e", "metadata": {}, "outputs": [], "source": [ "for key, val in artist_id_to_vox_paths.items():\n", " if val is not None and len(val) > 1:\n", " print(key)\n", " for v in val:\n", " print(v)\n", " break" ] }, { "cell_type": "code", "execution_count": null, "id": "4bfaf6a5", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import random\n", "from typing import List, Tuple\n", "\n", "def _fast_trim_mono(\n", " x: np.ndarray, # shape: (samples,), float32/64 in [-1, 1]\n", " sr: int, # sample rate (Hz)\n", " thresh_db_rel: float = -35, # keep where RMS > max_RMS + thresh (dB)\n", " win_ms: float = 20.0, # moving RMS window size (ms)\n", " pad_ms: float = 20.0, # pad around kept regions (ms)\n", " min_keep_ms: float = 40.0 # drop kept bits shorter than this (ms)\n", ") -> Tuple[np.ndarray, List[Tuple[int, int]]]:\n", " \"\"\"\n", " Ultra-fast silence trimmer for mono audio. No convolutions, all O(n).\n", " Returns (trimmed_audio, kept_spans) with kept_spans in original sample indices.\n", " \"\"\"\n", " assert x.ndim == 1, \"Expected mono waveform of shape (samples,)\"\n", " n = x.size\n", " if n == 0:\n", " return x[:0], []\n", "\n", " # --- Moving RMS via cumulative sums (box filter), O(n) ---\n", " # Compute moving average of power over a window, then sqrt.\n", " win = max(1, int(round(sr * win_ms / 1000.0)))\n", " if win > n:\n", " win = n\n", "\n", " # power and cumulative sum (use float64 for numeric safety)\n", " sq = x.astype(np.float64) ** 2\n", " csum = np.empty(n + 1, dtype=np.float64)\n", " csum[0] = 0.0\n", " np.cumsum(sq, out=csum[1:]) # csum[k] = sum_{i thresh_db_rel # True = keep\n", "\n", " # --- Turn mask into spans, expand by pad, merge, drop short ---\n", " pad = max(0, int(round(sr * pad_ms / 1000.0)))\n", " min_keep = max(1, int(round(sr * min_keep_ms / 1000.0)))\n", "\n", " # Find rising/falling edges\n", " m = mask.astype(np.int8)\n", " edges = np.flatnonzero(np.diff(m, prepend=0, append=0))\n", " # edges come in pairs [start0, end0, start1, end1, ...]\n", " starts = edges[::2]\n", " ends = edges[1::2]\n", "\n", " if starts.size == 0:\n", " return x[:0], []\n", "\n", " # Expand by pad and clamp\n", " starts = np.maximum(0, starts - pad)\n", " ends = np.minimum(n, ends + pad)\n", "\n", " # Merge overlaps and drop short spans\n", " spans: List[Tuple[int, int]] = []\n", " s_prev = int(starts[0])\n", " e_prev = int(ends[0])\n", " for s, e in zip(starts[1:], ends[1:]):\n", " s = int(s); e = int(e)\n", " if s <= e_prev: # overlap/adjacent -> merge\n", " e_prev = max(e_prev, e)\n", " else:\n", " if (e_prev - s_prev) >= min_keep:\n", " spans.append((s_prev, e_prev))\n", " s_prev, e_prev = s, e\n", " # last span\n", " if (e_prev - s_prev) >= min_keep:\n", " spans.append((s_prev, e_prev))\n", "\n", " if not spans:\n", " return x[:0], []\n", "\n", " # --- Concatenate kept spans (one pass) ---\n", " parts = [x[a:b] for (a, b) in spans]\n", " y = np.concatenate(parts, axis=0).astype(x.dtype)\n", " return y, spans\n", "\n", "def _read_opus_window(path: str, start_s: float, total_len_s: float, sample_rate: int = 48000) -> \"Audio\":\n", " buf_size = int(sample_rate * total_len_s)\n", " float_arr = OpusFile(path=path).read(\n", " buf_size=buf_size,\n", " float_samples=True,\n", " from_position=int(sample_rate * max(0.0, start_s)),\n", " )\n", " return Audio.from_array_float(float_arr, sample_rate=sample_rate, max_allowed_val=12)\n", "\n", "\n", "def _extract_vox_segment(\n", " path: str,\n", " vox_duration_s: float,\n", " vox_cond_duration_s: float = 30.0,\n", " min_segment_duration_s: float = 3.0,\n", " min_segments: int = 1,\n", " max_segments: int = 4,\n", " sample_rate: int = 48000,\n", ") -> \"Audio\":\n", " \"\"\"Extract vox segment from the file.\"\"\"\n", " # read the entire file\n", " vox_window = _read_opus_window(path, 0.0, vox_duration_s)\n", "\n", " mono_audio_array = (vox_window.array_float[0] + vox_window.array_float[1]) / 2\n", "\n", " # strip silence\n", " trimmed_audio, kept_spans = _fast_trim_mono(mono_audio_array, sample_rate)\n", "\n", " if len(trimmed_audio) == 0:\n", " # If no audio after trimming, return empty audio\n", " return None\n", "\n", " # Convert kept_spans from original indices to trimmed audio indices\n", " # Since _fast_trim_mono returns trimmed audio, we need to work with the trimmed length\n", " trimmed_duration_s = len(trimmed_audio) / sample_rate\n", "\n", " # select the number of random segments\n", " num_segments = random.randint(min_segments, max_segments)\n", " print(num_segments)\n", "\n", " # Calculate available duration for segments (min of vox_cond_duration_s and trimmed_duration_s)\n", " available_duration_s = min(vox_cond_duration_s, trimmed_duration_s)\n", "\n", " # Generate segment lengths that sum to available_duration_s\n", " # Each segment must be at least min_segment_duration_s\n", " segment_lengths = []\n", " remaining_duration = available_duration_s\n", "\n", " for i in range(num_segments):\n", " if i == num_segments - 1:\n", " # Last segment gets all remaining duration\n", " segment_lengths.append(remaining_duration)\n", " else:\n", " # Calculate max possible length for this segment\n", " # Need to leave room for remaining segments (each needs min_segment_duration_s)\n", " max_length = remaining_duration - (num_segments - i - 1) * min_segment_duration_s\n", " min_length = min_segment_duration_s\n", "\n", " if max_length <= min_length:\n", " # Not enough duration left, use minimum\n", " segment_lengths.append(min_length)\n", " else:\n", " # Random length between min and max\n", " segment_lengths.append(random.uniform(min_length, max_length))\n", "\n", " remaining_duration -= segment_lengths[-1]\n", "\n", " # Sample segments from the trimmed audio\n", " segments = []\n", " current_pos = 0\n", "\n", " for segment_length in segment_lengths:\n", " segment_samples = int(segment_length * sample_rate)\n", "\n", " # Ensure we don't go beyond the trimmed audio length\n", " if current_pos + segment_samples > len(trimmed_audio):\n", " segment_samples = len(trimmed_audio) - current_pos\n", "\n", " if segment_samples > 0:\n", " segment = trimmed_audio[current_pos : current_pos + segment_samples]\n", " segments.append(segment)\n", " current_pos += segment_samples\n", " else:\n", " break\n", "\n", " if not segments:\n", " # If no valid segments, return empty audio\n", " return None\n", "\n", " # concatenate the segments into one Audio object\n", " concatenated_audio = np.concatenate(segments)\n", " vox_window = Audio.from_array_float(concatenated_audio, sample_rate=sample_rate, max_allowed_val=12)\n", "\n", " return vox_window" ] }, { "cell_type": "code", "execution_count": null, "id": "d1447b5f", "metadata": {}, "outputs": [], "source": [ "# find the top 5 artist_ids with the most metas\n", "from operator import itemgetter\n", "\n", "# Get a list of tuples (artist_id, list_of_meta_indices), sorted by length of list_of_meta_indices descending\n", "top_artist_ids = sorted(artist_id_to_vox_paths.items(), key=lambda x: len(x[1]), reverse=True)[:5]\n", "\n", "for idx, (artist_id, meta_indices) in enumerate(top_artist_ids, 1):\n", " print(f\"Top {idx} artist_id: {artist_id}\")\n", " print(f\" Number of metas: {len(meta_indices)}\")\n", " print(f\" Meta indices: {meta_indices}\")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "61c27a48", "metadata": {}, "outputs": [], "source": [ "import time\n", "\n", "# vox should actually be mono\n", "vox_cond_duration_s = 30.0\n", "min_segment_duration_s = 3.0\n", "\n", "vox_filepaths = top_artist_ids[1][1]\n", "for vox_filepath in vox_filepaths:\n", " vox_audio = Audio.from_file(vox_filepath, n_channels=2)\n", " vox_duration_s = vox_audio.duration_s\n", " #start_time = time.time()\n", " #vox_trimmed = _fast_trim_mono(vox_audio.array_float[0], vox_audio.sample_rate)\n", " #elapsed = time.time() - start_time\n", "\n", " vox_segment = _extract_vox_segment(vox_filepath, vox_duration_s, vox_cond_duration_s, min_segment_duration_s)\n", "\n", " #audio_trimmed = Audio.from_array_float(vox_trimmed[0], vox_audio.sample_rate)\n", " vox_segment.play()\n", " #print(f\"Trim took {elapsed:.4f} seconds\")" ] }, { "cell_type": "code", "execution_count": null, "id": "9799e831", "metadata": {}, "outputs": [], "source": [ "for artist_id, meta_indices in artist_id_to_meta_idx.items():\n", " if len(meta_indices) > 1:\n", " # first check if any have stems\n", " has_stems = False\n", " for meta_idx in meta_indices:\n", " if \"stems\" in metas[meta_idx]:\n", " has_stems = True\n", " break\n", " if has_stems:\n", " print(\"artist_id\", artist_id, len(meta_indices))\n", " for meta_idx in meta_indices:\n", " print(metas[meta_idx][\"id\"])\n", " print(metas[meta_idx][\"local_filepath\"])\n", " print(\"stems\", metas[meta_idx].get(\"stems\", None))\n", " print()\n", " \n" ] }, { "cell_type": "code", "execution_count": null, "id": "ec22018b", "metadata": {}, "outputs": [], "source": [ "from suno_utils.audio import Audio\n", "\n", "audio1 = Audio.from_file(\"/app2/suno/data/raw_audio_opus_v0/RdGw36uLG7s.opus\")\n", "audio1_vox = Audio.from_file(\"/app2/suno/data/sft_stems_12_output_v11/RdGw36uLG7s_Vocals.opus\")\n", "\n", "audio2 = Audio.from_file(\"/app2/suno/data/raw_audio_opus_v0/6H6zGLP9Qos.opus\")\n", "audio2_vox = Audio.from_file(\"/app2/suno/data/sft_stems_12_output_v11/6H6zGLP9Qos_Vocals.opus\")\n", "\n", "audio1.play()\n", "audio1_vox.play()\n", "\n", "audio2.play()\n", "audio2_vox.play()\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "19a19ba8", "metadata": {}, "outputs": [], "source": [ "# count how many times each key occurs in the metas\n", "from collections import Counter\n", "key_counter = Counter()\n", "for meta in metas:\n", " for key in meta.keys():\n", " key_counter[key] += 1\n", "print(key_counter)\n", "\n", "for key, count in key_counter.items():\n", " if count > 1:\n", " print(key, count)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "61721e07", "metadata": {}, "outputs": [], "source": [ "for meta in metas:\n", " if \"stems\" in meta:\n", " print(meta[\"stems\"])\n" ] }, { "cell_type": "code", "execution_count": null, "id": "a9a7c487", "metadata": {}, "outputs": [], "source": [ "# lets count how many 30s chunks per second we can load\n", "import sys\n", "sys.path.append(\"/home/christian/code/neon/sunoDiff\")\n", "\n", "# combine text and tags into a single string, tags is a list of strings\n", "def augment_text_training(tags, text):\n", " return \",\".join(tags) + \" \" + text\n", "\n", "def prepare_text_inference(tags, text):\n", " return \",\".join(tags) + \" \" + text\n", "\n", "from tokenizers import Tokenizer\n", "from contextlib import contextmanager\n", "import tempfile\n", "import os\n", "\n", "\n", "from helpers import (\n", " dist_barrier,\n", " download_s3_file,\n", " get_filename,\n", " read_jsonl,\n", " write_jsonl,\n", " print_with_time_master,\n", " respell_random_words_in_text,\n", ")\n", "\n", "@contextmanager\n", "def _download_from_s3_if_needed(maybe_s3_filepath):\n", " tmp_filepath = maybe_s3_filepath\n", " if maybe_s3_filepath.startswith(\"s3://\"):\n", " temp_dir = tempfile.TemporaryDirectory()\n", " filename = get_filename(maybe_s3_filepath, keep_ext=True)\n", " tmp_filepath = os.path.join(temp_dir.name, filename)\n", " download_s3_file(maybe_s3_filepath, tmp_filepath)\n", " yield tmp_filepath\n", "\n", "\n", "def load_tokenizer(\n", " tokenizer_filepath=\"s3://suno-data/georg/models/tokenizers/tokenizer_60k.json\",\n", "):\n", " with _download_from_s3_if_needed(tokenizer_filepath) as tmp_fp:\n", " tokenizer = Tokenizer.from_file(tmp_fp)\n", " tokenizer.add_special_tokens([\"\\n\"])\n", " tokenizer.pad_idx = tokenizer.token_to_id(\"[PAD]\")\n", " return tokenizer\n", "\n", "SAMPLE_RATE = 48000\n", "\n", "def load_meta(meta, audio_chunk_s, audio_ctx_s):\n", " # 50% load 30s, 50% load 60s\n", " use_ctx = random.random() < 0.5\n", " if not use_ctx:\n", " buf_size_samples = int(SAMPLE_RATE * audio_chunk_s)\n", " else:\n", " buf_size_samples = int(SAMPLE_RATE * (audio_chunk_s + audio_ctx_s))\n", "\n", " full_audio_start_s = random.uniform(0, meta[\"duration_s\"] - buf_size_samples / SAMPLE_RATE)\n", " full_audio_end_s = full_audio_start_s + buf_size_samples / SAMPLE_RATE\n", " print(full_audio_start_s, full_audio_end_s)\n", "\n", " # read audio chunk\n", " full_audio = Audio.from_array_float(\n", " OpusFile(path=meta[\"local_filepath\"]).read(\n", " buf_size=buf_size_samples,\n", " float_samples=True,\n", " from_position=int(SAMPLE_RATE * full_audio_start_s),\n", " ),\n", " sample_rate=SAMPLE_RATE,\n", " max_allowed_val=12,\n", " )\n", "\n", " # crop the chunk if needed\n", " if not use_ctx:\n", " audio_target = full_audio\n", " audio_target_start_s = full_audio_start_s\n", " audio_target_end_s = full_audio_end_s\n", " audio_ctx = None\n", " else:\n", " # first half goes to the audio, second half goes to the ctx\n", " audio_target = full_audio.get_segment(full_audio_start_s, full_audio_start_s + audio_chunk_s)\n", " audio_target_start_s = full_audio_start_s\n", " audio_target_end_s = int(SAMPLE_RATE * (full_audio_start_s + audio_chunk_s))\n", " audio_ctx = full_audio.get_segment(full_audio_start_s + audio_chunk_s, full_audio_start_s + audio_chunk_s + audio_ctx_s)\n", "\n", " # get text (lyrics) and tags\n", " tags = meta.get(\"tags\", [])\n", " # check if we have algined lyrics\n", " if \"text_aligned\" in meta:\n", " text_aligned = meta[\"text_aligned\"]\n", " # get all the text for time from 0 to 30 seconds\n", " # the format is a list, in each element, (start_time, end_time, text)\n", " text_for_time = [text for start_time, end_time, text in text_aligned if start_time >= audio_target_start_s and end_time >= audio_target_end_s]\n", " text = \"\".join([text for text in text_for_time])\n", " else:\n", " text = None\n", "\n", " # check for None cases in lyrics and tags\n", " if text is None:\n", " text = \"\"\n", " if tags is None:\n", " tags = []\n", "\n", " if is_training:\n", " if random.random() <= 0.1:\n", " text = \"\"\n", " else:\n", " full_text = augment_text_training(tags, text)\n", " else:\n", " full_text = prepare_text_inference(tags, text)\n", "\n", " print(full_text)\n", "\n", " text_codes = tokenizer.encode(full_text).ids[: cond_text_len]\n", " text_codes = text_codes + [tokenizer.pad_idx] * max(0, cond_text_len - len(text_codes))\n", " text_codes = torch.tensor(text_codes).long()\n", "\n", " return audio_target, audio_ctx, text_codes\n", "\n", "import random\n", "class DummyDataset(torch.utils.data.IterableDataset):\n", " def __init__(\n", " self, \n", " metas, \n", " audio_chunk_s: float = 30.0, \n", " audio_ctx_s: float = 30.0, \n", " cond_text_len=1536, \n", " is_training=False\n", " ):\n", " self.metas = metas\n", " self.is_training = is_training\n", " self.cond_text_len = cond_text_len\n", " self.audio_chunk_s = audio_chunk_s\n", " self.audio_ctx_s = audio_ctx_s\n", " self.tokenizer = load_tokenizer()\n", "\n", " def __iter__(self):\n", " for meta in self.metas:\n", " try:\n", " audio_target, audio_ctx, text_codes = load_meta(meta)\n", " yield audio_target, audio_ctx, text_codes\n", " except Exception as e:\n", " print(f\"Error loading meta: {e}\")\n", " continue\n", " \n", "\n", "# we will need a special collate function to handle the different length audios\n", "def collate_fn(batch):\n", " audio_target_list = [item[0] for item in batch]\n", " audio_ctx_list = [item[1] for item in batch]\n", " text_codes_list = [item[2] for item in batch]\n", " return audio_target_list, audio_ctx_list, text_codes_list\n", "\n", "\n", "dataset = DummyDataset(metas)\n", "dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=False, collate_fn=collate_fn, num_workers=1)\n", "\n", "import time\n", "from tqdm import tqdm\n", "num_eval_batches = 100\n", "# compute it/s\n", "# lets count how many 30s chunks per second we can load \n", "time_start = time.time()\n", "for i, batch in enumerate(tqdm(dataloader)):\n", " audio_target_list, audio_ctx_list, text_codes_list = batch\n", "\n", " try:\n", " latents = codec_encode(audio_target_list, normalize_volume=False)\n", " except Exception as e:\n", " for audio_target in audio_target_list:\n", " print(audio_target)\n", " if i > num_eval_batches:\n", " break\n", "\n", "time_end = time.time()\n", "print(f\"Time taken: {time_end - time_start} seconds\")\n", "# time per batch \n", "time_per_batch = (time_end - time_start) / num_eval_batches\n", "print(f\"Time per batch: {time_per_batch} seconds\")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "0d2ffcee", "metadata": {}, "outputs": [], "source": [ "# -*- coding: utf-8 -*-\n", "\"\"\"\n", "Benchmark: how many 30s chunks/second can we load & encode\n", "\"\"\"\n", "\n", "from __future__ import annotations\n", "import os\n", "import time\n", "import math\n", "import json\n", "import random\n", "import tempfile\n", "from typing import Any, Dict, Iterable, List, Optional, Tuple\n", "from contextlib import contextmanager\n", "\n", "import torch\n", "from torch.utils.data import IterableDataset, DataLoader\n", "from tokenizers import Tokenizer\n", "\n", "# If these are in a local repo, keep your path append — but do it once, at top.\n", "import sys\n", "sys.path.append(\"/home/christian/code/neon/sunoDiff\")\n", "\n", "# Replace with your actual audio module imports\n", "# from your_audio_lib import Audio, OpusFile\n", "\n", "from helpers import (\n", " dist_barrier,\n", " download_s3_file,\n", " get_filename,\n", " read_jsonl,\n", " write_jsonl,\n", " print_with_time_master,\n", " respell_random_words_in_text,\n", ")\n", "\n", "# -----------------------------\n", "# Text prep\n", "# -----------------------------\n", "def augment_text_training(tags: List[str], text: str) -> str:\n", " \"\"\"Combine tags and text for training; tags may be empty.\"\"\"\n", " tags = tags or []\n", " text = text or \"\"\n", " return (\",\".join(tags) + \" \" + text).strip()\n", "\n", "def prepare_text_inference(tags: List[str], text: str) -> str:\n", " \"\"\"Combine tags and text for inference; tags may be empty.\"\"\"\n", " tags = tags or []\n", " text = text or \"\"\n", " return (\",\".join(tags) + \" \" + text).strip()\n", "\n", "# -----------------------------\n", "# Tokenizer\n", "# -----------------------------\n", "@contextmanager\n", "def _download_from_s3_if_needed(maybe_s3_filepath: str):\n", " \"\"\"If path is S3, download to a temp file and yield local path.\"\"\"\n", " temp_dir = None\n", " tmp_filepath = maybe_s3_filepath\n", " if maybe_s3_filepath.startswith(\"s3://\"):\n", " temp_dir = tempfile.TemporaryDirectory()\n", " filename = get_filename(maybe_s3_filepath, keep_ext=True)\n", " tmp_filepath = os.path.join(temp_dir.name, filename)\n", " download_s3_file(maybe_s3_filepath, tmp_filepath)\n", " try:\n", " yield tmp_filepath\n", " finally:\n", " if temp_dir is not None:\n", " temp_dir.cleanup()\n", "\n", "def load_tokenizer(\n", " tokenizer_filepath: str = \"s3://suno-data/georg/models/tokenizers/tokenizer_60k.json\",\n", " pad_token: str = \"[PAD]\",\n", " extra_specials: Optional[List[str]] = None,\n", ") -> Tokenizer:\n", " \"\"\"Load and minimally configure a Tokenizer.\"\"\"\n", " extra_specials = extra_specials or [\"\\n\"]\n", " with _download_from_s3_if_needed(tokenizer_filepath) as tmp_fp:\n", " tok = Tokenizer.from_file(tmp_fp)\n", " tok.add_special_tokens(extra_specials)\n", "\n", " # Ensure PAD exists; fall back to adding if missing\n", " pad_id = tok.token_to_id(pad_token)\n", " if pad_id is None:\n", " tok.add_special_tokens([pad_token])\n", " pad_id = tok.token_to_id(pad_token)\n", "\n", " # Attach convenient attributes\n", " tok.pad_token = pad_token\n", " tok.pad_idx = pad_id\n", " return tok\n", "\n", "# -----------------------------\n", "# Audio slicing + meta handling\n", "# -----------------------------\n", "SAMPLE_RATE = 48_000\n", "\n", "def _choose_window(\n", " duration_s: float,\n", " chunk_s: float,\n", " ctx_s: float,\n", " use_ctx: bool,\n", ") -> Tuple[float, float, float, float]:\n", " \"\"\"\n", " Returns:\n", " full_start, full_end, target_start, target_end\n", " \"\"\"\n", " full_len = chunk_s + (ctx_s if use_ctx else 0.0)\n", " buf_size_samples = int(SAMPLE_RATE * full_len)\n", " max_start = max(0.0, duration_s - buf_size_samples / SAMPLE_RATE)\n", " full_start = random.uniform(0.0, max_start)\n", " full_end = full_start + buf_size_samples / SAMPLE_RATE\n", "\n", " target_start = full_start\n", " target_end = target_start + chunk_s\n", " return full_start, full_end, target_start, target_end\n", "\n", "def _extract_lyrics_for_window(\n", " text_aligned: List[Tuple[float, float, str]],\n", " window_start: float,\n", " window_end: float,\n", ") -> str:\n", " \"\"\"\n", " text_aligned format: list of (start_time, end_time, text)\n", " We include items fully contained within [window_start, window_end].\n", " \"\"\"\n", " if not text_aligned:\n", " return \"\"\n", " parts = [\n", " text\n", " for (s, e, text) in text_aligned\n", " if (s >= window_start) and (e <= window_end)\n", " ]\n", " return \"\".join(parts).strip()\n", " \n", "def load_meta(\n", " meta: Dict[str, Any],\n", " tokenizer: Tokenizer,\n", " cond_text_len: int,\n", " audio_chunk_s: float,\n", " audio_ctx_s: float,\n", " is_training: bool,\n", " text_drop_prob: float = 0.1,\n", "):\n", " \"\"\" Assumes meta['local_filepath'] is an OPUS file.\n", " Guarantees:\n", " - audio_target is exactly `audio_chunk_s` seconds (pads with silence if needed)\n", " - If the file is shorter than `audio_chunk_s`, audio_ctx is None\n", " \"\"\"\n", " duration_s = float(meta[\"duration_s\"])\n", " path = meta[\"local_filepath\"]\n", "\n", " # Helper: read an opus window as float samples, then wrap as Audio\n", " def read_opus_window(start_s: float, total_len_s: float) -> \"Audio\":\n", " buf_size = int(SAMPLE_RATE * total_len_s)\n", " float_arr = OpusFile(path=path).read(\n", " buf_size=buf_size,\n", " float_samples=True,\n", " from_position=int(SAMPLE_RATE * max(0.0, start_s)),\n", " )\n", " # Wrap as Audio; from_array_float handles float32 -> int16 safely\n", " return Audio.from_array_float(float_arr, sample_rate=SAMPLE_RATE, max_allowed_val=12)\n", "\n", " # --- Short-file path: no context; pad target out to exact length.\n", " if duration_s < audio_chunk_s:\n", " tgt_start, tgt_end = 0.0, min(audio_chunk_s, duration_s)\n", " window = read_opus_window(0.0, tgt_end) # read what's available\n", " audio_target = window.pad_to_length(audio_chunk_s)\n", " audio_ctx = None\n", "\n", " else:\n", " use_ctx = (random.random() < 0.5)\n", " full_start, full_end, tgt_start, tgt_end = _choose_window(\n", " duration_s=duration_s,\n", " chunk_s=audio_chunk_s,\n", " ctx_s=audio_ctx_s,\n", " use_ctx=use_ctx,\n", " )\n", "\n", " full_len_s = audio_chunk_s + (audio_ctx_s if use_ctx else 0.0)\n", " window = read_opus_window(full_start, full_len_s)\n", "\n", " # First part → target, then (optionally) ctx. Pad target defensively to exact length.\n", " audio_target = window.get_segment(from_s=0.0, to_s=audio_chunk_s).pad_to_length(audio_chunk_s)\n", " audio_ctx = window.get_segment(from_s=audio_chunk_s, to_s=audio_chunk_s + audio_ctx_s) if use_ctx else None\n", "\n", " # ----- Text/tags (unchanged) -----\n", " tags = meta.get(\"tags\") or []\n", " if \"text_aligned\" in meta and meta[\"text_aligned\"]:\n", " text = _extract_lyrics_for_window(meta[\"text_aligned\"], tgt_start, tgt_end)\n", " else:\n", " text = \"\"\n", "\n", " if is_training:\n", " full_text = \"\" if random.random() <= text_drop_prob else augment_text_training(tags, text)\n", " else:\n", " full_text = prepare_text_inference(tags, text)\n", "\n", " # Tokenize & pad/truncate\n", " ids = tokenizer.encode(full_text).ids[:cond_text_len]\n", " if tokenizer.pad_idx is None:\n", " raise ValueError(\"Tokenizer must define a pad_idx.\")\n", " if len(ids) < cond_text_len:\n", " ids += [tokenizer.pad_idx] * (cond_text_len - len(ids))\n", " text_codes = torch.tensor(ids, dtype=torch.long)\n", "\n", " return audio_target, audio_ctx, text_codes\n", "\n", "\n", "# -----------------------------\n", "# Dataset & collate\n", "# -----------------------------\n", "class DummyDataset(IterableDataset):\n", " def __init__(\n", " self,\n", " metas: Iterable[Dict[str, Any]],\n", " *,\n", " audio_chunk_s: float = 30.0,\n", " audio_ctx_s: float = 30.0,\n", " cond_text_len: int = 1536,\n", " text_drop_prob: float = 0.1,\n", " is_training: bool = False,\n", " tokenizer: Optional[Tokenizer] = None,\n", " tokenizer_path: str = \"s3://suno-data/georg/models/tokenizers/tokenizer_60k.json\",\n", " ):\n", " self.metas = list(metas)\n", " self.audio_chunk_s = float(audio_chunk_s)\n", " self.audio_ctx_s = float(audio_ctx_s)\n", " self.cond_text_len = int(cond_text_len)\n", " self.is_training = bool(is_training)\n", " self.text_drop_prob = float(text_drop_prob)\n", " self.tokenizer = tokenizer or load_tokenizer(tokenizer_path)\n", "\n", " def __iter__(self):\n", " # IterableDataset must not rely on shuffle=True in DataLoader.\n", " for meta in self.metas:\n", " try:\n", " audio_target, audio_ctx, text_codes = load_meta(\n", " meta=meta,\n", " tokenizer=self.tokenizer,\n", " cond_text_len=self.cond_text_len,\n", " audio_chunk_s=self.audio_chunk_s,\n", " audio_ctx_s=self.audio_ctx_s,\n", " is_training=self.is_training,\n", " text_drop_prob=self.text_drop_prob,\n", " )\n", " yield audio_target, audio_ctx, text_codes\n", " except Exception as e:\n", " print(f\"[WARN] Error loading meta {meta.get('id', '')}: {e}\")\n", " continue\n", "\n", "def collate_fn(batch: List[Tuple[Any, Optional[Any], torch.Tensor]]):\n", " audio_target_list = [item[0] for item in batch]\n", " audio_ctx_list = [item[1] for item in batch]\n", " text_codes_list = [item[2] for item in batch]\n", " # If your encoder expects tensors, you can add padding/stacking here instead.\n", " return audio_target_list, audio_ctx_list, text_codes_list\n", "\n", "def _as_tensor_batch(x):\n", " \"\"\"Convert codec output to a torch tensor (batch-first).\"\"\"\n", " if isinstance(x, torch.Tensor):\n", " return x.detach().cpu()\n", " try:\n", " import numpy as np\n", " if isinstance(x, np.ndarray):\n", " return torch.from_numpy(x)\n", " except Exception:\n", " pass\n", " if isinstance(x, (list, tuple)):\n", " elems = [_as_tensor_batch(e) for e in x]\n", " return torch.stack(elems, dim=0)\n", " return torch.as_tensor(x)\n", "\n", "\n", "def benchmark(\n", " metas: Iterable[Dict[str, Any]],\n", " *,\n", " batch_size: int = 2,\n", " num_workers: int = 1,\n", " num_eval_batches: int = 100,\n", " codec_encode_fn=None,\n", " audio_chunk_s: float = 30.0,\n", " audio_ctx_s: float = 30.0,\n", " cond_text_len: int = 1536,\n", " is_training: bool = False,\n", " text_drop_prob: float = 0.1,\n", " n_vae_tokens: int = 750,\n", " vae_dim: int = 128,\n", "):\n", " \"\"\"\n", " codec_encode_fn: function(audio_list, normalize_volume=False) -> (B, T, D)\n", " Expected output shape: (B, n_vae_tokens, vae_dim)\n", " \n", " Builds:\n", " - vae_target: (B, n_vae_tokens, vae_dim)\n", " - audio_ctx_vae: (B, n_vae_tokens, vae_dim) (zeros where ctx=None)\n", " - audio_ctx_mask: (B, 1, 1) (1.0 where ctx exists, 0.0 otherwise)\n", " \"\"\"\n", " assert codec_encode_fn is not None, \"Please pass codec_encode_fn=your_encoder\"\n", "\n", " dataset = DummyDataset(\n", " metas=metas,\n", " audio_chunk_s=audio_chunk_s,\n", " audio_ctx_s=audio_ctx_s,\n", " cond_text_len=cond_text_len,\n", " is_training=is_training,\n", " text_drop_prob=text_drop_prob,\n", " )\n", " loader = DataLoader(\n", " dataset,\n", " batch_size=batch_size,\n", " shuffle=False, # IterableDataset must not be shuffled here\n", " collate_fn=collate_fn,\n", " num_workers=num_workers,\n", " pin_memory=False,\n", " )\n", "\n", " seen = 0\n", " t0 = time.time()\n", "\n", " for i, (audio_target_list, audio_ctx_list, text_codes_list) in enumerate(loader):\n", " # ----------------------------\n", " # Encode audio targets (always present)\n", " # ----------------------------\n", " vae_target = _as_tensor_batch(codec_encode_fn(audio_target_list, normalize_volume=False))\n", "\n", " if vae_target.ndim != 3:\n", " raise ValueError(f\"codec_encode(target) must return (B, T, D); got {vae_target.shape}\")\n", "\n", " B, T, D = vae_target.shape\n", " if T != n_vae_tokens or D != vae_dim:\n", " raise ValueError(\n", " f\"codec_encode(target) shape mismatch: expected (B, {n_vae_tokens}, {vae_dim}), \"\n", " f\"got (B, {T}, {D})\"\n", " )\n", "\n", " device, dtype = vae_target.device, vae_target.dtype\n", "\n", " # ----------------------------\n", " # Encode audio context (may be None)\n", " # ----------------------------\n", " audio_ctx_vae = torch.zeros((B, n_vae_tokens, vae_dim), dtype=dtype, device=device)\n", " audio_ctx_mask = torch.zeros((B, 1, 1), dtype=dtype, device=device)\n", "\n", " ctx_indices = [idx for idx, a in enumerate(audio_ctx_list) if a is not None]\n", " if ctx_indices:\n", " ctx_batch = [audio_ctx_list[idx] for idx in ctx_indices]\n", " encoded_ctx = _as_tensor_batch(codec_encode_fn(ctx_batch, normalize_volume=False))\n", "\n", " if encoded_ctx.ndim != 3:\n", " raise ValueError(f\"codec_encode(ctx) must return (b, T, D); got {encoded_ctx.shape}\")\n", "\n", " b_ctx, T_ctx, D_ctx = encoded_ctx.shape\n", " if T_ctx != n_vae_tokens or D_ctx != vae_dim:\n", " raise ValueError(\n", " f\"codec_encode(ctx) shape mismatch: expected (b, {n_vae_tokens}, {vae_dim}), \"\n", " f\"got (b, {T_ctx}, {D_ctx})\"\n", " )\n", "\n", " # Scatter ctx results into full batch tensors\n", " for k, global_idx in enumerate(ctx_indices):\n", " audio_ctx_vae[global_idx] = encoded_ctx[k]\n", " audio_ctx_mask[ctx_indices] = 1.0\n", "\n", " # ----------------------------\n", " # Here you can run downstream model / profiling / etc.\n", " # ----------------------------\n", " # Example: dummy use so these aren't optimized out\n", " _ = (vae_target, audio_ctx_vae, audio_ctx_mask, text_codes_list)\n", "\n", " seen += 1\n", " if seen >= num_eval_batches:\n", " break\n", "\n", " # ----------------------------\n", " # Timing summary\n", " # ----------------------------\n", " t1 = time.time()\n", " elapsed = t1 - t0\n", " time_per_batch = elapsed / max(1, seen)\n", " chunks_per_batch = batch_size\n", " chunks_per_second = chunks_per_batch / time_per_batch if time_per_batch > 0 else float(\"inf\")\n", "\n", " print(f\"Time taken: {elapsed:.3f}s for {seen} batches\")\n", " print(f\"Time per batch: {time_per_batch:.4f}s\")\n", " print(f\"~Chunks/sec (30s targets): {chunks_per_second:.2f}\")\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "a3c76953", "metadata": {}, "outputs": [], "source": [ "benchmark(metas, codec_encode_fn=codec_encode)" ] }, { "cell_type": "code", "execution_count": null, "id": "292f7017", "metadata": {}, "outputs": [], "source": [ "unique_keys = set()\n", "for meta in metas:\n", " for key in meta.keys():\n", " unique_keys.add(key)\n", "\n", "print(unique_keys)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "51f38434", "metadata": {}, "outputs": [], "source": [ "for meta in metas:\n", " if \"text_aligned\" in meta:\n", " text_aligned = meta[\"text_aligned\"]\n", " # get all the text for time from 0 to 30 seconds\n", " # the format is a list, in each element, (start_time, end_time, text)\n", " text_for_time = [text for start_time, end_time, text in text_aligned if start_time >= 0.0 and end_time >= 30.0]\n", " if len(text_for_time) > 0:\n", " text_str = \"\".join([text for text in text_for_time])\n", " print(text_str)\n", " print()\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "4482d04a", "metadata": {}, "outputs": [], "source": [ "# we will need to load the raw pytorch models for semantic and codec i believe" ] }, { "cell_type": "code", "execution_count": null, "id": "4e9851d4", "metadata": {}, "outputs": [], "source": [ "filepath = \"/app2/suno/data/dpo/diff2_v2_d5_metas/metas_tr_t14.jsonl\"\n", "\n", "metas = read_jsonl(filepath)\n", "print(len(metas))\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "9434de9a", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "from tqdm import tqdm\n", "for meta in tqdm(metas):\n", " # there are two ways that semantic codes are stored\n", " semantic_data = np.load(meta[\"semantic_codes_filepath\"])\n", "\n", " if \"semantic_codes\" in semantic_data.keys():\n", " semantic_codes = semantic_data[\"semantic_codes\"]\n", " else:\n", " # this may be prod data, so check for the old keys\n", " if \"v2.0_raw\" in semantic_data:\n", " semantic_codes = semantic_data[\"v2.0_raw\"]\n", " elif \"v3.0_raw\" in semantic_data:\n", " semantic_codes = semantic_data[\"v3.0_raw\"]\n", " elif \"v3.5_raw\" in semantic_data:\n", " semantic_codes = semantic_data[\"v3.5_raw\"]\n", " elif \"v4.0_raw\" in semantic_data:\n", " semantic_codes = semantic_data[\"v4.0_raw\"]\n", " elif \"v5.0_raw\" in semantic_data:\n", " semantic_codes = semantic_data[\"v5.0_raw\"]\n", " else:\n", " raise ValueError(\"No codes found\")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "20bf4600", "metadata": {}, "outputs": [], "source": [ "semantic_data\n" ] }, { "cell_type": "code", "execution_count": null, "id": "ad5c9eff", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "suno_env_fa2", "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": 5 }