{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.mert_25 import preload_models\n",
    "from suno_utils.diffusion import generation as diffusion_gen\n",
    "from suno_utils.tasks.upsample_engine import UpsampleEngine\n",
    "from suno_utils.tasks.data_loader import load_audio_opus\n",
    "\n",
    "from suno_utils.tasks.mert_25 import encode as mert_encode, SAMPLE_RATE, EMBEDDING_RATE\n",
    "from suno_utils.diffusion import generation as diffusion_gen\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import (\n",
    "    encode as vae_encode,\n",
    "    decode as vae_decode,\n",
    ")\n",
    "from suno_utils.tasks.upsample_engine import Request\n",
    "from suno_utils.audio.conversion import Audio\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "model_filepath = \"s3://suno-data/georg/models/semantic/mert_25.pt\"\n",
    "extra_filepath = \"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\"\n",
    "\n",
    "_ = diffusion_gen.preload_models(\n",
    "    tokenizer_filepath=\"s3://suno-data/georg/models/tokenizers/tokenizer_60k.json\",\n",
    "    semantic_model_filepath=\"s3://suno-data/georg/models/semantic/mert_25.pt\",\n",
    "    semantic_clusters_filepath=\"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\",\n",
    "    codec_filepath=\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\",\n",
    "    dit_model_filepath=\"s3://suno-data/sara/models/stems_41.pt\",\n",
    ")\n",
    "_ = preload_models(\n",
    "    checkpoint_filepath=model_filepath,\n",
    "    centroids_filepath=extra_filepath,\n",
    ")\n",
    "\n",
    "engine = UpsampleEngine(compile=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opus_file = \"/home/sara/glockenspiel/suno_utils/suno_utils/scripts/gpt/NrkrK3W1tEA.opus\"\n",
    "audio = load_audio_opus(\n",
    "    opus_file,\n",
    "    n_channels=2,\n",
    "    max_duration_s=120,\n",
    "    target_sample_rate=48_000,\n",
    "    normalize_volume=False,\n",
    "    debug=True,\n",
    ")\n",
    "audio = Audio.from_array_float(audio, sample_rate=48_000)\n",
    "audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tags = \"remove Lead Vocal, Backing Vocals\"\n",
    "gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n",
    "    tags=tags,\n",
    "    text_cfg_coef=1.0,\n",
    "    ctx_cfg_coef=1.0,\n",
    "    steps=12,\n",
    "    codec_scale_factor=0.4,\n",
    "    noise_ctx_level=0.3,\n",
    "    scale_ctx_vector=0.4,\n",
    "    noise_stem_ctx_level=0.0,\n",
    ")\n",
    "\n",
    "vae_latents = vae_encode(audio)\n",
    "request = Request(\n",
    "    id=\"dummy\",\n",
    "    generation_config=gen_cfg,\n",
    "    tokens=np.full(\n",
    "        (vae_latents.shape[0]), engine.dit_config[\"cond_semantic_n_vocab\"] - 1\n",
    "    ),\n",
    "    stem_ctx_latents=vae_latents,\n",
    "    input_tokens_finished=True,\n",
    ")\n",
    "\n",
    "result = engine.run_request(request)\n",
    "out_vae_latents = torch.concat(result.vae_latents)\n",
    "audio_stem = vae_decode(out_vae_latents)\n",
    "audio_stem.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": 2
}