{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!echo $HOSTNAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !aws s3 cp s3://suno-data/georg/tmp/model.pt /app/suno/tmp/2b.pt\n",
    "# !aws s3 cp s3://suno-data/georg/tmp/model_c25.pt /app/suno/tmp/2b_c25.pt\n",
    "# !aws s3 cp s3://suno-data/georg/tmp/model_s250c1.pt /app/suno/tmp/2b_s250c1.pt\n",
    "# !aws s3 cp s3://suno-data/georg/tmp/model_s250c5.pt /app/suno/tmp/2b_s250c5.pt\n",
    "# !aws s3 cp s3://suno-data/georg/tmp/model_s250c25.pt /app/suno/tmp/2b_s250c25.pt\n",
    "\n",
    "# !aws s3 cp s3://suno-data/georg/tmp/model_noSem_s250c25.pt /app/suno/tmp/2b_noSem_s250c25.pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ## fix 13b\n",
    "# from collections import OrderedDict\n",
    "# import os\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
    "# import torch\n",
    "# sd = torch.load(\"/app/suno/checkpoints/2024-05-09_14-19-24/last_ckpt_infer.pt\")\n",
    "# sd[\"model_args\"][\"embed_scale_factor\"] = 1.0\n",
    "# sd[\"model_args\"][\"activation_f\"] = \"gelu\"\n",
    "# sd[\"model_args\"][\"use_qk_norm\"] = False\n",
    "# del sd[\"model_args\"][\"attention_type\"]\n",
    "# sd[\"model\"] = OrderedDict({k: v.to(torch.bfloat16) for k, v in sd[\"model\"].items()})\n",
    "# torch.save(sd, \"/app/suno/checkpoints/2024-05-09_14-19-24/last_ckpt_infer__fix.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !aws s3 cp /app/suno/checkpoints/2024-05-09_14-19-24/last_ckpt_infer__fix.pt s3://suno-data/georg/tmp/model_13b.pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.gpt.chirp_v2_5 import (\n",
    "    preload_codec_models,\n",
    "    _get_model_if_needed,\n",
    "    GenerationConfig,\n",
    "    codec_decode_stream_to_full_audio,\n",
    "    \n",
    ")\n",
    "from suno_utils.gpt.generation_engine import (\n",
    "    make_prompt,\n",
    "    align_codes,\n",
    "    make_request,\n",
    ")\n",
    "from suno_utils.gpt.engine import Engine, Request, Job\n",
    "from suno_utils.gpt.prompt import Prompt\n",
    "\n",
    "N_BATCH = 2\n",
    "MAX_STREAMS = N_BATCH * 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = preload_codec_models(\"/app/suno/models/chirp_v2/dac_2c_25x12.pt\")\n",
    "gpt_ckpt_path = _get_model_if_needed(\n",
    "#     \"s3://suno-data/georg/checkpoints/chirp_v2_5/7b_ft.pt\"\n",
    "#     \"/app/suno/checkpoints/2024-05-09_14-19-24/last_ckpt_infer__fix.pt\"  # 13b\n",
    "    #\"/app/suno/checkpoints/2024-05-19_11-34-42/last_ckpt_infer.pt\"  # 30b\n",
    "  # \"/app/suno/checkpoints/2024-07-29_05-51-55/last_ckpt_infer.pt\"\n",
    "  #\"/app/suno/checkpoints/2024-08-11_18-28-59/last_ckpt_infer.pt\" # genius filtered finetuning\n",
    "  #\"/app/suno/checkpoints/2024-08-05_02-39-41/last_ckpt_infer.pt\" # v-12 dpo-ed model\n",
    "  \"/app/suno/data/dpo/models/model_30b_fix_ft2_20k.pt\" # base 30b\n",
    "  #\"/app/suno/checkpoints/2024-08-13_21-11-38/last_ckpt_infer.pt\" # 30b finetune datasetname control tag\n",
    "#     \"/app/suno/tmp/2b_plus.pt\"\n",
    "#     \"/app/suno/tmp/2b_musicfm.pt\"\n",
    "#     \"/app/suno/tmp/2b.pt\",\n",
    "#     \"/app/suno/tmp/2b_c25.pt\",\n",
    "#     \"/app/suno/tmp/2b_s250c1.pt\",\n",
    "#     \"/app/suno/tmp/2b_s250c5.pt\",\n",
    "#     \"/app/suno/tmp/2b_s250c25.pt\",\n",
    "#     \"/app/suno/tmp/2b_noSem_s250c25.pt\",\n",
    ")\n",
    "engine = Engine(\n",
    "    gpt_ckpt_path,\n",
    "    \"/app/suno/data/dpo/models/tokenizer_60k.json\",\n",
    "    max_sequences=MAX_STREAMS,\n",
    "    compile=False,\n",
    ")\n",
    "model = engine.model\n",
    "cfg = model.config\n",
    "tokenizer = engine.tokenizer\n",
    "\n",
    "print(cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\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",
    "text_love = \"\"\"\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",
    "\n",
    "text_train = \"\"\"\n",
    "[Verse 1]\n",
    "I have a demon in me\n",
    "It haunts me in my dreams\n",
    "I haven't slept in a week\n",
    "Could use some THC\n",
    "Forgive me, I'm impatient\n",
    "Nothing can calm me down\n",
    "These monsters think I'm famous\n",
    "I'm just a sad clown\n",
    "\n",
    "[Pre-Chorus]\n",
    "One more drop of poison for me\n",
    "Right into my chest\n",
    "Dirty little secrets I'll keep\n",
    "You can keep the rest\n",
    "\n",
    "[Chorus]\n",
    "I gotta get low just to get high\n",
    "It feels good, but it doesn't feel right\n",
    "I die slow just to know I'm alive\n",
    "I gotta get low just to get high (High)\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "#     text_tags=\"dark epic trailer music, violin, vocalizations, female\",\n",
    "     #text_tags=\"pop, power ballad, female\",\n",
    "     text_tags=\"indie, indie rock, alternative, rock, alternative rock\",\n",
    "   #text_tags=\"\",\n",
    "    #text_tags=\"female vocal, powerful, acapella, solo, sad, opera, ballad, vocalizations\",\n",
    "    #text_tags=\"bluegrass, female vocal, guitar, acoustic\",\n",
    "    #text_tags=\"pop, R&B, energetic, bass\",\n",
    "#     text_tags=\"sea shanty, sea shanty, sea shanty\",\n",
    "    \n",
    "    cfg_coef=1.3,\n",
    "    cfg_coef_tags=2.0,\n",
    "    n_repeat_tags=3,\n",
    " # start:0;vocals:intro\n",
    "    text_start_control_tags=\"{start:0;vocals:intro}\",\n",
    "    text_end_control_tags=\"{end}\",\n",
    "    #cfg_coef_neg_tags=-2,\n",
    "    #text_neg_tags=\"repetitive, loop, noise, noisy\",\n",
    "    \n",
    "     #temp_semantic=1.0,\n",
    "     #temp_coarse=1.0,\n",
    "#     top_k_semantic=None,\n",
    "#     top_k_coarse=None,\n",
    "     #top_p_semantic=0.95,\n",
    "     #top_p_coarse=0.98,\n",
    "    \n",
    "    n_batch=1,\n",
    "    #min_eos_p=0.1,\n",
    "#     allow_eos=False,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=60,\n",
    ")\n",
    "\n",
    "# request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "# in_arr = request.streams[0].prompt\n",
    "# prompt = Prompt(gconf, cfg, tokenizer)\n",
    "# prompt.visualize(in_arr)\n",
    "# for n in range(4):\n",
    "#     print(\"--\", n, \"--\")\n",
    "#     print(tokenizer.decode(request.streams[n].prompt[0][:-1].tolist()))\n",
    "#     print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "requests = [\n",
    "    make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "# print(requests[0])\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "\n",
    "for job in jobs:\n",
    "    stream = engine.token_generator(job)\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Covers generation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import uuid\n",
    "import IPython\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "from suno_utils.utils.text import read_jsonl\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",
    "#from suno_utils.tasks.dac_2c_12cb import (\n",
    "#    preload_models as preload_codec_models,\n",
    "#    encode as codec_encode,\n",
    "#    encode_files as codec_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",
    "#)\n",
    "\n",
    "#_ = preload_codec_models(\n",
    "#    \"s3://suno-data/georg/models/codec/dac_2c_25x12.pt\",\n",
    "#    device=\"cuda\",\n",
    "#)\n",
    "\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.gpt.chirp_v2_5 import (\n",
    "    codec_encode, codec_decode, preload_semantic_models, semantic_encode\n",
    ")\n",
    "_ = preload_semantic_models(\n",
    "    checkpoint_filepath=\"/app/suno/models/chirp_v2/mert_25.pt\",\n",
    "    centroids_filepath=\"/app/suno/models/chirp_v2/mert_25_2x4k.npy\",\n",
    "    device=\"cpu\",\n",
    ")\n",
    "\n",
    "def process_audio(audio):\n",
    "    #audio = audio.normalize_volume(-16)\n",
    "    sem_arr = semantic_encode(audio, device=\"cpu\")\n",
    "    coarse_arr = codec_encode(audio)\n",
    "    n_frames = min(sem_arr.shape[0], coarse_arr.shape[0])\n",
    "    sem_arr = sem_arr[:n_frames, : 1]\n",
    "    coarse_arr = coarse_arr[:n_frames, : 12]\n",
    "\n",
    "    a_arr = np.concatenate([sem_arr, coarse_arr], axis=-1)\n",
    "    return a_arr\n",
    "\n",
    "RATE_HZ = 25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first read in metadata \n",
    "metas = read_jsonl(\"/home/christian/code/christian/metadata/genius_hq_metas_filtered.jsonl\", progress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_covers = 97\n",
    "print(f\"Found {len(metas)} songs...\")\n",
    "# do only english songs\n",
    "metas = [meta for meta in metas if meta[\"lang\"] == \"en\"]\n",
    "print(f\"Found {len(metas)} english songs...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "covers_dir = \"covers-20240816\"\n",
    "os.makedirs(f\"/home/christian/code/christian/notebooks/outputs/{covers_dir}\", exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "for n in range(1):\n",
    "    uid = uuid.uuid4()\n",
    "    rand_meta_idx = np.random.randint(len(metas))\n",
    "    meta = metas[rand_meta_idx]\n",
    "    text = meta[\"lyrics\"]\n",
    "\n",
    "    # tags of the source song\n",
    "    text_tags = meta[\"tags_text\"]\n",
    "    text_tags = [tag.replace(\"Genius\", \"\") for tag in text_tags]\n",
    "    text_tags = [tag.strip() for tag in text_tags]\n",
    "    print(\"source tags:\", text_tags)\n",
    "\n",
    "    # download audio from s3\n",
    "    filename = os.path.basename(meta[\"audio_filepath\"])\n",
    "    os.system(f\"aws s3 cp {meta['audio_filepath']} /home/christian/code/christian/tmp2/{filename}\")\n",
    "    audio, sr = torchaudio.load(f\"/home/christian/code/christian/tmp2/{filename}\")\n",
    "\n",
    "    # crop first 30 sec\n",
    "    source_dur_sec = 60.0\n",
    "    start_idx = np.random.randint(0, int(sr*10))\n",
    "    end_idx = start_idx + int(sr*source_dur_sec)\n",
    "    start_sec = start_idx / sr\n",
    "    audio = audio[:, start_idx:end_idx]\n",
    "    IPython.display.display(IPython.display.Audio(data=audio, rate=sr))\n",
    "    torchaudio.save(f\"/home/christian/code/christian/notebooks/outputs/{covers_dir}/{uid}-source.mp3\", audio, sr)\n",
    "\n",
    "    # get random tags from a different song\n",
    "    rand_meta_idx = np.random.randint(len(metas))\n",
    "    target_meta = metas[rand_meta_idx]\n",
    "    target_tags = target_meta[\"tags_text\"]\n",
    "    target_tags = [tag.replace(\"Genius\", \"\") for tag in target_tags]\n",
    "    target_tags = [tag.strip() for tag in target_tags]\n",
    "    print(\"target tags:\", target_tags)\n",
    "\n",
    "    # encode the audio\n",
    "    cover_arr = process_audio(Audio.from_s3(meta[\"audio_filepath\"], n_channels=2))\n",
    "    print(cover_arr.shape)\n",
    "\n",
    "    num_tokens = int(RATE_HZ * source_dur_sec)\n",
    "    start_token = int(RATE_HZ * start_sec)\n",
    "    end_token = start_token + num_tokens\n",
    "    cover_arr = cover_arr[start_token:end_token,:]\n",
    "    print(cover_arr.shape)\n",
    "\n",
    "    # randomly sample from 1.0 to 6.0\n",
    "    cfg_coef = np.random.choice([1.3, 6.0])\n",
    "    cfg_coef_tags = np.random.choice([1.3, 6.0])\n",
    "\n",
    "    gconf = GenerationConfig(\n",
    "        text=text,\n",
    "        text_tags=target_tags,\n",
    "        cover_arr=cover_arr,\n",
    "        cfg_coef=1.3,\n",
    "        cfg_coef_tags=cfg_coef_tags,\n",
    "        n_repeat_tags=3,\n",
    "        text_start_control_tags=\"{start}\",\n",
    "        text_end_control_tags=\"{end}\",\n",
    "        n_batch=1,\n",
    "        min_eos_p=0.1,\n",
    "        min_text_offset=0,\n",
    "        eos_pad_duration_s=0,\n",
    "        max_gen_duration_s=60,\n",
    "    )\n",
    "\n",
    "    example_meta = {\n",
    "        \"uid\": str(uid),\n",
    "        \"source_start_sec\": start_sec,\n",
    "        \"source_end_sec\": start_sec + 60.0,\n",
    "        \"source_tags\" : text_tags,\n",
    "        \"source_s3_filepath\" : meta[\"audio_filepath\"],\n",
    "        \"target_tags\" : target_tags,\n",
    "        \"lyrics\" : text,\n",
    "        \"cfg_coef\" : cfg_coef,\n",
    "        \"cfg_coef_tags\" : cfg_coef_tags,\n",
    "    }\n",
    "\n",
    "    with open(f\"/home/christian/code/christian/notebooks/outputs/{covers_dir}/{uid}-meta.json\", \"w\") as f:\n",
    "        json.dump(example_meta, f)\n",
    "\n",
    "    requests = [\n",
    "        make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "        for i in range(N_BATCH)\n",
    "    ]\n",
    "    # print(requests[0])\n",
    "    jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "\n",
    "    for job_idx, job in enumerate(jobs):\n",
    "        stream = engine.token_generator(job)\n",
    "\n",
    "        # stack stream\n",
    "        #stream_arr = torch.stack(stream_arr)\n",
    "        #print(stream_arr.shape)\n",
    "\n",
    "        # convert to numpy and save as npz\n",
    "        #stream_arr = stream_arr.cpu().numpy()\n",
    "        #np.savez(f\"/home/christian/code/christian/notebooks/outputs/{covers_dir}/{uid}-cover-{job_idx}.npz\", stream_arr)\n",
    "        \n",
    "        audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "        audio.play()\n",
    "\n",
    "        audio.write_mp3(f\"/home/christian/code/christian/notebooks/outputs/{covers_dir}/{uid}-cover-{job_idx}.mp3\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def safe_iterate(generator):\n",
    "    try:\n",
    "        yield from generator\n",
    "    except RuntimeError as e:\n",
    "        if str(e) == \"generator raised StopIteration\":\n",
    "            return\n",
    "        else:\n",
    "            raise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Intro]\n",
    "I can feel the heat risin'\n",
    "Everything is on fire\n",
    "Today's a painful re-\n",
    "Minder of why\n",
    "It can only get brighter\n",
    "The further you put it behind ya\n",
    "But right now I'm on the\n",
    "Inside lookin' out, ‘cause—\n",
    "\n",
    "[Chorus]\n",
    "I'm standin' in the flames\n",
    "And it’s a beautiful kind of pain\n",
    "Settin' fire to yesterday\n",
    "To find the light, find the light, find the light\n",
    "Standin' in the flames\n",
    "And it’s a beautiful kind of pain\n",
    "Settin' fire to yesterday\n",
    "Find the light, find the light, find the light\n",
    "\n",
    "[Verse 1]\n",
    "Yesterday was the tornado warning, today's like the morning after\n",
    "Your world is torn in half, you wake in its wake\n",
    "To start the mourning process and rebuilding, you're still a work in progress\n",
    "Today's a whole new chapter, it's like an enormous asthma\n",
    "Thunderstorm has passed ya, you weathered it and poked its\n",
    "Eye out with the thornbush that ya used to smell the roses\n",
    "Stopped to inhale, can't even tell your nose is stuffed\n",
    "So focused on the bright side then you floor the gas pedal\n",
    "And hit the corner faster, more assertive, never looking back\n",
    "May hit the curb, but every day's a new learning curve as ya\n",
    "Steer through life, sometimes you might not wanna swerve but you have to\n",
    "To avert a disaster, lucky no permanent damage\n",
    "‘Cause they hurt you so ba\n",
    "\"\"\"\n",
    "\n",
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=\"Rap, In English, USA, Midwest Rap, Detroit Rap, Alternative, Hip-Hop\",\n",
    "\n",
    "    cfg_coef=1.3,\n",
    "    cfg_coef_tags=2,\n",
    "    n_repeat_tags=3,\n",
    "\n",
    "#     temp_semantic=0.8,\n",
    "#     temp_coarse=1,\n",
    "#     top_k_semantic=1000,\n",
    "#     top_k_coarse=100,\n",
    "#     top_p_semantic=None,\n",
    "#     top_p_coarse=None,\n",
    "    \n",
    "#     text_start_control_tags=\"{start;start:0;vocals:intro}\",\n",
    "#     text_end_control_tags=\"{end}\",\n",
    "#     cfg_coef_neg_tags=-2,\n",
    "#     text_neg_tags=\"repetitive, loop\",\n",
    "\n",
    "    n_batch=1,\n",
    "    min_eos_p=0.1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=60*2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "requests = [\n",
    "    make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "# print(requests[0])\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "\n",
    "for job in jobs:\n",
    "    stream = engine.token_generator(job)\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom Input Array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.gpt.chirp_v2_5 import (\n",
    "    codec_encode, codec_decode, preload_semantic_models, semantic_encode\n",
    ")\n",
    "_ = preload_semantic_models(\n",
    "    checkpoint_filepath=\"/app/suno/models/chirp_v2/mert_25.pt\",\n",
    "    centroids_filepath=\"/app/suno/models/chirp_v2/mert_25_2x4k.npy\",\n",
    "    device=\"cpu\",\n",
    ")\n",
    "\n",
    "# # MusicFM\n",
    "# from suno_utils.tasks.musicfm_v2 import (\n",
    "#     preload_models as preload_semantic_models, \n",
    "#     encode as semantic_encode\n",
    "# )\n",
    "# _ = preload_semantic_models(\n",
    "#     centroids_filepath=\"s3://suno-data/minz/models/musicfm_concat_centroids_2s_8192.npy\",\n",
    "# )\n",
    "\n",
    "def process_audio(audio, cfg):\n",
    "    audio = audio.normalize_volume(-16)\n",
    "    sem_arr = semantic_encode(audio, device=\"cpu\")\n",
    "    coarse_arr = codec_encode(audio)\n",
    "    n_frames = min(sem_arr.shape[0], coarse_arr.shape[0])\n",
    "    sem_arr = sem_arr[:n_frames, : cfg.semantic_n_codebooks]\n",
    "    coarse_arr = coarse_arr[:n_frames, : cfg.coarse_n_codebooks]\n",
    "\n",
    "    a_arr = np.concatenate([sem_arr, coarse_arr], axis=-1)\n",
    "    return a_arr\n",
    "\n",
    "def load_audio(fp):\n",
    "    return Audio.from_file(fp, n_channels=2, sample_rate=48_000, byte_width=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = Audio.from_file(\"tmp/50cent.wav\", n_channels=2, sample_rate=48_000, byte_width=2)  # 50 cent\n",
    "audio = audio.get_segment(from_s=24.99, to_s=45)\n",
    "# audio = Audio.from_file(\"tmp/eminem_doctor.wav\", n_channels=2, sample_rate=48_000, byte_width=2) # eminem doctor\n",
    "history_text = \"\"\"\n",
    "Welcome to the candy shop\n",
    "\n",
    "Yeah, uh-huh\n",
    "So seductive\n",
    "\n",
    "I take you to the candy shop\n",
    "I let you lick the lollipop\n",
    "Go 'head, girl, don't you stop\n",
    "\"\"\"\n",
    "text = \"\"\"\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",
    "Country roads, take me home\n",
    "To the place I belong\n",
    "West Virginia, mountain mama\n",
    "Take me home, country roads\n",
    "\"\"\"\n",
    "in_history_arr = process_audio(audio, cfg)\n",
    "audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# codec_decode(codec_encode(audio)).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "#     text_tags=\"bluegrass female\",\n",
    "    history_arr=in_history_arr,\n",
    "    history_text=history_text,\n",
    "    n_batch=1,\n",
    "    cfg_coef=1.2,\n",
    "    cfg_coef_tags=0,\n",
    "    max_gen_duration_s=40,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    ")\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "for job in engine.run_request(requests, tqdm_enabled=True):\n",
    "    stream = engine.token_generator(job)\n",
    "\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Oracle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# audio = load_audio(\"tmp/50cent.wav\").get_segment(from_s=24.99, to_s=45)\n",
    "# text = \"\"\"\n",
    "# Welcome to the candy shop\n",
    "\n",
    "# Yeah, uh-huh\n",
    "# So seductive\n",
    "\n",
    "# I take you to the candy shop\n",
    "# I let you lick the lollipop\n",
    "# Go 'head, girl, don't you stop\n",
    "# \"\"\"\n",
    "# in_oracle_arr = process_audio(audio, cfg)\n",
    "# audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = load_audio(\"tmp/halo.wav\").get_segment(from_s=49.99, to_s=70)\n",
    "text = \"\"\"\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",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "\"\"\"\n",
    "in_oracle_arr = process_audio(audio, cfg)\n",
    "# in_oracle_arr[:,:1] = cfg.semantic_mask_token\n",
    "audio.normalize_volume().play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# codec_decode(codec_encode(audio)[:,:6]).normalize_volume().play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "#     text_tags=\"bluegrass female\",\n",
    "    oracle_arr=in_oracle_arr[:,:1],\n",
    "    n_batch=1,\n",
    "    cfg_coef=1, #.3,\n",
    "    cfg_coef_tags=0,\n",
    "    cfg_coef_neg_tags=0,\n",
    "    max_gen_duration_s=20,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    ")\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "for job in engine.run_request(requests, tqdm_enabled=True):\n",
    "    stream = engine.token_generator(job)\n",
    "\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Covers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# audio = Audio.from_file(\"tmp/georg_country.mp3\", n_channels=2).get_segment(to_s=60.01)\n",
    "# audio = Audio.from_file(\"tmp/georg_bday.mp3\", n_channels=2).get_segment(to_s=60.01)\n",
    "audio = load_audio(\"tmp/halo.wav\").get_segment(from_s=49.99, to_s=70)\n",
    "audio.play()\n",
    "in_cover_arr = process_audio(audio, cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# text = \"\"\"\n",
    "# Happy birthday to you\n",
    "# Happy birthday to you\n",
    "# Happy birthday dear Margu\n",
    "# Happy birthday to you\n",
    "\n",
    "# Happy birthday to you\n",
    "# Happy birthday to you\n",
    "# Happy birthday dear Keenan\n",
    "# Happy birthday to you\n",
    "\n",
    "# Happy birthday to you\n",
    "# Happy birthday to you\n",
    "# Happy birthday dear Mikey\n",
    "# Happy birthday to you\n",
    "\n",
    "# Happy birthday to you\n",
    "# Happy birthday to you\n",
    "# Happy birthday dear Suno\n",
    "# Happy birthday to you\n",
    "\n",
    "# Happy birthday to you\n",
    "# Happy birthday to you\n",
    "# I don't know what I'm gonna do next\n",
    "# Happy birthday to you\n",
    "# \"\"\"\n",
    "\n",
    "# text = \"\"\"\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",
    "# [Verse 2]\n",
    "# All my memories gather 'round her\n",
    "# Miner's lady, stranger to blue water\n",
    "# Dark and dusty, painted on the sky\n",
    "# Misty taste of moonshine, teardrop in my eye\n",
    "# \"\"\"\n",
    "\n",
    "text = \"\"\"\n",
    "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",
    "\"\"\"\n",
    "\n",
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=\"bluegrass, female\",\n",
    "#     text_tags=\"pop, power ballad, female\",\n",
    "    cover_arr=in_cover_arr,\n",
    "    n_batch=1,\n",
    "    cfg_coef=1.3,\n",
    "    cfg_coef_tags=2,\n",
    "    max_gen_duration_s=40,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    ")\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "\n",
    "assert(N_BATCH * 4 <= 16)\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "for job in engine.run_request(requests, tqdm_enabled=True):\n",
    "    stream = engine.token_generator(job)\n",
    "\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Artist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = Audio.from_file(\"audios/georg_country.mp3\", n_channels=2).get_segment(to_s=5.01)\n",
    "# audio = load_audio(\"tmp/halo.wav\").get_segment(from_s=49.99, to_s=70)\n",
    "audio.play()\n",
    "in_artist_arr = process_audio(audio, cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# text = \"\"\"\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",
    "# [Verse 2]\n",
    "# All my memories gather 'round her\n",
    "# Miner's lady, stranger to blue water\n",
    "# Dark and dusty, painted on the sky\n",
    "# Misty taste of moonshine, teardrop in my eye\n",
    "# \"\"\"\n",
    "\n",
    "text = \"\"\"\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",
    "\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",
    "\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",
    "\"\"\"\n",
    "\n",
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=\"bluegrass\",\n",
    "    artist_arr=in_artist_arr,\n",
    "    n_batch=1,\n",
    "    cfg_coef=1.5,\n",
    "    cfg_coef_tags=2,\n",
    "    max_gen_duration_s=40,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    ")\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "\n",
    "assert(N_BATCH * 4 <= 16)\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "for job in engine.run_request(requests, tqdm_enabled=True):\n",
    "    stream = engine.token_generator(job)\n",
    "\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Infill"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# audio = Audio.from_file(\"tmp/georg_country.mp3\", n_channels=2).get_segment(to_s=5.01)\n",
    "audio = load_audio(\"audios/country_road.mp3\").get_segment(from_s=0, to_s=20)\n",
    "audio.play()\n",
    "in_infill_arr = process_audio(audio, cfg)\n",
    "in_history_arr = in_infill_arr[:25*5,:].copy()\n",
    "in_future_arr = in_infill_arr[-25*5:,:].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Verse 1]\n",
    "Younger than the mountains, growing like a breeze\n",
    "AH AH AH AH AH\n",
    "\n",
    "[Chorus]\n",
    "Test Test Test Test Test\n",
    "To the place I belong\n",
    "West Virginia, mountain mama\n",
    "Take me home, country roads\n",
    "\"\"\"\n",
    "\n",
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    history_arr=in_history_arr,\n",
    "    future_arr=in_future_arr,\n",
    "    n_batch=1,\n",
    "    cfg_coef=1.3,\n",
    "    cfg_coef_tags=0,\n",
    "    max_gen_duration_s=15,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    ")\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)\n",
    "in_arr = request.streams[0].prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "\n",
    "assert(N_BATCH * 4 <= 16)\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "for job in engine.run_request(requests, tqdm_enabled=True):\n",
    "    stream = engine.token_generator(job)\n",
    "    audio = codec_decode(np.concatenate([\n",
    "        in_history_arr[:,-12:],\n",
    "        torch.stack(list(align_codes(stream, cfg))).detach().cpu().numpy()[:,-12:],\n",
    "        in_future_arr[:,-12:],\n",
    "    ], axis=0))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import numpy as np\n",
    "# in_arr = requests[0].streams[0].prompt\n",
    "# generated = np.stack([t.cpu().numpy() for t in job.generated_tokens]).T\n",
    "# print(generated.shape)\n",
    "# arr = np.concatenate([in_arr, generated], axis=1)\n",
    "# prompt.visualize(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_jsonl(filepath, parse_idx_set=None):\n",
    "    data = []\n",
    "    with open(filepath) as f:\n",
    "        line_idx = 0\n",
    "        for line in f:\n",
    "            line = line.strip()\n",
    "            if len(line) == 0:\n",
    "                continue\n",
    "            if parse_idx_set is not None and line_idx not in parse_idx_set:\n",
    "                data.append(None)\n",
    "                line_idx += 1\n",
    "                continue\n",
    "            m = json.loads(line)\n",
    "            data.append(m)\n",
    "            line_idx += 1\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_loss = read_jsonl(\"/app/suno/data/dpo/2b_before_recode_v0/eval_val_loss/loss.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(eval_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_loss[1193]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[x for x in eval_loss if x[\"idx\"] == 311]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "c = Counter()\n",
    "for i in eval_loss:\n",
    "    c[i[\"idx\"]] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(c))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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": 4
}
