{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fae8f9c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:13:55.101018Z",
     "start_time": "2024-05-11T20:13:51.243900Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\"\n",
    "os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c27fb2bf-0bb3-4b5d-9930-22c1a6867a90",
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt_ckpt_path = \"/app/suno/checkpoints/2024-07-18_05-08-02/last_ckpt_infer.pt\"\n",
    "gpt_output_path = os.path.join(\"/home/tony/Work/gpt_samples/\", \"13b_2h_ft1_v2\")\n",
    "os.makedirs(gpt_output_path, exist_ok=True)\n",
    "MAX_STREAMS = 4 if \"30b\" not in gpt_output_path else 3\n",
    "MAX_SEQ = 40 if \"30b\" not in gpt_output_path else 12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a035f3d2-697c-401f-85ef-17289c9f8743",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.gpt.chirp_v2_5 import (\n",
    "    GenerationConfig,\n",
    "    decode_stream,\n",
    "    preload_models,\n",
    "    prep_gconf,\n",
    ")\n",
    "from suno_utils.gpt.generation_engine import (\n",
    "    align_codes,\n",
    "    make_request,\n",
    ")\n",
    "from suno_utils.gpt.engine import Engine\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "\n",
    "preload_models(load_gpt=False, load_semantic=False, load_codec_device=\"cuda\")\n",
    "\n",
    "engine = Engine(\n",
    "    gpt_ckpt_path,\n",
    "    \"/app/suno/data/dpo/models/tokenizer_60k.json\",\n",
    "    max_sequences=MAX_SEQ,\n",
    "    compile=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18bb8e5f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:14:14.500100Z",
     "start_time": "2024-05-11T20:14:14.497272Z"
    }
   },
   "outputs": [],
   "source": [
    "general_config = dict(\n",
    "    cfg_coef=1.0,  # no text cfg for dpo stream\n",
    "    min_eos_p=0.1,\n",
    "    eos_pad_duration_s=0,\n",
    "    # cfg_coef_tags=0.0,\n",
    "    # cfg_coef_tags_max_steps=None,  # collect the data for now\n",
    "    cfg_coef_tags=2,\n",
    "    cfg_coef_neg_tags=-1,\n",
    "    text_neg_tags=\"repetitive, loop, noisy, distorted\",\n",
    "    n_repeat_tags=1,\n",
    "    use_whisper=False,\n",
    "    text_start_control_tags=\"{start} \",\n",
    "    # text_end_control_tags=\"{end}\",\n",
    "    random_seed=42,\n",
    "    n_batch=1,\n",
    "    min_text_offset=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f57ea1b-4133-4a37-a50f-872ae74ccb08",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_audio_with_engine(text, text_tags, title):\n",
    "    cfg = GenerationConfig(\n",
    "        text=text, text_tags=text_tags, max_gen_duration_s=240, **general_config\n",
    "    )\n",
    "    preped_cfg = prep_gconf(cfg)\n",
    "    model_conf = engine.model.config\n",
    "    requests = []\n",
    "    for i in range(1, MAX_STREAMS + 1):\n",
    "        request = make_request(\n",
    "            f\"{i}\", preped_cfg, engine.model.config, engine.tokenizer\n",
    "        )\n",
    "        requests.append(request)\n",
    "\n",
    "    jobs = engine.run_request(requests)\n",
    "\n",
    "    for i, job in enumerate(jobs):\n",
    "        stream = engine.token_generator(job)\n",
    "        audios = []\n",
    "        for audio in decode_stream(align_codes(tqdm(stream), model_conf)):\n",
    "            audios.append(audio[0])\n",
    "        audio_continued = Audio.concatenate(audios)\n",
    "        # save the file to wav\n",
    "        audio_continued.to_hq_mp3(os.path.join(gpt_output_path, f\"{title}_{i}.mp3\"))\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa44f837",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:14:14.566214Z",
     "start_time": "2024-05-11T20:14:14.501592Z"
    }
   },
   "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",
    "text_tags=\"fast female vocalist, pop-punk, pop, melodic, hd stems, goes hard, 155bpm\"\n",
    "\n",
    "generate_audio_with_engine(text, text_tags, title=\"fast_female_walking\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "072adcda",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:15:43.758219Z",
     "start_time": "2024-05-11T20:15:43.756159Z"
    }
   },
   "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",
    "text_tags=\"pop woman\"\n",
    "\n",
    "generate_audio_with_engine(text, text_tags, title=\"chinese_pop_woman\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fa595a6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:18:36.675596Z",
     "start_time": "2024-05-11T20:18:36.616364Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\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",
    "text_tags=\"orchestral film epic\"\n",
    "\n",
    "generate_audio_with_engine(text, text_tags, title=\"film_epic_walking\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0fff5d8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:21:24.466761Z",
     "start_time": "2024-05-11T20:20:02.669267Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\n",
    "\n",
    "generate_audio_with_engine(text, text_tags, title=\"film_epic\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef2a197a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:21:28.557837Z",
     "start_time": "2024-05-11T20:21:28.555006Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Intro]\n",
    "[Verse]\n",
    "I used to think it's all a game\n",
    "That somethin' more was to blame\n",
    "But now I see, it's in the details, baby\n",
    "Just a little tweak here and there\n",
    "Can make our model beyond compare\n",
    "Yeah, it's the small things that drive it crazy (ooh-yeah)\n",
    "\n",
    "[Pre-Chorus]\n",
    "Needed to get rid of the rest,\n",
    "Cause baby you deserve what's best\n",
    "Gotta give you the cream of the crop\n",
    "So you can get us back to the top\n",
    "\n",
    "[Chorus]\n",
    "Fine tuning, baby, it's the key\n",
    "Turnin' up the signal for you and me\n",
    "When we're off, we find our way back again (oh-oh)\n",
    "Fine tuning, baby, it's the truth\n",
    "We make it work, we're bulletproof\n",
    "Together we train a model that's the best (oh-oh-oh-oh)\n",
    "\n",
    "[end]\n",
    "\"\"\"\n",
    "text_tags=\"k-pop female\"\n",
    "\n",
    "generate_audio_with_engine(text, text_tags, title=\"ft_kpop_female\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96c51840",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:22:54.784678Z",
     "start_time": "2024-05-11T20:22:54.781837Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Intro]\n",
    "\n",
    "\n",
    "[Verse 1]\n",
    "You're on the phone with your girlfriend, she's upset\n",
    "She's going off about something that you said\n",
    "'Cause she doesn't get your humor like I do\n",
    "I'm in the room, it's a typical Tuesday night\n",
    "I'm listening to the kind of music she doesn't like\n",
    "And she'll never know your story like I do\n",
    "\n",
    "[Pre-Chorus]\n",
    "But she wears short skirts, I wear T-shirts\n",
    "She's Cheer Captain and I'm on the bleachers\n",
    "Dreaming 'bout the day when you wake up and find\n",
    "That what you're looking for has been here the whole time\n",
    "\n",
    "[Chorus]\n",
    "If you could see that I'm the one who understands you\n",
    "Been here all along, so why can't you see?\n",
    "You belong with me, you belong with me\n",
    "\n",
    "[Verse 2]\n",
    "Walking the streets with you and your worn-out jeans\n",
    "I can't help thinking this is how it ought to be\n",
    "Laughing on a park bench, thinking to myself\n",
    "Hey, isn't this easy?\n",
    "And you've got a smile that could light up this whole town\n",
    "I haven't seen it in a while since she brought you down\n",
    "You say you're fine, I know you better than that\n",
    "Hey, whatcha doing with a girl like that?\n",
    "\n",
    "[Pre-Chorus]\n",
    "She wears high heels, I wear sneakers\n",
    "She's Cheer Captain and I'm on the bleachers\n",
    "Dreaming 'bout the day when you wake up and find\n",
    "That what you're looking for has been here the whole time\n",
    "\n",
    "[Chorus]\n",
    "If you could see that I'm the one who understands you\n",
    "Been here all along, so why can't you see?\n",
    "You belong with me\n",
    "Standing by and waiting at your back door\n",
    "All this time, how could you not know, baby?\n",
    "You belong with me, you belong with me\n",
    "\n",
    "[Bridge]\n",
    "Oh, I remember you driving to my house\n",
    "In the middle of the night\n",
    "I'm the one who makes you laugh\n",
    "When you know you're 'bout to cry\n",
    "And I know your favorite songs\n",
    "And you tell me 'bout your dreams\n",
    "Think I know where you belong\n",
    "Think I know it's with me\n",
    "\n",
    "[Chorus]\n",
    "Can't you see that I'm the one who understands you\n",
    "Been here all along, so why can't you see?\n",
    "You belong with me\n",
    "Standing by and waiting at your back door\n",
    "All this time, how could you not know, baby?\n",
    "You belong with me, you belong with me\n",
    "\n",
    "[Outro]\n",
    "You belong with me\n",
    "Have you ever thought just maybe\n",
    "You belong with me?\n",
    "You belong with me\n",
    "\"\"\"\n",
    "\n",
    "text_tags=\"pop country, female vocal, fast tempo, taylor swift\"\n",
    "\n",
    "generate_audio_with_engine(text, text_tags, title=\"ts_pop\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9592f21d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:24:21.959612Z",
     "start_time": "2024-05-11T20:24:21.957387Z"
    }
   },
   "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",
    "[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",
    "text_tags=\"gospel soulful choir\"\n",
    "generate_audio_with_engine(text, text_tags, title=\"yum_gospel\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "993e617a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:26:54.983332Z",
     "start_time": "2024-05-11T20:25:48.043415Z"
    }
   },
   "outputs": [],
   "source": [
    "text_tags=\"gregorian chant\"\n",
    "generate_audio_with_engine(text, text_tags, title=\"yum_chant\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8eba0a1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:26:57.993032Z",
     "start_time": "2024-05-11T20:26:57.991084Z"
    }
   },
   "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",
    "text_tags=\"bluegrass, female\"\n",
    "generate_audio_with_engine(text, text_tags, title=\"country_road_bluegrass_female\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82572f44",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:29:44.904731Z",
     "start_time": "2024-05-11T20:28:22.956074Z"
    }
   },
   "outputs": [],
   "source": [
    "generate_audio_with_engine(text + \"\\n [end]\", text_tags, title=\"country_road_bluegrass_female_end\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec07b363",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:29:49.112104Z",
     "start_time": "2024-05-11T20:29:49.043366Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Verse]\n",
    "I've created a monster\n",
    "'Cause nobody wants to see Marshall no more, they want Shady, I'm chopped liver\n",
    "Well, if you want Shady, this is what I'll give ya\n",
    "A little bit of weed mixed with some hard liquor\n",
    "Some vodka that'll jump-start my heart quicker\n",
    "Than a shock when I get shocked at the hospital\n",
    "By the doctor when I'm not cooperating\n",
    "When I'm rockin' the table while he's operating (Hey!)\n",
    "You waited this long, now stop debating\n",
    "'Cause I'm back, I'm on the rag and ovulating\n",
    "I know that you got a job, Ms. Cheney\n",
    "But your husband's heart problem's complicating\n",
    "So the FCC won't let me be\n",
    "Or let me be me, so let me see\n",
    "They tried to shut me down on MTV\n",
    "But it feels so empty without me\n",
    "So come on and dip, bum on your lips\n",
    "Fuck that, cum on your lips and some on your tits\n",
    "And get ready, 'cause this shit's about to get heavy\n",
    "I just settled all my lawsuits (Fuck you, Debbie!)\n",
    "\n",
    "[Chorus]\n",
    "Now, this looks like a job for me\n",
    "So everybody, just follow me\n",
    "'Cause we need a little controversy\n",
    "'Cause it feels so empty without me\n",
    "I said this looks like a job for me\n",
    "So everybody, just follow me\n",
    "'Cause we need a little controversy\n",
    "'Cause it feels so empty without me\n",
    "\n",
    "[end]\n",
    "\"\"\"\n",
    "\n",
    "text_tags=\"rap, female\"\n",
    "generate_audio_with_engine(text, text_tags, title=\"rap_female\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94ae33bf",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:37.896869Z",
     "start_time": "2024-05-11T20:31:15.871843Z"
    }
   },
   "outputs": [],
   "source": [
    "generate_audio_with_engine(\"\", text_tags=\"playful energetic synthwave\", title=\"synthwave\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a18c2238",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:35:18.465267Z",
     "start_time": "2024-05-11T20:33:56.569725Z"
    }
   },
   "outputs": [],
   "source": [
    "generate_audio_with_engine(\"\", text_tags=\"piano concerto beethoven\", title=\"beethoven_pc\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa3ab7e9-1ec1-4b49-9d5f-ba2adedccf98",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_audio_with_engine(\"\", text_tags=\"orchestral symphony beethoven\", title=\"beethoven_sym\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcea0c56",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.185484Z",
     "start_time": "2024-05-11T20:32:41.809855Z"
    }
   },
   "outputs": [],
   "source": [
    "BREAK"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33cd6562",
   "metadata": {},
   "source": [
    "# continue canon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f9e0366",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.190518Z",
     "start_time": "2024-05-11T20:32:42.190509Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio = Audio.from_file(\"audios/Pachelbel - Canon In D Major. Best version. [NlprozGcs80].wav\").get_segment(from_s=60, to_s=80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37d50b41",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.191304Z",
     "start_time": "2024-05-11T20:32:42.191295Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio = Audio.from_file(\"audios/Pachelbel - Canon In D Major. Best version. [NlprozGcs80].wav\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f9d3323",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.192059Z",
     "start_time": "2024-05-11T20:32:42.192050Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "001ec800",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.192894Z",
     "start_time": "2024-05-11T20:32:42.192884Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\"\",\n",
    "        text_tags=\"canon\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        # stream=False,\n",
    "    ),\n",
    "    history_audio=input_audio\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6309815d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.193501Z",
     "start_time": "2024-05-11T20:32:42.193493Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d1a8965",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.194403Z",
     "start_time": "2024-05-11T20:32:42.194393Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\"[start] \\n\" + \"I love you \\n\" * 5 + \"[end]\",\n",
    "        text_tags=\"\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        cfg_coef=1.8,\n",
    "        # stream=False,\n",
    "    ),\n",
    "    history_audio=input_audio\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c36e347c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.195297Z",
     "start_time": "2024-05-11T20:32:42.195287Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2270b189",
   "metadata": {},
   "source": [
    "# Oracle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "772a5d20",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.196090Z",
     "start_time": "2024-05-11T20:32:42.196081Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio = Audio.from_file(\"audios/test_swift.mp3\", n_channels=2)\n",
    "\n",
    "test_input_audio = input_audio.get_segment(from_s=180, to_s=200)\n",
    "test_input_audio.play(compress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4b5d48c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.196817Z",
     "start_time": "2024-05-11T20:32:42.196807Z"
    }
   },
   "outputs": [],
   "source": [
    "oracle_semantic_arr = semantic_encode(test_input_audio, n_codebooks=1)\n",
    "oracle_coarse_arr = codec_encode(test_input_audio,)[:oracle_semantic_arr.shape[0], :]\n",
    "print(oracle_semantic_arr.shape, oracle_coarse_arr.shape)\n",
    "oracle_arr = np.concatenate(\n",
    "    [oracle_semantic_arr, oracle_coarse_arr],\n",
    "    axis=-1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "808f5130",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.197671Z",
     "start_time": "2024-05-11T20:32:42.197661Z"
    }
   },
   "outputs": [],
   "source": [
    "input_oracle_arr = oracle_arr[:, :1]\n",
    "print(input_oracle_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "752ba22f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.198461Z",
     "start_time": "2024-05-11T20:32:42.198451Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\" \" * 20,\n",
    "        text_tags=\"\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        oracle_arr=input_oracle_arr,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bba0d8c9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.199207Z",
     "start_time": "2024-05-11T20:32:42.199197Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()\n",
    "    print(e.array_float[-100:].mean(axis=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad979c8b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.200057Z",
     "start_time": "2024-05-11T20:32:42.200047Z"
    }
   },
   "outputs": [],
   "source": [
    "input_pavarotti = Audio.from_file(\"audios/test_pavarotti.mp3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f0f8785",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.200671Z",
     "start_time": "2024-05-11T20:32:42.200662Z"
    }
   },
   "outputs": [],
   "source": [
    "input_pavarotti_seg = input_pavarotti.get_segment(from_s=152, to_s=172)\n",
    "input_pavarotti_seg.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e91f1d1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.201665Z",
     "start_time": "2024-05-11T20:32:42.201654Z"
    }
   },
   "outputs": [],
   "source": [
    "pavarotti_oracle_semantic_arr = semantic_encode(input_pavarotti_seg, n_codebooks=1)\n",
    "pavarotti_oracle_coarse_arr = codec_encode(input_pavarotti_seg)[\n",
    "    : pavarotti_oracle_semantic_arr.shape[0], :\n",
    "]\n",
    "print(pavarotti_oracle_semantic_arr.shape, pavarotti_oracle_coarse_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6502b74a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.202415Z",
     "start_time": "2024-05-11T20:32:42.202405Z"
    }
   },
   "outputs": [],
   "source": [
    "masked_oracle_semantic_arr = oracle_semantic_arr.copy().astype(np.int16)\n",
    "# masked_oracle_semantic_arr[0::5, :] = -1\n",
    "masked_oracle_coarse_arr = pavarotti_oracle_coarse_arr.copy().astype(np.int16)\n",
    "# masked_oracle_coarse_arr[20:, :] = -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b97481d6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.393795Z",
     "start_time": "2024-05-11T20:32:42.393783Z"
    }
   },
   "outputs": [],
   "source": [
    "mixed_oracle_arr =  np.concatenate(\n",
    "    [masked_oracle_semantic_arr], #, masked_oracle_coarse_arr[:, 0:1]],\n",
    "    axis=-1,\n",
    ")\n",
    "print(mixed_oracle_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76dcc73a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.394785Z",
     "start_time": "2024-05-11T20:32:42.394775Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\" \" * 20,\n",
    "        text_tags=\"\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        oracle_arr=mixed_oracle_arr,\n",
    "    ),\n",
    "    # history_audio=input_pavarotti_seg\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84e04fd4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.395542Z",
     "start_time": "2024-05-11T20:32:42.395533Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()\n",
    "    print(e.array_float[-100:].mean(axis=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edc87efa",
   "metadata": {},
   "source": [
    "# continuation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b6938c4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.396152Z",
     "start_time": "2024-05-11T20:32:42.396144Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio = Audio.from_file(\"audios/test_swift.mp3\", n_channels=2)\n",
    "test_input_audio = input_audio.get_segment(from_s=180, to_s=200)\n",
    "test_input_audio.play(compress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4cbb76b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.397111Z",
     "start_time": "2024-05-11T20:32:42.397102Z"
    }
   },
   "outputs": [],
   "source": [
    "hoot_encode([test_input_audio])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "984e8c9b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.397640Z",
     "start_time": "2024-05-11T20:32:42.397632Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\"[chorus] purple rain\" * 20,\n",
    "        text_tags=\"\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        # stream=False,\n",
    "    ),\n",
    "    history_audio=test_input_audio,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "296b34eb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.398394Z",
     "start_time": "2024-05-11T20:32:42.398385Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df83ed93",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.399468Z",
     "start_time": "2024-05-11T20:32:42.399459Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "On a vibrant summer morning, curiosity drew me to the innovative realm of Suno AI. Their latest\n",
    "masterpiece, Chirp, known for its unique ability to generate music from words, promised a bold\n",
    "step forward in music creation. I dived into their platform, typing in familiar lyrics. As I clicked\n",
    "'Create,' the ensuing symphony seamlessly blended human emotion and machine precision,\n",
    "heralding a new era in the convergence of music, words, and machine learning\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b5a2270",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-11T20:32:42.400260Z",
     "start_time": "2024-05-11T20:32:42.400249Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=text,\n",
    "        text_tags=\"hifi rap\",\n",
    "        n_batch=2,\n",
    "        max_gen_duration_s=120,\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4d197af-166d-4133-86bc-1a7bf04506ce",
   "metadata": {},
   "source": [
    "# Generate HTML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a344ca61-3b37-4e8c-9fd8-46c1c51df8bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "\n",
    "# # Define the directory containing the MP3 files\n",
    "# directory = gpt_output_path # Change this to your directory path\n",
    "\n",
    "# # Start the HTML content\n",
    "# html_content = f'''\n",
    "# <!DOCTYPE html>\n",
    "# <html lang=\"en\">\n",
    "# <head>\n",
    "#     <meta charset=\"UTF-8\">\n",
    "#     <meta name=\"viewport\" content=\"width=device-width, initial-scale=1.0\">\n",
    "#     <title>Audio Player</title>\n",
    "# </head>\n",
    "# <body>\n",
    "#     <h1>Audio Player for model {gpt_output_path} </h1>\n",
    "# '''\n",
    "\n",
    "# # Loop through all files in the directory\n",
    "# for filename in os.listdir(directory):\n",
    "#     if filename.endswith('.mp3'):\n",
    "#         file_path = os.path.join(directory, filename)\n",
    "#         # Add an audio player for each MP3 file\n",
    "#         html_content += f'''\n",
    "#     <div>\n",
    "#         <p>{filename}</p>\n",
    "#         <audio controls>\n",
    "#             <source src=\"{file_path}\" type=\"audio/mpeg\">\n",
    "#             Your browser does not support the audio element.\n",
    "#         </audio>\n",
    "#     </div>\n",
    "#     '''\n",
    "\n",
    "# # Close the HTML content\n",
    "# html_content += '''\n",
    "# </body>\n",
    "# </html>\n",
    "# '''\n",
    "\n",
    "# # Write the HTML content to a file\n",
    "# with open('mp3_player.html', 'w') as file:\n",
    "#     file.write(html_content)\n",
    "\n",
    "# print('HTML file generated successfully.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b437dbc-40d9-4dfb-9bd1-c4b1dd3b745a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from IPython.display import IFrame\n",
    "\n",
    "# # Path to the generated HTML file\n",
    "# html_file_path = 'mp3_player.html'  # Change this to your directory path\n",
    "\n",
    "# # Display the HTML file\n",
    "# IFrame(src=html_file_path, width='100%', height=600)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a841423-4079-498d-bc8f-4d9884e587e1",
   "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.14"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}