{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.audio import Audio\n",
    "from suno_utils.tasks.hoot import (\n",
    "    preload_models,\n",
    "    load_model_list,\n",
    ")\n",
    "\n",
    "_ = preload_models(\n",
    "    checkpoint_filepath=\"/home/tony/Data/checkpoints/hoot/2025-06-23_16-49-02/last_ckpt.pt\",\n",
    "    tokenizer_filepath=\"/home/tony/Work/tony/hoot/tokenizers/v5/tokenizer_spe_bpe_v20481/tokenizer.model\",\n",
    ")\n",
    "model_dict = load_model_list()[0]\n",
    "tokenizer = model_dict[\"tokenizer\"]\n",
    "model = model_dict[\"model\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_audio = Audio.from_s3(\n",
    "    \"s3://suno-data-uploads/studio/uploads/a5e2198a-f352-4abb-9a24-7f81b143ded3.mp3\"\n",
    ")\n",
    "\n",
    "prior_text = \"\"\"\n",
    "[verse]\n",
    "oh, my love\n",
    "My friend you know\n",
    "it's been a while\n",
    "Without thinking of you\n",
    "but the thought makes me smile\n",
    "\n",
    "[chorus]\n",
    "I'm so tired of wanting\n",
    "wanting more than this\n",
    "i know it but what am i to do\n",
    "i need some space to breathe,\n",
    "so give me some room\n",
    "\n",
    "[verse]\n",
    "oh, my love\n",
    "you have a heart of stone\n",
    "cause since i've come home\n",
    "i've never felt so alone\n",
    "but the thought makes me smile\n",
    "\n",
    "[verse]\n",
    "oh, my love\n",
    "you have a heart of stone\n",
    "cause since i've come home\n",
    "i've never felt so alone\n",
    "but the thought makes me smile\n",
    "\n",
    "[verse]\n",
    "oh, my love\n",
    "you're coming with me\n",
    "I know you'll leave me one day\n",
    "can you just stay till Monday\n",
    "the thought makes me smile\n",
    "\n",
    "[chorus]\n",
    "I'm so tired of wanting\n",
    "wanting more than this\n",
    "I know you'll leave me one day\n",
    "can you just stay till Monday\n",
    "the thought makes me smile\n",
    "\n",
    "[chorus]\n",
    "I'm so tired of wanting\n",
    "wanting more than this\n",
    "I know you'll leave me one day\n",
    "can you just stay till Monday\n",
    "the thought makes me smile\n",
    "\n",
    "[outro]\n",
    "So good\n",
    "Smile\n",
    "Smile\n",
    "Smile\n",
    "Smile\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.hoot import encode\n",
    "\n",
    "from suno_utils.tasks.lyrics_alignment.shortest_path_aligner import (\n",
    "    ShortestPathAligner,\n",
    "    ShortestPathHootV6Config,\n",
    ")\n",
    "\n",
    "spa = ShortestPathAligner.from_sentencepiece(tokenizer._tokenizer, ShortestPathHootV6Config)\n",
    "\n",
    "print(\"Greedy: \", encode(example_audio, batch_size=1))\n",
    "\n",
    "# try out debug=True\n",
    "alignment = spa.align(prior_text, encode(example_audio, return_logits=True, batch_size=1), debug=False)\n",
    "\n",
    "aligned_text = \"\".join([w[\"word\"] for w in alignment])\n",
    "print(\"Aligned (augment enabled):\", aligned_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alignment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_audio.play(alignment=alignment)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}