{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic Engine Usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\" # Set the gpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.gpt.chirp_v2_5 import (\n",
    "    preload_codec_models,\n",
    "    _get_model_if_needed,\n",
    "    GenerationConfig,\n",
    "    codec_decode_stream_to_full_audio,\n",
    "    GPT_CKPT_PATH,\n",
    ")\n",
    "from suno_utils.gpt.generation_engine import align_codes, make_request\n",
    "from suno_utils.gpt.engine import Engine\n",
    "from suno_utils.gpt.prompt import Prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preload_codec_models(\"/app/suno/models/chirp_v2/dac_2c_25x12.pt\")\n",
    "gpt_ckpt_path = _get_model_if_needed(GPT_CKPT_PATH)\n",
    "\n",
    "engine = Engine(\n",
    "    gpt_ckpt_path,\n",
    "    max_sequences=20,\n",
    "    max_length_s=240,\n",
    "    compile=False,  # set to True for faster inference, but wait for warmup\n",
    ")\n",
    "model = engine.model\n",
    "cfg = model.config\n",
    "tokenizer = engine.tokenizer\n",
    "print(cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=\"\"\"[intro]\n",
    "Maybe, Maybe, Maybe, maybe i had enough of all that (huh?)\n",
    "\n",
    "\n",
    "[verse]\n",
    "I could never be free,\n",
    "Maybe I'm not where I want to be\n",
    "lately it just feels like im going craaazy\n",
    "but once more ill open that door searching for solutions for my never ending game that my brain insists i play!\n",
    "\n",
    "[pre chorus]\n",
    "its just another of those days.. another one of those, days\n",
    "(what do i usually say?)\"\"\",\n",
    "    text_tags=\"bluegrass female\",\n",
    "    text_neg_tags=\"noise\",\n",
    "    n_batch=1,\n",
    "    cfg_coef=1.5,\n",
    "    max_gen_duration_s=90,\n",
    "    min_text_offset=0,\n",
    ")\n",
    "\n",
    "\n",
    "request = make_request(\"0\", gconf, cfg, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize prompt\n",
    "\n",
    "prompt = Prompt(gconf, cfg, tokenizer)\n",
    "\n",
    "for stream in request.streams:\n",
    "    prompt.visualize(stream.prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "requests = [make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer) for i in range(2)]\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "\n",
    "for job in jobs:\n",
    "    stream = engine.token_generator(job)\n",
    "\n",
    "    audio = codec_decode_stream_to_full_audio(align_codes(stream, cfg))\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# concat generated to prompt and visualize\n",
    "import numpy as np\n",
    "\n",
    "in_arr = requests[0].streams[0].prompt\n",
    "generated = np.stack([t.cpu().numpy() for t in job.generated_tokens]).T\n",
    "print(generated.shape)\n",
    "arr = np.concatenate([in_arr, generated], axis=1)\n",
    "prompt.visualize(arr)\n",
    "prompt.visualize(arr, compress=True)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}