{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.gpt.chirp_v2_5 import (\n",
    "    _get_model_if_needed,\n",
    ")\n",
    "from suno_utils.gpt.engine import Engine\n",
    "from suno_utils.tasks.dac_2c_12cb import (\n",
    "    preload_models as preload_codec_models,\n",
    "    encode as codec_encode,\n",
    ")\n",
    "from suno_utils.tasks.mert_25 import (\n",
    "    preload_models as preload_semantic_models,\n",
    "    encode as semantic_encode,\n",
    ")\n",
    "\n",
    "from suno_utils.tasks.upsample_engine import (\n",
    "    UpsampleEngine,\n",
    "    Request,\n",
    "    DiffusionGenerationConfig,\n",
    ")\n",
    "from suno_utils.diffusion import generation as diffusion_gen\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"5\"\n",
    "\n",
    "torch._logging.set_logs(recompiles=True, recompiles_verbose=True)  # , guards=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _interleave(semantic_arr, n_factor=1):\n",
    "    new_semantic_arr = (\n",
    "        np.zeros(\n",
    "            (semantic_arr.shape[0] * n_factor, semantic_arr.shape[-1]),\n",
    "            dtype=semantic_arr.dtype,\n",
    "        )\n",
    "        + cfg.semantic_vocab_size\n",
    "        - 1\n",
    "    )\n",
    "    new_semantic_arr[::n_factor] = semantic_arr\n",
    "    return new_semantic_arr\n",
    "\n",
    "\n",
    "def process_audio(audio, cfg, n_factor=1):\n",
    "    audio = audio.normalize_volume(-16)\n",
    "    sem_arr = semantic_encode(audio, device=\"cpu\")\n",
    "    if n_factor > 1:\n",
    "        sem_arr = _interleave(sem_arr, n_factor=n_factor)\n",
    "    coarse_arr = codec_encode(audio)\n",
    "    n_frames = min(sem_arr.shape[0], coarse_arr.shape[0])\n",
    "    sem_arr = sem_arr[:n_frames, : cfg.semantic_n_codebooks]\n",
    "    coarse_arr = coarse_arr[:n_frames, : cfg.coarse_n_codebooks]\n",
    "\n",
    "    a_arr = np.concatenate([sem_arr, coarse_arr], axis=-1)\n",
    "    return a_arr\n",
    "\n",
    "\n",
    "def load_audio(fp):\n",
    "    return Audio.from_file(fp, n_channels=2, sample_rate=48_000, byte_width=2)\n",
    "\n",
    "\n",
    "N_BATCH = 2\n",
    "MAX_STREAMS = N_BATCH * 4\n",
    "\n",
    "# preload codec\n",
    "_ = preload_codec_models(\"s3://suno-data/georg/models/codec/dac_2c_25x8.pt\")\n",
    "\n",
    "# preload mert\n",
    "_ = preload_semantic_models(\n",
    "    checkpoint_filepath=\"s3://suno-data/georg/models/semantic/mert_25.pt\",\n",
    "    centroids_filepath=\"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\",\n",
    "    device=\"cpu\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "USE_COMPILE = False\n",
    "gpt_ckpt_path = _get_model_if_needed(\n",
    "    \"/app/suno/checkpoints/2024-12-08_06-09-46/last_ckpt_infer.pt\"\n",
    ")\n",
    "engine = Engine(\n",
    "    gpt_ckpt_path,\n",
    "    \"s3://suno-data/georg/trained_models/chirp_v2/tokenizer_60k.json\",\n",
    "    max_sequences=MAX_STREAMS,\n",
    "    compile=USE_COMPILE,\n",
    ")\n",
    "model = engine.model\n",
    "cfg = model.config\n",
    "tokenizer = engine.tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = load_audio(\"/home/sara/samples/synthwave_fantasy.mp3\").get_segment(\n",
    "    from_s=45, to_s=105\n",
    ")\n",
    "audio.play()\n",
    "in_infill_arr = process_audio(audio, cfg)\n",
    "in_history_arr = in_infill_arr[: 25 * 15, :].copy()\n",
    "in_future_arr = in_infill_arr[-25 * 15 :, :].copy()\n",
    "in_infill_arr = in_infill_arr[25 * 15 : (-25 * 15), :].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diffusion_gen.preload_models(compile=False)\n",
    "engine = UpsampleEngine(min_chunk_size=25 * 30)\n",
    "diffusion_generation_config = DiffusionGenerationConfig()\n",
    "audio_semantic_codes = torch.from_numpy(in_infill_arr[:, 0]).reshape(1, -1).long()\n",
    "print(audio_semantic_codes.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "engine.run_request(\n",
    "    Request(\n",
    "        id=\"0\",\n",
    "        generation_config=diffusion_generation_config,\n",
    "        tokens=[c for c in audio_semantic_codes[0]],\n",
    "        input_tokens_finished=True,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}