{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import random\n",
    "import copy\n",
    "\n",
    "from suno_utils.utils.clip import SunoClip\n",
    "from suno_utils.tasks.audio_features.instrument import InstrumentExtractor\n",
    "from suno_utils.tasks.audio_features.vocal import VocalExtractor\n",
    "from suno_utils.tasks.hoot import (\n",
    "    encode,\n",
    "    encode_and_align,\n",
    "    get_word_timing_from_audio_and_lyrics,\n",
    ")\n",
    "from suno_utils.tasks.lyrics_alignment.shortest_path_aligner import (\n",
    "    ShortestPathAlignerConfig,\n",
    "    ShortestPathAligner,\n",
    ")\n",
    "\n",
    "from typing import Dict, Any, List\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_bracketed_content(text):\n",
    "    # This regex matches [section_name] or [section_name:description] with optional whitespace\n",
    "    pattern = r\"\\[\\s*([^:\\[\\]]+)(?:\\s*:\\s*([^:\\[\\]]+))?\\s*\\]\"\n",
    "\n",
    "    # Find all matches\n",
    "    matches = re.findall(pattern, text)\n",
    "\n",
    "    # Process results into a more usable format\n",
    "    results = []\n",
    "    for match in matches:\n",
    "        # Each match is a tuple of (section_name, description)\n",
    "        # If no description was found, the second element will be an empty string\n",
    "        section_name = match[0].strip()\n",
    "        description_part = match[1].strip() if match[1] else None\n",
    "\n",
    "        results.append({\"section_name\": section_name, \"description\": description_part})\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_hoot():\n",
    "    global tokenizer, model\n",
    "    from suno_utils.tasks.hoot import (\n",
    "        preload_models,\n",
    "        load_model_list,\n",
    "    )\n",
    "\n",
    "    _ = preload_models(\n",
    "        checkpoint_filepath=\"s3://suno-data/checkpoints/hoot_v3/hoot_ckpt.pt\",\n",
    "        tokenizer_filepath=\"s3://suno-data/checkpoints/hoot_v3/tokenizer.model\",\n",
    "    )\n",
    "    model_dict = load_model_list()[0]\n",
    "    return model_dict[\"tokenizer\"], model_dict[\"model\"]\n",
    "\n",
    "\n",
    "tokenizer, model = init_hoot()\n",
    "ie = InstrumentExtractor()\n",
    "ve = VocalExtractor()\n",
    "\n",
    "good_spa_config = ShortestPathAlignerConfig(\n",
    "    enable_jumps=True,\n",
    "    logit_skip_coef=13.2,\n",
    "    char_skip_coef=6.1,\n",
    "    spelling_error_coef=2.1,\n",
    "    section_skip_coef=27.2,\n",
    "    line_skip_coef=34.0,\n",
    "    char_epsilon=0.0007,\n",
    "    silence_threshold_p=0.9,\n",
    ")\n",
    "\n",
    "spa = ShortestPathAligner.from_sentencepiece(tokenizer._tokenizer, good_spa_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clip = SunoClip(\"75af6dcd-07b2-4f01-9f65-46747e639ed8\")\n",
    "vae_arr = clip.full_arr()\n",
    "example_audio = clip.audio()\n",
    "\n",
    "length_diff = (example_audio.duration_s) / (vae_arr.shape[0] / 25)\n",
    "print(length_diff)\n",
    "assert length_diff < 1.01 and length_diff > 0.99\n",
    "example_audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prior_text = \"\"\"\n",
    "[Verse 1]\n",
    "If you really loved me \n",
    "You wouldn't make me dance\n",
    "If you stepped into my mind \n",
    "You'd see we're not the same\n",
    "\n",
    "[Build Up]\n",
    "Dance\n",
    "no no no no no no no\n",
    "Dance\n",
    "won't dance won't dance\n",
    "\n",
    "[Chorus]\n",
    "No dance\n",
    "(No dance)\n",
    "Absence of rhythm.\n",
    "Lacking in movement.\n",
    "Won't dance\n",
    "(Won't Dance)\n",
    "\n",
    "[melodic interlude]\n",
    "\n",
    "[Verse 2]\n",
    "If you really loved me\n",
    "You wouldn't make me dance \n",
    "If you could see it through my eyes\n",
    "You'd see the alteration in my stance\n",
    "\n",
    "[Build Up]\n",
    "Dance\n",
    "no no no no no no no\n",
    "Dance\n",
    "won't dance won't dance\n",
    "\n",
    "[Chorus]\n",
    "No dance\n",
    "(No dance)\n",
    "Absence of rhythm,\n",
    "Lacking in movement.\n",
    "Won't dance\n",
    "(Won't dance)\n",
    "\n",
    "[breakdown]\n",
    "no dance no dance no dance\n",
    "no movement\n",
    "no dance\n",
    "\n",
    "[Final Chorus]\n",
    "No dance\n",
    "(No dance)\n",
    "Absence of rhythm,\n",
    "Lacking in movement.\n",
    "Won't dance\n",
    "(Won't dance)\n",
    "\n",
    "[outro]\n",
    "no dance no dance no dance\n",
    "no movement\n",
    "no dance\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = spa.align(prior_text, encode(example_audio, return_logits=True, batch_size=1))\n",
    "\n",
    "section_timings = []\n",
    "last_word_start_s = 0.0\n",
    "for idx, row in enumerate(output):\n",
    "    text = row[\"word\"]\n",
    "    headers = extract_bracketed_content(text)\n",
    "    if len(headers) > 0:\n",
    "        if idx > 0:\n",
    "            last_word_start_s = output[idx - 1][\"start_s\"]\n",
    "        for header in headers:\n",
    "            section_timings.append(\n",
    "                {\n",
    "                    \"section_name\": header[\"section_name\"],\n",
    "                    \"start_time_s\": row[\"start_s\"] if len(section_timings) > 0 else 0.0,\n",
    "                    \"end_time_s\": example_audio.duration_s,\n",
    "                    \"description\": header[\"description\"],\n",
    "                    \"last_start_s\": last_word_start_s,\n",
    "                }\n",
    "            )\n",
    "            if len(section_timings) > 1:\n",
    "                if row[\"start_s\"] == section_timings[-2][\"start_time_s\"]:\n",
    "                    section_timings[-2][\"start_time_s\"] = section_timings[-2][\n",
    "                        \"last_start_s\"\n",
    "                    ]\n",
    "                section_timings[-2][\"end_time_s\"] = last_word_start_s\n",
    "            last_word_start_s = row[\"start_s\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sections = []\n",
    "for idx, row in enumerate(section_timings):\n",
    "    from_token = int(row[\"start_time_s\"] * 25)\n",
    "    to_token = int(row[\"end_time_s\"] * 25)\n",
    "    name = row[\"section_name\"]\n",
    "    description = row[\"description\"]\n",
    "    segment_audio = example_audio.get_segment(\n",
    "        from_s=row[\"start_time_s\"], to_s=row[\"end_time_s\"]\n",
    "    )\n",
    "    data = vae_arr[from_token:to_token]\n",
    "    group_logits, instrument_logits, group_tags, instrument_tags = ie.extract(\n",
    "        data, threshold=0.6\n",
    "    )\n",
    "    gender_logits, gender = ve.extract(data, threshold=0.6)\n",
    "    print(\n",
    "        f\"{name}: {description}, gender: {gender}, found {group_tags}, {instrument_tags}\"\n",
    "    )\n",
    "    segment_audio.play(compress=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_clean",
   "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": 2
}