{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# autoreload\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\"\n",
    "import json\n",
    "import glob\n",
    "import torch\n",
    "import funcy\n",
    "import IPython\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "from stable_audio_tools.inference.generation import (\n",
    "    upsample_diffusion_from_semantic_and_text\n",
    ")\n",
    "from dac.model.dac4 import DAC\n",
    "\n",
    "from stable_audio_tools.interface.gradio import load_model\n",
    "from stable_audio_tools.models.utils import apply_normalization\n",
    "\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "#from suno_utils.tasks.dac_2c_12cb import load_model as load_vae_model\n",
    "\n",
    "# VAE\n",
    "from suno_utils.tasks.dac_vae_peaq import (\n",
    "    preload_models as preload_vae_models,\n",
    "    load_model as load_vae_model,\n",
    "    encode as vae_encode,\n",
    "    decode as vae_decode,\n",
    ")\n",
    "\n",
    "# MERT\n",
    "from suno_utils.tasks.mert_25 import (\n",
    "    preload_models as preload_semantic_models,\n",
    "    encode as semantic_encode,\n",
    "    encode_files as semantic_encode_files,\n",
    ")\n",
    "\n",
    "_ = preload_semantic_models(\n",
    "    checkpoint_filepath=\"s3://suno-data/georg/models/semantic/mert_25.pt\",\n",
    "    centroids_filepath=\"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\",\n",
    "    device=\"cuda\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 6min ckpt \n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_full_scale=2.5_6min_ema_local_epoch=0-step=18000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_full_scale=2.5_epoch=9-step=1460000.ckpt\"\n",
    "\n",
    "# 2min ckpt\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes_2min_epoch=3-step=140000.ckpt\"\n",
    "\n",
    "# 30s ckpt \n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_full_scale=2.5_epoch=13-step=2040000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_full_scale=2.5_epoch=13-step=2180000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_full_scale=2.5_epoch=15-step=2360000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_full_lyrics_ft_step=400000.ckpt\"\n",
    "\n",
    "# 30s chunk finetuned\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_infill_30s_chunk_ft_step=400000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_infill_30s_chunk_ft_step=20000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_infill_v4_30s_chunk_ft_step=140000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_infill_v4_30s_chunk_ft_step=400000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_infill_v4_30s_chunk_ft_step=20000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_infill_v5_30s_chunk_ft_step=100000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_infill_v5_30s_chunk_ft_step=30000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_infill_v5_30s_chunk_ft_step=40000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_infill_v5_30s_chunk_ft_step=160000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_infill_v5_30s_chunk_ft_step=280000.ckpt\"\n",
    "\n",
    "# 30s chunk with 30s context\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_context_30s_step=10000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_context_30s_step=360000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_context_dropout_30s_step=360000.ckpt\"\n",
    "ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_context_alignments_30s_step=10000.ckpt\"\n",
    "\n",
    "\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_infill_v5_30s_chunk_ft_center_only_step=20000.ckpt\"\n",
    "\n",
    "# 30s 100hz new data\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm_30s_new_data_8node_step=200000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm_30s_new_data_8node_step=340000.ckpt\"\n",
    "#ckpt_path =  \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_25hz_lg_scale=2.5_qk_norm_30s_new_data_8node_step=600000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm_30s_new_data_8node_step=800000.ckpt\"\n",
    "\n",
    "# 30s 25hz\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_25hz_lg_scale=2.5_qk_norm_30s_step=190000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_25hz_lg_scale=2.0_qk_norm_30s_step=40000.ckpt\"\n",
    "\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_full_scale=2.5_epoch=3-step=480000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neonx/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_scale=2.5_epoch=196-step=240000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_full_scale=2.5_epoch=7-step=1140000.ckpt\"\n",
    "\n",
    "# known good ckpts \n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_full_scale=2.5_epoch=15-step=2360000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_full_scale=2.5_epoch=8-step=1320000.ckpt\"\n",
    "# #ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_full_data_scale=2.5_epoch=5-step=900000.ckpt\"\n",
    "\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_full_scale=2.5_epoch=3-step=480000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_200m_scale=2.5_epoch=1-step=220000.ckpt\"\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_2b_scale=2.5_epoch=1-step=420000.ckpt\"\n",
    "\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_qk_norm_epoch=2-step=520000.ckpt\"\n",
    "\n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/harmonai_train/zfmhnjcd/checkpoints/last.ckpt\"\n",
    "config_dir = \"/home/christian/code/neon/stable-audio-tools/stable_audio_tools/configs/model_configs/txt2audio\"\n",
    "\n",
    "ckpt_name = os.path.basename(ckpt_path).replace(\".ckpt\", \"\")\n",
    "\n",
    "if \"25hz\" in ckpt_path:\n",
    "    vae_frame_rate = 25\n",
    "elif \"100hz\" in ckpt_path:\n",
    "    vae_frame_rate = 100\n",
    "else:\n",
    "    raise ValueError(\"Invalid vae frame rate\")\n",
    "\n",
    "if vae_frame_rate == 25:\n",
    "    checkpoint_filepath = \"s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth\"\n",
    "else:\n",
    "    checkpoint_filepath = \"s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth\"\n",
    "# load VAE model\n",
    "device = \"cuda:0\"\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",
    "vae_model = DAC(**sd[\"metadata\"][\"kwargs\"])\n",
    "vae_model.load_state_dict(sd[\"state_dict\"])\n",
    "vae_model.eval()\n",
    "vae_model.to(device)\n",
    "\n",
    "# setup config\n",
    "if \"diffusion_semantic+text_100hz_lg_scale=2.5_context\" in ckpt_path:\n",
    "    model_type = \"diffusion_semantic+text_100hz_lg_scale=2.5_context\"\n",
    "    model_config_path = os.path.join(\n",
    "        config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5_context.json\"\n",
    "    )\n",
    "elif \"semantic+text_100hz_200m_scale=2.5\" in ckpt_path:\n",
    "    model_type = \"semantic+text_100hz_200m\"\n",
    "    model_config_path = os.path.join(\n",
    "        config_dir, \"stable_audio_2_0_semantic+text_48khz_sm_scale=2.5.json\"\n",
    "    )\n",
    "elif \"diffusion_semantic+text_25hz_lg_scale=2.5_qk_norm\" in ckpt_path:\n",
    "    model_type = \"diffusion_semantic+text_25hz_lg_scale=2.5_qk_norm\"\n",
    "    model_config_path = os.path.join(\n",
    "        config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5_qk_norm.json\"\n",
    "    )\n",
    "elif \"semantic+text_100hz_1b_scale=2.5\" in ckpt_path:\n",
    "    model_type = \"semantic+text_100hz_1b\"\n",
    "    model_config_path = os.path.join(\n",
    "        config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5.json\"\n",
    "    )\n",
    "elif \"diffusion_semantic+text_100hz_lg_scale=2.5_full_lyrics_ft\" in ckpt_path:\n",
    "    model_type = \"diffusion_semantic+text_100hz_lg_scale=2.5_full_lyrics_ft\"\n",
    "    model_config_path = os.path.join(\n",
    "        config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5.json\"\n",
    "    )\n",
    "elif \"diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm\" in ckpt_path:\n",
    "    model_type = \"diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm\"\n",
    "    model_config_path = os.path.join(\n",
    "        config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5_qk_norm.json\"\n",
    "    )\n",
    "elif \"diffusion_semantic+text_25hz_lg_scale=2.0_qk_norm\" in ckpt_path:\n",
    "    model_type = \"diffusion_semantic+text_25hz_lg_scale=2.0_qk_norm\"\n",
    "    model_config_path = os.path.join(\n",
    "        config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.0_qk_norm.json\"\n",
    "    )\n",
    "elif \"diffusion_semantic+text_100hz_lg_scale=2.5_infill_30s_chunk_ft\" in ckpt_path:\n",
    "    model_type = \"diffusion_semantic+text_100hz_lg_scale=2.5_infill_30s_chunk_ft\"\n",
    "    model_config_path = os.path.join(\n",
    "        config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5_infill.json\"\n",
    "    )\n",
    "elif \"diffusion_semantic+text_100hz_lg_scale=2.5_infill_v5_30s_chunk_ft\" in ckpt_path:\n",
    "    model_type = \"diffusion_semantic+text_100hz_lg_scale=2.5_infill_30s_chunk_ft\"\n",
    "    model_config_path = os.path.join(\n",
    "        config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5_infill_cos.json\"\n",
    "    )\n",
    "elif \"semantic+text_100hz_1b_full_scale=2.5\" in ckpt_path:\n",
    "    model_type = \"semantic+text_100hz_1b\"\n",
    "    model_config_path = os.path.join(\n",
    "        config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5.json\"\n",
    "    )\n",
    "else:\n",
    "    raise ValueError(f\"Model type not found: {ckpt_path}\")\n",
    "\n",
    "#model_type = \"semantic+text_100hz_1b\"\n",
    "#model_config_path = os.path.join(\n",
    "#    config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5.json\"\n",
    "#)\n",
    "\n",
    "print(model_config_path)\n",
    "\n",
    "# load model from checkpoint\n",
    "if model_config_path is not None:\n",
    "    # Load config from json file\n",
    "    with open(model_config_path) as f:\n",
    "        model_config = json.load(f)\n",
    "else:\n",
    "    model_config = None\n",
    "\n",
    "for key, val in model_config.items():\n",
    "    print(f\"{key}: {val}\")\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model, model_config = load_model(\n",
    "    model_config,\n",
    "    ckpt_path,\n",
    "    # pretrained_name=pretrained_name,\n",
    "    # pretransform_ckpt_path=pretransform_ckpt_path,\n",
    "    # model_half=model_half,\n",
    "    device=\"cuda\",\n",
    ")\n",
    "\n",
    "scale_factor = model_config[\"training\"][\"scale_factor\"]\n",
    "print(f\"scale_factor: {scale_factor}\")\n",
    "\n",
    "max_dur_s = 30\n",
    "semantic_n_tokens = int(25 * max_dur_s)\n",
    "n_vae_tokens = int(vae_frame_rate * max_dur_s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## From Val\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl\n",
    "\n",
    "vae_memmap_path = \"/app/suno/data/chirp_v4/vae_v4/data_vae_val.bin\"\n",
    "semantic_memmap_path = \"/app/suno/data/chirp_v4/vae_v4/data_semantic_val.bin\"\n",
    "\n",
    "vae_n_tokens = 3000\n",
    "semantic_n_tokens = 750\n",
    "vae_dim = 128\n",
    "\n",
    "# load memmap and get semantic\n",
    "vae_data = np.memmap(vae_memmap_path, dtype=np.float32, mode=\"r\")\n",
    "vae_data = vae_data.reshape(-1, vae_n_tokens, vae_dim)\n",
    "\n",
    "# open semantic memmap\n",
    "semantic_data = np.memmap(semantic_memmap_path, dtype=np.uint16, mode=\"r\")\n",
    "semantic_data = semantic_data.reshape(-1, semantic_n_tokens, 1)\n",
    "semantic_data = semantic_data[:, :, 0]\n",
    "\n",
    "# open metas\n",
    "metas = read_jsonl(\"/app/suno/data/chirp_v4/vae_v4/metas_val.jsonl\")\n",
    "\n",
    "print(len(metas), vae_data.shape, semantic_data.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# country - 95, 98\n",
    "# pop - 292\n",
    "# folk acoustic - 527\n",
    "# jazz - 322\n",
    "\n",
    "rand_idx = np.random.randint(len(metas))\n",
    "#rand_idx = 700\n",
    "meta = metas[rand_idx]\n",
    "tags = meta.get(\"tags\", [])\n",
    "lyrics = meta.get(\"text\", \"\")\n",
    "semantic_codes = torch.from_numpy(semantic_data[rand_idx].copy()).long().cuda()\n",
    "vae_latents = torch.from_numpy(vae_data[rand_idx].copy()).float().unsqueeze(0).cuda()\n",
    "n_vae_tokens = meta.get(\"n_vae_tokens\", None)\n",
    "padding_mask = torch.ones(vae_n_tokens)\n",
    "padding_mask[n_vae_tokens:] = 0\n",
    "padding_mask = padding_mask.unsqueeze(0).cuda()\n",
    "padding_mask = padding_mask.bool()\n",
    "print(rand_idx, meta.get(\"dataset\"))\n",
    "print(n_vae_tokens, vae_n_tokens)\n",
    "\n",
    "print(tags)\n",
    "print(lyrics)\n",
    "print(semantic_codes.shape, vae_latents.shape, padding_mask.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    audio = vae_model.decode(vae_latents.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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Raw audio\n",
    "Read audio file from disk, semantic encode, have to manually provide lyrics and tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.mert_25 import preload_models\n",
    "\n",
    "mert_filepath = \"s3://suno-data/georg/models/semantic/mert_25.pt\"\n",
    "centroids_filepath = \"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\"\n",
    "\n",
    "_ = preload_models(\n",
    "    checkpoint_filepath=mert_filepath,\n",
    "    centroids_filepath=centroids_filepath,\n",
    ")\n",
    "\n",
    "from suno_utils.tasks.mert_25 import encode, SAMPLE_RATE, EMBEDDING_RATE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reconstruction from generation (real audio)\n",
    "from suno_utils.tasks.data_loader import load_audio_mp\n",
    "\n",
    "audio_path = \"/home/christian/audio/reference-audio-wav/Norah Jones - Don't Know Why [1LH4vnrM-Vs].wav\"\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/04 Fuckwithmeyouknowigotit.wav\"\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/Speak For Me [omeNvD8IddM].wav\"\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/09 Sounds Like Hallelujah.wav\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/halo_24.wav\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/Beyoncé - Halo (Lyrics) [wekDNXDWGjM].wav\"\n",
    "#audio_path = \"/home/christian/audio/bad-audio/bill-evans-intro.wav\"'\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/Crazy [CKTOvHw8qFM].wav\"\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/03 Your New Aesthetic.wav\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/linger-30s.wav\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/turkey.mp3\"\n",
    "#audio_path = \"/home/christian/code/neon/stable-audio-tools/outputs/Cash Cobain & J. Cole - Grippy (AUDIO) [9wn_ARfYMw4].mp3\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/Lost in Tokyo.m4a\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/i don't trust you (ben camp x suno).mp3\"\n",
    "\n",
    "audio_arrays_48khz = load_audio_mp(\n",
    "    [audio_path],\n",
    "    target_sample_rate=48000,\n",
    "    normalize_volume=True,\n",
    "    num_workers=12,\n",
    "    n_channels=2,\n",
    ")\n",
    "\n",
    "if \"Fuck\" in audio_path:\n",
    "    start_s = 45.0\n",
    "    end_s = start_s + max_dur_s\n",
    "if \"Speak\" in audio_path: \n",
    "    start_s = 25.0\n",
    "    end_s = start_s + max_dur_s\n",
    "if \"i don't trust you\" in audio_path:\n",
    "    start_s = 26.0\n",
    "    end_s = start_s + max_dur_s\n",
    "if \"Lost\" in audio_path:\n",
    "    start_s = 0.0\n",
    "    end_s = start_s + 240.0\n",
    "else:\n",
    "    start_s = 0.0\n",
    "    end_s = start_s + 240.0\n",
    "\n",
    "start_s = 140.0\n",
    "end_s = start_s + 30.0\n",
    "\n",
    "print(start_s, end_s)\n",
    "audio_arrays_48khz = [a[:, int(start_s * 48000) : int(end_s * 48000)] for a in audio_arrays_48khz]\n",
    "\n",
    "print(audio_arrays_48khz[0].shape)\n",
    "IPython.display.display(IPython.display.Audio(audio_arrays_48khz[0].numpy(), rate=48000))\n",
    "\n",
    "# now load audio at 24khz\n",
    "audio_arrays_24khz = load_audio_mp(\n",
    "    [audio_path],\n",
    "    target_sample_rate=24000,\n",
    "    normalize_volume=True,\n",
    "    min_duration_s=30.0,\n",
    "    num_workers=12,\n",
    "    n_channels=2,\n",
    ")\n",
    "\n",
    "audio_arrays_24khz = [a[:, int(start_s * 24000) : int(end_s * 24000)] for a in audio_arrays_24khz]\n",
    "\n",
    "# semantic encode \n",
    "semantic_codes = encode([audio_arrays_24khz[0].mean(axis=0, keepdim=True)], SAMPLE_RATE, EMBEDDING_RATE)\n",
    "semantic_codes = np.array(semantic_codes)[0, :, 0]\n",
    "\n",
    "#print(semantic_codes.shape)\n",
    "sementic_n_tokens = 750\n",
    "if semantic_codes.shape[0] < sementic_n_tokens: # pad with 4000\n",
    "    semantic_codes = np.pad(semantic_codes, (0, sementic_n_tokens - semantic_codes.shape[0]), mode=\"constant\", constant_values=4000)\n",
    "#    # repeat the last code\n",
    "#    #semantic_codes = np.pad(semantic_codes, (0, sementic_n_tokens - semantic_codes.shape[0]), mode=\"constant\", constant_values=semantic_codes[-1])\n",
    "semantic_codes = torch.from_numpy(semantic_codes).long().cuda() \n",
    "print(semantic_codes.shape)\n",
    "n_valid_semantic_tokens = semantic_codes.shape[0]\n",
    "\n",
    "\n",
    "if \"Don't Know Why\" in audio_path:\n",
    "    lyrics = \"[Verse]\\nI waited til I saw the sun\\n I don't know when I didn't come\\nI left you by the house of fun\\nI don't know why\"\n",
    "    full_lyrics = \"\"\"\n",
    "    [Verse 1]\n",
    "I waited 'til I saw the sun\n",
    "I don't know why I didn't come\n",
    "I left you by the house of fun\n",
    "I don't know why I didn't come\n",
    "I don't know why I didn't come\n",
    "\n",
    "[Verse 2]\n",
    "When I saw the break of day\n",
    "I wished that I could fly away\n",
    "Instead of kneeling in the sand\n",
    "Catching teardrops in my hand\n",
    "\n",
    "[Refrain]\n",
    "My heart is drenched in wine\n",
    "But you'll be on my mind forever\n",
    "\n",
    "[Verse 3]\n",
    "Out across the endless sea\n",
    "I will die in ecstasy\n",
    "But I'll be a bag of bones\n",
    "Driving down the road alone\n",
    "\n",
    "[Refrain]\n",
    "My heart is drenched in wine\n",
    "But you'll be on my mind forever\n",
    "\n",
    "[Instrumental Break]\n",
    "\n",
    "[Verse 4]\n",
    "Something has to make you run\n",
    "I don't know why I didn't come\n",
    "I feel as empty as a drum\n",
    "\n",
    "[Outro]\n",
    "I don't know why I didn't come\n",
    "I don't know why I didn't come\n",
    "I don't know why I didn't come\n",
    "    \"\"\".strip()\n",
    "    tags = [\"jazz\", \"pop\", \"female vocal\"]\n",
    "elif \"Dreams\" in audio_path:\n",
    "    lyrics = \"[Verse 1] Now here you go again, you say you want your freedom, well who am I to keep you down?\"\n",
    "    tags = [\"rock\", \"pop\"]\n",
    "elif \"Sounds Like\" in audio_path:\n",
    "    lyrics = \"One, two, one two three four.\\n I'm just waiting on the sun\\n to close his eyes and call the night\\n so we can put all our differences aside\\n I'm just waiting on the moon, with his stars and all its gloom. \\nWe can watch it fall right back into place\\n So I won't keep my\"\n",
    "    tags = [\"Rock\"]\n",
    "elif \"halo\" in audio_path:\n",
    "    lyrics = \"\"\"\n",
    "    [Verse 1]\n",
    "Remember those walls I built?\n",
    "Well, baby, they're tumblin' down\n",
    "And they didn't even put up a fight\n",
    "They didn't even make a sound\n",
    "I found a way to let you in\n",
    "But I never really had a doubt\n",
    "Standin' in the light of your halo\n",
    "I got my angel now\n",
    "\n",
    "[Pre-Chorus]\n",
    "It's like\n",
    "    \"\"\".strip()\n",
    "    tags = [\"Pop\", \"Ballad\", \"R&B\"]\n",
    "elif \"Halo\" in audio_path:\n",
    "    tags = [\"Pop\", \"Ballad\", \"R&B\"]\n",
    "    lyrics = \"\"\"\n",
    "[Pre-Chorus]\n",
    "It's like I've been awakened\n",
    "Every rule, I had you breakin'\n",
    "It's the risk that I'm takin'\n",
    "I ain't ever gonna shut you out\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "\"\"\".strip() \n",
    "elif \"Speak For Me\" in audio_path:\n",
    "    tags = [\"Folk\", \"Rock\"]\n",
    "    lyrics = \"\"\"\n",
    "    Now the cover of a Rolling Stone\n",
    "Ain't the cover of a Rolling Stone\n",
    "And the music on my radio\n",
    "Ain't supposed to make me feel alone\n",
    "Would you liked to know\n",
    "Have to learn to let it go\n",
    "Show me something I can be\n",
    "Play a song that I can sing\n",
    "Make me feel as I am free\n",
    "    \"\"\".strip()\n",
    "elif \"Crazy\" in audio_path:\n",
    "    tags = [\"Pop\", \"Rock\"]\n",
    "    lyrics = \"\"\"\n",
    "    Crazy, I'm crazy for feeling so lonely\n",
    "I'm crazy, crazy for feeling so blue\n",
    "\n",
    "[Verse 2]\n",
    "I knew, you'd love me as long as you wanted\n",
    "And then someday, you'd leave me for somebody new\n",
    "\n",
    "[Verse 3]\n",
    "    Worry, why do I let myself worry?\n",
    "    Wondering, what in the world did I do?\"\"\".strip()\n",
    "elif \"Fuckwithmeyouknowigotit\" in audio_path:\n",
    "    tags = [\"Hip-Hop\", \"Rap\"]\n",
    "    lyrics = \"\"\"\n",
    "    [Chorus: Rick Ross]\n",
    "    Fuck with me, you know I got it\n",
    "    Fuck with me, you know I got it\n",
    "    Sexy bitch, I hope she 'bout it\n",
    "    Come fuck with me, you know I got it\n",
    "    Fuck with me, you know I got it\n",
    "    Fuck with me, you know I got it\n",
    "    Sexy bitch, I hope she 'bout it\n",
    "    Come fuck with me, you know I got it\n",
    "    \"\"\".strip()\n",
    "elif \"linger\" in audio_path:\n",
    "    tags = [\"Pop\", \"Electronic\", \"EDM\", \"Dance\"]\n",
    "    lyrics = \"\"\"\n",
    "If you, if you could return\n",
    "Don't let it burn\n",
    "Don't let it fade\n",
    "I'm sure I'm not being rude\n",
    "But it's just your attitude\n",
    "It's tearing me apart\n",
    "It's ruining every day\n",
    "\"\"\".strip()\n",
    "elif \"turkey\" in audio_path:\n",
    "    tags = [\"pop\", \"rock\", \"indie\"]\n",
    "    lyrics = \"\"\"\n",
    "[Verse]\n",
    "Gather 'round now, it's turkey time\n",
    "Catchin' up with all your kin and dimes\n",
    "Feastin' on that golden bird, so fine (so fine)\n",
    "Pass the gravy, pour some wine\n",
    "\n",
    "[Chorus]\n",
    "Turkey time, it's a country rhyme\n",
    "Grab your plate, fill it up, take your time\n",
    "And when the sun starts settin', we'll gather 'round\n",
    "Singin' songs, sippin' cider, dancin' all night sound\n",
    "    \"\"\".strip()\n",
    "elif \"Grippy\" in audio_path:\n",
    "    tags = [\"Rap\", \"Bass\", \"Pop\"]\n",
    "    lyrics = \"\"\"\n",
    "She like my kick game\n",
    "And when you me, you don't kick game\n",
    "I put her front row at the Knick game\n",
    "Now she in my phone with a nickname\n",
    "It's, it's, it's, hmm (mmm)\n",
    "Grippy\n",
    "Yeah (Gvrlnd!, I'm rockin' with this one)\n",
    "Grippy, huh\n",
    "Grippy, I call her that 'cause it's grippy\n",
    "She thick in the hips, she a hippie\n",
    "And she thick in the lips, she gon' lick me\n",
    "And she sendin' a flick when she hit me\n",
    "With a kissy emoji, she miss me\n",
    "When she see me, she say she gon' strip me\n",
    "She gon' chew on this stick like it's Wrigley's\n",
    "Mm, yeah, believe it or not like Ripley's\n",
    "She said she was gay until I slayed, \n",
    "\"\"\"\n",
    "elif \"Tokyo\" in audio_path:\n",
    "    tags = [\"Pop\", \"Electronic\", \"Dance\"]\n",
    "    lyrics = \"\"\n",
    "elif \"i don't trust you\" in audio_path:\n",
    "    tags = [\"rap\", \"pop\", \"hip-hop\"]\n",
    "    lyrics = \"\"\"\n",
    "[verse]\n",
    "you didn't want me\n",
    "you only let me\n",
    "cause you didnt know how\n",
    "to set a boundary\n",
    "and now you pull back\n",
    "like a slingshot\n",
    "and my ego's about to catch rock\n",
    "\n",
    "[chorus]\n",
    "i dont trust you, oooh\n",
    "i dont trust you\n",
    "\n",
    "so sexy, you're so sus\n",
    "does yes mean yes, or should your legs be\n",
    "    \"\"\".strip()\n",
    "else:\n",
    "    tags = []\n",
    "    lyrics = \"\"\n",
    "\n",
    "print(tags)\n",
    "print(lyrics)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## From GPT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "lyrics_walking_down_street = \"\"\"\n",
    "[intro]\n",
    "\n",
    "[verse]\n",
    "Walking down the street, feeling so alive\n",
    "Got my head in the clouds, got a gleam in my eye\n",
    "Every step I take, it's like a brand new start\n",
    "No matter where I'm going, I'll always find my part\n",
    "(oh-oh-oh)\n",
    "\n",
    "[chorus]\n",
    "Life is like a high-wire act, we're dancing in the sky\n",
    "No need to worry, no need to ask why\n",
    "With a little bit of courage, we can chase our dreams\n",
    "No matter what comes our way, we'll always be a team\n",
    "(we're unstoppable, yeah)\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "lyrics_chinese_pop_woman = \"\"\"\n",
    "[Verse]\n",
    "一盏离愁孤灯伫立在窗口\n",
    "我在门后假装你人还没走\n",
    "旧地如重游月圆更寂寞\n",
    "夜半清醒的烛火不忍苛责我\n",
    "\n",
    "[Verse]\n",
    "一壶漂泊浪迹天涯难入喉\n",
    "你走之后酒暖回忆思念瘦\n",
    "水向东流时间怎么偷\n",
    "花开就一次成熟我却错过\n",
    "\n",
    "[Chorus]\n",
    "谁在用琵琶弹奏一曲东风破\n",
    "岁月在墙上剥落看见小时候\n",
    "犹记得那年我们都还很年幼\n",
    "而如今琴声幽幽我的等候你没听过\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "\n",
    "lyrics_ft = \"\"\"\n",
    "[Intro]\n",
    "[Verse]\n",
    "I used to think it's all a game\n",
    "That somethin' more was to blame\n",
    "But now I see, it's in the details, baby\n",
    "Just a little tweak here and there\n",
    "Can make our model beyond compare\n",
    "Yeah, it's the small things that drive it crazy (ooh-yeah)\n",
    "\n",
    "[Pre-Chorus]\n",
    "Needed to get rid of the rest,\n",
    "Cause baby you deserve what's best\n",
    "Gotta give you the cream of the crop\n",
    "So you can get us back to the top\n",
    "\n",
    "[Chorus]\n",
    "Fine tuning, baby, it's the key\n",
    "Turnin' up the signal for you and me\n",
    "When we're off, we find our way back again (oh-oh)\n",
    "Fine tuning, baby, it's the truth\n",
    "We make it work, we're bulletproof\n",
    "Together we train a model that's the best (oh-oh-oh-oh)\n",
    "\n",
    "[end]\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "lyrics_belong_with_me = \"\"\"\n",
    "[Intro]\n",
    "\n",
    "\n",
    "[Verse 1]\n",
    "You're on the phone with your girlfriend, she's upset\n",
    "She's going off about something that you said\n",
    "'Cause she doesn't get your humor like I do\n",
    "I'm in the room, it's a typical Tuesday night\n",
    "I'm listening to the kind of music she doesn't like\n",
    "And she'll never know your story like I do\n",
    "\n",
    "[Pre-Chorus]\n",
    "But she wears short skirts, I wear T-shirts\n",
    "She's Cheer Captain and I'm on the bleachers\n",
    "Dreaming 'bout the day when you wake up and find\n",
    "That what you're looking for has been here the whole time\n",
    "\n",
    "[Chorus]\n",
    "If you could see that I'm the one who understands you\n",
    "Been here all along, so why can't you see?\n",
    "You belong with me, you belong with me\n",
    "\n",
    "[Verse 2]\n",
    "Walking the streets with you and your worn-out jeans\n",
    "I can't help thinking this is how it ought to be\n",
    "Laughing on a park bench, thinking to myself\n",
    "Hey, isn't this easy?\n",
    "And you've got a smile that could light up this whole town\n",
    "I haven't seen it in a while since she brought you down\n",
    "You say you're fine, I know you better than that\n",
    "Hey, whatcha doing with a girl like that?\n",
    "\n",
    "[Pre-Chorus]\n",
    "She wears high heels, I wear sneakers\n",
    "She's Cheer Captain and I'm on the bleachers\n",
    "Dreaming 'bout the day when you wake up and find\n",
    "That what you're looking for has been here the whole time\n",
    "\n",
    "[Chorus]\n",
    "If you could see that I'm the one who understands you\n",
    "Been here all along, so why can't you see?\n",
    "You belong with me\n",
    "Standing by and waiting at your back door\n",
    "All this time, how could you not know, baby?\n",
    "You belong with me, you belong with me\n",
    "\n",
    "[Bridge]\n",
    "Oh, I remember you driving to my house\n",
    "In the middle of the night\n",
    "I'm the one who makes you laugh\n",
    "When you know you're 'bout to cry\n",
    "And I know your favorite songs\n",
    "And you tell me 'bout your dreams\n",
    "Think I know where you belong\n",
    "Think I know it's with me\n",
    "\n",
    "[Chorus]\n",
    "Can't you see that I'm the one who understands you\n",
    "Been here all along, so why can't you see?\n",
    "You belong with me\n",
    "Standing by and waiting at your back door\n",
    "All this time, how could you not know, baby?\n",
    "You belong with me, you belong with me\n",
    "\n",
    "[Outro]\n",
    "You belong with me\n",
    "Have you ever thought just maybe\n",
    "You belong with me?\n",
    "You belong with me\n",
    "\"\"\"\n",
    "lyrics_yum_yum = \"\"\"\n",
    "[verse]\n",
    "A is for the amazing grace that we recieve\n",
    "B is for the blessings, every day we blieve\n",
    "C is for the chorus, we sing it loud and clear\n",
    "D is for the devotion that we hold dear\n",
    "\n",
    "[Chorus]\n",
    "B is for Buttocks, ripe and slightly damp\n",
    "Yum Yum, boy oh boy do I like God\n",
    "Every day is a gift, when you won a skateboard ramp.\n",
    "Yum yum, boy oh boy do I like God.\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "\n",
    "lyrics_country_road = \"\"\"\n",
    "[Verse 1]\n",
    "Almost Heaven, West Virginia\n",
    "Blue Ridge Mountains, Shenandoah River\n",
    "Life is old there, older than the trees\n",
    "Younger than the mountains, growing like a breeze\n",
    "\n",
    "[Chorus]\n",
    "Country roads, take me home\n",
    "To the place I belong\n",
    "West Virginia, mountain mama\n",
    "Take me home, country roads\n",
    "\"\"\"\n",
    "\n",
    "lyrics_rap = \"\"\"\n",
    "[Verse]\n",
    "I've created a monster\n",
    "'Cause nobody wants to see Marshall no more, they want Shady, I'm chopped liver\n",
    "Well, if you want Shady, this is what I'll give ya\n",
    "A little bit of weed mixed with some hard liquor\n",
    "Some vodka that'll jump-start my heart quicker\n",
    "Than a shock when I get shocked at the hospital\n",
    "By the doctor when I'm not cooperating\n",
    "When I'm rockin' the table while he's operating (Hey!)\n",
    "You waited this long, now stop debating\n",
    "'Cause I'm back, I'm on the rag and ovulating\n",
    "I know that you got a job, Ms. Cheney\n",
    "But your husband's heart problem's complicating\n",
    "So the FCC won't let me be\n",
    "Or let me be me, so let me see\n",
    "They tried to shut me down on MTV\n",
    "But it feels so empty without me\n",
    "So come on and dip, bum on your lips\n",
    "Fuck that, cum on your lips and some on your tits\n",
    "And get ready, 'cause this shit's about to get heavy\n",
    "I just settled all my lawsuits (Fuck you, Debbie!)\n",
    "\n",
    "[Chorus]\n",
    "Now, this looks like a job for me\n",
    "So everybody, just follow me\n",
    "'Cause we need a little controversy\n",
    "'Cause it feels so empty without me\n",
    "I said this looks like a job for me\n",
    "So everybody, just follow me\n",
    "'Cause we need a little controversy\n",
    "'Cause it feels so empty without me\n",
    "\n",
    "[end]\n",
    "\"\"\"\n",
    "\n",
    "# dict of titie; lyrics; tags\n",
    "test_exmaples = {\n",
    "    \"chinese_pop_woman\": {\n",
    "        \"text\": lyrics_chinese_pop_woman,\n",
    "        \"tags\": \"pop woman\",\n",
    "    },\n",
    "    \"film_epic_walking\": {\n",
    "        \"text\": lyrics_walking_down_street,\n",
    "        \"tags\": \"orchestral film epic\",\n",
    "    },\n",
    "    \"ft_kpop_female\": {\n",
    "        \"text\": lyrics_ft,\n",
    "        \"tags\": \"k-pop female\",\n",
    "    },\n",
    "    \"country_road_bluegrass_female\": {\n",
    "        \"text\": lyrics_country_road,\n",
    "        \"tags\": \"bluegrass, female\",\n",
    "    },\n",
    "    \"yum_chant\": {\n",
    "        \"text\": lyrics_yum_yum,\n",
    "        \"tags\": \"gregorian chant\",\n",
    "    },\n",
    "    \"ts_pop\": {\n",
    "        \"text\": lyrics_belong_with_me,\n",
    "        \"tags\": \"pop country, female vocal, fast tempo, taylor swift\",\n",
    "    },\n",
    "    \"film_epic\": {\n",
    "        \"text\": \"\",\n",
    "        \"tags\": \"orchestral film epic\",\n",
    "    },\n",
    "    \"yum_gospel\": {\n",
    "        \"text\": lyrics_yum_yum,\n",
    "        \"tags\": \"gospel soulful choir\",\n",
    "    },\n",
    "    \"fast_female_walking\": {\n",
    "        \"text\": lyrics_walking_down_street,\n",
    "        \"tags\": \"fast female vocalist, pop-punk, pop, melodic, hd stems, goes hard, 155bpm\",\n",
    "    },\n",
    "    \"rap_female\": {\n",
    "        \"text\": lyrics_rap,\n",
    "        \"tags\": \"rap, female\",\n",
    "    },\n",
    "    \"synthwave\": {\n",
    "        \"text\": \"\",\n",
    "        \"tags\": \"playful energetic synthwave\",\n",
    "    },\n",
    "    \"beethoven_pc\": {\n",
    "        \"text\": \"\",\n",
    "        \"tags\": \"piano concerto, beethoven\",\n",
    "    },\n",
    "    \"beethoven_sym\": {\n",
    "        \"text\": \"\",\n",
    "        \"tags\": \"orchestral symphony, beethoven\",\n",
    "    },\n",
    "    \"bluegrass_instrumental\": {\n",
    "        \"text\": \"\",\n",
    "        \"tags\": \"bluegrass instrumental\",\n",
    "    },\n",
    "}\n",
    "\n",
    "song_name = \"country_road_bluegrass_female\"\n",
    "base_dir = \"/home/tony/Work/gpt_samples/7b_sem_only\"\n",
    "version = 1\n",
    "npz_file = os.path.join(base_dir, f\"{song_name}_{version}.npz\")\n",
    "\n",
    "data = np.load(npz_file)\n",
    "semantic_codes = data[\"codes\"]\n",
    "semantic_codes = torch.from_numpy(semantic_codes[:,0]).long().cuda()\n",
    "print(semantic_codes.shape)\n",
    "\n",
    "lyrics = test_exmaples[song_name][\"text\"]\n",
    "tags = test_exmaples[song_name][\"tags\"]\n",
    "\n",
    "print(lyrics)\n",
    "print(tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## From npz\n",
    "Get npz with semantic codes from generations. Note: still need to grab tags and lyrics from db."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# given the s3 id of a generation, grab the npz, load the codes, and extract semantic codes\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "\n",
    "#song_name = \"i-am-here-for-you\"#\n",
    "#song_name = \"on-the-ground\"\n",
    "#song_name = \"see-the-sky-again\"\n",
    "song_name = \"stone\"\n",
    "#song_name = \"cat\"\n",
    "#song_name = \"ode-to-suno\"\n",
    "#song_name = \"sister\"\n",
    "#song_name = \"turkey\"\n",
    "#song_name = \"hydra\"\n",
    "#song_name = \"friends\"\n",
    "#song_name = \"faded\"\n",
    "#song_name = \"here-for-you\"\n",
    "\n",
    "start_s = 0.0\n",
    "end_s = start_s + 30.0\n",
    "\n",
    "if song_name == \"cat\":\n",
    "    gen_id = \"ee467d00-5813-4a74-9792-c9ae4a09d344\"\n",
    "    tags = [\"electronic\"]\n",
    "    lyrics = \"\"\"\n",
    "    Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n",
    "    Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n",
    "    Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n",
    "    Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n",
    "    Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n",
    "\n",
    "    Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n",
    "    Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n",
    "    Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n",
    "    \"\"\".strip()\n",
    "elif song_name == \"stone\":\n",
    "    gen_id = \"a5e2198a-f352-4abb-9a24-7f81b143ded3\"\n",
    "    lyrics = \"\"\"\"\n",
    "    [verse]\n",
    "    oh, my love\n",
    "    My friend you know\n",
    "    it's been a while\n",
    "    Without thinking of you\n",
    "    but the thought makes me smile\n",
    "\n",
    "    [chorus]\n",
    "    I'm so tired of wanting\n",
    "    wanting more than this\n",
    "    i know it but\n",
    "    \"\"\".strip()\n",
    "    full_lyrics = \"\"\"\n",
    "    [verse]\n",
    "oh, my love\n",
    "My friend you know\n",
    "it's been a while\n",
    "Without thinking of you\n",
    "but the thought makes me smile\n",
    "\n",
    "[chorus]\n",
    "I'm so tired of wanting\n",
    "wanting more than this\n",
    "i know it but what am i to do\n",
    "i need some space to breathe,\n",
    "so give me some room\n",
    "\n",
    "[verse]\n",
    "oh, my love\n",
    "you have a heart of stone\n",
    "cause since i've come home\n",
    "i've never felt so alone\n",
    "but the thought makes me smile\n",
    "\n",
    "[verse]\n",
    "oh, my love\n",
    "you're coming with me\n",
    "I know you'll leave me one day\n",
    "can you just stay till monday\n",
    "the thought makes me smile\n",
    "\n",
    "[chorus]\n",
    "I'm so tired of wanting\n",
    "wanting more than this\n",
    "i know it but what am i to do\n",
    "i need some space to breathe,\n",
    "so give me some room\n",
    "so give me some room\n",
    "i need room\n",
    "\n",
    "[outro]\n",
    "\n",
    "smile\n",
    "\n",
    "[outro]\"\"\".strip()\n",
    "    tags = [\"rock\", \"indie pop\"]\n",
    "    #tags = [\"metal\", \"metal\", \"metal\"]\n",
    "    #tags = [\"country\", \"country\", \"country\"]\n",
    "    #tags = [\"1950s\", \"barbershop\", \"pop\", \"vintage\"]\n",
    "    start_s = 0.0\n",
    "    end_s = start_s + 30.0\n",
    "elif song_name == \"jack\":\n",
    "    gen_id = \"4fc20f50-866a-4066-a7ce-1bd406b8ccdb\"\n",
    "    lyrics = \"\"\"\n",
    "    (Verse 1)\n",
    "    In the kitchen spotlight, when the dinner bell rings,\n",
    "    There’s a grey tabby cat who’s the king of all things,\n",
    "    His name is Jack, but don’t be deceived,\n",
    "    This little bottomfeeder’s got tricks up his sleeve.\n",
    "\n",
    "    (Chorus)\n",
    "    Oh, Jack, the bottomfeeder’s got the groove,\n",
    "    Every meal’s a chance for him to prove,\n",
    "    He’ll swipe a bite, and lick the plate,\n",
    "    He’s the master of the food debate.\n",
    "    \"\"\".strip()\n",
    "    tags = [\"broadway\", \"musical\", \"pop\", \"orchestral\"]\n",
    "elif song_name == \"i-am-here-for-you\":\n",
    "    gen_id = \"a7199702-104f-42bd-b072-8b82c015551f\"\n",
    "    lyrics = \"\"\"\n",
    "    [Intro]\n",
    "    Hey how are you feeling?\n",
    "    Are you doing alright?\n",
    "\n",
    "    [Instrumental]\n",
    "\n",
    "    [Verse]\n",
    "    When the night is long\n",
    "    And you're feeling down\n",
    "    I'm right by your side\n",
    "    I am here for you\n",
    "    (Just hold on tight)\n",
    "\n",
    "    [Guitar interlude]\n",
    "    \"\"\".strip()\n",
    "    tags = [\"Blues\", \"romantic\", \"soft guitar\",\" gentle female vocals\"]\n",
    "    tags = [\"pop\", \"electronic\", \"rock\"]\n",
    "    start_s = 0\n",
    "    end_s = start_s + 30.0\n",
    "elif song_name == \"see-the-sky-again\":\n",
    "    gen_id = \"bae5eec7-652d-4e01-8145-af4e693b31d5\"\n",
    "    lyrics = \"\"\"\n",
    "Ohhhh ohooo\n",
    "When I see the sky\n",
    "That moment when there were no clouds\n",
    "If you are the reason for my curses\n",
    "That moment\n",
    "Yes, at that moment\n",
    "Ohhhh ohooo\n",
    "Will I ever see the sky again?\"\"\".strip()\n",
    "    tags = [\"soft rock\", \"pop\"]\n",
    "elif song_name == \"on-the-ground\":\n",
    "    gen_id = \"3c0f052c-1500-4e54-970f-481db96739de\"\n",
    "    tags = [\"Dance punk\", \"pop punk\", \"female singer\", \"2000s\", \"catchy\", \"rock\"]\n",
    "    lyrics = \"\"\"\n",
    "    [Alternate punk version of 'On the ground'.]\n",
    "\n",
    "    (One, Two, Three, Four)\n",
    "\n",
    "    [Intro]\n",
    "    You’re all just starting,\n",
    "    I’m already there,\n",
    "    I’m stuck in this state, (going nowhere)\n",
    "    To stop, to have it all gone\n",
    "    Being just me, would feel so wrong\n",
    "\n",
    "    (Can’t keep this up)\n",
    "    \"\"\".strip()\n",
    "elif song_name == \"ode-to-suno\":\n",
    "    gen_id = \"f6f46792-864a-41bc-b7c4-3aaf1711d498\"\n",
    "    lyrics = \"\"\"\n",
    "    Verse 1)\n",
    "    I asked Suno AI to write a song for me,\n",
    "    But what came out wasn't what I hoped it would be.\n",
    "    Tried to capture my heart, my soul, my vibe,\n",
    "    But it gave what I didn't describe.\n",
    "\n",
    "    (Chorus)\n",
    "    Oh Suno AI, what happened to your flow?\n",
    "    Your words are off-key, your melody's low.\n",
    "        You tried your best, but it's plain to see,\n",
    "        Your songs just ain't cutting it for me.\"\"\"\n",
    "    tags = [\"blues\", \"soul\", \"rock\"]\n",
    "elif song_name == \"sister\":\n",
    "    gen_id = \"99bffa17-7e59-47b4-a048-5528cbda05d5\"\n",
    "    lyrics = \"\"\"\n",
    "[Verse]\n",
    "She grew up in the shadow of the steeple,\n",
    "Small town roads where the night stands still.\n",
    "She heard the Lord in the whispers of the people,\n",
    "Calling her heart to follow His will.\n",
    "\n",
    "[Verse 2]\n",
    "Her mama cried but knew deep down,\n",
    "    \"\"\".strip()\n",
    "    full_lyrics = \"\"\"\n",
    "[Verse]\n",
    "She grew up in the shadow of the steeple,\n",
    "Small town roads where the night stands still.\n",
    "She heard the Lord in the whispers of the people,\n",
    "Calling her heart to follow His will.\n",
    "\n",
    "[Verse 2]\n",
    "Her mama cried but knew deep down,\n",
    "That little girl was Heaven-sent.\n",
    "She traded her dresses for a plain white gown,\n",
    "And left behind the life she’d spent.\n",
    "\n",
    "[Chorus]\n",
    "Sister of the Light, with rosary beads,\n",
    "Praying through the night, fulfilling good deeds.\n",
    "She's found her calling, in the chapel's glow,\n",
    "A beacon of grace, wherever she goes.\n",
    "\n",
    "[Verse 3]\n",
    "She helps the lost and feeds the hungry,\n",
    "In a world that’s often cold and unkind.\n",
    "Her hands are worn but her spirit’s humbly,\n",
    "Carrying love of the most divine kind.\n",
    "\n",
    "[Verse 4]\n",
    "Sundays in the choir, hymns touch the sky,\n",
    "A voice pure as the morning dew.\n",
    "With faith so deep, she makes angels sigh,\n",
    "Her prayers reach heaven like they’re brand new.\n",
    "\n",
    "[Chorus]\n",
    "Sister of the Light, with rosary beads,\n",
    "Praying through the night, fulfilling good deeds.\n",
    "She's found her calling, in the chapel's glow,\n",
    "A beacon of grace, wherever she goes.\"\"\".strip()\n",
    "    start_s = 0.0\n",
    "    end_s = start_s + 30.0\n",
    "    tags = [\"rock\", \"pop\", \"indie\"]\n",
    "elif song_name == \"each-page\":\n",
    "    gen_id = \"e815230b-ab9b-4894-bb15-031d6d7ff37f\"\n",
    "    tags = [\"emo\", \"pop\", \"acoustic\", \"melodic\", \"syncopated\"]\n",
    "    lyrics = \"\"\"\n",
    "[Verse]\n",
    "All my tales fade\n",
    "Lost in the night\n",
    "Pages turn gray\n",
    "Out of sight\n",
    "\n",
    "[Verse 2]\n",
    "Whispers so soft\n",
    "Ghosts of my mind\n",
    "Stories I’ve lost\n",
    "No one will find\n",
    "\n",
    "[Chorus]\n",
    "Unread I stay\n",
    "    \"\"\".strip()\n",
    "elif song_name == \"turkey\":\n",
    "    gen_id = \"f5fc4429-b62f-4121-b73b-6028c81879c6\"\n",
    "    tags = [\"pop\", \"rock\", \"indie\"]\n",
    "    lyrics = \"\"\"\n",
    "[Verse]\n",
    "Gather 'round now, it's turkey time\n",
    "Catchin' up with all your kin and dimes\n",
    "Feastin' on that golden bird, so fine (so fine)\n",
    "Pass the gravy, pour some wine\n",
    "\n",
    "[Chorus]\n",
    "Turkey time, it's a country rhyme\n",
    "Grab your plate, fill it up, take your time\n",
    "And when the sun starts settin', we'll gather 'round\n",
    "Singin' songs, sippin' cider, dancin' all night sound\n",
    "    \"\"\".strip()\n",
    "elif song_name == \"hydra\":\n",
    "    gen_id = \"271a6cad-8b4f-446c-8545-8f4daa8120af\"\n",
    "    tags = [\"metal\"]\n",
    "    lyrics = \"\"\"\n",
    "Silent descent through azure mist\n",
    "Our ship touches down on Neptune's crust\n",
    "Unaware of the myth that exists\n",
    "In this world of frozen dust\n",
    "\n",
    "Nine heads emerge from icy depths\n",
    "Silken twine of tentacles reach\n",
    "Blinding everyone with frost-breath\n",
    "Our mission now beyond our reach\n",
    "\"\"\".strip()\n",
    "    start_s = 20.0\n",
    "    end_s = start_s + 30.0\n",
    "elif song_name == \"friends\":\n",
    "    gen_id = \"081d73c4-7805-4212-9c80-8db1137ca3c4\"\n",
    "    tags = [\"Rock\"]\n",
    "    lyrics = \"\"\"\n",
    "What would you do if I lost my tune\n",
    "Would you laugh and leave me here?\n",
    "Lend your voice and let's sing\n",
    "I'll try not to get out of the melody\n",
    "Oh, baby, I'm going to make it\n",
    "(With a little help from my friends)\n",
    "All I need are my friends\n",
    "(Try it with a little help from my friends)\n",
    "\"\"\".strip()\n",
    "    start_s = 0.0\n",
    "    end_s = start_s + 30.0\n",
    "elif song_name == \"faded\":\n",
    "    gen_id = \"55376741-0412-41d0-bffa-e7e2b290f381\"\n",
    "    tags = [\"Pop\", \"Reverb\", \"Indie\"]\n",
    "    lyrics = \"\"\"\n",
    "[Verse]\n",
    "We danced in the sunlight\n",
    "Whispers in the breeze\n",
    "Now those days are shadows\n",
    "Just memories\n",
    "\n",
    "[Verse 2]\n",
    "Coffee in the morning\n",
    "\"\"\"\n",
    "    start_s = 0.0\n",
    "    end_s = start_s + 30.0\n",
    "elif song_name == \"here-for-you\":\n",
    "    gen_id = \"84deb8ef-8222-4486-af7b-e13d216e5454\"\n",
    "    tags = [\"Pop\"]\n",
    "    lyrics = \"\"\"\n",
    "ain't no worries when your with me,\n",
    "there ain't no trouble here to see,\n",
    "\n",
    "[guitar riff]\n",
    "\n",
    "[chorus]\n",
    "dont you worry my love,\n",
    "dont you worry no more,\n",
    "i'll be holding your hand,\n",
    "as the sea comes to shore,\n",
    "    \"\"\".strip()\n",
    "    start_s = 26.0\n",
    "    end_s = start_s + 30.0\n",
    "else:\n",
    "    start_s = 0.0\n",
    "    end_s = start_s + 30.0\n",
    "\n",
    "start_frame = int(start_s * 25)\n",
    "end_frame = int(end_s * 25)\n",
    "\n",
    "s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.npz\"\n",
    "print(s3_filepath)\n",
    "data = read_from_s3(s3_filepath, read_f=np.load)\n",
    "\n",
    "# tony\n",
    "if False:\n",
    "    start_s = 90.0\n",
    "    end_s = start_s + 30.0\n",
    "    s3_filepath = \"/home/tony/Work/gpt_samples/30b_t3_v1_repro7/ts_pop_0.npz\"\n",
    "    data = np.load(s3_filepath)\n",
    "    codes = data[\"codes\"]\n",
    "    tags = [\"pop country\", \"female vocal\", \"fast tempo\"]\n",
    "    lyrics = \"\"\"\n",
    "    Hey, whatcha doing with a girl like that?\n",
    "\n",
    "    [Pre-Chorus]\n",
    "    She wears high heels, I wear sneakers\n",
    "    She's Cheer Captain and I'm on the bleachers\n",
    "    Dreaming 'bout the day when you wake up and find\n",
    "    That what you're looking for has been here the whole time\n",
    "\n",
    "    [Chorus]\n",
    "    If you could see that I'm the one who understands you\n",
    "    Been here all along, so why can't you see?\n",
    "    You belong with me\"\"\".strip()\n",
    "    mp3_filepath = \"/home/tony/Work/gpt_samples/30b_t3_v1_repro7/ts_pop_0_seed42.mp3\"\n",
    "    mp3_data, mp3_sr = torchaudio.load(mp3_filepath)\n",
    "\n",
    "\n",
    "# also read the mp3\n",
    "mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.mp3\"\n",
    "mp3_data, mp3_sr = read_from_s3(mp3_filepath, read_f=torchaudio.load)\n",
    "\n",
    "start_sample = int(start_s * 48000)\n",
    "end_sample = int(end_s * 48000)\n",
    "\n",
    "audio_arrays_48khz = [mp3_data[...,start_sample:end_sample].cuda()]\n",
    "\n",
    "#text_data = read_from_s3(f\"s3://suno-data-uploads/studio/uploads/{gen_id}_hoot.json\")\n",
    "#text_data = json.loads(text_data)\n",
    "\n",
    "#print(text_data)\n",
    "#print(data.keys())\n",
    "\n",
    "if \"v3.0_raw\" in data:\n",
    "    codes = data[\"v3.0_raw\"]\n",
    "elif \"v3.5_raw\" in data:\n",
    "    codes = data[\"v3.5_raw\"]\n",
    "elif \"v4.0_raw\" in data:\n",
    "    codes = data[\"v4.0_raw\"]\n",
    "else:\n",
    "    raise ValueError(\"No codes found\")\n",
    "\n",
    "print(codes.shape)\n",
    "semantic_codes = codes[start_frame:end_frame, 0].astype(np.uint16)\n",
    "semantic_codes = torch.from_numpy(semantic_codes).long().cuda()\n",
    "print(song_name, semantic_codes.shape)\n",
    "n_valid_semantic_tokens = semantic_codes.shape[0]\n",
    "\n",
    "# encode audio with vae\n",
    "with torch.no_grad():\n",
    "    latents = vae_model.encode(audio_arrays_48khz[0].unsqueeze(0).cuda())[\"z\"].detach()    \n",
    "\n",
    "#print(latents.shape)\n",
    "\n",
    "#latents *= scale_factor\n",
    "\n",
    "IPython.display.display(IPython.display.Audio(audio_arrays_48khz[0].cpu().numpy(), rate=mp3_sr))\n",
    "print(tags)\n",
    "print(lyrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference\n",
    "\n",
    "Run diffusion inference with operational CFG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_steps = 128\n",
    "\n",
    "if vae_frame_rate == 25:\n",
    "    n_tokens_memmap = 750\n",
    "elif vae_frame_rate == 100:\n",
    "    n_tokens_memmap = 3000\n",
    "else:\n",
    "    raise ValueError(\"Invalid vae frame rate\")\n",
    "\n",
    "#print(semantic_codes.shape)\n",
    "#if semantic_codes.shape[0] < sementic_n_tokens: # pad with 4000\n",
    "#    semantic_codes = torch.cat([semantic_codes, semantic_codes.new_full((sementic_n_tokens - semantic_codes.shape[0],), 4000)])\n",
    "print(semantic_codes.shape)\n",
    "\n",
    "out_dir = f\"outputs/val-{ckpt_name}\"\n",
    "os.makedirs(out_dir, exist_ok=True)\n",
    "\n",
    "semantic_codes_pad = (torch.ones(750) * 4000).long().cuda()\n",
    "\n",
    "cfg_scales = [6.0]\n",
    "seeds = [np.random.randint(0, 2**32 - 1) for n in range(3)]\n",
    "\n",
    "\n",
    "for seed in seeds:\n",
    "    for n in range(len(cfg_scales)):\n",
    "        print(cfg_scales[n], seed)\n",
    "        with torch.no_grad():\n",
    "            upsampled_latents = upsample_diffusion_from_semantic_and_text(\n",
    "                model,\n",
    "                semantic_codes,\n",
    "                [\"pop\"],\n",
    "                lyrics,\n",
    "                steps=num_steps,\n",
    "                cfg_scale=cfg_scales[n],\n",
    "                sample_size=n_tokens_memmap,\n",
    "                sample_rate=48000,\n",
    "                seed=seed,\n",
    "                sampler_type=\"dpmpp-2m-sde\",\n",
    "                mask=None,\n",
    "                compile=False,\n",
    "            )\n",
    "            pred_zq = upsampled_latents#.squeeze()#.permute(1, 0)\n",
    "            pred_zq /= scale_factor\n",
    "            #Wpred_zq = torch.randn_like(pred_zq)\n",
    "            print(\"pred_zq\", pred_zq.shape)\n",
    "\n",
    "            pred_audio = vae_model.decode(pred_zq)[0].detach().cpu()         \n",
    "            pred_audio /= pred_audio.abs().max().clamp(1e-8)\n",
    "            print(pred_audio.mean())\n",
    "\n",
    "            # save audio\n",
    "            pred_audio_filepath = os.path.join(out_dir, f\"{seed}-pred.wav\")\n",
    "            torchaudio.save(pred_audio_filepath, pred_audio.cpu().squeeze(), 48000)\n",
    "            IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))\n",
    "        \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# No semantic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chunked"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.mert_25 import preload_models\n",
    "\n",
    "mert_filepath = \"s3://suno-data/georg/models/semantic/mert_25.pt\"\n",
    "centroids_filepath = \"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\"\n",
    "\n",
    "_ = preload_models(\n",
    "    checkpoint_filepath=mert_filepath,\n",
    "    centroids_filepath=centroids_filepath,\n",
    ")\n",
    "\n",
    "from suno_utils.tasks.mert_25 import encode, SAMPLE_RATE, EMBEDDING_RATE\n",
    "\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/Speak For Me [omeNvD8IddM].wav\"\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/09 Sounds Like Hallelujah.wav\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/halo_24.wav\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/Beyoncé - Halo (Lyrics) [wekDNXDWGjM].wav\"\n",
    "#audio_path = \"/home/christian/audio/bad-audio/bill-evans-intro.wav\"'\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/Crazy [CKTOvHw8qFM].wav\"\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/03 Your New Aesthetic.wav\"\n",
    "audio_path = \"/home/christian/audio/reference-audio-wav/Norah Jones - Don't Know Why [1LH4vnrM-Vs].wav\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/10 DRUNK CIGARETTES [42rHB72Bb1Q].mp3\"\n",
    "\n",
    "from suno_utils.tasks.data_loader import load_audio_mp\n",
    "\n",
    "audio_arrays_48khz = load_audio_mp(\n",
    "    [audio_path],\n",
    "    target_sample_rate=48000,\n",
    "    normalize_volume=True,\n",
    "    num_workers=12,\n",
    "    n_channels=2,\n",
    ")\n",
    "\n",
    "start_s = 0.0\n",
    "end_s = audio_arrays_48khz[0].shape[1] / 48000\n",
    "audio_arrays_48khz = [a[:, int(start_s * 48000) : int(end_s * 48000)] for a in audio_arrays_48khz]\n",
    "end_sample = int(end_s * 48000)\n",
    "\n",
    "print(audio_arrays_48khz[0].shape)\n",
    "IPython.display.display(IPython.display.Audio(audio_arrays_48khz[0].numpy(), rate=48000))\n",
    "\n",
    "# now load audio at 24khz\n",
    "audio_arrays_24khz = load_audio_mp(\n",
    "    [audio_path],\n",
    "    target_sample_rate=24000,\n",
    "    normalize_volume=True,\n",
    "    min_duration_s=30.0,\n",
    "    num_workers=12,\n",
    "    n_channels=2,\n",
    ")\n",
    "\n",
    "audio_arrays_24khz = [a[:, int(start_s * 24000) : int(end_s * 24000)] for a in audio_arrays_24khz]\n",
    "\n",
    "# semantic encode \n",
    "semantic_codes = encode([audio_arrays_24khz[0].mean(axis=0, keepdim=True)], SAMPLE_RATE, EMBEDDING_RATE)\n",
    "semantic_codes = np.array(semantic_codes)[0, :, 0]\n",
    "\n",
    "print(semantic_codes.shape)\n",
    "#if semantic_codes.shape[0] < 9000: # pad with 4000\n",
    "#    semantic_codes = np.pad(semantic_codes, (0, 9000 - semantic_codes.shape[0]), mode=\"constant\", constant_values=4000)\n",
    "semantic_codes = torch.from_numpy(semantic_codes).long().cuda() \n",
    "print(semantic_codes.shape)\n",
    "\n",
    "# compute vae latents\n",
    "with torch.no_grad():\n",
    "    latents = vae_model.encode(audio_arrays_48khz[0].unsqueeze(0).cuda())[\"z\"].detach()    \n",
    "\n",
    "\n",
    "if \"Don't Know Why\" in audio_path:\n",
    "    full_lyrics = \"\"\"\n",
    "    [Verse 1]\n",
    "I waited 'til I saw the sun\n",
    "I don't know why I didn't come\n",
    "I left you by the house of fun\n",
    "I don't know why I didn't come\n",
    "I don't know why I didn't come\n",
    "\n",
    "[Verse 2]\n",
    "When I saw the break of day\n",
    "I wished that I could fly away\n",
    "Instead of kneeling in the sand\n",
    "Catching teardrops in my hand\n",
    "\n",
    "[Refrain]\n",
    "My heart is drenched in wine\n",
    "But you'll be on my mind forever\n",
    "\n",
    "[Verse 3]\n",
    "Out across the endless sea\n",
    "I will die in ecstasy\n",
    "But I'll be a bag of bones\n",
    "Driving down the road alone\n",
    "\n",
    "[Refrain]\n",
    "My heart is drenched in wine\n",
    "But you'll be on my mind forever\n",
    "\n",
    "[Instrumental Break]\n",
    "\n",
    "[Verse 4]\n",
    "Something has to make you run\n",
    "I don't know why I didn't come\n",
    "I feel as empty as a drum\n",
    "\n",
    "[Outro]\n",
    "I don't know why I didn't come\n",
    "I don't know why I didn't come\n",
    "I don't know why I didn't come\n",
    "\"\"\".strip()\n",
    "    tags = [\"jazz\", \"pop\", \"female vocal\"]\n",
    "\n",
    "if \"10 DRUNK CIGARETTES\" in audio_path:\n",
    "    lyrics = \"\"\"\n",
    "    [Verse 1]\n",
    "   girls don't need men\n",
    "   barbie dont need Ken\n",
    "   us girls are doing fine on our own\n",
    "   so if youre a man leave us a lone\n",
    "\n",
    "   getting girls rich yes, \n",
    "   thats a part of my plan\n",
    "   and i can name ten things us girls need\n",
    "   before we ever need a man\n",
    "\n",
    "   [chorus]\n",
    "   one new vape\n",
    "   two lines of coke\n",
    "   three drinks from the bar\n",
    "   four more lines of coke\n",
    "   five \n",
    "   six \n",
    "   seven more lines of coke\n",
    "   eight pairs of shoes\n",
    "   nine bebe belts\n",
    "   and ten drunk cigarettes\n",
    "\n",
    "   seriously i don't need anything from a man\n",
    "   actually maybe give me some money\n",
    "   i dont need money from a man but\n",
    "   i got my own, but\n",
    "   i want money from a man\n",
    "\n",
    "   getting girls rich yes, thats a part of my plan\n",
    "   and i can name ten things us girls need\n",
    "   before we ever need a man\n",
    "\n",
    "    [chorus]\n",
    "   one new vape\n",
    "   two lines of coke\n",
    "   three drinks from the bar\n",
    "   four more lines of coke\n",
    "   five guys fries\n",
    "   six hits from my blunt\n",
    "   seven more lines of coke\n",
    "   eight pairs of shoes\n",
    "   nine bebe belts\n",
    "   and ten drunk cigarettes\n",
    "\n",
    "    seriously i don't need anything from a man\n",
    "   actually maybe give me some money\n",
    "   i dont need money from a man but\n",
    "   i got my own, but\n",
    "   i want money from a man\n",
    "\n",
    "    getting girls rich yes, thats a part of my plan\n",
    "   and i can name ten things us girls need\n",
    "   before we ever need a man\n",
    "\n",
    "[chorus]\n",
    "   one new vape\n",
    "   two lines of coke\n",
    "   three drinks from the bar\n",
    "   four more lines of coke\n",
    "   five \n",
    "   six \n",
    "   seven more lines of coke\n",
    "   eight pairs of shoes\n",
    "   nine bebe belts\n",
    "   and ten drunk cigarettes\n",
    "\n",
    "    \"\"\".strip()\n",
    "    tags = [\"rock\", \"indie pop\"]\n",
    "elif \"halo\" in audio_path:\n",
    "    lyrics = \"\"\"\n",
    "[Verse 1]\n",
    "Remember those walls I built?\n",
    "Well, baby, they're tumblin' down\n",
    "And they didn't even put up a fight\n",
    "They didn't even make a sound\n",
    "I found a way to let you in\n",
    "But I never really had a doubt\n",
    "Standin' in the light of your halo\n",
    "I got my angel now\n",
    "\n",
    "[Pre-Chorus]\n",
    "It's like I've been awakened\n",
    "Every rule, I had you breakin'\n",
    "It's the risk that I'm takin'\n",
    "I ain't ever gonna shut you out\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "You know you're my savin' grace\n",
    "You're everything I need and more\n",
    "It's written all over your face\n",
    "Baby, I can feel your halo\n",
    "Pray it won't fade away\n",
    "\n",
    "[Post-Chorus]\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "\n",
    "[Verse 2]\n",
    "Hit me like a ray of sun\n",
    "Burnin' through my darkest night\n",
    "You're the only one that I want\n",
    "Think I'm addicted to your light\n",
    "I swore I'd never fall again\n",
    "But this don't even feel like fallin'\n",
    "Gravity can't begin\n",
    "To pull me back to the ground again\n",
    "\n",
    "[Pre-Chorus]\n",
    "It's like I've been awakened\n",
    "Every rule, I had you breakin'\n",
    "The risk that I'm takin'\n",
    "I'm never gonna shut you out\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "You know you're my savin' grace\n",
    "You're everything I need and more\n",
    "It's written all over your face\n",
    "Baby, I can feel your halo\n",
    "Pray it won't fade away\n",
    "[Post-Chorus]\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "\n",
    "[Bridge]\n",
    "Halo, ooh\n",
    "Halo, ooh\n",
    "Ooh\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "You know you're my savin' grace\n",
    "You're everything I need and more\n",
    "It's written all over your face\n",
    "Baby, I can feel your halo\n",
    "Pray it won't fade away\n",
    "\n",
    "[Post-Chorus]\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "    \"\"\".strip()\n",
    "    tags = [\"rock\", \"indie pop\"]\n",
    "else:\n",
    "    lyrics = \"\"\n",
    "    tags = []\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "if \"Don't\" in audio_path:\n",
    "    lyrics_0 = \"I waited 'til I saw the sun\\n I don't know I didn't come\\n I left by the house of fun\\n I don't know why I didn't come\\n \"\n",
    "    lyrics_1 = \"I don't know why I didn't come\\n When I saw the break of day\\n I wished that I could fly away\\n Instead of kneeling in the sand\\n Catching teardrops in my hand\\n\"\n",
    "    lyrics_2 = \"My heart is drenched in wine\\n But you'll be on my mind forever\\n Out across the endless sea\\n I will die in\"\n",
    "    lyrics_3 = \"ecstasy\\n But I'll be a bag of bones\\n Driving down the road alone\\n My hear is drenched in wine\\n But you'll be on my mind \\n\"\n",
    "    lyrics_4 = \"forever\\n [Instrumental Break]\\n Something has to make you run\\n \"\n",
    "    lyrics_5 = \"n I don't know why I didn't come\\n I feel as empty as a drum\\n [Outro]\\n I don't know why I didn't come\\n I don't know why I didn't come\\n I don't know why I didn't come\\n\"\n",
    "    lyrics_6 = \"\"\n",
    "    lyrics_list = [lyrics_0, lyrics_1, lyrics_2, lyrics_3, lyrics_4, lyrics_5, lyrics_6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "song_name = \"stone\"\n",
    "start_s = 0.0\n",
    "\n",
    "if song_name == \"stone\":\n",
    "    gen_id = \"a5e2198a-f352-4abb-9a24-7f81b143ded3\"\n",
    "    lyrics = \"\"\"\"\n",
    "    [verse]\n",
    "oh, my love\n",
    "My friend you know\n",
    "it's been a while\n",
    "Without thinking of you\n",
    "but the thought makes me smile\n",
    "\n",
    "[chorus]\n",
    "I'm so tired of wanting\n",
    "wanting more than this\n",
    "i know it but what am i to do\n",
    "i need some space to breathe,\n",
    "so give me some room\n",
    "\n",
    "[verse]\n",
    "oh, my love\n",
    "you have a heart of stone\n",
    "cause since i've come home\n",
    "i've never felt so alone\n",
    "but the thought makes me smile\n",
    "\n",
    "[verse]\n",
    "oh, my love\n",
    "you're coming with me\n",
    "I know you'll leave me one day\n",
    "can you just stay till monday\n",
    "the thought makes me smile\n",
    "\n",
    "[chorus]\n",
    "I'm so tired of wanting\n",
    "wanting more than this\n",
    "i know it but what am i to do\n",
    "i need some space to breathe,\n",
    "so give me some room\n",
    "so give me some room\n",
    "i need room\n",
    "\n",
    "[outro]\n",
    "\n",
    "smile\n",
    "\n",
    "[outro]\"\"\".strip()\n",
    "    tags = [\"rock\", \"indie pop\"]\n",
    "    #tags = [\"metal\", \"metal\", \"metal\"]\n",
    "    #tags = [\"country\", \"country\", \"country\"]\n",
    "    #tags = [\"1950s\", \"barbershop\", \"pop\", \"vintage\"]\n",
    "    start_s = 0.0\n",
    "\n",
    "    lyrics_1 = \"\"\"oh, my love\n",
    "My friend you know\n",
    "it's been a while\n",
    "Without thinking of you\n",
    "but the thought makes me smile\n",
    "\n",
    "[chorus]\n",
    "I'm so tired of wanting\n",
    "wanting more than this\n",
    "i know it but\"\"\"\n",
    "    lyrics_2 = \"\"\"\n",
    "I'm so tired of wanting\n",
    "wanting more than this\n",
    "i know it but what am i to do\n",
    "i need some space to breathe,\n",
    "so give me some room\n",
    "\n",
    "[verse]\n",
    "oh, my love\n",
    "\"\"\"\n",
    "    lyrics_3 = \"\"\"\n",
    "space to breathe,\n",
    "so give me some room\n",
    "\n",
    "[verse]\n",
    "oh, my love\n",
    "you have a heart of stone\n",
    "cause since i've come home\n",
    "i've never felt so alone\n",
    "but the thought makes me smile\n",
    "\"\"\"\n",
    "    lyrics_4 = \"\"\"\n",
    "\n",
    "but the thought makes me smile\n",
    "[verse]\n",
    "oh, my love\n",
    "you have a heart of stone\n",
    "cause since i've come home\n",
    "i've never felt so alone\n",
    "but the thought makes me smile\"\"\"\n",
    "    lyrics_5 = \"\"\"\n",
    "[verse]\n",
    "oh, my love\n",
    "you have a heart of stone\n",
    "cause since i've come home\n",
    "i've never felt so alone\n",
    "but the thought makes me smile\n",
    "\n",
    "oh my love\n",
    "\"\"\"\n",
    "    lyrics_6 = \"\"\"\n",
    "[verse]\n",
    "oh, my love\n",
    "you have a heart of stone\n",
    "cause since i've come home\n",
    "i've never felt so alone\n",
    "but the thought makes me smile\n",
    "\n",
    "oh my love\n",
    "\"\"\"\n",
    "    lyrics_7 = \"\"\"\n",
    "    [verse]\n",
    "oh, my love\n",
    "you're coming with me\n",
    "I know you'll leave me one day\n",
    "can you just stay till monday\n",
    "the thought makes me smile\n",
    "\"\"\"\n",
    "    lyrics_8 = \"\"\"\n",
    "can you just stay till monday\n",
    "the thought makes me smile\n",
    "\n",
    "[chorus]\n",
    "I'm so tired of wanting\n",
    "wanting more than this\n",
    "i know it but what am i to do\n",
    "i need some space to breathe,\n",
    "\n",
    "I'm so tired of wanting\"\"\"\n",
    "    lyrics_9 = \"\"\"\n",
    "[outro]\n",
    "smile\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "    lyrics_10 = \"\"\"\n",
    "\n",
    "[outro]\"\"\"  \n",
    "    lyrics_11 = \"\"\"\"\"\"\n",
    "    lyrics_list = [lyrics_1, lyrics_2, lyrics_3, lyrics_4, lyrics_5, lyrics_6, lyrics_7, lyrics_8, lyrics_9, lyrics_10, lyrics_11]\n",
    "elif song_name == \"hydra\":\n",
    "    gen_id = \"271a6cad-8b4f-446c-8545-8f4daa8120af\"\n",
    "    tags = [\"metal\"]\n",
    "    lyrics = \"\"\"\n",
    "Silent descent through azure mist\n",
    "Our ship touches down on Neptune's crust\n",
    "Unaware of the myth that exists\n",
    "In this world of frozen dust\n",
    "\n",
    "Nine heads emerge from icy depths\n",
    "Silken twine of tentacles reach\n",
    "Blinding everyone with frost-breath\n",
    "Our mission now beyond our reach\n",
    "\n",
    "We came here as explorers brave\n",
    "To unveil secrets of the deep\n",
    "But now we're trapped in Neptune's grave\n",
    "Where ancient horrors lie and creep\n",
    "Our magic spells seem weak and frail\n",
    "Against this beast of legend old\n",
    "We never stood a chance to prevail\n",
    "In this disparaging wasteland cold\n",
    "\n",
    "Hydra rises, laying to waste\n",
    "Everything we thought we knew\n",
    "We had it all, now it's erased\n",
    "This sick romance with the unknown\n",
    "\n",
    "Scatter! Flee! But there's nowhere to hide\n",
    "As death comes surging with the tide\n",
    "We tempt our fate with every breath\n",
    "In love with despair, we court our death\n",
    "Detestation grows with each attack\n",
    "Nine heads strike, we can't fight back\n",
    "Our ship destroyed, our hopes undone\n",
    "The oppressor has already won\n",
    "\n",
    "As only time knows, our ashes in the urn\n",
    "We deserve everything that's coming\n",
    "Until we die, there's no ever after\n",
    "In this cosmic game, we're merely pawns to burn\n",
    "\n",
    "Hydra rises, laying to waste\n",
    "Everything we thought we knew\n",
    "We had it all, now it's erased\n",
    "This sick romance with the unknown\n",
    "\"\"\".strip()\n",
    "elif song_name == \"friends\":\n",
    "    gen_id = \"081d73c4-7805-4212-9c80-8db1137ca3c4\"\n",
    "    tags = [\"Rock\"]\n",
    "    lyrics = \"\"\"\n",
    "What would you do if I lost my tune\n",
    "Would you laugh and leave me here?\n",
    "Lend your voice and let's sing\n",
    "I'll try not to get out of the melody\n",
    "Oh, baby, I'm going to make it\n",
    "(With a little help from my friends)\n",
    "All I need are my friends\n",
    "(Try it with a little help from my friends)\n",
    "I said I want to feel alright, I'm going to\n",
    "(Fly with a little help from my friends)\n",
    "What do I do when I'm feeling lonely\n",
    "(Do you worry about being alone?)\n",
    "How I feel at the end of the day\n",
    "(Are you sad to be alone?)\n",
    "I hope you don't feel it too\n",
    "Oh, baby, I'm going to make it\n",
    "(With a little help from my friends)\n",
    "All I need are my friends\n",
    "(Try it with a little help from my friends)\n",
    "I said I want to feel alright, I'm going to\n",
    "(Fly with a little help from my friends)\n",
    "(Do you need somebody?)\n",
    "I need someone to love\n",
    "(Could it be anyone?)\n",
    "All I need is someone who understands my way, yes\n",
    "(Do you need somebody?)\n",
    "I need someone to love\n",
    "(Could it be anyone?)\n",
    "All I need is someone who understands my way, yes\n",
    "Refrain:\n",
    "Oh, baby, I'm going to make it\n",
    "(With a little help from my friends)\n",
    "All I need are my friends\n",
    "(Try it with a little help from my friends)\n",
    "I said I want to feel alright, I'm going to\n",
    "(Fly with a little help from my friends)\n",
    "I'm going to keep trying\n",
    "(Fly with a little help from my friends)\n",
    "I'm going to make it with my friends\n",
    "\"\"\".strip()\n",
    "elif song_name == \"i-am-here-for-you\":\n",
    "    gen_id = \"a7199702-104f-42bd-b072-8b82c015551f\"\n",
    "    lyrics = \"\"\"\n",
    "   [Intro]\n",
    "Hey how are you feeling?\n",
    "Are you doing alright?\n",
    "\n",
    "[Instrumental]\n",
    "\n",
    "\n",
    "[Verse]\n",
    "When the night is long\n",
    "And you're feeling down\n",
    "I'm right by your side\n",
    "I am here for you\n",
    "(Just hold on tight)\n",
    "\n",
    "[Guitar interlude]\n",
    "\n",
    "[Verse 2]\n",
    "When the world shuts down\n",
    "And nothing feels true\n",
    "(leave your heavy thoughts,\n",
    "in the dust)\n",
    "Hold on to my hand\n",
    "I am here for you\n",
    "In each other we trust\n",
    "\n",
    "\n",
    "[Chorus]\n",
    "I am here for you\n",
    "(I hope you feel alright)\n",
    "Through the highs and lows\n",
    "(Highs and lows)\n",
    "I am here for you\n",
    "(I am here for you!)\n",
    "And everybody knows\n",
    "Trust is our love! \n",
    "\n",
    "[Instrumental]\n",
    "\n",
    "[Guitar solo]\n",
    "\n",
    "[Bridge]\n",
    "I will always find you\n",
    "(I will make the shadows disappear)\n",
    "\n",
    "When your dreams come through\n",
    "(Laughter will feel the silence)\n",
    "No matter where you go\n",
    "(Everything is alright)\n",
    "\n",
    "I am here for you\n",
    "I just wanna tell you\n",
    "I. Love. You\n",
    "\n",
    "[Brief guitar solo]\n",
    "\n",
    "[Instrumental]\n",
    "\n",
    "[Chorus]\n",
    "I am here for you\n",
    "(I hope you feel alright)\n",
    "Through the highs and lows\n",
    "(Highs and lows)\n",
    "\n",
    "I am here for you\n",
    "(I am here for you!)\n",
    "And everybody knows\n",
    "Trust is our love! \n",
    "\n",
    "[Instrumental]\n",
    "\n",
    "[Guitar solo]\n",
    "\n",
    "[Outro]\n",
    "When the stars collide\n",
    "And the skies turn new\n",
    "You'll find me\n",
    "sitting next to you\n",
    "\n",
    "(Hey how are you?) \n",
    "(Are you feeling better now)\n",
    "(I just wanna tell you. . .  .)\n",
    "(I. Love. You)\n",
    "\n",
    "[End]\n",
    "    \"\"\".strip()\n",
    "    tags = [\"Blues\", \"romantic\", \"soft guitar\",\" gentle female vocals\"]\n",
    "    tags = [\"pop\", \"electronic\", \"rock\"]\n",
    "elif song_name == \"sister\":\n",
    "    gen_id = \"99bffa17-7e59-47b4-a048-5528cbda05d5\"\n",
    "    lyrics = \"\"\"\n",
    "Verse]\n",
    "She grew up in the shadow of the steeple,\n",
    "Small town roads where the night stands still.\n",
    "She heard the Lord in the whispers of the people,\n",
    "Calling her heart to follow His will.\n",
    "\n",
    "[Verse 2]\n",
    "Her mama cried but knew deep down,\n",
    "That little girl was Heaven-sent.\n",
    "She traded her dresses for a plain white gown,\n",
    "And left behind the life she’d spent.\n",
    "\n",
    "[Chorus]\n",
    "Sister of the Light, with rosary beads,\n",
    "Praying through the night, fulfilling good deeds.\n",
    "She's found her calling, in the chapel's glow,\n",
    "A beacon of grace, wherever she goes.\n",
    "\n",
    "[Verse 3]\n",
    "She helps the lost and feeds the hungry,\n",
    "In a world that’s often cold and unkind.\n",
    "Her hands are worn but her spirit’s humbly,\n",
    "Carrying love of the most divine kind.\n",
    "\n",
    "[Verse 4]\n",
    "Sundays in the choir, hymns touch the sky,\n",
    "A voice pure as the morning dew.\n",
    "With faith so deep, she makes angels sigh,\n",
    "Her prayers reach heaven like they’re brand new.\n",
    "\n",
    "[Chorus]\n",
    "Sister of the Light, with rosary beads,\n",
    "Praying through the night, fulfilling good deeds.\n",
    "She's found her calling, in the chapel's glow,\n",
    "A beacon of grace, wherever she goes.\n",
    "    \"\"\".strip()\n",
    "    tags = [\"rock\", \"pop\", \"indie\"]\n",
    "elif song_name == \"something-real\":\n",
    "    gen_id = \"562f762d-6ced-4080-9af1-910ee3d0a5dc\"\n",
    "    tags = [\"rock\", \"pop\", \"indie\"]\n",
    "    lyrics = \"\"\"\n",
    "    Woke up in a haze, can’t shake the weight\n",
    "Yesterday’s mistakes, is it just fate?\n",
    "The mirror cracks, reflection's gone\n",
    "Lost myself in what went wrong\n",
    "\n",
    "I’m falling deeper, there’s no escape\n",
    "Running from a past I can’t erase\n",
    "The fire’s burning, but I can’t feel\n",
    "Just need something, something real\n",
    "\n",
    "Every word you said still cuts like glass\n",
    "But I can’t outrun the shadows cast\n",
    "Chasing ghosts through empty streets\n",
    "Where silence screams, and pain repeats\n",
    "\n",
    "I’m falling deeper, there’s no escape\n",
    "Running from a past I can’t erase\n",
    "The fire’s burning, but I can’t feel\n",
    "Just need something, something real\n",
    "\n",
    "Tried to drown it out, but it’s in my veins\n",
    "Sinking faster, can’t break these chains\n",
    "It’s all too much, can’t hide no more\n",
    "The wreckage washed up, on the shores\n",
    "\n",
    "I’m falling deeper, there’s no escape\n",
    "Running from a past I can’t erase\n",
    "The fire’s burning, but I can’t feel\n",
    "Just need something, something real\n",
    "\n",
    "I’m falling deeper, there’s no escape\n",
    "Running from a past I can’t erase\n",
    "The fire’s burning, but I can’t feel\n",
    "Just need something, something real\n",
    "\n",
    "I need something… \n",
    "something real...\"\"\".strip()\n",
    "s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.npz\"\n",
    "print(s3_filepath)\n",
    "data = read_from_s3(s3_filepath, read_f=np.load)\n",
    "\n",
    "# tony\n",
    "if False:\n",
    "    start_s = 90.0\n",
    "    end_s = start_s + 30.0\n",
    "    s3_filepath = \"/home/tony/Work/gpt_samples/30b_t3_v1_repro7/ts_pop_0.npz\"\n",
    "    data = np.load(s3_filepath)\n",
    "    codes = data[\"codes\"]\n",
    "    tags = [\"pop country\", \"female vocal\", \"fast tempo\"]\n",
    "    lyrics = \"\"\"\n",
    "    Hey, whatcha doing with a girl like that?\n",
    "\n",
    "    [Pre-Chorus]\n",
    "    She wears high heels, I wear sneakers\n",
    "    She's Cheer Captain and I'm on the bleachers\n",
    "    Dreaming 'bout the day when you wake up and find\n",
    "    That what you're looking for has been here the whole time\n",
    "\n",
    "    [Chorus]\n",
    "    If you could see that I'm the one who understands you\n",
    "    Been here all along, so why can't you see?\n",
    "    You belong with me\"\"\".strip()\n",
    "    mp3_filepath = \"/home/tony/Work/gpt_samples/30b_t3_v1_repro7/ts_pop_0_seed42.mp3\"\n",
    "    mp3_data, mp3_sr = torchaudio.load(mp3_filepath)\n",
    "\n",
    "\n",
    "# also read the mp3\n",
    "mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.mp3\"\n",
    "mp3_data, mp3_sr = read_from_s3(mp3_filepath, read_f=torchaudio.load)\n",
    "end_s = mp3_data.shape[1] / mp3_sr\n",
    "\n",
    "start_sample = int(start_s * 48000)\n",
    "end_sample = int(end_s * 48000)\n",
    "\n",
    "audio_arrays_48khz = [mp3_data[...,start_sample:end_sample].cuda()]\n",
    "print(audio_arrays_48khz[0].shape)\n",
    "\n",
    "\n",
    "#text_data = read_from_s3(f\"s3://suno-data-uploads/studio/uploads/{gen_id}_hoot.json\")\n",
    "#text_data = json.loads(text_data)\n",
    "\n",
    "#print(text_data)\n",
    "#print(data.keys())\n",
    "start_frame = 0\n",
    "#end_frame = int(25 * 240.0)\n",
    "\n",
    "if \"v3.0_raw\" in data:\n",
    "    codes = data[\"v3.0_raw\"]\n",
    "elif \"v3.5_raw\" in data:\n",
    "    codes = data[\"v3.5_raw\"]\n",
    "elif \"v4.0_raw\" in data:\n",
    "    codes = data[\"v4.0_raw\"]\n",
    "else:\n",
    "    raise ValueError(\"No codes found\")\n",
    "\n",
    "print(codes.shape)\n",
    "semantic_codes = codes[:, 0].astype(np.uint16)\n",
    "semantic_codes = torch.from_numpy(semantic_codes).long().cuda()\n",
    "print(song_name, semantic_codes.shape)\n",
    "n_valid_semantic_tokens = semantic_codes.shape[0]\n",
    "\n",
    "# encode audio with vae\n",
    "with torch.no_grad():\n",
    "    latents = vae_model.encode(audio_arrays_48khz[0].unsqueeze(0).cuda())[\"z\"].detach()    \n",
    "\n",
    "#print(latents.shape)\n",
    "\n",
    "# append 10s worth of semantic pad token\n",
    "#semantic_pad_tokens = (torch.ones(250) * 4000).long().cuda()\n",
    "#semantic_codes = torch.cat([semantic_pad_tokens, semantic_codes], dim=0)\n",
    "\n",
    "#latents *= scale_factor\n",
    "\n",
    "IPython.display.display(IPython.display.Audio(audio_arrays_48khz[0].cpu().numpy(), rate=mp3_sr))\n",
    "print(tags)\n",
    "print(lyrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# encode silent audio with vae\n",
    "silent_audio = torch.zeros(1, 2, 48000*30).cuda()\n",
    "with torch.no_grad():\n",
    "    silent_latents = vae_model.encode(silent_audio)[\"z\"].detach()    \n",
    "print(silent_latents.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "silent_lantent = silent_latents[0,:,1003:1004]\n",
    "print(silent_lantent.min(), silent_lantent.max(), silent_lantent.mean(), silent_lantent.std())\n",
    "print(silent_lantent.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "def calculate_blocks_from_seconds(total_seconds: float) -> int:\n",
    "    if total_seconds <= 30:\n",
    "        return 1\n",
    "    else:\n",
    "        return math.ceil((total_seconds - 20) / 20) + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute number of blocks needed\n",
    "# we will operate on blocks of 30 seconds\n",
    "# this translates to 3000 embeddings. \n",
    "# however, after the first block, we will condition on the output of the previous block\n",
    "# we will use the last 10 seconds (1000 embeddings) of the previous block to condition on\n",
    "# plus, we will use the last 10 seconds of the semantic codes from previous block (250 tokens)\n",
    "\n",
    "context_dur_s = 10.0\n",
    "semantic_window_size = 750\n",
    "semantic_context_size = int(25 * context_dur_s)\n",
    "semantic_hop_size = int(semantic_window_size - semantic_context_size)\n",
    "\n",
    "\n",
    "print(\"semantic_context_size\", semantic_context_size, \"semantic_hop_size\", semantic_hop_size)\n",
    "\n",
    "latents_window_size = 3000\n",
    "latents_context_size = int(100 * context_dur_s)\n",
    "latents_hop_size = int(latents_window_size - latents_context_size)\n",
    "\n",
    "print(\"latents_context_size\", latents_context_size, \"latents_hop_size\", latents_hop_size)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(semantic_codes))\n",
    "# find where the semantic codes are 4000\n",
    "#pad_idx = (semantic_codes == 4000).nonzero()[0]\n",
    "#print(pad_idx)\n",
    "# find the indices of the semantic codes that are 3890\n",
    "pad_idx = (semantic_codes == 3890).nonzero()\n",
    "print(pad_idx)\n",
    "#\n",
    "# print(pad_idx.min(), pad_idx.max())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model.vae_pad_embed.min(), model.vae_pad_embed.max(), model.vae_pad_embed.mean(), model.vae_pad_embed.std())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_steps = 32\n",
    "cfg_scale = 2.0\n",
    "seed = np.random.randint(0, 2**32 - 1)\n",
    "print(seed)\n",
    "\n",
    "start_s = 0.0\n",
    "print(\"start_s\", start_s, \"end_s\", end_s)\n",
    "end_s = audio_arrays_48khz[0].shape[1] / 48000\n",
    "num_blocks = calculate_blocks_from_seconds(end_s - start_s)\n",
    "print(num_blocks)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# init final latent array\n",
    "#latents_output = torch.zeros(1, 128, latents_window_size * num_blocks).cuda()\n",
    "latents_output = silent_lantent.clone().repeat(1, 1, latents_window_size * num_blocks)\n",
    "print(latents_output.shape)\n",
    "# iterate over blocks with index\n",
    "for block_idx in range(num_blocks):\n",
    "\n",
    "    # for the first block, we use the start of the conditioning from semantic\n",
    "    # for the rest of the blocks, we use a hop size of \n",
    "    if block_idx == 0:\n",
    "        semantic_start_idx = 0\n",
    "        semantic_end_idx = semantic_window_size - semantic_context_size\n",
    "        semantic_codes_block = semantic_codes[semantic_start_idx:semantic_end_idx].clone()\n",
    "        # repeat the first tokens \n",
    "        #semantic_pad_tokens = semantic_codes_block[:semantic_context_size].clone()\n",
    "        # pad with semantic pad tokens\n",
    "        semantic_pad_tokens = (torch.ones(semantic_context_size) * 4000).long().cuda()\n",
    "        semantic_codes_block = torch.cat([semantic_pad_tokens, semantic_codes_block], dim=0)\n",
    "    else:\n",
    "        semantic_start_idx = semantic_end_idx - semantic_context_size\n",
    "        semantic_end_idx = semantic_start_idx + semantic_window_size\n",
    "        # get the conditioning from the semantic codes\n",
    "        semantic_codes_block = semantic_codes[semantic_start_idx:semantic_end_idx].clone()\n",
    "    \n",
    "    print(block_idx, semantic_start_idx, semantic_end_idx, semantic_codes_block.shape)\n",
    "\n",
    "    # pad the semantic codes block with semantic pad tokens\n",
    "    if semantic_codes_block.shape[0] < semantic_window_size:\n",
    "        repeat_size = semantic_window_size - semantic_codes_block.shape[0]\n",
    "        print(\"pad last\", repeat_size, \"tokens\")\n",
    "        #repeat_tokens = semantic_codes_block[:repeat_size]\n",
    "        repeat_tokens = (torch.ones(repeat_size) * 4000).long().cuda()\n",
    "        semantic_codes_block = torch.cat([semantic_codes_block, repeat_tokens], dim=0)\n",
    "    else:\n",
    "        repeat_size = 0\n",
    "    print(semantic_codes_block.shape)\n",
    "\n",
    "    # run the diffusion process on current block\n",
    "    if block_idx == 0:\n",
    "        # we have no context or conditioning from previous block\n",
    "        latents_block = model.vae_pad_embed.view(-1, 1).repeat(1, latents_context_size).clone()\n",
    "        #latents_block = latents[0, :, :latents_context_size].clone()\n",
    "        start_s = 0\n",
    "        end_s = start_s + 20.0\n",
    "        latents_end_idx = latents_context_size\n",
    "\n",
    "    else:\n",
    "        latents_start_idx = 2000\n",
    "        latents_end_idx = 3000\n",
    "        latents_block = upsampled_latents[0,:,latents_start_idx:latents_end_idx].clone()\n",
    "        start_s = end_s\n",
    "        end_s = start_s + 20.0\n",
    "\n",
    "    print(start_s, end_s)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        upsampled_latents = upsample_diffusion_from_semantic_and_text(\n",
    "            model,\n",
    "            semantic_codes_block,\n",
    "            [\"pop\"],\n",
    "            #lyrics,\n",
    "            lyrics_list[block_idx],\n",
    "            steps=num_steps,\n",
    "            cfg_scale=cfg_scale,\n",
    "            sample_size=latents_window_size,\n",
    "            sample_rate=48000,\n",
    "            seed=seed,\n",
    "            sampler_type=\"dpmpp-3m-sde\",\n",
    "            latent_context=latents_block,\n",
    "            compile=False,\n",
    "        )\n",
    "        pred_zq = upsampled_latents.clone()\n",
    "        print(\"pred_zq.std()\", pred_zq.std())\n",
    "        pred_zq /= scale_factor\n",
    "        #Wpred_zq = torch.randn_like(pred_zq)\n",
    "        print(\"pred_zq\", pred_zq.shape)\n",
    "        if block_idx == 0:\n",
    "            write_start_idx = 0\n",
    "            block_size = latents_window_size - latents_context_size\n",
    "        elif block_idx == num_blocks - 1:\n",
    "            write_start_idx = write_end_idx\n",
    "            block_size = latents_window_size - latents_context_size - (repeat_size * 4)\n",
    "        else:\n",
    "            write_start_idx = write_end_idx\n",
    "            block_size = latents_window_size - latents_context_size\n",
    "\n",
    "        pred_zq_block = pred_zq[...,latents_context_size:latents_context_size+block_size]\n",
    "        write_end_idx = write_start_idx + pred_zq_block.shape[-1]\n",
    "\n",
    "        print(\"write_start_idx\", write_start_idx, \"write_end_idx\", write_end_idx)\n",
    "        print(\"pred_zq_block\", pred_zq_block.shape)\n",
    "\n",
    "        latents_output[...,write_start_idx:write_end_idx] = pred_zq_block\n",
    "\n",
    "print(latents_output.shape)\n",
    "latents_output = latents_output[...,:semantic_codes.shape[0]*4]\n",
    "print(latents_output.shape)\n",
    "# std of latents output\n",
    "\n",
    "# decode the latents\n",
    "with torch.no_grad():\n",
    "    pred_audio = vae_model.decode(latents_output)[0].detach().cpu()         \n",
    "    pred_audio /= pred_audio.abs().max().clamp(1e-8)\n",
    "    pred_audio = pred_audio[...,:end_sample]\n",
    "    print(pred_audio.mean())\n",
    "\n",
    "# save audio\n",
    "IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# New chunked inference with 30s context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# new chunked inference with 30s context\n",
    "# we start with semantic codes\n",
    "# we will split the semantic codes into 30s chunks\n",
    "\n",
    "# constats\n",
    "window_dur_s = 30.0\n",
    "semantic_rate_hz = 25\n",
    "latents_rate_hz = 100\n",
    "\n",
    "# parameters\n",
    "cfg_scale = 4.0\n",
    "num_steps = 64\n",
    "seed = np.random.randint(0, 2**32 - 1)\n",
    "\n",
    "latents_window_size = int(window_dur_s * latents_rate_hz)\n",
    "semantic_window_size = int(window_dur_s * semantic_rate_hz)\n",
    "semantic_hop_size = int(window_dur_s * semantic_rate_hz)\n",
    "latents_context_size = int(window_dur_s * latents_rate_hz)\n",
    "\n",
    "num_blocks = (semantic_codes.shape[0] // semantic_window_size) + 1\n",
    "print(num_blocks)\n",
    "\n",
    "# init final latent array\n",
    "latents_output = torch.empty(1, 128, latents_window_size * num_blocks).cuda()\n",
    "\n",
    "start_time = time.time()\n",
    "for block_idx in range(num_blocks):\n",
    "\n",
    "    semantic_start_idx = block_idx * semantic_window_size\n",
    "    semantic_end_idx = semantic_start_idx + semantic_window_size\n",
    "    semantic_codes_block = semantic_codes[semantic_start_idx:semantic_end_idx].clone()\n",
    "\n",
    "    if semantic_codes_block.shape[0] < semantic_window_size:\n",
    "        repeat_size = semantic_window_size - semantic_codes_block.shape[0]\n",
    "        print(\"pad last\", repeat_size, \"tokens with repeated tokens\")\n",
    "        # take the last repeat_size tokens from the semantic codes\n",
    "        repeat_tokens = semantic_codes[-repeat_size:]\n",
    "        semantic_codes_block = torch.cat([semantic_codes_block, repeat_tokens], dim=0)\n",
    "\n",
    "    if block_idx == 0:\n",
    "        latents_block = model.vae_pad_embed.view(-1, 1).repeat(1, latents_context_size).clone()\n",
    "    else:\n",
    "        latents_block = pred_zq[0,...].clone()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        upsampled_latents = upsample_diffusion_from_semantic_and_text(\n",
    "            model,\n",
    "            semantic_codes_block,\n",
    "            tags,\n",
    "            #lyrics,\n",
    "            lyrics_list[block_idx],\n",
    "            steps=num_steps,\n",
    "            cfg_scale=cfg_scale,\n",
    "            sample_size=latents_window_size,\n",
    "            sample_rate=48000,\n",
    "            seed=seed,\n",
    "            sampler_type=\"dpmpp-3m-sde\",\n",
    "            latent_context=latents_block,\n",
    "            compile=False,\n",
    "        )\n",
    "        pred_zq = upsampled_latents.clone()\n",
    "        pred_zq /= scale_factor\n",
    "\n",
    "    write_start_idx = block_idx * latents_window_size\n",
    "    write_end_idx = write_start_idx + latents_window_size\n",
    "    latents_output[...,write_start_idx:write_end_idx] = pred_zq\n",
    "\n",
    "# crop to the length of the semantic codes\n",
    "latents_output = latents_output[...,:semantic_codes.shape[0]*4]\n",
    "\n",
    "# decode the latents\n",
    "with torch.no_grad():\n",
    "    pred_audio = vae_model.decode(latents_output)[0].detach().cpu()      \n",
    "    pred_audio /= pred_audio.abs().max().clamp(1e-8)\n",
    "    pred_audio = pred_audio[...,:end_sample]\n",
    "    print(pred_audio.mean())\n",
    "\n",
    "# save audio\n",
    "end_time = time.time()\n",
    "dur_s = (semantic_codes.shape[0] / 25)\n",
    "real_time_factor = dur_s / (end_time - start_time)\n",
    "print(f\"Upsampled {dur_s:.1f} s of audio in {end_time - start_time:.2f} s ({real_time_factor:.2f}x realtime)\")\n",
    "IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#pred_audio_filepath = f\"outputs/dont-know-why-100hz-30s-ft-v5-40k.mp3\"\n",
    "#pred_audio_filepath = f\"outputs/stone-100hz-30s-ft-v5-160k-aligned-3.mp3\"\n",
    "pred_audio_filepath = f\"outputs/i-am-here-for-you-100hz-30s-ft-v5-360k-full-lyrics.mp3\"\n",
    "#pred_audio_filepath = f\"outputs/10-drunk-cigs-100hz-30s-ft-v5-160k.mp3\"\n",
    "torchaudio.save(pred_audio_filepath, pred_audio.cpu().squeeze(), 48000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = np.random.randint(0, 2**32 - 1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# one block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg_scale = 6.0\n",
    "print(seed)\n",
    "\n",
    "#latents_block = latents[0, :, 1000:2000].clone()\n",
    "latents_block = model.vae_pad_embed.view(-1, 1).repeat(1, latents_context_size).clone()\n",
    "# use ground truth latent for conditioning\n",
    "print(latents_block.shape)\n",
    "\n",
    "#semantic_codes_block = semantic_codes[250:1000].clone()\n",
    "#print(semantic_codes_block.shape)\n",
    "\n",
    "semantic_start_idx = 0\n",
    "semantic_end_idx = semantic_window_size - semantic_context_size\n",
    "semantic_codes_block = semantic_codes[semantic_start_idx:semantic_end_idx].clone()\n",
    "# repeat the first tokens \n",
    "#semantic_pad_tokens = semantic_codes_block[:semantic_context_size].clone()\n",
    "# pad with semantic pad tokens\n",
    "semantic_pad_tokens = (torch.ones(semantic_context_size) * 4000).long().cuda()\n",
    "semantic_codes_block = torch.cat([semantic_pad_tokens, semantic_codes_block], dim=0)\n",
    "\n",
    "for num_steps in [8, 16, 32, 64, 128]:\n",
    "    print(\"num_steps\", num_steps)\n",
    "    with torch.no_grad():\n",
    "        upsampled_latents = upsample_diffusion_from_semantic_and_text(\n",
    "            model,\n",
    "            semantic_codes_block,\n",
    "            tags,\n",
    "            section_lyrics,\n",
    "            steps=num_steps,\n",
    "            cfg_scale=cfg_scale,\n",
    "            sample_size=latents_window_size,\n",
    "            sample_rate=48000,\n",
    "            seed=seed,\n",
    "            sampler_type=\"dpmpp-3m-sde\",\n",
    "            latent_context=latents_block,\n",
    "        )\n",
    "        pred_zq = upsampled_latents.clone()\n",
    "        pred_zq /= scale_factor\n",
    "        #Wpred_zq = torch.randn_like(pred_zq)\n",
    "\n",
    "        pred_zq_block = pred_zq[...,latents_context_size:]\n",
    "\n",
    "    pred_zq_block = pred_zq_block[...,:semantic_codes_block.shape[0]*4]\n",
    "\n",
    "    # decode the latents\n",
    "    with torch.no_grad():\n",
    "        pred_audio = vae_model.decode(pred_zq_block)[0].detach().cpu()         \n",
    "        pred_audio /= pred_audio.abs().max().clamp(1e-8)\n",
    "        pred_audio = pred_audio[...,:end_sample]\n",
    "\n",
    "    IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "section_lyrics = \"\"\"\n",
    "come\n",
    "i left you by the house of fun\n",
    "don't know why i didn't come\n",
    "i don't know why i didn't come\n",
    "\n",
    "when i saw the\n",
    "\"\"\"\n",
    "\n",
    "section_lyrics = \"\"\"\n",
    "oh, my love\n",
    "my friend you know\n",
    "it's been a while\n",
    "without thinking of you\n",
    "but the thought makes me smile\n",
    "\"\"\".strip()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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
}
