{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9c62b25",
   "metadata": {},
   "source": [
    "## GPT engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c7d0311",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.gpt.generation import GenerationConfig\n",
    "from suno_utils.gpt.generation_engine import make_request\n",
    "from tqdm import tqdm\n",
    "from suno_utils.diffusion import generation as diffusion_gen\n",
    "from suno_utils.tasks.upsample_engine import UpsampleEngine, Request\n",
    "from suno_utils.audio import Audio\n",
    "import torch\n",
    "\n",
    "N_BATCH = 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06f57a07",
   "metadata": {},
   "outputs": [],
   "source": [
    "dit_model_filepath = \"/app/suno/checkpoints/2025-04-27_00-05-49_s1686/last_ckpt.pt\"  # base\n",
    "diffusion_gen.preload_models(\n",
    "    dit_model_filepath=dit_model_filepath,\n",
    "    codec_filepath=\"/app/suno/data/dpo/models/dac_vae_tuned_25hz.pth\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b21561e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.dac_vae_fixed_25hz import preload_models as preload_codec_models, decode, encode\n",
    "\n",
    "preload_codec_models(\"/app/suno/data/dpo/models/dac_vae_tuned_25hz.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93d9d9e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "diffusion_engine = UpsampleEngine(compile=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c5014e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.gpt.rpc_zmq import start_service_processes\n",
    "\n",
    "client = start_service_processes(\n",
    "    \"/app2/suno/checkpoints/2025-08-24_08-31-23/last_ckpt_infer.pt\",\n",
    "    max_sequences=40,\n",
    "    tokenizer_path=\"s3://suno-data/georg/models/tokenizers/tokenizer_60k.json\",\n",
    "    compile=False,\n",
    "    port=18865,\n",
    ")\n",
    "cfg = client.get_model_cfg()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2e7e428",
   "metadata": {},
   "outputs": [],
   "source": [
    "lyrics = \"\"\"\n",
    "[verse]\n",
    "oh, my love\n",
    "My friend you know\n",
    "it's been a while\n",
    "Without thinking of you\n",
    "but the thought makes me smile\n",
    "\n",
    "[chorus]\n",
    "I'm so tired of wanting\n",
    "wanting more than this\n",
    "i know it but what am i to do\n",
    "i need some space to breathe,\n",
    "so give me some room\n",
    "\n",
    "[verse]\n",
    "oh, my love\n",
    "you have a heart of stone\n",
    "cause since i've come home\n",
    "i've never felt so alone\n",
    "but the thought makes me smile\n",
    "\"\"\"\n",
    "gconf = GenerationConfig(\n",
    "    text=lyrics,\n",
    "    text_tags=\"\",\n",
    "    cfg_coef=1.0,\n",
    "    # cfg_coef_tags=1.5,\n",
    "    n_batch=1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=90,\n",
    ")\n",
    "requests = [make_request(f\"_{i}\", gconf, cfg, client.get_tokenizer()) for i in range(4)]\n",
    "jobs = []\n",
    "for request in tqdm(requests):\n",
    "    jobs.append(client.add_request(request))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cd2d82f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "while any(client.get_job_state(job)[\"completed\"] == False for job in jobs):\n",
    "    time.sleep(1)\n",
    "client.get_job_state(jobs[0])\n",
    "\n",
    "for job in jobs:\n",
    "    print(client.get_job_state(job))\n",
    "\n",
    "codes = client.get_generated_codes(jobs[0])\n",
    "print(codes.shape)\n",
    "assert codes.shape[0] > 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8628f293",
   "metadata": {},
   "outputs": [],
   "source": [
    "codes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01101b29",
   "metadata": {},
   "outputs": [],
   "source": [
    "codes.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3021a427",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n",
    "    # audio=vae_latents,\n",
    "    # tags=\"extract Lead Vocal\",\n",
    "    lyrics=lyrics,\n",
    "    text_cfg_coef=2.0,\n",
    "    ctx_cfg_coef=1.0,\n",
    "    steps=10,\n",
    "    codec_scale_factor=0.4,\n",
    "    scale_ctx_vector=True,\n",
    "    noise_ctx_level=0.75,\n",
    "    noise_ctx_pad_len=0,\n",
    ")\n",
    "request = Request(\n",
    "    id=\"dummy\",\n",
    "    generation_config=gen_cfg,\n",
    "    tokens=codes[:, 0],\n",
    "    input_tokens_finished=True,\n",
    ")\n",
    "result = diffusion_engine.run_request(request)\n",
    "vae_latents = torch.concat(result.vae_latents)\n",
    "audio = decode(vae_latents)\n",
    "audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdce0448",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1656dc31",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8219262",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}