{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0d83737",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !nvidia-smi\n",
    "# !echo $HOSTNAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74d60c8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from suno_utils.gpt.generation import GenerationConfig\n",
    "from suno_utils.gpt.engine import Engine\n",
    "from suno_utils.gpt.generation_engine import make_request\n",
    "\n",
    "N_BATCH = 4\n",
    "\n",
    "engine = Engine(\n",
    "    #     \"/app/suno/checkpoints/2025-01-26_01-35-19/last_ckpt_infer.pt\",  # skip\n",
    "    # \"/app/suno/checkpoints/2025-01-30_14-26-14/last_ckpt_infer.pt\",  # base long\n",
    "    \"/app/suno/checkpoints/2025-02-04_16-45-51/last_ckpt_infer.pt\", # 13b dpo\n",
    "    #     \"/app/suno/checkpoints/2025-01-18_18-03-13/last_ckpt_infer.pt\",  # base\n",
    "    #     \"/app/suno/checkpoints/2024-12-05_22-01-59/last_ckpt_infer.pt\",  # (old chef)\n",
    "    \"/app/suno/models/chirp_v2/tokenizer_60k.json\",\n",
    "    max_sequences=4 * N_BATCH,\n",
    "    compile=False,\n",
    ")\n",
    "cfg = engine.model.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6a9f401",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "sys.path.insert(0, \"/home/tony/Work/neon/sunoDiff/\")\n",
    "\n",
    "diff_model_fp = \"/app/suno/data/dpo/models/diff_vae_25_peaq_v4_jan28.pt\"  # prod\n",
    "# diff_model_fp = \"/app/suno/checkpoints/2025-01-19_05-10-13_s1983/step_3000_infer.pt\"  # best\n",
    "\n",
    "from generation import preload_models as preload_diff_models, generate\n",
    "\n",
    "_ = preload_diff_models(\n",
    "    tokenizer_filepath=\"/app/suno/data/dpo/models/tokenizer_60k.json\",\n",
    "    semantic_model_filepath=\"/app/suno/data/dpo/models/mert_25.pt\",\n",
    "    semantic_clusters_filepath=\"/app/suno/data/dpo/models/mert_25_2x4k.npy\",\n",
    "    codec_filepath=\"/app/suno/data/dpo/models/25hz_vae_peaq_kl_0.005.pth\",\n",
    "    dit_model_filepath=diff_model_fp,\n",
    "    model_type=\"prefix\",\n",
    "    weights_precision=torch.bfloat16,\n",
    "    compile=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ec28b37",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # [Verse 1: female vocals, soft, beautiful]\n",
    "# # [Pre-Chorus: male, aggressive, shouting]\n",
    "# # [Chorus: female & male, together, harmonies]\n",
    "\n",
    "# text = \"\"\"\n",
    "# [Verse 1]\n",
    "# You are my fire\n",
    "# The one desire\n",
    "# Believe when I say\n",
    "# I want it that way\n",
    "\n",
    "# [Verse 2]\n",
    "# But we are two worlds apart\n",
    "# Can't reach to your heart\n",
    "# When you say\n",
    "# That I want it that way\n",
    "\n",
    "# [Chorus]\n",
    "# Tell me why\n",
    "# Ain't nothin' but a heartache\n",
    "# Tell me why\n",
    "# Ain't nothing but a mistake\n",
    "# Tell me why\n",
    "# I never wanna hear you say\n",
    "# I want it that way\n",
    "\n",
    "# [Verse 3]\n",
    "# Am I your fire?\n",
    "# Your one desire\n",
    "# Yes, I know it's too late\n",
    "# But I want it that way\n",
    "\n",
    "# [Chorus]\n",
    "# Tell me why\n",
    "# Ain't nothin' but a heartache\n",
    "# Tell me why\n",
    "# Ain't nothin' but a mistake\n",
    "# Tell me why\n",
    "# I never wanna hear you say\n",
    "# I want it that way\n",
    "\n",
    "# [Bridge]\n",
    "# Now I can see that we've fallen apart\n",
    "# From the way that it used to be, yeah\n",
    "# No matter the distance, I want you to know\n",
    "# That deep down inside of me\n",
    "\n",
    "# [Verse 4]\n",
    "# You are my fire\n",
    "# The one desire\n",
    "# You are (You are, you are, you are)\n",
    "# Don't wanna hear you say\n",
    "\n",
    "# [Break]\n",
    "# Ain't nothin' but a heartache (Hey, yeah)\n",
    "# Ain't nothin' but a mistake (Don't wanna hear you say)\n",
    "# I never wanna hear you say (Oh, yeah)\n",
    "# I want it that way\n",
    "# [Chorus: Brian, All, Nick, AJ]\n",
    "# Tell me why\n",
    "# Ain't nothin' but a heartache\n",
    "# Tell me why\n",
    "# Ain't nothin' but a mistake\n",
    "# Tell me why\n",
    "# I never wanna hear you say (Don't wanna hear you say)\n",
    "# I want it that way\n",
    "# Tell me why\n",
    "# Ain't nothin' but a heartache\n",
    "# Ain't nothin' but a mistake\n",
    "# Tell me why\n",
    "# I never wanna hear you say (Never wanna hear you say, yeah)\n",
    "# I want it that way\n",
    "\n",
    "# [Outro]\n",
    "# 'Cause I want it that way\n",
    "# \"\"\"\n",
    "\n",
    "text = \"\"\"\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",
    "\n",
    "# tags = \"grunge alternative rock\"\n",
    "# tags = \"guitar violin drums\"\n",
    "# tags = \"A melancholic country song featuring prominent pedal steel guitar.\"\n",
    "\n",
    "# tags = \"Pop\"\n",
    "# tags = \"Gregorian Chant, Hymns, Christian, ethereal, female vocalist, medieval, atmospheric, spiritual, melodic, choral, peaceful, hypnotic, hymn, meditative, minimalistic, soothing, religious, soft, longing, dense, melancholic, mysterious, repetitive, acoustic\"\n",
    "# tags = \"Two-track instrumental featuring ambient nature sounds and rhythmic clanging, followed by sinister synthesizer music and buzzing electric circuit.\"\n",
    "# tags = \"Country, southern accent, deep south, raspy vocals\"\n",
    "# tags = \"Gospel\"\n",
    "# tags = \"Americana, Country, Contemporary Folk, lonely, sombre, introspective, longing, acoustic, melodic, bittersweet, pastoral, poetic, melancholic, male vocalist, depressive, pessimistic, dark, passionate\"\n",
    "# tags = \"Contemporary R&B, Pop Rap, Alternative R&B, Trap, nocturnal, female vocalist, aquatic, alcohol, male vocalist, sexual\"\n",
    "# tags = \"Pop Soul, Neo-Soul, Jazz Pop, Brill Building, Twist, drugs, female vocalist, rebellious, alcohol, rhythmic, depressive, introspective, passionate, sarcastic, melodic, hedonistic, anthemic\"\n",
    "# tags = \"power ballad, female vocals, male vocals, pop, rap, hip-hop, emotive, powerful, haunting\"\n",
    "# tags = \"Melodic Techno, Future Rave, Melodic House, rhythmic, atmospheric\"\n",
    "# tags = \"Dance-Pop, Electro House, Festival Progressive House, Electropop, Dutch House, passionate, female vocalist, triumphant, introspective, energetic, anthemic, rhythmic, sentimental\"\n",
    "# tags = \"passionate, raspy vocals, powerful delivery, female, pop\"\n",
    "# tags = \"Pop, female vocals, orchestra\"\n",
    "# tags = \"female vocals; This track start off quiet and slow featuring string instruments including a violin. It then builds into an energetic chorus with drums and trumpets.\"\n",
    "\n",
    "# text = \"\"\n",
    "# tags = \"orchestra, symphony\"\n",
    "# tags = \"orchestra, powerful trailer music, epic, dark, violin, female vocalizations\"\n",
    "# tags = \"The piece that opens with a soft, flute-led melody gently awakening the senses, conjuring the warmth of a sunrise over tranquil hills. The orchestra gradually swells with smooth strings and light percussion, evoking a blossoming sense of optimism and renewed energy. A graceful interplay of woodwinds and strings concludes on a delicate, uplifting note that captures the fresh promise of morning.\"\n",
    "# tags = \"A stirring, festive march led by bold brass fanfares and triumphant rhythms that command attention right from the start. Let snappy snare drums and buoyant woodwinds propel the music forward with an infectious energy. The piece culminates in a grand, celebratory finale that invites enthusiastic clapping and cheering.\"\n",
    "tags = \"epic film orchestral\"\n",
    "n_skip_semantic = 1\n",
    "# n_skip_semantic = 4\n",
    "\n",
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=tags,\n",
    "    cfg_coef=1.0,\n",
    "    cfg_coef_tags=2.0,\n",
    "    # cfg_coef_max_steps=None,\n",
    "    cfg_coef_tags_max_steps=250,\n",
    "    n_repeat_tags=1,\n",
    "    n_skip_semantic=n_skip_semantic,\n",
    "    #     text_start_control_tags=\"{duration:210}\",\n",
    "    #     cfg_coef_neg_tags=-1,\n",
    "    #     text_neg_tags=\"repetitive, loop\",\n",
    "    #     temp_semantic=0.92,\n",
    "    #     temp_coarse=0.9,\n",
    "    #     top_k_semantic=None,\n",
    "    #     top_k_coarse=None,\n",
    "    #     top_p_semantic=None,\n",
    "    #     top_p_coarse=None,\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=int(4 * 60 / n_skip_semantic),\n",
    "    random_seed=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3133ec0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "requests = [\n",
    "    make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "out_gpt = []\n",
    "for n, job in enumerate(jobs):\n",
    "    stream = engine.token_generator(job)\n",
    "    arr = torch.stack(list(stream))[:, 1]\n",
    "    if arr[-1] == 4000:\n",
    "        arr = arr[:-1]\n",
    "    print(f\"{round(arr.shape[-1]/25*n_skip_semantic)}s for track {n}\")\n",
    "    # do stuff incase skip\n",
    "    arr2 = (\n",
    "        torch.zeros(arr.shape[0] * n_skip_semantic, dtype=arr.dtype)\n",
    "        + cfg.semantic_pad_token\n",
    "    )\n",
    "    arr2[::n_skip_semantic] = arr\n",
    "    # add\n",
    "    out_gpt.append(arr2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19788eb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "for in_sem_arr in out_gpt:\n",
    "    generate(\n",
    "        in_sem_arr,\n",
    "        lyrics=text,\n",
    "        tags=tags,\n",
    "        text_cfg_coef=2.0,\n",
    "        steps=10,\n",
    "        seed=0,\n",
    "    ).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f64ec057",
   "metadata": {},
   "outputs": [],
   "source": [
    "tags = \"gregorian chant\"\n",
    "n_skip_semantic = 1\n",
    "# n_skip_semantic = 4\n",
    "\n",
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=tags,\n",
    "    cfg_coef=1.0,\n",
    "    cfg_coef_tags=2.0,\n",
    "    # cfg_coef_max_steps=None,\n",
    "    cfg_coef_tags_max_steps=250,\n",
    "    n_repeat_tags=1,\n",
    "    n_skip_semantic=n_skip_semantic,\n",
    "    #     text_start_control_tags=\"{duration:210}\",\n",
    "    #     cfg_coef_neg_tags=-1,\n",
    "    #     text_neg_tags=\"repetitive, loop\",\n",
    "    #     temp_semantic=0.92,\n",
    "    #     temp_coarse=0.9,\n",
    "    #     top_k_semantic=None,\n",
    "    #     top_k_coarse=None,\n",
    "    #     top_p_semantic=None,\n",
    "    #     top_p_coarse=None,\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=int(4 * 60 / n_skip_semantic),\n",
    "    random_seed=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cb96fc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "requests = [\n",
    "    make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "out_gpt = []\n",
    "for n, job in enumerate(jobs):\n",
    "    stream = engine.token_generator(job)\n",
    "    arr = torch.stack(list(stream))[:, 1]\n",
    "    if arr[-1] == 4000:\n",
    "        arr = arr[:-1]\n",
    "    print(f\"{round(arr.shape[-1]/25*n_skip_semantic)}s for track {n}\")\n",
    "    # do stuff incase skip\n",
    "    arr2 = (\n",
    "        torch.zeros(arr.shape[0] * n_skip_semantic, dtype=arr.dtype)\n",
    "        + cfg.semantic_pad_token\n",
    "    )\n",
    "    arr2[::n_skip_semantic] = arr\n",
    "    # add\n",
    "    out_gpt.append(arr2)\n",
    "for in_sem_arr in out_gpt:\n",
    "    generate(\n",
    "        in_sem_arr,\n",
    "        lyrics=text,\n",
    "        tags=tags,\n",
    "        text_cfg_coef=2.0,\n",
    "        steps=10,\n",
    "        seed=0,\n",
    "    ).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43513761",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "fbe4f87d",
   "metadata": {},
   "source": [
    "### Over/Underpaint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6566827e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from suno_utils.audio import Audio\n",
    "from generation import encode_semantic\n",
    "from suno_utils.tasks.demucs import split_vocals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaf94996",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "\n",
    "mm = np.memmap(\"/app/suno/data/chirp_v5/v1/data_val.bin\", dtype=np.uint16, mode=\"r\")\n",
    "metas = read_jsonl(\"/app/suno/data/chirp_v5/v1/metas_val.jsonl\")\n",
    "overpaint_metas = [\n",
    "    m\n",
    "    for m in metas\n",
    "    if m[\"dataset\"] == \"musdb_stems_overpaint\" and m[\"type\"] == \"instrumental\"\n",
    "]\n",
    "underpaint_metas = [\n",
    "    m\n",
    "    for m in metas\n",
    "    if m[\"dataset\"] == \"musdb_stems_underpaint\" and m[\"type\"] == \"vocals\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f98dcf7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "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",
    "\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",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa3cd8c8",
   "metadata": {},
   "source": [
    "#### Underpaint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28bf850c",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = Audio.from_file(\"../samples/halo.wav\", sample_rate=44_100, n_channels=2)\n",
    "# # a = Audio.from_file(\"../samples/calm_down.mp3\", sample_rate=44_100, n_channels=2)\n",
    "\n",
    "a_vocals, a_other = split_vocals(a.convert(44_100, 2, 2))\n",
    "vocals_arr = encode_semantic(a_vocals.convert(44_100, 2, 2).normalize_volume())[:, :1]\n",
    "# instrumental_arr = encode_semantic(a_other.convert(44_100, 2, 2).normalize_volume())[:,:1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10e6d9b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# m = random.choice(underpaint_metas)\n",
    "# vocals_arr = mm[m[\"offset_idx\"]:m[\"offset_idx\"]+m[\"n_tokens\"]][:60*4][None].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1023d4b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "tags = \"Pop\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fa2fa71",
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=tags,\n",
    "    underpaint_arr=vocals_arr,\n",
    "    cfg_coef=1.0,\n",
    "    cfg_coef_tags=1.0,\n",
    "    n_repeat_tags=1,\n",
    "    n_batch=1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=2 * 60,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "020e0ac5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "requests = [\n",
    "    make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "out_gpt = []\n",
    "for n, job in enumerate(jobs):\n",
    "    stream = engine.token_generator(job)\n",
    "    arr = torch.stack(list(stream))[:, 1]\n",
    "    if arr[-1] == 4000:\n",
    "        arr = arr[:-1]\n",
    "    print(f\"{round(arr.shape[-1]/25)}s for track {n}\")\n",
    "    out_gpt.append(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf653ddf",
   "metadata": {},
   "outputs": [],
   "source": [
    "for in_sem_arr in out_gpt:\n",
    "    generate(\n",
    "        in_sem_arr,\n",
    "        lyrics=text,\n",
    "        tags=tags,\n",
    "        text_cfg_coef=1.0,\n",
    "        steps=16,\n",
    "        seed=0,\n",
    "    ).play()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f687e4c",
   "metadata": {},
   "source": [
    "#### Overpaint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a779ace0",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = Audio.from_file(\"../samples/halo.wav\", sample_rate=44_100, n_channels=2)\n",
    "# a = Audio.from_file(\"../samples/calm_down.mp3\", sample_rate=44_100, n_channels=2)\n",
    "\n",
    "a_vocals, a_other = split_vocals(a.convert(44_100, 2, 2))\n",
    "vocals_arr = encode_semantic(a_vocals.convert(44_100, 2, 2).normalize_volume())[:, :1]\n",
    "instrumental_arr = encode_semantic(a_other.convert(44_100, 2, 2).normalize_volume())[\n",
    "    :, :1\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "233a72dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# a_vocals = Audio.from_file(\"../samples/martin_vocals_up.mp3\", sample_rate=44_100, n_channels=2)\n",
    "# vocals_arr = encode_semantic(a_vocals.convert(44_100, 2, 2).normalize_volume())[:,:1]\n",
    "\n",
    "# a_instrumental = Audio.from_file(\"../samples/martin_instrumental.m4a\", sample_rate=44_100, n_channels=2)\n",
    "# instrumental_arr = encode_semantic(a_instrumental.convert(44_100, 2, 2).normalize_volume())[:,:1]\n",
    "\n",
    "# m = random.choice(overpaint_metas)\n",
    "# instrumental_arr = mm[m[\"offset_idx\"]:m[\"offset_idx\"]+m[\"n_tokens\"]][:60*4][None].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fc2a900",
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    #     text_tags=\"pop, sentimental, folk\",\n",
    "    overpaint_arr=instrumental_arr,\n",
    "    cfg_coef=1.2,\n",
    "    #     cfg_coef_tags=2.0,\n",
    "    n_repeat_tags=1,\n",
    "    n_batch=1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=2 * 60,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4462f43",
   "metadata": {},
   "outputs": [],
   "source": [
    "requests = [\n",
    "    make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "out_gpt = []\n",
    "for n, job in enumerate(jobs):\n",
    "    stream = engine.token_generator(job)\n",
    "    arr = torch.stack(list(stream))[:, 1]\n",
    "    if arr[-1] == 4000:\n",
    "        arr = arr[:-1]\n",
    "    print(f\"{round(arr.shape[-1]/25)}s for track {n}\")\n",
    "    out_gpt.append(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c4e154e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for in_sem_arr in out_gpt:\n",
    "    generate(\n",
    "        in_sem_arr,\n",
    "        lyrics=text,\n",
    "        #         tags=tags,\n",
    "        text_cfg_coef=1.0,\n",
    "        steps=16,\n",
    "        seed=0,\n",
    "    ).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b1d7b0b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6267ad78",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abc1d47c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c61e6914",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "c95fb777",
   "metadata": {},
   "source": [
    "### Cover"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f21c1f65",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.audio import Audio\n",
    "from generation import encode_semantic\n",
    "from suno_utils.tasks.demucs import split_vocals\n",
    "\n",
    "# audio_filepath = \"../samples/martin2.m4a\"\n",
    "audio_filepath = \"../samples/halo.wav\"\n",
    "# audio_filepath = \"../samples/corn_chase.mp3\"\n",
    "# audio_filepath = \"../samples/martin_vocals.m4a\"\n",
    "\n",
    "a = Audio.from_file(audio_filepath, sample_rate=44_100, n_channels=2)\n",
    "cover_arr = encode_semantic(a.normalize_volume())[:, :1]\n",
    "a.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75677dcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# text = \"\"\"\n",
    "# [Verse]\n",
    "# Raindrops falling away\n",
    "# Dancing in the grey\n",
    "# Clouds are here to stay\n",
    "# Love on a rainy day\n",
    "\n",
    "# [Verse 2]\n",
    "# Puddles on the ground\n",
    "# Feet splashing 'round\n",
    "# You and me we found\n",
    "# Paradise in the sound\n",
    "\n",
    "# [Chorus]\n",
    "# Love on a rainy day\n",
    "# Washing doubts away\n",
    "# Hold me and let's sway\n",
    "# It's our perfect cliche\n",
    "\n",
    "# Love on a rainy day\n",
    "# Washing doubts away\n",
    "# Hold me and let's sway\n",
    "# It's our perfect cliche\n",
    "# \"\"\"\n",
    "\n",
    "# text = \"\"\"\n",
    "# Well, you only need the light when it's burning low\n",
    "# Only miss the sun when it starts to snow\n",
    "# Only know you love her when you let her go\n",
    "# Only know you've been high when you're feeling low\n",
    "# Only hate the road when you're missing home\n",
    "# Only know you love her when you let her go\n",
    "\n",
    "# And you let her go\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",
    "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",
    "\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",
    "\"\"\"\n",
    "\n",
    "tags = \"folk music, haunting, sad, female vocals, emotive, guitar\"\n",
    "# tags = \"R&B, Pop, Contemporary R&B, Ballad, male vocals\"\n",
    "\n",
    "# text = \"\"\n",
    "# tags = \"Big Room Trance\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99745ae8",
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=tags,\n",
    "    cover_arr=cover_arr,\n",
    "    cfg_coef=1.2,\n",
    "    cfg_coef_tags=2.5,\n",
    "    n_repeat_tags=1,\n",
    "    n_batch=1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=2 * 60,\n",
    "    random_seed=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2a70da1",
   "metadata": {},
   "outputs": [],
   "source": [
    "requests = [\n",
    "    make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "out_gpt = []\n",
    "for n, job in enumerate(jobs):\n",
    "    stream = engine.token_generator(job)\n",
    "    arr = torch.stack(list(stream))[:, 1]\n",
    "    if arr[-1] == 4000:\n",
    "        arr = arr[:-1]\n",
    "    print(f\"{round(arr.shape[-1]/25)}s for track {n}\")\n",
    "    out_gpt.append(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d44152cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "for in_sem_arr in out_gpt:\n",
    "    generate(\n",
    "        in_sem_arr,\n",
    "        lyrics=text,\n",
    "        tags=tags,\n",
    "        text_cfg_coef=1.0,\n",
    "        steps=16,\n",
    "        seed=0,\n",
    "    ).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3452380a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be956547",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "19570e8d",
   "metadata": {},
   "source": [
    "### Visualize prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e5ff457",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.gpt.prompt import Prompt\n",
    "\n",
    "in_arr = requests[0].streams[0].prompt\n",
    "prompt = Prompt(\"\", engine.model.config)\n",
    "prompt.visualize(in_arr, compress=True)\n",
    "prompt.visualize(in_arr, compress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d042392a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "699bda31",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "312a749f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7f5cce1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03a473d7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbb3d211",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "addcbe70",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b290ea0b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "242acd6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "def convert_to_type(path):\n",
    "    d = torch.load(path)\n",
    "    d[\"best_val_loss\"] = float(d[\"best_val_loss\"])\n",
    "    torch.save(d, \"converted.pt\")\n",
    "    print(\"DONE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abd3f654-4903-4a5e-b6ef-fe2013c912d8",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
