{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from suno_utils.tasks.hoot import preload_models, load_model_list, encode\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "UnpicklingError",
     "evalue": "Weights only load failed. This file can still be loaded, to do so you have two options, \u001b[1mdo those steps only if you trust the source of the checkpoint\u001b[0m. \n\t(1) In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.\n\t(2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.\n\tWeightsUnpickler error: Unsupported global: GLOBAL numpy.core.multiarray.scalar was not an allowed global by default. Please use `torch.serialization.add_safe_globals([scalar])` or the `torch.serialization.safe_globals([scalar])` context manager to allowlist this global if you trust this class/function.\n\nCheck the documentation of torch.load to learn more about types accepted by default with weights_only https://pytorch.org/docs/stable/generated/torch.load.html.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mUnpicklingError\u001b[0m                           Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[3], line 3\u001b[0m\n\u001b[1;32m      1\u001b[0m HOOT_CKPT_PATH \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ms3://suno-data/checkpoints/hoot_v3/hoot_ckpt.pt\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m      2\u001b[0m HOOT_TOKENIZER_PATH \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ms3://suno-data/checkpoints/hoot_v3/tokenizer.model\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m----> 3\u001b[0m _ \u001b[38;5;241m=\u001b[39m \u001b[43mpreload_models\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m      4\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcheckpoint_filepath\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mos\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpath\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mjoin\u001b[49m\u001b[43m(\u001b[49m\u001b[43mHOOT_CKPT_PATH\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      5\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtokenizer_filepath\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mHOOT_TOKENIZER_PATH\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      6\u001b[0m \u001b[43m)\u001b[49m\n\u001b[1;32m      7\u001b[0m tokenizer \u001b[38;5;241m=\u001b[39m load_model_list()[\u001b[38;5;241m0\u001b[39m][\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtokenizer\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n",
      "File \u001b[0;32m~/code/glockenspiel/suno_utils/suno_utils/tasks/hoot.py:174\u001b[0m, in \u001b[0;36mpreload_models\u001b[0;34m(checkpoint_filepath, tokenizer_filepath, n_gpus)\u001b[0m\n\u001b[1;32m    173\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mpreload_models\u001b[39m(checkpoint_filepath\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, tokenizer_filepath\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, n_gpus\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m--> 174\u001b[0m     model_list \u001b[38;5;241m=\u001b[39m \u001b[43mload_model_list\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    175\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcheckpoint_filepath\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcheckpoint_filepath\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    176\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtokenizer_filepath\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtokenizer_filepath\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    177\u001b[0m \u001b[43m        \u001b[49m\u001b[43mn_gpus\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mn_gpus\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    178\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    179\u001b[0m     \u001b[38;5;66;03m# do dummy preds to preinitialize stuff\u001b[39;00m\n\u001b[1;32m    180\u001b[0m     _ \u001b[38;5;241m=\u001b[39m encode(\n\u001b[1;32m    181\u001b[0m         [torch\u001b[38;5;241m.\u001b[39mzeros((\u001b[38;5;241m1\u001b[39m, SPLIT_PREDICT_ARRAY_LEN), dtype\u001b[38;5;241m=\u001b[39mtorch\u001b[38;5;241m.\u001b[39mfloat32)] \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mlen\u001b[39m(model_list),\n\u001b[1;32m    182\u001b[0m         batch_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m,\n\u001b[1;32m    183\u001b[0m     )\n",
      "File \u001b[0;32m~/code/glockenspiel/suno_utils/suno_utils/tasks/hoot.py:164\u001b[0m, in \u001b[0;36mload_model_list\u001b[0;34m(checkpoint_filepath, tokenizer_filepath, n_gpus)\u001b[0m\n\u001b[1;32m    161\u001b[0m model_list \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m    162\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m n \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(n_gpus):\n\u001b[1;32m    163\u001b[0m     model_list\u001b[38;5;241m.\u001b[39mappend(\n\u001b[0;32m--> 164\u001b[0m         \u001b[43mload_model\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    165\u001b[0m \u001b[43m            \u001b[49m\u001b[43mcheckpoint_filepath\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcheckpoint_filepath\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    166\u001b[0m \u001b[43m            \u001b[49m\u001b[43mtokenizer_filepath\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtokenizer_filepath\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    167\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mcuda:\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mn\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    168\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    169\u001b[0m     )\n\u001b[1;32m    170\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m model_list\n",
      "File \u001b[0;32m~/code/glockenspiel/suno_utils/suno_utils/tasks/hoot.py:141\u001b[0m, in \u001b[0;36mload_model\u001b[0;34m(checkpoint_filepath, tokenizer_filepath, device)\u001b[0m\n\u001b[1;32m    139\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m checkpoint_filepath \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mor\u001b[39;00m tokenizer_filepath \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    140\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel not loaded, need to specify checkpoint\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 141\u001b[0m     model, tokenizer \u001b[38;5;241m=\u001b[39m \u001b[43m_load_model\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcheckpoint_filepath\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtokenizer_filepath\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    142\u001b[0m     models[model_key] \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m    143\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel\u001b[39m\u001b[38;5;124m\"\u001b[39m: model,\n\u001b[1;32m    144\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtokenizer\u001b[39m\u001b[38;5;124m\"\u001b[39m: tokenizer,\n\u001b[1;32m    145\u001b[0m     }\n\u001b[1;32m    146\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m models[model_key]\n",
      "File \u001b[0;32m~/code/glockenspiel/suno_utils/suno_utils/tasks/hoot.py:120\u001b[0m, in \u001b[0;36m_load_model\u001b[0;34m(checkpoint_filepath, tokenizer_filepath, device)\u001b[0m\n\u001b[1;32m    118\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_load_model\u001b[39m(checkpoint_filepath, tokenizer_filepath, device):\n\u001b[1;32m    119\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m checkpoint_filepath\u001b[38;5;241m.\u001b[39mstartswith(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ms3://\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m--> 120\u001b[0m         model \u001b[38;5;241m=\u001b[39m \u001b[43mread_from_s3\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcheckpoint_filepath\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mread_f\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mHoot\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_checkpoint\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    121\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    122\u001b[0m         model \u001b[38;5;241m=\u001b[39m Hoot\u001b[38;5;241m.\u001b[39mfrom_checkpoint(checkpoint_filepath)\n",
      "File \u001b[0;32m~/code/glockenspiel/suno_utils/suno_utils/utils/s3.py:325\u001b[0m, in \u001b[0;36mread_from_s3\u001b[0;34m(filepath, read_f, client_config)\u001b[0m\n\u001b[1;32m    318\u001b[0m     confirmed_download \u001b[38;5;241m=\u001b[39m download_s3_files(\n\u001b[1;32m    319\u001b[0m         filepath,\n\u001b[1;32m    320\u001b[0m         tmp_filepath,\n\u001b[1;32m    321\u001b[0m         silent\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m    322\u001b[0m         client_config\u001b[38;5;241m=\u001b[39mclient_config,\n\u001b[1;32m    323\u001b[0m     )\n\u001b[1;32m    324\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m confirmed_download\n\u001b[0;32m--> 325\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[43mread_f\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtmp_filepath\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    326\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m data\n",
      "File \u001b[0;32m~/code/hoot/hoot/model.py:186\u001b[0m, in \u001b[0;36mHoot.from_checkpoint\u001b[0;34m(cls, checkpoint_fp)\u001b[0m\n\u001b[1;32m    184\u001b[0m \u001b[38;5;129m@classmethod\u001b[39m\n\u001b[1;32m    185\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mfrom_checkpoint\u001b[39m(\u001b[38;5;28mcls\u001b[39m, checkpoint_fp):\n\u001b[0;32m--> 186\u001b[0m     d \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcheckpoint_fp\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmap_location\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mcpu\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    187\u001b[0m     model \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mcls\u001b[39m(HootConfig(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39md[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel_args\u001b[39m\u001b[38;5;124m\"\u001b[39m]))\n\u001b[1;32m    188\u001b[0m     model\u001b[38;5;241m.\u001b[39mload_state_dict(d[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel\u001b[39m\u001b[38;5;124m\"\u001b[39m])\n",
      "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/serialization.py:1470\u001b[0m, in \u001b[0;36mload\u001b[0;34m(f, map_location, pickle_module, weights_only, mmap, **pickle_load_args)\u001b[0m\n\u001b[1;32m   1462\u001b[0m                 \u001b[38;5;28;01mreturn\u001b[39;00m _load(\n\u001b[1;32m   1463\u001b[0m                     opened_zipfile,\n\u001b[1;32m   1464\u001b[0m                     map_location,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1467\u001b[0m                     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mpickle_load_args,\n\u001b[1;32m   1468\u001b[0m                 )\n\u001b[1;32m   1469\u001b[0m             \u001b[38;5;28;01mexcept\u001b[39;00m pickle\u001b[38;5;241m.\u001b[39mUnpicklingError \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[0;32m-> 1470\u001b[0m                 \u001b[38;5;28;01mraise\u001b[39;00m pickle\u001b[38;5;241m.\u001b[39mUnpicklingError(_get_wo_message(\u001b[38;5;28mstr\u001b[39m(e))) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m   1471\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m _load(\n\u001b[1;32m   1472\u001b[0m             opened_zipfile,\n\u001b[1;32m   1473\u001b[0m             map_location,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1476\u001b[0m             \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mpickle_load_args,\n\u001b[1;32m   1477\u001b[0m         )\n\u001b[1;32m   1478\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m mmap:\n",
      "\u001b[0;31mUnpicklingError\u001b[0m: Weights only load failed. This file can still be loaded, to do so you have two options, \u001b[1mdo those steps only if you trust the source of the checkpoint\u001b[0m. \n\t(1) In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.\n\t(2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.\n\tWeightsUnpickler error: Unsupported global: GLOBAL numpy.core.multiarray.scalar was not an allowed global by default. Please use `torch.serialization.add_safe_globals([scalar])` or the `torch.serialization.safe_globals([scalar])` context manager to allowlist this global if you trust this class/function.\n\nCheck the documentation of torch.load to learn more about types accepted by default with weights_only https://pytorch.org/docs/stable/generated/torch.load.html."
     ]
    }
   ],
   "source": [
    "HOOT_CKPT_PATH = \"s3://suno-data/checkpoints/hoot_v3/hoot_ckpt.pt\"\n",
    "HOOT_TOKENIZER_PATH = \"s3://suno-data/checkpoints/hoot_v3/tokenizer.model\"\n",
    "_ = preload_models(\n",
    "    checkpoint_filepath=os.path.join(HOOT_CKPT_PATH),\n",
    "    tokenizer_filepath=HOOT_TOKENIZER_PATH,\n",
    ")\n",
    "tokenizer = load_model_list()[0][\"tokenizer\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}