{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0d83737",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !nvidia-smi\n",
    "# !echo $HOSTNAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "sys.path.append(\"..\")\n",
    "\n",
    "# Check if the path was added correctly\n",
    "print(f\"Current path: {sys.path}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd1bb050",
   "metadata": {},
   "outputs": [],
   "source": [
    "dit_model_filepath = (\n",
    "    \"/app/suno/checkpoints/2025-03-20_00-27-17_s8851/last_ckpt_infer.pt\"  # with infill, noisy\n",
    ")\n",
    "dit_model_filepath = \"/app/suno/checkpoints/2025-03-22_13-17-52_s9785/last_ckpt_infer.pt\"  # prefix\n",
    "dit_model_filepath = \"/app/suno/checkpoints/2025-03-22_22-19-43_s2078/last_ckpt_infer.pt\"  # prefix fixes\n",
    "dit_model_filepath = \"/app/suno/checkpoints/2025-03-23_04-09-35_s4633/last_ckpt_infer.pt\"  # prefix fixes\n",
    "# # multi-instrument\n",
    "dit_model_filepath = \"/app/suno/checkpoints/2025-03-24_02-52-32_s786/last_ckpt_infer.pt\"\n",
    "# # single instrument\n",
    "# dit_model_filepath = \"/app/suno/checkpoints/2025-03-24_19-40-25_s9528/last_ckpt_infer.pt\"\n",
    "\n",
    "# dit_model_filepath = \"/app/suno/checkpoints/2025-03-26_13-24-59_s1608/step_100000_infer.pt\"\n",
    "\n",
    "# old arch\n",
    "# dit_model_filepath = \"/app/suno/checkpoints/2025-03-06_23-22-48_s4786/step_40000_infer.pt\"\n",
    "# transfer old arch to new\n",
    "dit_model_filepath = \"/app/suno/checkpoints/2025-03-28_13-45-42_s1953/step_40000_infer.pt\"\n",
    "# grouped\n",
    "dit_model_filepath = \"/app/suno/checkpoints/2025-03-29_03-49-41_s6886/step_40000_infer.pt\"\n",
    "dit_model_filepath = \"/app/suno/checkpoints/2025-04-24_03-51-55_s1736/last_ckpt_infer.pt\"  # main\n",
    "dit_model_filepath = \"/app/suno/checkpoints/2025-04-24_13-06-58_s8444/last_ckpt_infer.pt\"  # otf\n",
    "dit_model_filepath = (\n",
    "    \"/app/suno/checkpoints/2025-04-25_02-47-42_s734/last_ckpt_infer.pt\"  # otf extract only\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22f198cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "\n",
    "# !ls sunoDiff/stem_type_map.json\n",
    "\n",
    "with open(\"../stem_type_map.json\", \"r\") as f:\n",
    "    stem_type_map = json.load(f)\n",
    "stem_df = pd.Series(stem_type_map).to_frame(\"stem_id\")\n",
    "stem_df.index = [f\"{e[1]['stem_id']} - {e[0]}\" for e in stem_df.iterrows()]\n",
    "# Print all stem types and their IDs\n",
    "with pd.option_context(\"display.max_rows\", None, \"display.max_columns\", None):\n",
    "    print(stem_df)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9522b25",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from suno_utils.audio import Audio\n",
    "from generation import preload_models, generate, _retrieve_models\n",
    "\n",
    "_ = preload_models(\n",
    "    codec_filepath=\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\",  # 25hz\n",
    "    dit_model_filepath=dit_model_filepath,\n",
    "    weights_precision=torch.bfloat16,\n",
    "    compile=True,\n",
    "    codec_scale_factor=0.4,\n",
    ")\n",
    "models = _retrieve_models()\n",
    "duration_s = 60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2a7b3c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "offset_s = 0\n",
    "\n",
    "# audio_fp = \"/home/georg/notebooks/samples/walking_down.mp3\"\n",
    "# tags = \"orchestral, orchestra, classical, epic movie\"\n",
    "# if duration_s <= 30.1:\n",
    "#     lyrics = \"Walking down the street\\nFeeling so alive\\nGot my head in the clouds\\nGot a gleam in my eye\\nEvery step that i take\\nIs like a brand new start\\nNo matter where I'm going\\nI'll always find my part\\n\\nOooo\\n\\nLife is like a high\"\n",
    "# else:\n",
    "#     lyrics = \"Walking down the street\\nFeeling so alive\\nGot my head in the clouds\\nGot a gleam in my eye\\nEvery step that i take\\nIs like a brand new start\\nNo matter where I'm going\\nI'll always find my part\\n\\nOooo\\n\\nLife is like a high wire act\\nwe're dancing in the sky\\nIt's a worry no need to ask why\\nWith a little bit or courage\\nWe can chase our dreams\\nNo matter what comes away\\nWe'll always be a team\\n\\nDon't stop oooo, yeah\"\n",
    "\n",
    "# audio_fp = \"/home/georg/notebooks/samples/instrumental.mp3\"\n",
    "# tags = \"orchestral, orchestra, classical, epic movie, dark trailer\"\n",
    "# offset_s = 20\n",
    "\n",
    "# audio_fp = \"/home/georg/notebooks/samples/md_mirror_baby.mp3\"\n",
    "# tags = \"reggae uplifting playful\"\n",
    "# lyrics = \"[Verse]\\nLittle Leo looks so fine\\nMirror baby by his side\\nThinks he’s seeing someone new\\nEvery time he smiles too\\n\\n[Chorus]\\nMirror baby best friend\\nReflections that don't end\\nLittle Leo's world so bright\\nEvery day a new delight\\n\\n[Verse 2]\\nChubby cheeks and laughing eyes\\nMirror baby never lies\\nEvery nod and every wink\\nLeo sees the bond they link\\n\\n[Chorus]\\nMirror baby best friend\\nReflections that don't end\\nLittle Leo's world so bright\\nEvery day a new delight\\n\\n[Bridge]\\nEvery giggle every glance\\nTiny feet that love to dance\\nMirror cheers him all day long\\nWith a smile and happy song\\n\\n[Verse 3]\\nCribs and rattles filled with glee\\nMirror baby laughs with me\\nLeo finds a friend so true\\nIn the mirror just like new\"\n",
    "\n",
    "# audio_fp = \"/home/georg/notebooks/samples/norah.wav\"\n",
    "# tags = \"jazz, pop, female vocal\"\n",
    "# lyrics = \"[Verse 1]\\nI waited 'til I saw the sun\\nI don't know why I didn't come\\nI left you by the house of fun\\nI don't know why I didn't come\\nI don't know why I didn't come\\n\\n[Verse 2]\\nWhen I saw the break of day\\nI wished that I could fly away\\nInstead of kneeling in the sand\\nCatching teardrops in my hand\\n\\n[Refrain]\\nMy heart is drenched in wine\\nBut you'll be on my mind forever\\n\\n[Verse 3]\\nOut across the endless sea\\nI will die in ecstasy\\nBut I'll be a bag of bones\\nDriving down the road alone\\n\\n[Refrain]\\nMy heart is drenched in wine\\nBut you'll be on my mind\"\n",
    "\n",
    "# audio_fp = \"/home/georg/notebooks/samples/turkey_time.mp3\"\n",
    "# tags = \"Country\"\n",
    "# lyrics = \"[Verse]\\nGather 'round now, it's turkey time\\nCatchin' up with all your kin and dimes\\nFeastin' on that golden bird, so fine (so fine)\\nPass the gravy, pour some wine\\n\\n[Chorus]\\nTurkey time, it's a country rhyme\\nGrab your plate, fill it up, take your time\\nAnd when the sun starts settin', we'll gather 'round\\nSingin' songs, sippin' cider, dancin' all night sound (oooh)\\nTurkey time, it's a country rhyme\\n\\n[Verse]\\nWell, it's that time of year, with family near\\nGatherin' 'round the table, filled with cheer\\nTurkey's roasting, gravy's flowin'\\nMashed potatoes piled high, it's Thanksiving time\\n\\n[Chorus]\\nTurkey time, it's a country rhyme\\nGrab your plate, fill it up, take your time\\nAnd when the sun starts settin', we'll gather 'round\\nSingin' songs, sippin' cider, dancin' all night sound (oooh)\\nTurkey time, it's a country rhyme\\n\\n[Bridge]  \\nTurkey time is here again\\nTurkey time is here again\\nTurkey time is here again\\nTurkey tiiiiime!\"\n",
    "\n",
    "# audio_fp = \"/home/georg/notebooks/samples/mistletoe.mp3\"\n",
    "# tags = \"Country\"\n",
    "# lyrics = \"[intro]\\nooooh\\n(what is life, what is love?)\\n[Verse]\\nAnother year has come and gone, and I'm not where I long to be\\nSnowfall on the city streets, but it ain't Christmas to me\"\n",
    "\n",
    "\n",
    "# audio_fp = \"/home/georg/notebooks/samples/blau2.mp3\"\n",
    "\n",
    "# audio_fp = \"/home/georg/notebooks/samples/halo.wav\"\n",
    "# offset_s = 20\n",
    "\n",
    "audio = Audio.from_file(audio_fp, n_channels=2, sample_rate=48000).get_segment(\n",
    "    from_s=offset_s, to_s=offset_s + duration_s + 0.01\n",
    ")\n",
    "audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "286abe63",
   "metadata": {},
   "outputs": [],
   "source": [
    "# gen_id = \"a5e2198a-f352-4abb-9a24-7f81b143ded3\"  # stone\n",
    "# # gen_id = \"37d0ac56-ca62-4ffc-bd91-351aa77ca9a7\"  # marc example\n",
    "gen_id = \"f966588f-d92f-46d0-8c6c-fef0a92d8764\"  # 20xx\n",
    "# # gen_id = \"053e61f7-2d4a-4f7a-ac92-e25c6d84c089\"  # sam ex\n",
    "mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.mp3\"\n",
    "audio = Audio.from_s3(mp3_filepath, n_channels=2, sample_rate=48000)\n",
    "audio.play()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9686747e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # wetdry\n",
    "# steps = 20\n",
    "# gen_id = \"f8a50e25-d705-49fa-b83a-82521014cb63\"  # vocal\n",
    "# # gen_id = \"0a3245dd-ceb9-46e1-a580-f7f8ca1ac8ed\"  # drums\n",
    "# # gen_id = \"14ce9bab-8b70-41a8-8e53-576479f4843a\" # piano\n",
    "# # gen_id = \"bdf764c9-cbc4-4069-bcb4-c1e7bc023e72\" # synth\n",
    "# # gen_id = \"9aea8ca8-37c2-48bf-8409-cdc8f3475e1b\" # backing vocal\n",
    "# # gen_id = \"b8ab993a-6934-4de2-a2ae-2764249b4c6b\" # lead vocal\n",
    "# lead_vocal = Audio.from_s3(gen_id, n_channels=2, sample_rate=48000)\n",
    "# lead_vocal = lead_vocal.get_segment(from_s=0, to_s=90)\n",
    "# # lead_vocal = generate(\n",
    "# #     audio,\n",
    "# #     steps=steps,\n",
    "# #     seed=0,\n",
    "# #     tags=\"extract Lead Vocal\",\n",
    "# #     downscale_ctx_vector=False,\n",
    "# # )\n",
    "# lead_vocal.play()\n",
    "\n",
    "# dry_audio = generate(\n",
    "#     lead_vocal,\n",
    "#     steps=steps,\n",
    "#     seed=0,\n",
    "#     tags=\"dry\",\n",
    "#     downscale_ctx_vector=False,\n",
    "#     # tag_cfg_coef=3.0,\n",
    "# )\n",
    "# dry_audio.play()\n",
    "\n",
    "# wet_audio = generate(\n",
    "#     dry_audio,\n",
    "#     steps=steps,\n",
    "#     seed=0,\n",
    "#     tags=\"wet\",\n",
    "#     downscale_ctx_vector=False,\n",
    "#     # tag_cfg_coef=3.0,\n",
    "# )\n",
    "# wet_audio.play()\n",
    "# # break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a43ae4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "steps = 12\n",
    "\n",
    "\n",
    "audios = []\n",
    "\n",
    "\n",
    "tags = [\n",
    "    \"extract [Guitar]\",\n",
    "    \"extract [Bass]\",\n",
    "    \"extract [Piano]\",\n",
    "    # \"remove Guitar\",\n",
    "    \"add [Flute]\",\n",
    "    \"extract [Lead Vocal]\",\n",
    "    # \"remove Lead Vocal, Backing Vocals\",\n",
    "    # \"extract Lead Vocal, Backing Vocals\",\n",
    "    # \"extract Lead Vocal, Backing Vocals, Drum Kit\",\n",
    "    # \"remove Lead Vocal\",\n",
    "    \"extract [Drum Kit]\",\n",
    "    # \"single_extract Drum Kit\",\n",
    "    \"remove [Drum Kit]\",\n",
    "    # \"extract Bass\",\n",
    "    # \"add Flute\",\n",
    "]\n",
    "for tag in tags:\n",
    "    print(f\"Generating {tag}\")\n",
    "    pred_audio, latents = generate(\n",
    "        audio,\n",
    "        steps=steps,\n",
    "        seed=3,\n",
    "        # tags=tag,\n",
    "        lyrics=tag,\n",
    "        downscale_ctx_vector=False,\n",
    "        # tag_cfg_coef=3.0,\n",
    "        text_cfg_coef=1.0,\n",
    "        # ctx_vector_scale=0.4,\n",
    "        return_latents=True,\n",
    "        normalize_volume=False,\n",
    "        noise_ctx_vector=0.3,\n",
    "    )\n",
    "    pred_audio.play()\n",
    "    # # Crop arrays to the shortest length to avoid broadcasting error\n",
    "    # min_length = min(audio.array_float.shape[1], pred_audio.array_float.shape[1])\n",
    "    # audio_cropped = audio.array_float[:, :min_length]\n",
    "    # pred_audio_cropped = pred_audio.array_float[:, :min_length]\n",
    "    # mixed_arr = audio_cropped + pred_audio_cropped\n",
    "    # mixed = Audio.from_array_float(mixed_arr, sample_rate=audio.sample_rate, max_allowed_val=10)\n",
    "    # mixed.play()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37c1188e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !aws s3 cp /app/suno/checkpoints/2025-03-06_23-22-48_s4786/step_40000_infer.pt s3://suno-data/victor/checkpoints/diffusion/stems_v1.pt\n",
    "# !aws s3 cp /app/suno/checkpoints/2025-03-09_16-56-18_s5489/last_ckpt_infer.pt s3://suno-data/victor/checkpoints/diffusion/stems_v1p1.pt\n",
    "# !aws s3 cp /app/suno/checkpoints/2025-03-14_21-02-44_s1097/last_ckpt_infer.pt s3://suno-data/victor/checkpoints/diffusion/stems_v1p2.pt\n",
    "# !aws s3 cp /app/suno/checkpoints/2025-03-26_13-24-59_s1608/last_ckpt_infer.pt s3://suno-data/victor/checkpoints/diffusion/stems_v1p3.pt\n",
    "# !aws s3 cp /app/suno/checkpoints/2025-03-28_13-45-42_s1953/step_40000_infer.pt s3://suno-data/victor/checkpoints/diffusion/stems_v1p3.pt\n",
    "\n",
    "# !cp /app/suno/checkpoints/2025-03-28_13-45-42_s1953/step_40000_infer.pt /mnt/localdisk/models/victor/checkpoints/diffusion/stems_multi.pt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a9f59bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !python -m suno_utils.worker.modal_manager_model_store upload --path ssd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c815c682",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}