{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "from suno_utils.utils.s3 import read_from_s3\n", "from suno_utils.utils.text import read_jsonl" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "semantic_n_tokens = 750\n", "pad_token = 4000\n", "import random\n", "\n", "semantic_codes = torch.ones((semantic_n_tokens,), dtype=torch.long)\n", "print(semantic_codes)\n", "\n", "p = 0.1 #random.random()\n", "print(p)\n", "# select randomize token indicies to corrupt\n", "mask = torch.bernoulli(torch.full((semantic_n_tokens,), p))\n", "# select random token to replace with\n", "random_tokens = torch.randint(0, pad_token, (semantic_n_tokens,))\n", "semantic_codes[mask == 1] = random_tokens[mask == 1]\n", "\n", "print(semantic_codes)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "covers_metas = \"/home/christian/code/christian/metadata/v4/covers_metas.jsonl\"\n", "covers_metas = read_jsonl(covers_metas)\n", "print(len(covers_metas))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "covers_metas[0]" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# let's save out a list of covers with is_reliable = True\n", "reliable_covers = [cover for cover in covers_metas if cover[\"is_reliable\"]]\n", "print(len(reliable_covers))\n", "\n", "# get the ids\n", "reliable_cover_ids = [cover[\"id\"] for cover in reliable_covers]\n", "\n", "# save out the ids\n", "with open(\"/home/christian/code/christian/metadata/v4/reliable_cover_ids.txt\", \"w\") as f:\n", " for cover_id in reliable_cover_ids:\n", " f.write(f\"{cover_id}\\n\")\n", "\n", "# save out the ids\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import polars as pl\n", "\n", "df = pl.read_ndjson(\"/home/christian/code/christian/metadata/v4/covers_metas.jsonl\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "for item in df.iter_rows(named=True):\n", " print(item)\n", " break\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# call generate_chunk directly to use infill model for normal upsample\n", "\n", "main_window_size = 750\n", "infill_window_size = 250\n", "output_window_size = main_window_size - infill_window_size\n", "\n", "n_chunks = int(np.ceil(codes.shape[0] / output_window_size))\n", "print(f\"n_chunks: {n_chunks}\")\n", "\n", "# this is global\n", "generation_config = diffusion_gen.DiffusionGenerationConfig(\n", " audio=audio,\n", " steps=32,\n", " lyrics=lyrics,\n", " tags=\"metal\",\n", " text_cfg_coef=2.0,\n", " codec_scale_factor=0.4,\n", " scale_ctx_vector=True,\n", " noise_ctx_level=0.5,\n", " semantic_skip_factor=1,\n", ")\n", "\n", "latents = torch.zeros(codes.shape[0], 128)\n", "for chunk_idx in range(n_chunks):\n", " if chunk_idx == 0:\n", " end_idx = main_window_size\n", " start_idx = 0\n", " else:\n", " start_idx = (chunk_idx * output_window_size) - infill_window_size\n", " end_idx = start_idx + main_window_size\n", " semantic_chunk = torch.from_numpy(codes[start_idx:end_idx, 0]).long().unsqueeze(0)\n", "\n", " # if semantic_chunk is less than main_window_size, repeat pad\n", " pad_len = 0\n", " if semantic_chunk.shape[1] < main_window_size:\n", " pad_len = main_window_size - semantic_chunk.shape[1]\n", " pad_chunk = semantic_chunk[:, :pad_len]\n", " semantic_chunk = torch.cat([semantic_chunk, pad_chunk], dim=1)\n", "\n", " print(f\"semantic chunk {chunk_idx} from {start_idx} to {end_idx} with shape {semantic_chunk.shape}\")\n", "\n", " if chunk_idx > 0:\n", " ctx_start_idx = start_idx\n", " ctx_end_idx = start_idx + infill_window_size\n", " print(f\"ctx start idx {ctx_start_idx} to {ctx_end_idx}\")\n", " infill_prefix_latents = latents[ctx_start_idx:ctx_end_idx, :]\n", " else:\n", " infill_prefix_latents = None\n", "\n", " vae_latents = diffusion_gen.generate_chunk(semantic_chunk, generation_config, infill_prefix_latents=infill_prefix_latents)\n", " print(vae_latents.shape)\n", " \n", " if pad_len > 0:\n", " print(f\"cropping {vae_latents.shape[1]} to {vae_latents.shape[1] - pad_len}\")\n", " vae_latents = vae_latents[:, :-pad_len]\n", "\n", " # insert into latents\n", " if chunk_idx > 0:\n", " vae_chunk = vae_latents[:, infill_window_size:, :]\n", " write_start_idx = chunk_idx * output_window_size\n", " if vae_chunk.shape[1] > output_window_size:\n", " write_end_idx = write_start_idx + output_window_size\n", " else:\n", " write_end_idx = write_start_idx + vae_chunk.shape[1]\n", " print(f\"writing {vae_chunk.shape[0]} latents to {write_start_idx} to {write_end_idx}\")\n", " latents[write_start_idx:write_end_idx, :] = vae_chunk\n", " else:\n", " write_start_idx = 0\n", " write_end_idx = output_window_size\n", " print(f\"writing {vae_latents.shape[0]} latents to {write_start_idx} to {write_end_idx}\")\n", " \n", " latents[write_start_idx:write_end_idx, :] = vae_latents[:, :output_window_size, :]\n", "\n", " print()\n", "\n", "decode_stream_to_full_audio(latents).play()\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ctx_mask = torch.tensor([[1, 1, 1, 1], [0, 0, 0, 0], [0, 0, 0, 0]])\n", "infill_ctx_mask = torch.tensor([[1, 1, 1, 1], [1, 1, 1, 1], [0, 0, 0, 0]])\n", "# bs, seq_len\n", "print(ctx_mask.shape)\n", "print(infill_ctx_mask.shape)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "default_embed_mask = ((ctx_mask == 0) & (infill_ctx_mask == 0)).float()\n", "print(default_embed_mask)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from suno_utils.utils.s3 import read_from_s3\n", "from suno_utils.utils.text import read_jsonl\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "work_items = read_from_s3(\n", " \"s3://suno-data/christian/sft/pos_interesting_clips_up_u_1_20241201_full.jsonl\",\n", " read_f=read_jsonl,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "for work_item in work_items:\n", " tags = work_item[\"tags\"]\n", " if tags is None:\n", " print(work_item)\n", " break\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "import json\n", "import random\n", "from suno_utils.utils.text import read_jsonl\n", "\n", "\n", "def _get_rand_int(min_val, max_val):\n", " if max_val <= min_val:\n", " return min_val\n", " return random.randint(min_val, max_val)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "base_dir = \"/app/suno/data/chirp_v5_ft/v2/metas_val.jsonl\"\n", "metas = read_jsonl(base_dir)\n", "\n", "semantic_rate_hz = 25\n", "\n", "for meta in metas:\n", " if \"text_lines\" in meta:\n", " start_text_lines = meta[\"text_lines\"]\n", " break" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "t = int(round(start_text_lines[-1][\"end_s\"] * semantic_rate_hz))\n", "print(t)\n", "\n", "\n", "left_idx = random.randint(0, len(start_text_lines) - 1)\n", "right_idx = random.randint(left_idx + 1, len(start_text_lines))\n", "text_lines = start_text_lines[left_idx:right_idx]\n", "sample_text = \"\\n\".join([m[\"text\"] for m in text_lines]) # more tex\n", "print(\"text_lines\", left_idx, right_idx)\n", "print(\"sample_text\", sample_text)\n", "print()\n", "a_right_idx = int(round(text_lines[0][\"start_s\"] * semantic_rate_hz))\n", "b_right_idx = int(round(text_lines[-1][\"end_s\"] * semantic_rate_hz))\n", "if random.random() <= 0.9: # half the time, we add extra lines \n", " # Calculate max values based on remaining lines on either side\n", " # Use all available lines instead of limiting to 10\n", " max_left_shift = left_idx # Maximum we can shift left is up to the beginning\n", " max_right_shift = len(start_text_lines) - right_idx # Maximum we can shift right is up to the end\n", " \n", " left_idx = left_idx - _get_rand_int(0, max_left_shift)\n", " right_idx = right_idx + _get_rand_int(0, max_right_shift)\n", " print(\"left_idx\", left_idx)\n", " print(\"right_idx\", right_idx)\n", " text_lines = start_text_lines[left_idx:right_idx]\n", " sample_text = \"\\n\".join([m[\"text\"] for m in text_lines])\n", " print(\"extra_text\", sample_text)\n", " print()\n", " #a_right_idx = int(round(text_lines[0][\"start_s\"] * semantic_rate_hz))\n", " #b_right_idx = int(round(text_lines[-1][\"end_s\"] * semantic_rate_hz))\n", "print(\"a_right_idx\", a_right_idx)\n", "print(\"b_right_idx\", b_right_idx)\n", "if random.random() <= 0.1:\n", " a_left_idx = a_right_idx\n", "else:\n", " a_left_idx = _get_rand_int(0, a_right_idx)\n", "if random.random() <= 0.1:\n", " c_right_idx = b_right_idx\n", "else:\n", " c_right_idx = _get_rand_int(b_right_idx, t)\n", "if (\n", " 0 <= a_right_idx\n", " and a_left_idx <= a_right_idx\n", " and a_right_idx < b_right_idx\n", " and b_right_idx <= c_right_idx\n", " and c_right_idx <= t\n", "):\n", " delta_infill = True\n", "\n", "print(\"a_arr\", a_left_idx, a_right_idx)\n", "print(\"c_arr\", a_right_idx, b_right_idx)\n", "print(\"b_arr\", b_right_idx, c_right_idx)\n", "print(\"delta_infill\", delta_infill)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# lets load some random vae data and try to normalize it\n", "import numpy as np\n", "from suno_utils.utils.s3 import read_from_s3\n", "from suno_utils.utils.text import read_jsonl\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/diffusion_mix/dac_vae_tuned_25hz\"\n", "#base_dir = \"/app/suno/data/diffusion_v5/v0\"\n", "#base_dir = \"/mnt/localdisk/cjs_shards/\"\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))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "mean_ = vae_data[3].mean()\n", "std_ = vae_data[4].astype(np.float32).std()\n", "print(mean_, std_)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "x = torch.randn(750, 128) * 2.0 * 0.4\n", "print(x.mean(), x.std())\n", "xn = torch.randn(750, 128) * 0.7\n", "y = xn + x\n", "print(y.mean(), y.std())\n", "\n", "\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# i have vae data of shape (examples, seq_len, feature_dim)\n", "# i want to normalize each feature_dim with l2 norm \n", "\n", "# get the l2 norm of each feature_dim using torch\n", "import torch\n", "\n", "# Convert numpy arrays to PyTorch tensors\n", "vae_data_tensor = torch.from_numpy(vae_data) * 0.4 # (apply training scale factor)\n", "\n", "scale = torch.linalg.norm(vae_data_tensor, dim=1, keepdim=True)\n", "scale = scale.mean(dim=0)\n", "\n", "scale = scale + 1e-6\n", "\n", "# take std over feature dim\n", "scale2 = torch.std(vae_data_tensor, dim=1).std(dim=0)\n", "print(scale2.shape)\n", "\n", "# Apply normalization using the computed factors\n", "vae_data_tensor_norm = vae_data_tensor / scale\n", "\n", "# Now we can use feature_norms to normalize new data\n", "# Example of how to normalize new data:\n", "# new_data_normalized = new_data_tensor / feature_norms\n", "print(vae_data_tensor_norm[100])" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print(scale.shape)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "for scal_val in scale.numpy().flatten():\n", " print(f\"{scal_val:0.2f}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "plt.bar(np.arange(VAE_DIM), scale2.numpy().flatten())\n", "print(min(scale.numpy().flatten()), max(scale.numpy().flatten()))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "new_scale = []\n", "for scale_val in scale2.numpy().flatten():\n", " print(f\"{(1/scale_val)*2:0.2f}\")\n", " new_scale.append((1/scale_val)*2)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# divide by new scale\n", "vae_data_tensor_norm2 = vae_data_tensor / torch.tensor(new_scale)\n", "\n", "print(vae_data_tensor_norm2[100])\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "\n", "# make a histogram of the values before and after normalization\n", "fig, axs = plt.subplots(3, 1, figsize=(10, 15))\n", "axs[0].hist(vae_data_tensor.numpy().flatten()[:100000], bins=1000, label='Before normalization', alpha=0.5)\n", "axs[1].hist(vae_data_tensor_norm.numpy().flatten()[:100000], bins=1000, label='After normalization', alpha=0.5)\n", "axs[2].hist(vae_data_tensor_norm2.numpy().flatten()[:100000], bins=1000, label='After normalization', alpha=0.5)\n", "axs[0].legend()\n", "axs[1].legend()\n", "axs[2].legend()\n", "plt.show()\n", "\n", "#print(min(vae_data_tensor_norm2.numpy().flatten()), max(vae_data_tensor_norm.numpy().flatten()))\n", "#print(min(vae_data_tensor.numpy().flatten()), max(vae_data_tensor.numpy().flatten()))\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "norm_layer = torch.nn.LayerNorm(128)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ctx_vector = torch.from_numpy(vae_data[0:10,...])\n", "print(ctx_vector.mean(), ctx_vector.std())\n", "ctx_vector_norm = norm_layer(ctx_vector)\n", "print(ctx_vector_norm.mean(), ctx_vector_norm.std())\n", "\n", "\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from suno_utils.utils.text import read_jsonl\n", "\n", "work_items = read_jsonl(\n", " \"/home/christian/code/christian/metadata/discogs_subset_sampled_metas.jsonl\"\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "work_items[0]" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# lets set the seed based on the id of the work item\n", "# item_id is a string\n", "# We need to handle non-numeric characters in the ID\n", "item_id = work_items[0][\"id\"]\n", "# Convert the string to a numeric value by using the ord() of each character\n", "diffusion_seed = 0\n", "for char in item_id:\n", " diffusion_seed = (diffusion_seed * 31 + ord(char)) % 1000000\n", "print(f\"Generated seed from '{item_id}': {diffusion_seed}\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# copy the audio from s3 to local for the first 128 work items\n", "import torchaudio\n", "import os\n", "from tqdm import tqdm\n", "from suno_utils.utils.s3 import read_from_s3\n", "from joblib import Parallel, delayed\n", "\n", "output_dir = \"/home/christian/audio/real-audio-test\"\n", "os.makedirs(output_dir, exist_ok=True)\n", "\n", "def process_item(item):\n", " try:\n", " s3_filepath = item[\"s3_filepath\"]\n", " item_output_dir = f\"{output_dir}/{item['id']}\"\n", " os.makedirs(item_output_dir, exist_ok=True)\n", " output_filepath = f\"{item_output_dir}/{item['id']}_reference.mp3\"\n", " \n", " audio, sr = read_from_s3(s3_filepath, read_f=torchaudio.load)\n", " torchaudio.save(output_filepath, audio, sr)\n", " return True\n", " except Exception as e:\n", " print(f\"Error processing {item['id']}: {e}\")\n", " return False\n", "\n", "items_to_process = work_items[:128]\n", "results = Parallel(n_jobs=-1)(delayed(process_item)(item) for item in tqdm(items_to_process))\n", "print(f\"Successfully processed {sum(results)} out of {len(items_to_process)} items\")\n" ] }, { "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 }