{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fae8f9c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:09.763933Z",
     "start_time": "2024-01-12T17:47:06.618228Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\"\n",
    "\n",
    "import numpy as np\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.gpt.chirp_v2_5 import (\n",
    "    generate_audio,\n",
    "    generate_audio_stream,\n",
    "    preload_models,\n",
    "    codec_decode,\n",
    "    GenerationConfig,\n",
    "    semantic_encode,\n",
    "    codec_encode,\n",
    "    hoot_encode_and_align,\n",
    ")\n",
    "from suno_utils.tasks.hoot import encode as hoot_encode\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9108708b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:21.807412Z",
     "start_time": "2024-01-12T17:47:09.766720Z"
    }
   },
   "outputs": [],
   "source": [
    "_ = preload_models(\n",
    "    gpt_ckpt_path=\"/app/suno/tony/v3/7b_ft.pt\",\n",
    "    tokenizer_path=\"/home/victor/data/models/chirp_v2/tokenizer_60k.json\",\n",
    "    semantic_ckpt_path=\"/home/victor/data/models/chirp_v2/mert_25.pt\",\n",
    "    semantic_centroids_path=\"/home/victor/data/models/chirp_v2/mert_25_2x4k.npy\",\n",
    "    codec_ckpt_path=\"/app/suno/tony/v3/dac_2c_25x12.pt\",\n",
    "    fasttext_ckpt_path=\"/home/victor/data/models/lid.176.bin\",\n",
    "    load_whisper=False,\n",
    "    use_gpu=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edc87efa",
   "metadata": {},
   "source": [
    "# continuation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00e1c66e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:21.811254Z",
     "start_time": "2024-01-12T17:47:21.809384Z"
    }
   },
   "outputs": [],
   "source": [
    "# # input_audio = Audio.from_file(\"audios/test_jingle.mp3\")\n",
    "# input_audio = Audio.from_file(\"audios/test_jingle_martin.wav\")\n",
    "# test_input_audio = input_audio # input_audio.get_segment(from_s=0, to_s=18)\n",
    "# test_input_audio.play(compress=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf2f70e8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-11T18:39:07.116160Z",
     "start_time": "2024-01-11T18:39:07.113475Z"
    }
   },
   "source": [
    "GT:\n",
    "Been here all along\n",
    "So, why can't you see?\n",
    "You belong with me\n",
    "[chorus]\n",
    "Standing by and waiting at your backdoor\n",
    "All this time how could you not know, baby?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bd670c6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:22.893833Z",
     "start_time": "2024-01-12T17:47:21.812277Z"
    }
   },
   "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.1)\n",
    "test_input_audio.play(compress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c155d20",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:22.897233Z",
     "start_time": "2024-01-12T17:47:22.895768Z"
    }
   },
   "outputs": [],
   "source": [
    "# input_audio = Audio.from_file(\"audios/test_hbd.mp3\")\n",
    "# test_input_audio = input_audio.get_segment(from_s=6, to_s=23)\n",
    "# test_input_audio.play(compress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1302205",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:22.972724Z",
     "start_time": "2024-01-12T17:47:22.898242Z"
    }
   },
   "outputs": [],
   "source": [
    "real_history_text = \"\"\"\n",
    "[chorus]\n",
    "Can't you see that I'm the one\n",
    "Who understands you?\n",
    "Been here all along\n",
    "So, why can't you see?\n",
    "You belong with me\n",
    "[chorus]\n",
    "Standing by and waiting at your backdoor\n",
    "All this time how could you not know, baby?\n",
    "You belong with me\n",
    "[chorus]\n",
    "You belong with me\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9550518",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:23.376309Z",
     "start_time": "2024-01-12T17:47:22.973825Z"
    }
   },
   "outputs": [],
   "source": [
    "print(hoot_encode([test_input_audio]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63134503",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:23.471729Z",
     "start_time": "2024-01-12T17:47:23.377840Z"
    }
   },
   "outputs": [],
   "source": [
    "aligned_lyrics = hoot_encode_and_align(test_input_audio, prior_texts=real_history_text)\n",
    "print(aligned_lyrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9be653a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:23.474466Z",
     "start_time": "2024-01-12T17:47:23.472920Z"
    }
   },
   "outputs": [],
   "source": [
    "# prompt_text = \"\"\"\n",
    "# [Intro]\n",
    "\n",
    "# [Verse]\n",
    "# In a stable, on a silent night\n",
    "# A star shining bright, a holy sight (ooh-yeah)\n",
    "# Wrapped in swaddling clothes, the Savior sleeps\n",
    "# Angels rejoice, the world in awe at His feet\n",
    "\n",
    "# [Chorus]\n",
    "# Star of Wonder, shining in the sky (shining bright)\n",
    "# Bringing hope and joy, the reason we testify (testify)\n",
    "# Celebrate the birth of Jesus, our King (our King)\n",
    "# With love and peace, let our voices sing (oh-oh-oh-oh-oh-oh-oh)\n",
    "\n",
    "# [Outro]\n",
    "# \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "984e8c9b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:23.533086Z",
     "start_time": "2024-01-12T17:47:23.475437Z"
    }
   },
   "outputs": [],
   "source": [
    "# out_audio = generate_audio(\n",
    "#     GenerationConfig(\n",
    "#         text=\" [verse] purple rain \" * 4 + \" [chorus] purple rain \" * 4,\n",
    "#         n_batch=4,\n",
    "#         # cfg_coef=1.5,\n",
    "#         max_gen_duration_s=30,\n",
    "#         # stream=False,\n",
    "#     ),\n",
    "#     history_audio=test_input_audio,\n",
    "#     history_text=real_history_text,\n",
    "#     max_history_duration_s=20.1,\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c04f09a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:23.589597Z",
     "start_time": "2024-01-12T17:47:23.534131Z"
    }
   },
   "outputs": [],
   "source": [
    "# for e in out_audio:\n",
    "#     e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fca409ab",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:23.655281Z",
     "start_time": "2024-01-12T17:47:23.590968Z"
    }
   },
   "outputs": [],
   "source": [
    "# out_audio = generate_audio(\n",
    "#     GenerationConfig(\n",
    "#         text=\" [verse] purple rain \" * 4 + \" [chorus] purple rain \" * 4,\n",
    "#         n_batch=4,\n",
    "#         # cfg_coef=1.5,\n",
    "#         max_gen_duration_s=30,\n",
    "#         # stream=False,\n",
    "#         use_whisper=False,\n",
    "#     ),\n",
    "#     history_audio=test_input_audio,\n",
    "#     history_text=real_history_text,\n",
    "#     max_history_duration_s=20.1,\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4a4e7e2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:54.243890Z",
     "start_time": "2024-01-12T17:47:23.656290Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\" [chorus] Victor pondering quantum kv caches \" * 4,\n",
    "        n_batch=4,\n",
    "        # cfg_coef=1.5,\n",
    "        max_gen_duration_s=30,\n",
    "        # stream=False,\n",
    "        use_whisper=False,\n",
    "    ),\n",
    "    history_audio=test_input_audio,\n",
    "    history_text=real_history_text,\n",
    "    max_history_duration_s=20.1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1663a82c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-12T17:47:56.147921Z",
     "start_time": "2024-01-12T17:47:54.246335Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04a1bfad",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1f4118f",
   "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.12"
  },
  "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
}