{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6537dc61",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "\n",
    "from suno_utils.tasks.mert_25 import preload_models as preload_mert_models_\n",
    "from suno_utils.tasks.musicfm_v3 import preload_models as preload_musicfm_models_\n",
    "from suno_utils.tasks.mert_25 import encode_both as encode_mert\n",
    "from suno_utils.tasks.musicfm_v3 import encode as encode_musicfm\n",
    "\n",
    "# preload_mert_models_(\n",
    "#     checkpoint_filepath=\"s3://suno-data/georg/models/semantic/mert_25.pt\",\n",
    "#     centroids_filepath=\"s3://suno-data/minz/models/mert_64residual_kmeans_centroids.npy\",\n",
    "# )\n",
    "\n",
    "preload_musicfm_models_(\n",
    "    checkpoint_filepath=\"/app2/suno/checkpoints/2025-10-23_23-53-15/last_ckpt_infer.pt\",\n",
    "    centroids_filepath=\"/home/minz/temp/musicfm3_1024d_nov17_centroids_4000_50.npy\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f68ee861",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from suno_utils.audio import Audio\n",
    "\n",
    "temp_audio = Audio.from_file(\n",
    "    \"/home/minz/temp/country_road.mp3\", sample_rate=48000, n_channels=2\n",
    ")\n",
    "inp_16 = torch.from_numpy(temp_audio.convert(16000, 2, 2).array_float)\n",
    "inp_24 = torch.from_numpy(temp_audio.convert(24000, 2, 1).array_float).unsqueeze(0)\n",
    "print(inp_16.shape, inp_24.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc23ec76",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%timeit\n",
    "# out = encode_mert([inp_24])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50ffab99",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "out = encode_musicfm([inp_16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2420b9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.profiling.profiler import ScheduledProfiler, ProfilerConfig\n",
    "\n",
    "s3_profile_bucket = \"suno-data\"\n",
    "s3_profile_prefix = \"traces/musicfm3\"\n",
    "config = ProfilerConfig(\n",
    "    wait_steps=0,\n",
    "    warmup_steps=1,\n",
    "    active_steps=3,\n",
    "    repeat=1,\n",
    "    export_backend=\"s3\",\n",
    "    s3_bucket=s3_profile_bucket,\n",
    "    s3_prefix=s3_profile_prefix,\n",
    "    local_directory=\"./profile_logs\",\n",
    "    project_name=\"musicfm_encode\",\n",
    "    with_stack=True,\n",
    ")\n",
    "\n",
    "profiler = ScheduledProfiler(config)\n",
    "\n",
    "with profiler:\n",
    "    for step in range(5):\n",
    "        out = encode_musicfm([inp_16])\n",
    "        profiler.step()\n",
    "        print(f\"Step {step} completed\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6369714",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88aee4b3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26a061f7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}