{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\"\n",
    "\n",
    "procid = int(os.environ.get(\"SLURM_PROCID\", 0))\n",
    "localid = int(os.environ.get(\"SLURM_LOCALID\", 0))\n",
    "world_size = int(os.environ.get(\"SLURM_JOB_NUM_NODES\", 1)) * int(\n",
    "    os.environ.get(\"SLURM_NTASKS_PER_NODE\", 1)\n",
    ")\n",
    "\n",
    "assert world_size > 0, \"WORLD_SIZE is 0\"\n",
    "\n",
    "print(f\"PROCID: {procid}, LOCALID: {localid}, WORLD_SIZE: {world_size}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c046f7b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.diffusion import generation as diffusion_gen\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, encode, decode\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from suno_utils.audio import Audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14151510",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "id_path = \"/home/christian/code/christian/metadata/v45_splits/ids_keep_sets_v11.json\"\n",
    "all_ids = []\n",
    "\n",
    "with open(id_path, \"r\") as f:\n",
    "    ids_keep_sets = json.load(f)\n",
    "for k, v in ids_keep_sets.items():\n",
    "    print(k, len(v))\n",
    "    all_ids.extend(v)\n",
    "# print an example\n",
    "print(ids_keep_sets[\"discogs_subset\"][0])\n",
    "print(len(all_ids))\n",
    "\n",
    "all_ids = all_ids[procid::world_size]\n",
    "\n",
    "\n",
    "def get_audio_path(id):\n",
    "    path = f\"/app2/suno/data/raw_audio_opus_v0/{id}.opus\"\n",
    "    if os.path.exists(path):\n",
    "        return path\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "\n",
    "# !ls /app/suno/data/auk_v0\n",
    "# !head /app/suno/data/auk_v0/metas_v2_val.jsonl\n",
    "\n",
    "# !ls /app2/suno/data/raw_audio_opus_v0/ZrGWefub8RU.opus\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a718fb2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "diffusion_gen.preload_models(\n",
    "    dit_model_filepath=\"/app/suno/checkpoints/2025-05-26_22-49-34_s8646/last_ckpt_infer.pt\",  # 12 stems\n",
    "    codec_filepath=\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\",\n",
    "    compile=True,\n",
    ")\n",
    "\n",
    "engine = UpsampleEngine(min_chunk_size=25 * 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2197af65",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_stem(\n",
    "    audio: Audio,\n",
    "    stem_type_cfg_scale=1.0,\n",
    "    tags=\"extract [split_karaoke]\",\n",
    "    steps=4,\n",
    "    seed=3,\n",
    "    codec_scale_factor=0.4,\n",
    "    scale_ctx_vector=True,\n",
    "    noise_ctx_level=0.0,\n",
    "    infill_prefix_latents=None,\n",
    "    infill_suffix_latents=None,\n",
    "):\n",
    "    vae = encode(audio)\n",
    "    gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n",
    "        lyrics=tags,\n",
    "        steps=steps,\n",
    "        seed=seed,\n",
    "        codec_scale_factor=codec_scale_factor,\n",
    "        scale_ctx_vector=scale_ctx_vector,\n",
    "        noise_ctx_level=noise_ctx_level,\n",
    "        text_cfg_coef=stem_type_cfg_scale,\n",
    "        infill_prefix_latents=infill_prefix_latents,\n",
    "        infill_suffix_latents=infill_suffix_latents,\n",
    "        drop_semantic_tokens=True,\n",
    "    )\n",
    "    print(gen_cfg)\n",
    "\n",
    "    request = Request(\n",
    "        id=\"dummy\",\n",
    "        generation_config=gen_cfg,\n",
    "        tokens=np.zeros((vae.shape[0], 1)),\n",
    "        input_tokens_finished=True,\n",
    "        stem_ctx_latents=vae,\n",
    "    )\n",
    "\n",
    "    result = engine.run_request(request)\n",
    "    vae_latents = torch.concat(result.vae_latents)\n",
    "    print(f\"vae_latents: {vae_latents.shape}\")\n",
    "    audios = []\n",
    "    for i in tqdm(range(vae_latents.shape[1]), desc=\"Decoding stems\"):\n",
    "        # audios.append(decode(vae_latents[:, i]))\n",
    "        audios.append(decode_stream_to_full_audio(vae_latents[:, i], n_stride_tokens=25 * 120))\n",
    "    return audios\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b28b1f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_audio = get_audio_path(\"ZrGWefub8RU\")\n",
    "test_audio_audio = Audio.from_file(test_audio, n_channels=2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f6bc7c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "audios = gen_stem(\n",
    "    test_audio_audio,\n",
    "    tags=\"extract [split_karaoke]\",\n",
    ")\n",
    "# mix = Audio.sum(audios)\n",
    "# mix.play()\n",
    "\n",
    "categories = [\n",
    "    \"Vocals\",\n",
    "    \"Backing_Vocals\",\n",
    "    \"Drums\",\n",
    "    \"Bass\",\n",
    "    \"Guitar\",\n",
    "    \"Keyboard\",\n",
    "    \"Percussion\",\n",
    "    \"Strings\",\n",
    "    \"Synth\",\n",
    "    \"FX\",\n",
    "    \"Brass\",\n",
    "    \"Woodwinds\",\n",
    "]\n",
    "\n",
    "for category, audio in zip(categories, audios):\n",
    "    print(category, audio.loudness)\n",
    "    if audio.loudness < -45:\n",
    "        continue\n",
    "    # audio.play()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b88fa2a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "OUT_DIR = \"/app2/suno/data/sft_stems_12_output_v11\"\n",
    "\n",
    "os.makedirs(OUT_DIR, exist_ok=True)\n",
    "\n",
    "\n",
    "def write_stem(args):\n",
    "    id, category, stem = args\n",
    "    if stem.loudness < -45:\n",
    "        return None\n",
    "    out_path = f\"{OUT_DIR}/{id}_{category}.opus\"\n",
    "    stem.write_opus(out_path)\n",
    "    return category\n",
    "\n",
    "\n",
    "import threading\n",
    "\n",
    "\n",
    "def process_id(id):\n",
    "    audio = get_audio_path(id)\n",
    "    if audio is None:\n",
    "        return None\n",
    "    audio = Audio.from_file(audio, n_channels=2)\n",
    "    stems = gen_stem(audio, steps=8)\n",
    "\n",
    "    # Prepare arguments for parallel processing\n",
    "    write_args = [(id, category, stem) for category, stem in zip(categories, stems)]\n",
    "\n",
    "    # Use threading to write stems in parallel\n",
    "    results = [None] * len(write_args)\n",
    "    threads = []\n",
    "\n",
    "    def write_stem_thread(i, args):\n",
    "        results[i] = write_stem(args)\n",
    "\n",
    "    for i, args in enumerate(write_args):\n",
    "        thread = threading.Thread(target=write_stem_thread, args=(i, args))\n",
    "        threads.append(thread)\n",
    "        thread.start()\n",
    "\n",
    "    for thread in threads:\n",
    "        thread.join()\n",
    "\n",
    "    # Filter out None results\n",
    "    found_categories = [result for result in results if result is not None]\n",
    "    return found_categories\n",
    "\n",
    "\n",
    "process_id(\"ZrGWefub8RU\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c29e64f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for id in tqdm(all_ids):\n",
    "    process_id(id)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}