{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0e09c52f-da2a-4a21-b659-b58383cba6f4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-27T04:39:31.263938Z",
     "iopub.status.busy": "2024-10-27T04:39:31.263804Z",
     "iopub.status.idle": "2024-10-27T04:39:31.268228Z",
     "shell.execute_reply": "2024-10-27T04:39:31.267834Z",
     "shell.execute_reply.started": "2024-10-27T04:39:31.263924Z"
    }
   },
   "outputs": [],
   "source": [
    "# !nvidia-smi\n",
    "# !echo $HOSTNAME\n",
    "\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\"\n",
    "\n",
    "import sys\n",
    "#sys.path.insert(0, \"/home/georg/code/neon/sunoDiff/\")\n",
    "sys.path.insert(0, \"/home/tony/Work/neon/sunoDiff/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c046f7b4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-27T04:39:31.268867Z",
     "iopub.status.busy": "2024-10-27T04:39:31.268739Z",
     "iopub.status.idle": "2024-10-27T04:39:40.373472Z",
     "shell.execute_reply": "2024-10-27T04:39:40.372413Z",
     "shell.execute_reply.started": "2024-10-27T04:39:31.268854Z"
    }
   },
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'torch.nn.attention.flex_attention'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[2], line 5\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mgeneration\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m preload_models, generate, _retrieve_models\n\u001b[1;32m      4\u001b[0m CKPT_DIR \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/home/christian/code/neon/stable-audio-tools/harmonai_train/\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m----> 5\u001b[0m _ \u001b[38;5;241m=\u001b[39m \u001b[43mpreload_models\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m      6\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtokenizer_filepath\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m/app/suno/data/dpo/models/tokenizer_60k.json\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m      7\u001b[0m \u001b[43m    \u001b[49m\u001b[43msemantic_model_filepath\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m/app/suno/data/dpo/models/mert_25.pt\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m      8\u001b[0m \u001b[43m    \u001b[49m\u001b[43msemantic_clusters_filepath\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m/app/suno/data/dpo/models/mert_25_2x4k.npy\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m      9\u001b[0m \u001b[43m    \u001b[49m\u001b[43mweights_precision\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbfloat16\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     10\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mcompile\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     11\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m### 25hz ###\u001b[39;49;00m\n\u001b[1;32m     12\u001b[0m \u001b[38;5;66;43;03m#     codec_filepath=\"/home/georg/notebooks/gpu_nb/tmp/25hz_vae_peaq_kl_0.005.pth\",\u001b[39;49;00m\n\u001b[1;32m     13\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m### 100hz ###\u001b[39;49;00m\n\u001b[1;32m     14\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcodec_filepath\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m/home/georg/notebooks/gpu_nb/tmp/100hz_vae_peaq_kl_0.005.pth\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     15\u001b[0m \u001b[43m    \u001b[49m\n\u001b[1;32m     16\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m# 30s (overlap newest)\u001b[39;49;00m\n\u001b[1;32m     17\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m# dit_model_filepath=\"/app/suno/checkpoints/2024-10-17_04-15-19_s4176/last_ckpt_infer.pt\",\u001b[39;49;00m\n\u001b[1;32m     18\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m# dit_model_filepath=\"/app/suno/checkpoints/2024-10-11_16-53-57_s8050/last_ckpt_infer.pt\", # default\u001b[39;49;00m\n\u001b[1;32m     19\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdit_model_filepath\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m/app/suno/checkpoints/2024-10-27_00-05-50_s2159/last_ckpt_infer.pt\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     20\u001b[0m \u001b[43m)\u001b[49m\n\u001b[1;32m     21\u001b[0m models \u001b[38;5;241m=\u001b[39m _retrieve_models()\n\u001b[1;32m     22\u001b[0m model_duration_s \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m30\u001b[39m\n",
      "File \u001b[0;32m~/Work/neon/sunoDiff/generation.py:145\u001b[0m, in \u001b[0;36mpreload_models\u001b[0;34m(tokenizer_filepath, semantic_model_filepath, semantic_clusters_filepath, codec_filepath, dit_model_filepath, dit_model_2_filepath, use_ema_if_exists, model_type, compile, weights_precision)\u001b[0m\n\u001b[1;32m    143\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmodel\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m DiffusionTransformer\n\u001b[1;32m    144\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m model_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mprefix\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m--> 145\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mprefix_model\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmodel\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m DiffusionTransformer\n\u001b[1;32m    146\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    147\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mNotImplementedError\u001b[39;00m()\n",
      "File \u001b[0;32m~/Work/neon/sunoDiff/prefix_model/model.py:9\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdistributed\u001b[39;00m\n\u001b[0;32m----> 9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbase\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ContinuousTransformer\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mFourierFeatures\u001b[39;00m(nn\u001b[38;5;241m.\u001b[39mModule):\n\u001b[1;32m     13\u001b[0m     \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, in_features, out_features):\n",
      "File \u001b[0;32m~/Work/neon/sunoDiff/prefix_model/base.py:14\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnn\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnn\u001b[39;00m\n\u001b[1;32m     13\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnn\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfunctional\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mF\u001b[39;00m\n\u001b[0;32m---> 14\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnn\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mattention\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mflex_attention\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m flex_attention, create_block_mask\n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# drastic speedup\u001b[39;00m\n\u001b[1;32m     18\u001b[0m flex_attention \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mcompile(flex_attention)\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'torch.nn.attention.flex_attention'"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from suno_utils.audio import Audio\n",
    "from generation import preload_models, generate, _retrieve_models\n",
    "CKPT_DIR = \"/home/christian/code/neon/stable-audio-tools/harmonai_train/\"\n",
    "_ = preload_models(\n",
    "    tokenizer_filepath=\"/app/suno/data/dpo/models/tokenizer_60k.json\",\n",
    "    semantic_model_filepath=\"/app/suno/data/dpo/models/mert_25.pt\",\n",
    "    semantic_clusters_filepath=\"/app/suno/data/dpo/models/mert_25_2x4k.npy\",\n",
    "    weights_precision=torch.bfloat16,\n",
    "    compile=True,\n",
    "    ### 25hz ###\n",
    "#     codec_filepath=\"/home/georg/notebooks/gpu_nb/tmp/25hz_vae_peaq_kl_0.005.pth\",\n",
    "    ### 100hz ###\n",
    "    codec_filepath=\"/home/georg/notebooks/gpu_nb/tmp/100hz_vae_peaq_kl_0.005.pth\",\n",
    "    \n",
    "    # 30s (overlap newest)\n",
    "    # dit_model_filepath=\"/app/suno/checkpoints/2024-10-17_04-15-19_s4176/last_ckpt_infer.pt\",\n",
    "    # dit_model_filepath=\"/app/suno/checkpoints/2024-10-11_16-53-57_s8050/last_ckpt_infer.pt\", # default\n",
    "    dit_model_filepath=\"/app/suno/checkpoints/2024-10-27_00-05-50_s2159/last_ckpt_infer.pt\",\n",
    ")\n",
    "models = _retrieve_models()\n",
    "model_duration_s = 30\n",
    "if models[\"dit_model\"].ctx_len is not None:\n",
    "    model_duration_s = 6 * 60\n",
    "else:\n",
    "    model_duration_s = models[\"dit_model\"].block_size // models[\"dit_model\"].io_hz\n",
    "duration_s = 2*60 if model_duration_s >= 2*60 else 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf696991",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-10-27T04:39:40.373980Z",
     "iopub.status.idle": "2024-10-27T04:39:40.374160Z",
     "shell.execute_reply": "2024-10-27T04:39:40.374075Z",
     "shell.execute_reply.started": "2024-10-27T04:39:40.374067Z"
    }
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "\n",
    "#gen_id = \"99bffa17-7e59-47b4-a048-5528cbda05d5\" # sister\n",
    "#gen_id = \"081d73c4-7805-4212-9c80-8db1137ca3c4\" # friends\n",
    "#gen_id = \"562f762d-6ced-4080-9af1-910ee3d0a5dc\" # something real\n",
    "#gen_id = \"23c15c62-494d-422d-8a60-8b0454044322\" # rubber duck\n",
    "#gen_id = \"4b140a9e-964b-422c-85b5-5861ad1a9d38\" # once\n",
    "#gen_id = \"7b214347-fa38-4e9b-96f4-f7ec65adea45\" # rock n roll\n",
    "# gen_id = \"a5e2198a-f352-4abb-9a24-7f81b143ded3\" # stone\n",
    "#gen_id = \"7f5e7819-b7a1-449d-a76c-9c38950c4274\" # canada\n",
    "#gen_id = \"27135940-957f-4152-9e40-72e5c10a46bd\" # drunk\n",
    "#gen_id = \"caf1d7ae-ce81-4c75-b737-fd41e47a38ba\" # linger\n",
    "#gen_id = \"ac5aa29a-f18f-406f-a904-c47493f1beec\" # linger vox\n",
    "#gen_id = \"b04719c3-012a-4a4a-9b02-3cb2875d1744\" # courtney\n",
    "#gen_id = \"7407c7b3-8aa2-4f0a-82e8-fa19933f6e94\" # opera\n",
    "#gen_id = \"5b9624ef-b416-4095-a5f3-4945f09bba86\" # butterflies\n",
    "#gen_id = \"e757aa19-5c33-4945-aaa1-1519b7890f0b\" # bubblegum\n",
    "#gen_id = \"7f774078-1672-4858-a37f-acad373c5a84\" # golden\n",
    "gen_id = \"1cee79db-a0ee-44e4-9646-d312cd620997\" # chinese\n",
    "\n",
    "s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.npz\"\n",
    "mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.mp3\"\n",
    "print(s3_filepath)\n",
    "data = read_from_s3(s3_filepath, read_f=np.load)\n",
    "\n",
    "audio = Audio.from_s3(mp3_filepath, n_channels=2)\n",
    "\n",
    "if \"v3.0_raw\" in data:\n",
    "    codes = data[\"v3.0_raw\"]\n",
    "elif \"v3.5_raw\" in data:\n",
    "    codes = data[\"v3.5_raw\"]\n",
    "elif \"v4.0_raw\" in data:\n",
    "    codes = data[\"v4.0_raw\"]\n",
    "else:\n",
    "    raise ValueError(\"No codes found\")\n",
    "\n",
    "text_data = read_from_s3(f\"s3://suno-data-uploads/studio/uploads/{gen_id}_hoot.json\")\n",
    "aligned_lyrics = json.loads(text_data)\n",
    "# print(aligned_lyrics)\n",
    "audio.normalize_volume().play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b6e8953",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-10-27T04:39:40.374802Z",
     "iopub.status.idle": "2024-10-27T04:39:40.374960Z",
     "shell.execute_reply": "2024-10-27T04:39:40.374887Z",
     "shell.execute_reply.started": "2024-10-27T04:39:40.374879Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "seeds = np.random.randint(0, 2**32 - 1, 3)\n",
    "print(seeds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0c34cd3",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-10-27T04:39:40.375432Z",
     "iopub.status.idle": "2024-10-27T04:39:40.375582Z",
     "shell.execute_reply": "2024-10-27T04:39:40.375512Z",
     "shell.execute_reply.started": "2024-10-27T04:39:40.375505Z"
    }
   },
   "outputs": [],
   "source": [
    "steps = 32\n",
    "text_cfg_coef = 1.0\n",
    "ctx_cfg_coef = 2.0\n",
    "\n",
    "audios = []\n",
    "for n in seeds:\n",
    "    pred_audio = generate(\n",
    "        audio, \n",
    "#         lyrics=lyrics,\n",
    "        aligned_lyrics=aligned_lyrics,\n",
    "        tags='chinese pop', \n",
    "        text_cfg_coef=text_cfg_coef, \n",
    "        ctx_cfg_coef=ctx_cfg_coef,\n",
    "        steps=steps,\n",
    "        seed=n,\n",
    "    ).normalize_volume()\n",
    "    audios.append(pred_audio)\n",
    "    pred_audio.play()\n",
    "#     break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00e3d45c",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-10-27T04:39:40.376132Z",
     "iopub.status.idle": "2024-10-27T04:39:40.376281Z",
     "shell.execute_reply": "2024-10-27T04:39:40.376211Z",
     "shell.execute_reply.started": "2024-10-27T04:39:40.376204Z"
    }
   },
   "outputs": [],
   "source": [
    "aligned_lyrics = [m for m in aligned_lyrics if \"word\" in m]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "311039b6",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-10-27T04:39:40.376793Z",
     "iopub.status.idle": "2024-10-27T04:39:40.376942Z",
     "shell.execute_reply": "2024-10-27T04:39:40.376872Z",
     "shell.execute_reply.started": "2024-10-27T04:39:40.376865Z"
    }
   },
   "outputs": [],
   "source": [
    "for word in aligned_lyrics:\n",
    "    print(word) #[\"word\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77725223",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-10-27T04:39:40.377409Z",
     "iopub.status.idle": "2024-10-27T04:39:40.377572Z",
     "shell.execute_reply": "2024-10-27T04:39:40.377498Z",
     "shell.execute_reply.started": "2024-10-27T04:39:40.377490Z"
    }
   },
   "outputs": [],
   "source": [
    "chunk_window_size = 3000\n",
    "\n",
    "for n_chunk in range(8):\n",
    "    idx_pad = 5\n",
    "    chunk_start_s = n_chunk * chunk_window_size / models[\"dit_model\"].io_hz\n",
    "    chunk_end_s = (n_chunk + 1) * chunk_window_size / models[\"dit_model\"].io_hz\n",
    "    print(chunk_start_s, chunk_end_s)\n",
    "\n",
    "    chunk_start_indices = [\n",
    "        idx\n",
    "        for idx, m in enumerate(aligned_lyrics)\n",
    "        if \"start_s\" in m and m[\"start_s\"] >= chunk_start_s\n",
    "    ]\n",
    "    chunk_end_indices = [\n",
    "        idx for idx, m in enumerate(aligned_lyrics) if \"end_s\" in m and m[\"end_s\"] <= chunk_end_s\n",
    "    ]\n",
    "\n",
    "\n",
    "    if len(chunk_start_indices) > 0 and len(chunk_end_indices) > 0:\n",
    "        chunk_start_idx = max(0, chunk_start_indices[0] - idx_pad)\n",
    "        chunk_end_idx = min(len(aligned_lyrics), chunk_end_indices[-1] + idx_pad)\n",
    "        print(chunk_start_idx, chunk_end_idx)\n",
    "        lyrics_chunk = \"\".join(\n",
    "            [m[\"word\"] for m in aligned_lyrics[chunk_start_idx:chunk_end_idx]]\n",
    "        )\n",
    "        #lyrics_chunk = simplify_whitespace(lyrics_chunk, retain_newlines=True)\n",
    "    else:\n",
    "        lyrics_chunk = \"\"\n",
    "    print(lyrics_chunk)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "441ef260",
   "metadata": {},
   "source": [
    "# Testbench"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2115a8e4",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-10-27T04:39:40.377925Z",
     "iopub.status.idle": "2024-10-27T04:39:40.378070Z",
     "shell.execute_reply": "2024-10-27T04:39:40.378003Z",
     "shell.execute_reply.started": "2024-10-27T04:39:40.377996Z"
    }
   },
   "outputs": [],
   "source": [
    "# import json\n",
    "# import numpy as np\n",
    "# from suno_utils.utils.s3 import read_from_s3\n",
    "# import shutil\n",
    "# output_dir = \"outputs/custom_infer_stable_audio-10102024\"\n",
    "# if os.path.exists(output_dir):\n",
    "#     shutil.rmtree(output_dir)\n",
    "# os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "# gen_ids = [\n",
    "#     (\"99bffa17-7e59-47b4-a048-5528cbda05d5\", \"country\"), # sister\n",
    "#     (\"081d73c4-7805-4212-9c80-8db1137ca3c4\", \"rock\"), # friends\n",
    "#     (\"562f762d-6ced-4080-9af1-910ee3d0a5dc\", \"rock\"), # something real\n",
    "#     (\"23c15c62-494d-422d-8a60-8b0454044322\", \"shanty\"), # rubber duck\n",
    "#     (\"4b140a9e-964b-422c-85b5-5861ad1a9d38\", \"indie rock\"), # once\n",
    "#     (\"7b214347-fa38-4e9b-96f4-f7ec65adea45\", \"rock n roll\"), # rock n roll\n",
    "#     (\"7f5e7819-b7a1-449d-a76c-9c38950c4274\", \"pop\"), # canada\n",
    "#     (\"27135940-957f-4152-9e40-72e5c10a46bd\", \"pop\"), # drunk\n",
    "#     (\"caf1d7ae-ce81-4c75-b737-fd41e47a38ba\", \"pop\"), # linger\n",
    "#     (\"ac5aa29a-f18f-406f-a904-c47493f1beec\", \"pop\"), # linger vox\n",
    "#     (\"b04719c3-012a-4a4a-9b02-3cb2875d1744\", \"pop\"), # courtney\n",
    "#     (\"7407c7b3-8aa2-4f0a-82e8-fa19933f6e94\", \"opera\"), # opera\n",
    "#     (\"5b9624ef-b416-4095-a5f3-4945f09bba86\", \"math rock\"), # butterflies\n",
    "# ]\n",
    "\n",
    "# for gen_id, tags in gen_ids:\n",
    "#     mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.mp3\"\n",
    "#     print(tags, mp3_filepath)\n",
    "#     audio = Audio.from_s3(mp3_filepath, n_channels=2)\n",
    "#     audio.normalize_volume()\n",
    "#     text_data = read_from_s3(f\"s3://suno-data-uploads/studio/uploads/{gen_id}_hoot.json\")\n",
    "#     aligned_lyrics = json.loads(text_data)\n",
    "\n",
    "#     steps = 32\n",
    "#     text_cfg_coef = 4.0\n",
    "#     ctx_cfg_coef = 2.0\n",
    "#     seeds = np.random.randint(0, 2**32 - 1, 3)\n",
    "\n",
    "#     for n in seeds:\n",
    "#         pred_audio = generate(\n",
    "#             audio, \n",
    "#             aligned_lyrics=aligned_lyrics,\n",
    "#             tags=tags, \n",
    "#             text_cfg_coef=text_cfg_coef, \n",
    "#             ctx_cfg_coef=ctx_cfg_coef,\n",
    "#             steps=steps,\n",
    "#             seed=n,\n",
    "#         ).normalize_volume()\n",
    "\n",
    "#         # save audio to disk\n",
    "#         pred_audio.write_mp3(f\"{output_dir}/pred_{gen_id}_{n}.mp3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb494e8a",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ce62b93",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-10-27T04:39:40.378615Z",
     "iopub.status.idle": "2024-10-27T04:39:40.378763Z",
     "shell.execute_reply": "2024-10-27T04:39:40.378693Z",
     "shell.execute_reply.started": "2024-10-27T04:39:40.378687Z"
    }
   },
   "outputs": [],
   "source": [
    "# import json\n",
    "# with open(\"walking_aligned_lyrics.json\") as f:\n",
    "#     aligned_lyrics = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75675270",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-10-27T04:39:40.379286Z",
     "iopub.status.idle": "2024-10-27T04:39:40.379429Z",
     "shell.execute_reply": "2024-10-27T04:39:40.379361Z",
     "shell.execute_reply.started": "2024-10-27T04:39:40.379354Z"
    }
   },
   "outputs": [],
   "source": [
    "# # TODO: pad with silence for integer multiple of sliding\n",
    "# if models[\"dit_model\"].ctx_len is not None:\n",
    "#     window_s = int(round((models[\"dit_model\"].block_size - models[\"dit_model\"].ctx_len) / models[\"dit_model\"].io_hz))\n",
    "#     if audio.duration_s % window_s > 0.1:\n",
    "#         audio = Audio.concatenate([\n",
    "#             audio, \n",
    "#             Audio.from_silence(window_s - audio.duration_s % window_s + 0.01, sample_rate=audio.sample_rate, n_channels=audio.n_channels)\n",
    "#         ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2d70822",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-10-27T04:39:40.379797Z",
     "iopub.status.idle": "2024-10-27T04:39:40.379944Z",
     "shell.execute_reply": "2024-10-27T04:39:40.379877Z",
     "shell.execute_reply.started": "2024-10-27T04:39:40.379869Z"
    }
   },
   "outputs": [],
   "source": [
    "# # get latents for preconditioning\n",
    "# import torch\n",
    "# from suno_utils.tasks.dac_vae_100hz_peaq import encode as codec_encode\n",
    "# # latent_audio = audio\n",
    "# latent_audio = Audio.from_file(\"test.mp3\", sample_rate=48_000, n_channels=2).get_segment(from_s=offset_s, to_s=offset_s+duration_s)\n",
    "# init_latents = torch.from_numpy(codec_encode(latent_audio).T[None]) * 2.5\n",
    "# latent_audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c946bc7",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-10-27T04:39:40.380445Z",
     "iopub.status.idle": "2024-10-27T04:39:40.380590Z",
     "shell.execute_reply": "2024-10-27T04:39:40.380523Z",
     "shell.execute_reply.started": "2024-10-27T04:39:40.380516Z"
    }
   },
   "outputs": [],
   "source": [
    "# import torch\n",
    "# models = _retrieve_models()\n",
    "# extra_args = torch.load(\"extra_args.pt\", weights_only=True)\n",
    "# torch.manual_seed(0)\n",
    "# x = torch.randn([1, 128, 100*30], device=\"cuda\")*50\n",
    "# with torch.no_grad():\n",
    "#     out = models[\"dit_model\"].forward_inference(\n",
    "#         x,\n",
    "#         x.new_ones([x.shape[0]])*50,\n",
    "#         cross_attn_cond=extra_args[\"cross_attn_cond\"],\n",
    "#         empty_cross_attn_cond=extra_args[\"empty_cross_attn_cond\"],\n",
    "#         cfg_scale=extra_args[\"cfg_scale\"]\n",
    "#     )\n",
    "# print(out.shape)\n",
    "# out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "818cf6a9",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-10-27T04:39:40.381072Z",
     "iopub.status.idle": "2024-10-27T04:39:40.381216Z",
     "shell.execute_reply": "2024-10-27T04:39:40.381148Z",
     "shell.execute_reply.started": "2024-10-27T04:39:40.381141Z"
    }
   },
   "outputs": [],
   "source": [
    "# PAG\n",
    "#  https://arxiv.org/pdf/2403.17377\n",
    "#  https://huggingface.co/hyoungwoncho/sd_perturbed_attention_guidance/blob/main/pipeline.py\n",
    "#  https://github.com/KU-CVLAB/Perturbed-Attention-Guidance?tab=readme-ov-file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abe70d19",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56142e0c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d56f346",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b12de9ec",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acda5fe8",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}