{
 "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\"] = \"5\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0d2ac4c-cba0-4a1d-be09-7036a0b78d52",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from suno_utils.diffusion import generation as diffusion_gen\n",
    "from suno_utils.tasks.upsample_engine import UpsampleEngine, Request\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import (\n",
    "    preload_models as preload_codec_models,\n",
    "    decode,\n",
    ")\n",
    "\n",
    "# Load codec/decoder model\n",
    "preload_codec_models(\"/app/suno/data/dpo/models/dac_vae_tuned_25hz.pth\")\n",
    "\n",
    "# dit_model_filepath = \"/app2/suno/checkpoints/2025-11-02_15-41-25_s2081/last_ckpt.pt\"\n",
    "dit_model_filepath = (\n",
    "    \"/app2/suno/checkpoints/2025-11-11_20-38-05_s6962/last_ckpt.pt\"  #  newer run\n",
    ")\n",
    "diffusion_gen.preload_tokenizer(\"/app/suno/models/chirp_v2/tokenizer_60k.json\")\n",
    "diffusion_gen.preload_dit_model(\n",
    "    dit_model_filepath=dit_model_filepath,\n",
    "    use_ema_if_exists=True,\n",
    "    compile=False,\n",
    "    weights_precision=torch.bfloat16,\n",
    "    use_vox=True,\n",
    ")\n",
    "\n",
    "diffusion_engine = UpsampleEngine(compile=False)"
   ]
  },
  {
   "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",
    "    # \"/app2/suno/data/dpo/models/sem_4rvq_6b_11141927_infer.pt\",  # d67\n",
    "    \"/app2/suno/checkpoints/2025-11-15_13-12-30/last_ckpt_infer.pt\",  # +8 to bluejay\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": "caa2446f-9692-42bb-8d60-db11478ed3c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_clips(\n",
    "    text, tags, neg_tags=None, cfg_coef_tags_max_steps=None, control_tags=None\n",
    "):\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=5,\n",
    "        cfg_coef_tags_max_steps=25 * 60\n",
    "        if cfg_coef_tags_max_steps is None\n",
    "        else cfg_coef_tags_max_steps,\n",
    "        n_repeat_tags=1,\n",
    "        n_skip_semantic=n_skip_semantic,\n",
    "        text_start_control_tags=f\"{control_tags}\" if control_tags is not None else None,\n",
    "        cfg_coef_neg_tags=-1,\n",
    "        text_neg_tags=\"repetitive, loop\" if neg_tags is None else neg_tags,\n",
    "        temp_semantic=0.90,\n",
    "        #     temp_coarse=0.9,\n",
    "        #     top_k_semantic=None,\n",
    "        #     top_k_coarse=None,\n",
    "        min_p_semantic=0.005,\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(8 * 60 / n_skip_semantic),\n",
    "        random_seed=0,\n",
    "    )\n",
    "    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",
    "        # Engine generates all streams: text (stream 0) + 4 RVQ codebooks (streams 1-4)\n",
    "        stream = engine.token_generator(job)\n",
    "        arr = torch.stack(list(stream))  # Shape: (T, 5) - all streams\n",
    "        print(f\"{round(arr.shape[0]/25*n_skip_semantic)}s for track {n}\")\n",
    "        out_gpt.append(arr)\n",
    "\n",
    "    for full_arr in out_gpt:\n",
    "        # Manual unshifting: extract semantic codebooks (streams 1-4) and align them\n",
    "        shift_amount = 5  # delay between RVQ codebooks\n",
    "        n_rvq = 4\n",
    "        max_shift = (n_rvq - 1) * shift_amount  # 15 frames\n",
    "\n",
    "        # Calculate the aligned length (shortest codebook after unshifting)\n",
    "        final_length = full_arr.shape[0] - max_shift\n",
    "\n",
    "        # Unshift each codebook by removing its delay from the start\n",
    "        unshifted = torch.zeros((n_rvq, final_length), dtype=full_arr.dtype)\n",
    "        for i in range(n_rvq):\n",
    "            shift = i * shift_amount\n",
    "            # Extract codebook i (stream i+1) and remove the delay\n",
    "            unshifted[i] = full_arr[shift : shift + final_length, i + 1]\n",
    "\n",
    "        unshifted[unshifted > 4000] = 4000\n",
    "        gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n",
    "            lyrics=text,\n",
    "            text_cfg_coef=2.0,\n",
    "            ctx_cfg_coef=1.0,\n",
    "            steps=12,\n",
    "            codec_scale_factor=0.4,\n",
    "            scale_ctx_vector=True,\n",
    "            objective=\"rectified_flow\",\n",
    "        )\n",
    "        request = Request(\n",
    "            id=\"dummy\",\n",
    "            generation_config=gen_cfg,\n",
    "            tokens=unshifted.T,\n",
    "            input_tokens_finished=True,\n",
    "        )\n",
    "        result = diffusion_engine.run_request(request)\n",
    "\n",
    "        vae_latents = torch.concat(result.vae_latents)\n",
    "        audio = decode(vae_latents)\n",
    "        audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9b4a2be-1505-4ef6-8198-bc4f6865ea96",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(\"\", tags=\"JS bach, JSbach, fugue, piano, baroque\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3133ec0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "generate_clips(text, tags=\"epic film orchestral\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2064fd87-541a-465c-8a47-8510e9123f47",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate_clips(text, tags=\"k-pop\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2517468-63bb-451a-9cad-12fabe5a8902",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[verse]\n",
    "A is for the amazing grace that we recieve\n",
    "B is for the blessings, every day we blieve\n",
    "C is for the chorus, we sing it loud and clear\n",
    "D is for the devotion that we hold dear\n",
    "\n",
    "[verse]\n",
    "E is for the everlasting love that sets us free\n",
    "F is for the faith that guides us on this journey\n",
    "G is for the goodness that we share each day\n",
    "H is for the hope that never fades away\n",
    "\n",
    "[Chorus]\n",
    "B is for Buttocks, ripe and slightly damp\n",
    "Yum Yum, boy oh boy do I like God\n",
    "Every day is a gift, when you won a skateboard ramp.\n",
    "Yum yum, boy oh boy do I like God.\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "generate_clips(text, tags=\"gregorian chant\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c044be24-be21-4810-82b1-bd30d05ffc31",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[verse]\n",
    "A is for the amazing grace that we recieve\n",
    "B is for the blessings, every day we believe\n",
    "C is for the chorus, we sing it loud and clear\n",
    "D is for the devotion that we hold dear\n",
    "\n",
    "[verse]\n",
    "E is for the everlasting love that sets us free\n",
    "F is for the faith that guides us on this journey\n",
    "G is for the goodness that we share each day\n",
    "H is for the hope that never fades away\n",
    "\n",
    "[Chorus]\n",
    "B is for Buttocks, ripe and slightly damp\n",
    "Yum Yum, boy oh boy do I like God\n",
    "Every day is a gift, when you won a skateboard ramp.\n",
    "Yum yum, boy oh boy do I like God.\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "generate_clips(text, tags=\"gregorian chant\", control_tags=\"tag_strength:1;\")\n",
    "# generate_clips(text, tags=\"neo soul\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4707f57c-3269-43fd-a189-eb8b51898046",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[verse]\n",
    "A is for the amazing grace that we recieve\n",
    "B is for the blessings, every day we believe\n",
    "C is for the chorus, we sing it loud and clear\n",
    "D is for the devotion that we hold dear\n",
    "\n",
    "[verse]\n",
    "E is for the everlasting love that sets us free\n",
    "F is for the faith that guides us on this journey\n",
    "G is for the goodness that we share each day\n",
    "H is for the hope that never fades away\n",
    "\n",
    "[Chorus]\n",
    "B is for Buttocks, ripe and slightly damp\n",
    "Yum Yum, boy oh boy do I like God\n",
    "Every day is a gift, when you won a skateboard ramp.\n",
    "Yum yum, boy oh boy do I like God.\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "generate_clips(text, tags=\"gregorian chant\", control_tags=\"tag_strength:9\")\n",
    "# generate_clips(text, tags=\"neo soul\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c9652cc-7ec3-405e-80fc-ffb2b58e5186",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[verse]\n",
    "A is for the amazing grace that we recieve\n",
    "B is for the blessings, every day we believe\n",
    "C is for the chorus, we sing it loud and clear\n",
    "D is for the devotion that we hold dear\n",
    "\n",
    "[verse]\n",
    "E is for the everlasting love that sets us free\n",
    "F is for the faith that guides us on this journey\n",
    "G is for the goodness that we share each day\n",
    "H is for the hope that never fades away\n",
    "\n",
    "[Chorus]\n",
    "B is for Buttocks, ripe and slightly damp\n",
    "Yum Yum, boy oh boy do I like God\n",
    "Every day is a gift, when you won a skateboard ramp.\n",
    "Yum yum, boy oh boy do I like God.\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "generate_clips(\n",
    "    text, tags=\"gregorian chant, no instrumental\", control_tags=\"tag_strength:9\"\n",
    ")\n",
    "# generate_clips(text, tags=\"neo soul\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b0c144b-96f7-4b8b-ab9b-8db235eb78d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[verse]\n",
    "A is for the amazing grace that we recieve\n",
    "B is for the blessings, every day we blieve\n",
    "C is for the chorus, we sing it loud and clear\n",
    "D is for the devotion that we hold dear\n",
    "\n",
    "[verse]\n",
    "E is for the everlasting love that sets us free\n",
    "F is for the faith that guides us on this journey\n",
    "G is for the goodness that we share each day\n",
    "H is for the hope that never fades away\n",
    "\n",
    "[Chorus]\n",
    "B is for Buttocks, ripe and slightly damp\n",
    "Yum Yum, boy oh boy do I like God\n",
    "Every day is a gift, when you won a skateboard ramp.\n",
    "Yum yum, boy oh boy do I like God.\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "generate_clips(text, tags=\"gregorian chant, gregorian chant, gregorian chant\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ad05205-272a-4fc6-bd8e-6193111cb141",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[bagpipes, verse]\n",
    "A is for the amazing grace that we recieve\n",
    "B is for the blessings, every day we blieve\n",
    "C is for the chorus, we sing it loud and clear\n",
    "D is for the devotion that we hold dear\n",
    "\n",
    "[verse]\n",
    "E is for the everlasting love that sets us free\n",
    "F is for the faith that guides us on this journey\n",
    "G is for the goodness that we share each day\n",
    "H is for the hope that never fades away\n",
    "\n",
    "[Chorus]\n",
    "B is for Buttocks, ripe and slightly damp\n",
    "Yum Yum, boy oh boy do I like God\n",
    "Every day is a gift, when you won a skateboard ramp.\n",
    "Yum yum, boy oh boy do I like God.\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "generate_clips(text, tags=\"bagpipes, scottish funeral march\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f64ec057",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Verse]\n",
    "一盏离愁孤灯伫立在窗口\n",
    "我在门后假装你人还没走\n",
    "旧地如重游月圆更寂寞\n",
    "夜半清醒的烛火不忍苛责我\n",
    "\n",
    "[Verse]\n",
    "一壶漂泊浪迹天涯难入喉\n",
    "你走之后酒暖回忆思念瘦\n",
    "水向东流时间怎么偷\n",
    "花开就一次成熟我却错过\n",
    "\n",
    "[Chorus]\n",
    "谁在用琵琶弹奏一曲东风破\n",
    "岁月在墙上剥落看见小时候\n",
    "犹记得那年我们都还很年幼\n",
    "而如今琴声幽幽我的等候你没听过\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "generate_clips(text, tags=\"chinese pop woman\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cb96fc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(\n",
    "    \"\", tags=\"late 19th century austrian romantic symphony, dramatic, large orchestra\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "741261a7-70bb-46a1-94f5-2547c17e6484",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(\n",
    "    \"[Dramatic slow opening, drums, strong strings, emotional, intense] [Woodwinds comes in with constrast, delicate, sad, sorrow] [Back strings, fast tempo, echo the theme]\",\n",
    "    tags=\"\"\"\n",
    "beethoven, beethoven, beethoven, ludwig van beethoven, ludwig van beethoven as the composer,\n",
    "dg, philharmonic orchestra, DG, symphony 1st movement\"\"\",\n",
    "    neg_tags=\"vocals, vocals, pianos, classical, baroque, pop, mozart, opera, mozart, haydn, brahms, chopin, overture, overture, march, march\",\n",
    "    cfg_coef_tags_max_steps=25 * 4 * 60,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c278ed9-e9df-47f3-ad5e-be9f225d4ace",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(\"\", tags=\"JS bach, JSbach, fugue, piano, baroque\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16e662b1-6515-4e5c-90b0-4d429600edfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(\"\", tags=\"fugue, piano, baroque, slow tempo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9769cd6-0b69-4dcb-bb24-cc72108fd0c5",
   "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",
    "[end][end][end]\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc47d650-4bec-4a91-b1ba-def5716e9f27",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(text, tags=\"female bluegrass, fast tempo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8b8ea55-f30f-4a82-b5d6-4f6f696e0808",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(\"\", tags=\"JS bach, JS bach, fugue, piano, baroque\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb7dda71-0320-458f-b58a-f9cd90eeb4cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(\"plastic 1, plastic 2, plastic 3, plastic 4, plastic 5\", tags=\"pop\")"
   ]
  },
  {
   "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": [
    "import torch\n",
    "\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": "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": []
  }
 ],
 "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
}
