{
 "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\"] = \"5\"\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",
    "    upsample_diffusion_from_semantic_and_text_with_phonemes\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": [
    "# 4min ckpt \n",
    "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_20_scale=2.5_qk_norm_4min_step=280000.ckpt\"\n",
    "ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_20_scale=2.5_qk_norm_4min_step=380000.ckpt\"\n",
    "\n",
    "\n",
    "if not os.path.isfile(ckpt_path):\n",
    "    raise ValueError(f\"Checkpoint not found: {ckpt_path}\")\n",
    "\n",
    "\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",
    "vae_frame_rate = 100\n",
    "device = \"cuda:0\"\n",
    "\n",
    "if vae_frame_rate == 25:\n",
    "    # load VAE model\n",
    "    checkpoint_filepath = \"s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth\"\n",
    "elif vae_frame_rate == 100:\n",
    "    checkpoint_filepath = \"s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth\"\n",
    "else:\n",
    "    raise ValueError(f\"VAE frame rate not supported: {vae_frame_rate}\")\n",
    "\n",
    "#checkpoint_filepath = \"/home/christian/code/christian/checkpoints/25hz_vae_peaq_kl_0.005.pth\"\n",
    "load_f = funcy.partial(torch.load, map_location=\"cpu\")\n",
    "\n",
    "if checkpoint_filepath.startswith(\"s3://\"):\n",
    "    sd = read_from_s3(checkpoint_filepath, read_f=load_f)\n",
    "else:\n",
    "    sd = load_f(checkpoint_filepath)\n",
    "\n",
    "sd[\"metadata\"][\"kwargs\"] = {\n",
    "    k: v\n",
    "    for k, v in sd[\"metadata\"][\"kwargs\"].items()\n",
    "    if k in DAC.__init__.__code__.co_varnames\n",
    "}\n",
    "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",
    "max_dur_s = 240\n",
    "semantic_n_tokens = int(25 * max_dur_s)\n",
    "n_vae_tokens = int(vae_frame_rate * max_dur_s)\n",
    "\n",
    "# setup config\n",
    "if \"diffusion_semantic+text_100hz_20_scale=2.5_qk_norm\" in ckpt_path:\n",
    "    model_type = \"diffusion_semantic+text_100hz_20_scale=2.5_qk_norm\"\n",
    "    model_config_path = os.path.join(\n",
    "        config_dir, \"stable_audio_2_0_semantic+text_48khz_20_scale=2.5_qk_norm.json\"\n",
    "    )\n",
    "elif \"diffusion_semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes_6min\" in ckpt_path:\n",
    "    model_type = \"diffusion_semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes_6mins\"\n",
    "    model_config_path = os.path.join(\n",
    "        config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes_mask.json\"\n",
    "    )\n",
    "else:\n",
    "    raise ValueError(f\"Model type not found: {ckpt_path}\")\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}\")"
   ]
  },
  {
   "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",
    "else:\n",
    "    start_s = 0.0\n",
    "    end_s = start_s + max_dur_s\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",
    "n_valid_semanitc_codes = semantic_codes.shape[0]\n",
    "#print(semantic_codes.shape)\n",
    "sementic_n_tokens = 9000\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",
    "\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",
    "    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 (oooh)\n",
    "Turkey time, it's a country rhyme\n",
    "\n",
    "[Verse]\n",
    "Well, it's that time of year, with family near\n",
    "Gatherin' 'round the table, filled with cheer\n",
    "Turkey's roasting, gravy's flowin'\n",
    "Mashed potatoes piled high, it's Thanksiving time\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 (oooh)\n",
    "Turkey time, it's a country rhyme\n",
    "\n",
    "[Bridge]  \n",
    "Turkey time is here again\n",
    "Turkey time is here again\n",
    "Turkey time is here again\n",
    "Turkey tiiiiime!\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 (oooh)\n",
    "Turkey time, it's a country rhyme\n",
    "\n",
    "Turkey time, it's a country rhyme\n",
    "\n",
    "[guitar, outro]\n",
    "\n",
    "[outro]\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 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 = \"cat-instrumental\"\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 + max_dur_s\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",
    "\n",
    "Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n",
    "Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat CatCat Cat Cat Cat Cat Cat Cat Cat Cat Cat CatCat 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",
    "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",
    "\n",
    "Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n",
    "Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat CatCat Cat Cat Cat Cat Cat Cat Cat Cat Cat CatCat 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 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",
    "[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'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]\n",
    "    \"\"\".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 == \"cat-instrumental\":\n",
    "    gen_id = \"5fc87a9f-03a5-44c6-8ff4-ffd64772e582\"\n",
    "    tags = [\"jazz\"]\n",
    "    lyrics = \"\"\n",
    "    start_s = 0.0\n",
    "    end_s = start_s + max_dur_s\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",
    "    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",
    "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\"\"\".strip()\n",
    "    start_s = 0.0\n",
    "    end_s = start_s + max_dur_s\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_s = 0.0\n",
    "end_s = 360.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 = 256\n",
    "n_tokens_memmap_used = n_valid_semantic_tokens * 4\n",
    "n_tokens_memmap = 24000\n",
    "#mask = torch.ones(n_tokens_memmap).bool().cuda()\n",
    "#mask[n_tokens_memmap_used:] = False\n",
    "\n",
    "print(semantic_codes.shape)\n",
    "if semantic_codes.shape[0] < semantic_n_tokens: # pad with 4000\n",
    "    semantic_codes = torch.cat([semantic_codes, semantic_codes.new_full((semantic_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(n_tokens_memmap) * 4000).long().cuda()\n",
    "\n",
    "cfg_scales = [4.0]\n",
    "seeds = [np.random.randint(0, 2**32 - 1) for n in range(3)]\n",
    "\n",
    "#tags = [\"Jazz\", \"Piano\", \"High-fidelity\"]\n",
    "#tags = [\"drums\", \"bass\", \"synth\", \"edm\", \"hard\"]#\n",
    "#tags = [\"piano\"]\n",
    "\n",
    "lyric_phonemes = \"\"\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[:semantic_n_tokens],\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(\"full pred_zq\", pred_zq.shape)\n",
    "            pred_zq = pred_zq[...,:n_tokens_memmap_used]\n",
    "            print(\"valid 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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_steps = 100\n",
    "n_tokens_memmap = 36000\n",
    "\n",
    "out_dir = f\"outputs/val-{ckpt_name}\"\n",
    "os.makedirs(out_dir, exist_ok=True)\n",
    "\n",
    "semantic_codes_pad = (torch.ones(9000) * 4000).long().cuda()\n",
    "lyrics_empty = \"\"\"\"\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",
    "#lyrics_empty = \"Country roads, take me home, to the place I belong, West Virginia, mountain mama, take me home, country roads.\"\n",
    "#tags_empty = [\"country\"]\n",
    "#tags_empty = [\"jazz\", \"pop\", \"female vocal\", \"ballad\"]\n",
    "#tags_empty = [\"rap\", \"hip-hop\", \"bass\"]\n",
    "tags_empty = [\"pop\", \"r&b\", \"indie\"]\n",
    "\n",
    "#yrics_empty = lyrics\n",
    "#tags_empty = tags\n",
    "print(tags_empty)\n",
    "print(lyrics_empty)\n",
    "\n",
    "#lyrics_empty = \"\"\n",
    "#tags_empty = []\n",
    "\n",
    "#lyrics_empty = \"\"\"\"\"\"\n",
    "#tags_empty = [\"southern rock\", \"country\", \"folk\"]\n",
    "\n",
    "num_seeds = 1\n",
    "cfg_scales = [6.0]\n",
    "seeds = [np.random.randint(0, 2**32 - 1) for _ in range(num_seeds)]\n",
    "\n",
    "for n in range(len(cfg_scales)):\n",
    "    for seed in seeds:\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_pad,\n",
    "                tags_empty,\n",
    "                lyrics_empty,\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-3m-sde\",\n",
    "                #init_audio=latents,\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"
   ]
  }
 ],
 "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
}
