{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37556f9b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-25T19:21:44.048855Z",
     "start_time": "2024-04-25T19:21:43.788713Z"
    },
    "execution": {
     "iopub.execute_input": "2024-08-27T18:12:08.112673Z",
     "iopub.status.busy": "2024-08-27T18:12:08.112533Z",
     "iopub.status.idle": "2024-08-27T18:12:08.444995Z",
     "shell.execute_reply": "2024-08-27T18:12:08.444542Z",
     "shell.execute_reply.started": "2024-08-27T18:12:08.112658Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "import numpy as np\n",
    "\n",
    "from suno_utils.audio import Audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4bb89de",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-25T19:21:48.454673Z",
     "start_time": "2024-04-25T19:21:44.050943Z"
    },
    "execution": {
     "iopub.execute_input": "2024-08-27T18:12:08.445728Z",
     "iopub.status.busy": "2024-08-27T18:12:08.445548Z",
     "iopub.status.idle": "2024-08-27T18:12:13.599688Z",
     "shell.execute_reply": "2024-08-27T18:12:13.599098Z",
     "shell.execute_reply.started": "2024-08-27T18:12:08.445713Z"
    }
   },
   "outputs": [],
   "source": [
    "from suno_utils.tasks.dac_2c_12cb import preload_models as preload_codec_models\n",
    "from suno_utils.tasks.dac_2c_12cb import (\n",
    "    encode as codec_encode,\n",
    "    decode_stream_to_full_audio as codec_decode,\n",
    "    EMBEDDING_RATE as CODEC_EMBEDDING_RATE,\n",
    "    decode as decode\n",
    ")\n",
    "from suno_utils.worker.settings import s3_client\n",
    "_ = preload_codec_models(\"/app/suno/data/dpo/models/dac_2c_25x12.pt\", device=\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f620d310",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tempfile\n",
    "f_path = \"audios/cbf8bfea-5f0d-400d-9b51-66c0136dd9c7.npz\"\n",
    "with open(f_path, \"wb\") as f:\n",
    "    s3_client.download_fileobj(\"suno-data-uploads\", \"studio/uploads/cbf8bfea-5f0d-400d-9b51-66c0136dd9c7.npz\", f)\n",
    "test_npz = np.load(f_path)\n",
    "for key in test_npz.keys():\n",
    "    print(key)\n",
    "codec_labels = test_npz.get(\"v4.0_raw\")[:, 1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26c75445",
   "metadata": {},
   "outputs": [],
   "source": [
    "codec_labels = test_npz.get(\"v4.0_raw\")[:, 1:]\n",
    "audio = codec_decode(codec_labels, n_stride_tokens=60, n_overlap_tokens=20)\n",
    "print(audio.duration_s, codec_labels.shape)\n",
    "audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78b4ff45",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-25T20:32:46.932412Z",
     "start_time": "2024-04-25T20:32:46.930237Z"
    },
    "execution": {
     "iopub.execute_input": "2024-08-27T18:12:13.600692Z",
     "iopub.status.busy": "2024-08-27T18:12:13.600455Z",
     "iopub.status.idle": "2024-08-27T18:12:13.603139Z",
     "shell.execute_reply": "2024-08-27T18:12:13.602739Z",
     "shell.execute_reply.started": "2024-08-27T18:12:13.600675Z"
    }
   },
   "outputs": [],
   "source": [
    "test_id = \"dc333f7e-862d-4209-9c89-6802ed9b018c\"\n",
    "test_id_part_1 = \"a372fddb-6894-45d5-82bf-c7da216c384b\"\n",
    "test_id_part_2 = \"d556816c-2f68-4f1d-8aee-c050f47ab3be\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af648485",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-25T20:32:52.600063Z",
     "start_time": "2024-04-25T20:32:47.705568Z"
    },
    "execution": {
     "iopub.execute_input": "2024-08-27T18:12:13.603957Z",
     "iopub.status.busy": "2024-08-27T18:12:13.603822Z",
     "iopub.status.idle": "2024-08-27T18:12:14.283948Z",
     "shell.execute_reply": "2024-08-27T18:12:14.283422Z",
     "shell.execute_reply.started": "2024-08-27T18:12:13.603943Z"
    }
   },
   "outputs": [],
   "source": [
    "# upsampled\n",
    "new_npz = np.load(f\"/home/tony/Data/test_npz/{test_id}.npz\")\n",
    "for key in new_npz.keys():\n",
    "    print(key)\n",
    "codec_labels = new_npz.get(\"v4.0_raw\")[:, 1:]\n",
    "audio = codec_decode(codec_labels)\n",
    "print(audio.duration_s, codec_labels.shape)\n",
    "audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfece04f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-25T20:33:39.004908Z",
     "start_time": "2024-04-25T20:33:34.118379Z"
    },
    "execution": {
     "iopub.execute_input": "2024-08-27T18:12:14.284728Z",
     "iopub.status.busy": "2024-08-27T18:12:14.284581Z",
     "iopub.status.idle": "2024-08-27T18:12:14.835964Z",
     "shell.execute_reply": "2024-08-27T18:12:14.835526Z",
     "shell.execute_reply.started": "2024-08-27T18:12:14.284714Z"
    }
   },
   "outputs": [],
   "source": [
    "# recycled\n",
    "new_npz = np.load(f\"/home/tony/Data/test_npz/{test_id_part_1}.npz\")\n",
    "for key in new_npz.keys():\n",
    "    print(key)\n",
    "codec_labels_1 = new_npz.get(\"v3.0_raw\")[:, 1:]\n",
    "audio_1 = codec_decode(codec_labels_1)\n",
    "print(audio_1.duration_s, codec_labels_1.shape)\n",
    "audio_1.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca42bef2-de4c-4a4f-bf34-0942131fb478",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-08-27T18:12:14.837337Z",
     "iopub.status.busy": "2024-08-27T18:12:14.837182Z",
     "iopub.status.idle": "2024-08-27T18:12:15.384157Z",
     "shell.execute_reply": "2024-08-27T18:12:15.383736Z",
     "shell.execute_reply.started": "2024-08-27T18:12:14.837323Z"
    }
   },
   "outputs": [],
   "source": [
    "new_npz = np.load(f\"/home/tony/Data/test_npz/{test_id_part_2}.npz\")\n",
    "for key in new_npz.keys():\n",
    "    print(key)\n",
    "print(new_npz.get(\"v4.0_raw\").shape)\n",
    "print(new_npz.get(\"full_arr\").shape)\n",
    "print(new_npz.get(\"history_start_index\"))\n",
    "print(new_npz.get(\"generated_start_index\"))\n",
    "print(new_npz.get(\"future_start_index\"))\n",
    "codec_labels_2 = new_npz.get(\"v4.0_raw\")[:, 1:]\n",
    "audio_2 = codec_decode(codec_labels_2)\n",
    "audio_2.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f87b2b59-321a-4918-b914-e227c1040415",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from typing import Tuple\n",
    "from suno_utils.audio import Audio\n",
    "\n",
    "def concat(audio_a: Audio, code_a: np.ndarray, audio_b: Audio, code_b: np.ndarray, \n",
    "           a_overlap: int = 25, b_overlap: int = 25, overlap_chunk: int = 5) -> Tuple[Audio, np.ndarray]:\n",
    "    \"\"\"\n",
    "    Concatenate two audio segments and their corresponding codec labels with overlap.\n",
    "\n",
    "    Args:\n",
    "        audio_a (Audio): First audio segment\n",
    "        code_a (np.ndarray): Codec labels for first audio segment\n",
    "        audio_b (Audio): Second audio segment\n",
    "        code_b (np.ndarray): Codec labels for second audio segment\n",
    "        a_overlap (int): Number of frames to overlap from the end of first segment\n",
    "        b_overlap (int): Number of frames to overlap from the start of second segment\n",
    "        overlap_chunk (int): Number of frames to use for smooth transition\n",
    "\n",
    "    Returns:\n",
    "        Tuple[Audio, np.ndarray]: Concatenated audio and codec labels\n",
    "    \"\"\"\n",
    "    codec_freq = 25  # Codec frequency in Hz\n",
    "\n",
    "    # Extract overlap regions from codec labels\n",
    "    code_b_overlap = code_b[:b_overlap, :]\n",
    "    code_a_head, code_a_overlap = code_a[:-a_overlap, ], code_a[-a_overlap:, :]\n",
    "    print(f\"Overlap shapes: {code_b_overlap.shape}, {code_a_overlap.shape}\")\n",
    "\n",
    "    # Merge overlapping codec labels\n",
    "    code_merge = np.vstack([code_a_overlap, code_b_overlap])\n",
    "    print(f\"Merged shape: {code_a_head.shape}, {code_merge.shape}\")\n",
    "\n",
    "    # Decode merged codec labels\n",
    "    audio_overlap = codec_decode(code_merge)\n",
    "    print(f\"Overlap duration: {audio_overlap.duration_s}\")\n",
    "\n",
    "    # Extract and merge audio segments\n",
    "    audio_a_head = audio_a.get_segment(0, (code_a_head.shape[0] + overlap_chunk) / codec_freq)\n",
    "    audio_overlap = audio_overlap.get_segment(overlap_chunk / codec_freq, (a_overlap + b_overlap - overlap_chunk) / codec_freq)\n",
    "    new_audio_a = audio_a_head.append(audio_overlap)\n",
    "    audio_b_end = audio_b.get_segment(from_s=(b_overlap - overlap_chunk) / codec_freq)\n",
    "    print(f\"Segment durations: {audio_a_head.duration_s}, {audio_overlap.duration_s}, {audio_b_end.duration_s}\")\n",
    "\n",
    "    # Concatenate final audio and codec labels\n",
    "    total_audio = new_audio_a.append(audio_b_end)\n",
    "    total_code = np.vstack([code_a, code_b])\n",
    "    print(f\"Total audio duration: {total_audio.duration_s}, {total_code.shape}\")\n",
    "    return total_audio, total_code\n",
    "\n",
    "audio_a = audio_1.get_segment(0, 5)\n",
    "codec_labels_a = codec_labels_1[:5 *25, :].copy()\n",
    "audio_b = audio_2.get_segment(0, 5)\n",
    "codec_labels_b = codec_labels_2.copy()\n",
    "test_output_audio_1, test_output_code_1 = concat(audio_a, codec_labels_a, audio_b, codec_labels_b)\n",
    "test_output_audio_1.play()\n",
    "print(test_output_audio_1.duration_s, test_output_audio_1.array_float.shape)\n",
    "audio_c = audio_1.get_segment(10, 20)\n",
    "codec_labels_c = codec_labels_1[10 * 25:, :].copy()\n",
    "test_output_audio_2, test_output_code_2 = concat(test_output_audio_1, test_output_code_1, audio_c, codec_labels_c)\n",
    "test_output_audio_2.play()\n",
    "print(test_output_audio_2.duration_s, test_output_audio_2.array_float.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4119543a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-25T19:29:05.685302Z",
     "start_time": "2024-04-25T19:29:05.667647Z"
    },
    "execution": {
     "iopub.execute_input": "2024-08-27T18:12:15.384858Z",
     "iopub.status.busy": "2024-08-27T18:12:15.384718Z",
     "iopub.status.idle": "2024-08-27T18:12:15.386754Z",
     "shell.execute_reply": "2024-08-27T18:12:15.386365Z",
     "shell.execute_reply.started": "2024-08-27T18:12:15.384844Z"
    }
   },
   "outputs": [],
   "source": [
    "# diffs = []\n",
    "# for i in range(min(codec_labels.shape[0], new_codec_labels.shape[0])):\n",
    "#     diff = codec_labels[i, :] == new_codec_labels[i, :]\n",
    "#     diffs.append(12 - sum(diff))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d22a710",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-25T19:30:05.548530Z",
     "start_time": "2024-04-25T19:30:05.437000Z"
    },
    "execution": {
     "iopub.execute_input": "2024-08-27T18:12:15.387418Z",
     "iopub.status.busy": "2024-08-27T18:12:15.387286Z",
     "iopub.status.idle": "2024-08-27T18:12:15.923613Z",
     "shell.execute_reply": "2024-08-27T18:12:15.923101Z",
     "shell.execute_reply.started": "2024-08-27T18:12:15.387405Z"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "test_output_audio_array = test_output_audio_2.array_float\n",
    "test_output_audio_array_length = test_output_audio_array.shape[1]\n",
    "diffs = (audio.array_float[:, :test_output_audio_array_length] - test_output_audio_array)\n",
    "print(f\"Mean difference: {np.mean(diffs):.6f}\")\n",
    "print(f\"Standard deviation of difference: {np.std(diffs):.6f}\")\n",
    "print(f\"Shape of difference array: {diffs.shape}\")\n",
    "\n",
    "# Calculate time axis in seconds\n",
    "time_axis = np.arange(test_output_audio_array_length) / audio.sample_rate\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.plot(time_axis, diffs[0], alpha=0.5, label='Channel 1')\n",
    "plt.plot(time_axis, diffs[1], alpha=0.5, label='Channel 2')\n",
    "plt.xlabel('Time (seconds)')\n",
    "plt.ylabel('Difference')\n",
    "plt.title('Difference between original and concatenated audio')\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93bb75e0",
   "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"
  },
  "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
}