{
 "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",
    "\n",
    "from suno_utils.diffusion.generation import preload_models as preload_diff_models, DiffusionGenerationConfig\n",
    "from suno_utils.diffusion.generation import encode_semantic\n",
    "from suno_utils.tasks.upsample_engine import UpsampleEngine, Request\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import decode_stream_to_full_audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "diff_model_fp = \"/app/suno/checkpoints/2025-08-14_14-08-33_s8943/last_ckpt_infer.pt\"# seeds model\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 * 15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "semantic_codes = torch.ones((375,1), dtype=torch.int32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_diffusion_seeds(diffusion_seed, sem_data, tag):\n",
    "    diffusion_steps = 32\n",
    "\n",
    "    diffusion_text_cfg_coef = 4.0\n",
    "    diffusion_ctx_cfg_coef = 1.0\n",
    "\n",
    "    noise_ctx_level = 0.5\n",
    "    noise_ctx_pad_len = 0\n",
    "    CODEC_SCALE_FACTOR = 0.4\n",
    "    gen_cfg = DiffusionGenerationConfig(\n",
    "        steps=diffusion_steps,\n",
    "        lyrics=\"\",\n",
    "        tags=tag,\n",
    "        text_cfg_coef=diffusion_text_cfg_coef,\n",
    "        ctx_cfg_coef=diffusion_ctx_cfg_coef,\n",
    "        codec_scale_factor=CODEC_SCALE_FACTOR,\n",
    "        scale_ctx_vector=True,\n",
    "        noise_ctx_level=noise_ctx_level,\n",
    "        noise_ctx_pad_len=noise_ctx_pad_len,\n",
    "        drop_semantic_tokens=True,\n",
    "        seed=diffusion_seed,\n",
    "        rho=1.0,\n",
    "        sigma_min=0.5,\n",
    "        sigma_max=50.0,\n",
    "        objective=\"rectified_flow\",\n",
    "    )\n",
    "\n",
    "    request = Request(\n",
    "        id=\"dummy\",\n",
    "        generation_config=gen_cfg,\n",
    "        tokens=sem_data[:375],\n",
    "        input_tokens_finished=True,\n",
    "    )\n",
    "\n",
    "    result = up_engine.run_request(request)\n",
    "\n",
    "    vae_latents = []\n",
    "    for vae_latent in result.vae_latents:\n",
    "        print(vae_latent.shape)\n",
    "        mean = vae_latent.mean()\n",
    "        std = vae_latent.std()\n",
    "        print(f\"mean: {mean}, std: {std}\")\n",
    "        vae_latents.append(vae_latent)\n",
    "\n",
    "    vae_latents = torch.concat(vae_latents)\n",
    "    print(vae_latents.shape)\n",
    "    upsampled_audio = decode_stream_to_full_audio(vae_latents)\n",
    "    return upsampled_audio\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_diffusion_prod(diffusion_seed, sem_data, tag):\n",
    "    gen_cfg = DiffusionGenerationConfig(\n",
    "        seed=diffusion_seed,\n",
    "        lyrics=\"\",\n",
    "        tags=tag,\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",
    "    )\n",
    "\n",
    "    request = Request(\n",
    "        id=\"dummy\",\n",
    "        generation_config=gen_cfg,\n",
    "        tokens=sem_data[:375],\n",
    "        input_tokens_finished=True,\n",
    "    )\n",
    "\n",
    "    result = up_engine.run_request(request)\n",
    "\n",
    "    vae_latents = []\n",
    "    for vae_latent in result.vae_latents:\n",
    "        print(vae_latent.shape)\n",
    "        mean = vae_latent.mean()\n",
    "        std = vae_latent.std()\n",
    "        print(f\"mean: {mean}, std: {std}\")\n",
    "        vae_latents.append(vae_latent)\n",
    "\n",
    "    vae_latents = torch.concat(vae_latents)\n",
    "    print(vae_latents.shape)\n",
    "    upsampled_audio = decode_stream_to_full_audio(vae_latents)\n",
    "    return upsampled_audio\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "diffusion_seeds = np.random.randint(0, 1000000, size=3)\n",
    "tag = \"electronic drum loop\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "semantics = []\n",
    "for idx, diffusion_seed in enumerate(diffusion_seeds):\n",
    "    upsampled_audio = run_diffusion_seeds(diffusion_seed, semantic_codes, tag)\n",
    "    upsampled_audio.play()\n",
    "    semantic = encode_semantic(upsampled_audio)[:,:1]\n",
    "    semantics.append(semantic)\n",
    "    from_semantic = run_diffusion_seeds(diffusion_seed, semantic, tag)\n",
    "    from_semantic.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "# 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 * 15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx, diffusion_seed in enumerate(diffusion_seeds):\n",
    "    semantic = semantics[idx]\n",
    "    from_semantic = run_diffusion_prod(diffusion_seed, semantic, tag)\n",
    "    from_semantic.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "## Training Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl\n",
    "import random\n",
    "import pandas as pd\n",
    "from suno_utils.audio import Audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "def summarize_meta(ds_path):\n",
    "    metas = read_jsonl(ds_path)\n",
    "    print(f\"{len(metas):,} tracks with {sum([m['duration_s'] for m in metas])/60/60:,.1f}h total\")\n",
    "    print(metas[0].keys())\n",
    "    return metas\n",
    "\n",
    "def sample_meta(metas):\n",
    "    random_index = random.randint(0, len(metas ) - 1)\n",
    "    sample = metas[random_index]\n",
    "    for k,v in sample.items():\n",
    "        print(f\"{k}: {v}\")\n",
    "\n",
    "    return sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = read_jsonl(\"/app2/suno/data/diffusion/sfx/v1/metas_filter_v1_add_key_bpm.jsonl\")\n",
    "df = pd.DataFrame(train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "silent_semantic = Audio.from_silence(duration_s=15, sample_rate=48000)\n",
    "semantic_codes = encode_semantic(silent_semantic)[:,:1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "test = sample_meta(train_data)\n",
    "audio = Audio.from_s3(test[\"s3_filepath\"])\n",
    "tags = \", \".join(test[\"tags\"])\n",
    "audio.play()\n",
    "diffusion_seeds = np.random.randint(0, 1000000, size=1)\n",
    "for seed in diffusion_seeds:\n",
    "    silent_semantic = Audio.from_silence(duration_s=15, sample_rate=48000)\n",
    "    semantic_codes = encode_semantic(silent_semantic)[:,:1]\n",
    "    semantic = torch.from_numpy(encode_semantic(audio)[:,:1])\n",
    "    semantic_codes[:(semantic.shape[0])] = semantic\n",
    "    print(semantic_codes.shape)\n",
    "    from_semantic = run_diffusion_prod(diffusion_seed, semantic_codes, tags)\n",
    "    from_semantic.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}