{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.gpt.generation import GenerationConfig, CfgGenerationConfig\n",
    "from suno_utils.gpt.engine import Engine\n",
    "from suno_utils.gpt.generation_engine import make_request\n",
    "\n",
    "from suno_utils.diffusion.generation import preload_models as preload_diff_models, generate, DiffusionGenerationConfig\n",
    "from suno_utils.diffusion.generation import encode_semantic\n",
    "from suno_utils.tasks.upsample_engine import UpsampleEngine, Request, Job\n",
    "\n",
    "gpt_model_fp = \"/app/suno/modal/models/tony/sem/model_45_6b_bjay_sft_0519_t1_r8_d20.pt\"\n",
    "diff_model_fp = \"/app/suno/modal/models/tony/tmp/diff/v45_2b_step_2mil_ft_8k_infill_apr21_d3_v23.pt\" # production model (double check this)\n",
    "\n",
    "\n",
    "# Diffusion\n",
    "_ = preload_diff_models(\n",
    "    tokenizer_filepath=\"/home/georg/notebooks/gpu_nb/tmp/tokenizer_60k.json\",\n",
    "    semantic_model_filepath=\"/home/georg/notebooks/gpu_nb/tmp/mert_25.pt\",\n",
    "    semantic_clusters_filepath=\"/home/georg/notebooks/gpu_nb/tmp/mert_25_2x4k.npy\",\n",
    "    codec_filepath=\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\",\n",
    "    dit_model_filepath=diff_model_fp,\n",
    "    weights_precision=torch.bfloat16,\n",
    "    compile=False,\n",
    ")\n",
    "\n",
    "up_engine = UpsampleEngine(min_chunk_size=25 * 30)\n",
    "\n",
    "# GPT\n",
    "N_BATCH = 2\n",
    "n_skip_semantic = 1\n",
    "engine = Engine(\n",
    "    gpt_model_fp,\n",
    "    \"/app/suno/models/chirp_v2/tokenizer_60k.json\",\n",
    "    max_sequences=8*N_BATCH,\n",
    "    compile=False,\n",
    ")\n",
    "cfg = engine.model.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "history_lyrics = \"\"\"\n",
    "[instrumental]\n",
    "destroy\n",
    "[instrumental]\n",
    "\"\"\"\n",
    "\n",
    "lyrics = \"\"\"\n",
    "[instrumental outro, gradually fading and filtering down elements and melody to end the song without immediately ending]\n",
    "\"\"\"\n",
    "\n",
    "tags = \"techno melodic house deep\\nA high-energy electronic dance track with a driving beat and prominent synth melodies, The song features a consistent four-on-the-floor kick drum pattern, a pulsating bassline, and layered synth pads creating a rich harmonic texture, Percussion includes a crisp snare drum on the backbeat and a closed hi-hat pattern providing rhythmic propulsion, A distinctive synth lead plays a catchy, repetitive melodic phrase throughout the track, often accompanied by arpeggiated synth elements, The overall production is clean and polished, with a strong emphasis on rhythmic clarity and melodic hooks, The tempo is fast, contributing to the energetic feel, experimental sound design and ear candy is used throughout\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.clip import SunoClip\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import decode, encode\n",
    "from suno_utils.tasks.mert_25 import encode as encode_mert\n",
    "from suno_utils.tasks.audio_features.downbeats_data_prep.augment import stretch_audio\n",
    "\n",
    "clip = SunoClip('15edc54b-a734-40ef-b796-c28d81fd8d7a')\n",
    "\n",
    "\n",
    "#clip2 = SunoClip('312d0c32-d1a4-46b8-bd9a-e6636a179116') # extend of clip\n",
    "#arr_extend = clip2._get_npz()['v5.0_raw']\n",
    "#vae_extend = torch.from_numpy(clip2._get_vae_v4()['vae_latents'])\n",
    "\n",
    "\n",
    "arr = clip._get_npz()['v5.0_raw']\n",
    "vae_base = torch.from_numpy(clip._get_vae_v4()['vae_latents'])\n",
    "#vae_base = torch.from_numpy(np.load('/home/m4burns/extend_tests.npy'))\n",
    "\n",
    "PREFIX_LEN = 95\n",
    "\n",
    "arr_sfx = arr[int(PREFIX_LEN * 25):]\n",
    "vae_pfx = vae_base[:int(PREFIX_LEN * 25)]\n",
    "\n",
    "# cycle vae\n",
    "\n",
    "vae_decoded = decode(vae_base)\n",
    "# simulate gain adjustment\n",
    "vae_decoded = vae_decoded.apply_gain(1.28)[0]\n",
    "# simulate studio\n",
    "vae_decoded.write_opus('vae_decoded.opus')\n",
    "vae_decoded = Audio.from_file('vae_decoded.opus', n_channels=vae_decoded.n_channels, sample_rate=vae_decoded.sample_rate)\n",
    "#vae_decoded = stretch_audio(vae_decoded, [(0, 0), (1, 1)])\n",
    "vae_pfx_cycled = torch.from_numpy(encode(vae_decoded)[:int(PREFIX_LEN * 25)])\n",
    "\n",
    "decode(vae_pfx_cycled).play()\n",
    "\n",
    "# cycle mert\n",
    "\n",
    "arr_pfx_cycled = encode_mert(vae_decoded)[:int(PREFIX_LEN * 25), 0]\n",
    "\n",
    "gconf = GenerationConfig(\n",
    "    text=lyrics,\n",
    "    text_tags=tags,\n",
    "    text_neg_tags=\"repetitive, loop\",\n",
    "    cfg_coef=1.0,  # no text cfgc\n",
    "    cfg_coef_tags=1.0,\n",
    "    cfg_coef_neg_tags=-1.0,\n",
    "    cfg_coef_tags_max_steps=25 * 120,\n",
    "    min_eos_p=0.1,  # this is okay\n",
    "    use_whisper=False,\n",
    "    min_text_offset=0,\n",
    "    n_repeat_tags=1,\n",
    "    temp_semantic=0.90,\n",
    "    top_k_semantic=1500,\n",
    "    top_p_semantic=None,\n",
    "    min_p_semantic=0.005,\n",
    "    eos_pad_duration_s=0,\n",
    "    text_start_control_tags = None,\n",
    "    n_batch=1,\n",
    "    #random_seed=0,\n",
    "    history_arr=arr_pfx_cycled.reshape(-1, 1),\n",
    ")\n",
    "\n",
    "gpt_request = make_request(\"foo\", gconf, engine.model.config, engine.tokenizer)\n",
    "\n",
    "job = engine.run_request(gpt_request, tqdm_enabled=True)\n",
    "stream = engine.token_generator(job)\n",
    "out_gpt = torch.stack(list(stream))[:,1]\n",
    "if out_gpt[-1] == 4000:\n",
    "    out_gpt = out_gpt[:-1]\n",
    "\n",
    "\n",
    "print(arr.shape)\n",
    "print(vae_base.shape)\n",
    "print(out_gpt.shape)\n",
    "\n",
    "gen_cfg = DiffusionGenerationConfig(\n",
    "    lyrics=lyrics,\n",
    "    history_lyrics=history_lyrics,\n",
    "    tags=tags,\n",
    "    text_cfg_coef=4.0,\n",
    "    steps=10,\n",
    "    codec_scale_factor=0.4,\n",
    "    scale_ctx_vector=True,\n",
    "    noise_ctx_level=0.75,\n",
    "    drop_semantic_tokens=False,\n",
    "    generation_history_latents=vae_pfx_cycled[-25*30:],\n",
    "    #seed=3,\n",
    ")\n",
    "\n",
    "request = Request(\n",
    "    id=\"dummy\",\n",
    "    generation_config=gen_cfg,\n",
    "    tokens=out_gpt,\n",
    "    input_tokens_finished=True,\n",
    ")\n",
    "\n",
    "result = up_engine.run_request(request).vae_latents\n",
    "extended_audio = decode(*result)\n",
    "\n",
    "#extended_audio = decode(torch.cat([vae_pfx.cuda(), *result], dim=0))\n",
    "extended_audio.get_segment(from_s=PREFIX_LEN - 3, to_s=PREFIX_LEN + 5).play()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "extended_audio.play()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_clean",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}