{
 "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\"] = \"6\"  # 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": [
    "from suno_utils.utils.s3 import read_from_s3, _read_npz\n",
    "\n",
    "npz = read_from_s3(\n",
    "    \"s3://suno-data-uploads/studio/uploads/30f189a4-9bb6-4bab-8afd-465f11c598ad.npz\", read_f=_read_npz\n",
    ")\n",
    "full_arr = npz[\"full_arr\"]\n",
    "full_arr.shape"
   ]
  },
  {
   "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(\"s3://suno-data/tony/tmp/model_30b_fix_ft2_20k.pt\")\n",
    "\n",
    "engine = Engine(\n",
    "    gpt_ckpt_path,\n",
    "    max_sequences=8,\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=\"\"\"[Verse]\n",
    "Chicken on the dance floor\n",
    "Feathers flying we want more\n",
    "Cluckin' to the beat feel the heat\n",
    "Move your feet let's get sweet\n",
    "\n",
    "[Verse 2]\n",
    "Pecking round and round we go\n",
    "Winging it we're stealing the show\n",
    "Shake your tail let's set sail\n",
    "In this coop we won't fail\n",
    "\n",
    "[Chorus]\n",
    "Do the chicken dance\n",
    "Hop and prance\n",
    "In a clucking trance\n",
    "Shake those wings take a chance\n",
    "\n",
    "[Verse 3]\n",
    "Feathers ruffling to the right\n",
    "Dancing through the night\n",
    "Golden beaks so chic\n",
    "Feeling free hearts beat\n",
    "\n",
    "[Bridge]\n",
    "Cluck cluck feeling fine\n",
    "Riding on the chicken line\n",
    "Groove and jive stay alive\n",
    "In our world we thrive\n",
    "\n",
    "[Verse 4]\n",
    "Underneath the bright moonlight\n",
    "Chickens groove it's such a sight\n",
    "Cackle loud sing it proud\n",
    "Cluckin’ crowd like a cloud\"\"\",\n",
    "    text_tags=\"bluegrass female\",\n",
    "    # text_neg_tags=\"bluegrass\",\n",
    "    n_batch=1,\n",
    "    cfg_coef=1.3,\n",
    "    cfg_coef_tags=1,\n",
    "    cfg_coef_neg_tags=3,\n",
    "    max_gen_duration_s=30,\n",
    "    min_text_offset=0,\n",
    "    cover_arr=full_arr,\n",
    "    n_repeat_tags=1,\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(4)]\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_arr[:20, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generated[1, :20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for job in jobs:\n",
    "    generated = np.stack([t.cpu().numpy() for t in job.generated_tokens]).T\n",
    "    print(generated[2, 50:100])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}