{ "cells": [ { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import boto3\n", "import IPython\n", "import numpy as np\n", "\n", "from suno_utils.utils.s3 import read_from_s3\n", "from suno_utils.utils.text import read_jsonl" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "\n", "\n", "def get_s3_files(bucket_name, prefix, max_keys: int = 100000):\n", " all_files = []\n", " continuation_token = None\n", "\n", " while True:\n", " # Prepare the arguments for the request\n", " list_kwargs = {\n", " \"Bucket\": bucket_name,\n", " \"Prefix\": prefix, # List objects under this prefix, or leave blank for all objects\n", " }\n", "\n", " if continuation_token:\n", " list_kwargs[\"ContinuationToken\"] = continuation_token\n", "\n", " # Make the request to list objects\n", " response = s3.list_objects_v2(**list_kwargs)\n", "\n", " # Collect the file keys\n", " all_files += [obj[\"Key\"] for obj in response.get(\"Contents\", [])]\n", "\n", " # Check if more results are available\n", " if response.get(\"IsTruncated\"): # True if there are more results to fetch\n", " continuation_token = response[\"NextContinuationToken\"]\n", " else:\n", " break # No more results to fetch\n", "\n", " return all_files\n", "\n", "\n", "def list_s3_directories(bucket_name: str, prefix: str = \"\"):\n", " \"\"\"\n", " List all directories (prefixes) in an S3 bucket.\n", "\n", " Args:\n", " bucket_name (str): Name of the S3 bucket\n", " prefix (str): Optional prefix to filter results (like a directory path)\n", "\n", " Returns:\n", " List[str]: List of directory paths (prefixes)\n", " \"\"\"\n", " s3_client = boto3.client(\"s3\")\n", " directories = set()\n", "\n", " # Use paginator to handle buckets with many objects\n", " paginator = s3_client.get_paginator(\"list_objects_v2\")\n", " page_iterator = paginator.paginate(Bucket=bucket_name, Prefix=prefix, Delimiter=\"/\")\n", "\n", " # Collect all prefixes (directories)\n", " for page in page_iterator:\n", " # Get common prefixes (directories)\n", " if \"CommonPrefixes\" in page:\n", " for prefix_obj in page[\"CommonPrefixes\"]:\n", " directories.add(prefix_obj[\"Prefix\"])\n", "\n", " # Also check Contents for any directory-like objects\n", " if \"Contents\" in page:\n", " for obj in page[\"Contents\"]:\n", " key = obj[\"Key\"]\n", " # If the key contains a slash, add the directory part\n", " if \"/\" in key:\n", " directory = key.rsplit(\"/\", 1)[0] + \"/\"\n", " directories.add(directory)\n", "\n", " return sorted(list(directories))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "prefix = \"suno-data\"\n", "test_dir = \"christian/data/genius_hq_filtered_20k/25hz_20241115_v1/\"\n", "\n", "# get files in dir\n", "s3 = boto3.client(\"s3\")\n", "directories = list_s3_directories(prefix, test_dir)\n", "print(len(directories))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "directory_idx = 7\n", "files = get_s3_files(prefix, directories[directory_idx])\n", "print(len(files))\n", "mp3_filepath = [f for f in files if f.endswith(\".mp3\")][0]\n", "npz_filepath = mp3_filepath.replace(\".mp3\", \".npz\")\n", "full_npz_filepath = f\"s3://{prefix}/{npz_filepath}\"\n", "\n", "# read the npz file\n", "npz_file = read_from_s3(full_npz_filepath, read_f=np.load)\n", "original_latents = npz_file[\"original_latents\"]\n", "upsampled_latents = npz_file[\"upsampled_latents\"]\n", "semantic_codes = npz_file[\"semantic_codes\"]\n", "print(original_latents.shape, upsampled_latents.shape, semantic_codes.shape)\n", "\n", "# decode the original latents\n", "original_latents = torch.from_numpy(original_latents[:750, :]).float().to(device)\n", "with torch.no_grad():\n", " decoded_latents = model_25hz.decode(original_latents.permute(1, 0).unsqueeze(0))\n", "print(decoded_latents.shape)\n", "\n", "# decode the upsampled latents\n", "upsampled_latents = torch.from_numpy(upsampled_latents[:750, :]).float().to(device)\n", "with torch.no_grad():\n", " decoded_upsampled_latents = model_25hz.decode(upsampled_latents.permute(0, 2, 1))\n", "print(decoded_upsampled_latents.shape)\n", "\n", "# play the decoded audio\n", "IPython.display.display(IPython.display.Audio(decoded_latents.squeeze().cpu().numpy(), rate=48000))\n", "IPython.display.display(IPython.display.Audio(decoded_upsampled_latents.squeeze().cpu().numpy(), rate=48000))\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import funcy\n", "import torch\n", "from dac.model.dac4 import DAC\n", "\n", "# load VAE\n", "device = \"cuda:0\"\n", "#device = \"cpu\"\n", "# checkpoint_filepath = \"/app/suno/christian/checkpoints/dac/100hz_vae_peaq_kl_0.005/best/dac/weights.pth\"\n", "#checkpoint_filepath = \"s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth\"\n", "checkpoint_filepath = \"s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth\"\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_25hz = DAC(**sd[\"metadata\"][\"kwargs\"])\n", "model_25hz.load_state_dict(sd[\"state_dict\"])\n", "model_25hz.eval()\n", "model_25hz.to(device)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load the memmaps and check the shapes\n", "import os\n", "VAE_DIM = 128\n", "VAE_RATE_HZ = 25\n", "VAE_MEMMAP_SIZE = 750\n", "SEMANTIC_VOCAB_SIZE = 4000\n", "SEMANTIC_MEMMAP_SIZE = 750\n", "\n", "out_dir = \"/app/suno/data/diffusion_ft/genius_hq_filtered_20k_25hz_20241115_v1\"\n", "metas_tr = read_jsonl(os.path.join(out_dir, f\"metas_val.jsonl\"))\n", "print(len(metas_tr)/ 2)\n", "mm_semantic_tr = np.memmap(os.path.join(out_dir, f\"data_semantic_val.bin\"), dtype=np.uint16, mode=\"r\")\n", "mm_vae_tr = np.memmap(os.path.join(out_dir, f\"data_vae_val.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": [ "rand_idx = 108\n", "print(metas_tr[rand_idx])\n", "vae_seq = mm_vae_tr[rand_idx]\n", "vae_seq = torch.from_numpy(vae_seq.copy()).unsqueeze(0).float().cuda()\n", "\n", "with torch.no_grad():\n", " audio = model_25hz.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))\n", "\n", "\n", "vae_seq = mm_vae_tr[rand_idx + 1]\n", "vae_seq = torch.from_numpy(vae_seq.copy()).unsqueeze(0).float().cuda()\n", "\n", "with torch.no_grad():\n", " audio = model_25hz.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": [] } ], "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 }