{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.audio import Audio\n",
    "from suno_utils.audio.midi import Midi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_samples = {\n",
    "    \"stone_vocals\": {\n",
    "        \"audio_path\": \"s3://suno-data-uploads/studio/uploads/45f04b32-8df3-4428-ab7a-47e561ff3847.mp3\",\n",
    "        \"instrument\": 53,\n",
    "    },\n",
    "    \"bass\": {\n",
    "        \"audio_path\": \"s3://suno-data-uploads/studio/uploads/ec626427-771d-4b2e-8cd0-7bd2a7f1cc19.mp3\",\n",
    "        \"instrument\": 32,\n",
    "    },\n",
    "    \"piano\": {\n",
    "        \"audio_path\": \"s3://suno-data-uploads/studio/uploads/fefebefd-a98b-4167-83ae-0722f8891824.mp3\",\n",
    "        \"instrument\": 1,\n",
    "    },\n",
    "    \"electric_guitar\": {\n",
    "        \"audio_path\": \"s3://suno-data-uploads/studio/uploads/bac51035-7732-47b7-b180-7a759edb6bd0.mp3\",\n",
    "        \"instrument\": 25,\n",
    "    },\n",
    "    \"flute\": {\n",
    "        \"audio_path\": \"s3://suno-data-uploads/studio/uploads/16dcfe3e-7760-40f5-9d6c-e598010383c9.mp3\",\n",
    "        \"instrument\": 73,\n",
    "    },\n",
    "    \"violin\": {\n",
    "        \"audio_path\": \"s3://suno-data-uploads/studio/uploads/6c957b14-66f1-4219-8b83-6001860055d7.mp3\",\n",
    "        \"instrument\": 40,\n",
    "    },\n",
    "}\n",
    "\n",
    "\n",
    "def audio_path_to_midi(audio_path: str):\n",
    "    id = audio_path.split(\"/\")[-1].split(\".\")[0]\n",
    "    midi_path = f\"s3://suno-data-uploads/studio/uploads/{id}.mid\"\n",
    "    return Midi.from_s3(midi_path)\n",
    "\n",
    "\n",
    "midi = audio_path_to_midi(\n",
    "    \"s3://suno-data-uploads/studio/uploads/632f8788-225b-4127-8974-57afa0b7a71c.mp3\"\n",
    ")\n",
    "midi.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(midi.num_instruments)\n",
    "print(midi.num_notes)\n",
    "midi.to_audio().play()\n",
    "# midi.write(\"test.mid\")\n",
    "midi = Midi.from_scale(12)\n",
    "midi.simplify_time().to_audio().play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from miditok import REMI, TokenizerConfig\n",
    "\n",
    "config = TokenizerConfig(\n",
    "    pitch_range=(0, 127),\n",
    "    use_rests=False,\n",
    "    use_programs=True,\n",
    "    one_token_stream_for_programs=True,\n",
    "    use_tempos=False,\n",
    "    num_velocities=32,\n",
    ")\n",
    "max_bar_embedding = 60 // 2 * 30  # 30 min max. 2s per bar\n",
    "midi_tokenizer = REMI(config, max_bar_embedding=60 // 2 * 30)\n",
    "midi = Midi.from_scale(12)\n",
    "midi.to_tokens(midi_tokenizer).tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "midi_tokenizer = Midi.load_tokenizer()\n",
    "tokens = midi.to_tokens(midi_tokenizer).ids\n",
    "print(tokens)\n",
    "midi = Midi.from_tokens(tokens, midi_tokenizer)\n",
    "midi.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "midi.make_stereo_comparison(midi.to_audio()).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, sample in test_samples.items():\n",
    "    print(f\"Transcribing {name}...\")\n",
    "    print(sample[\"audio_path\"])\n",
    "    midi = audio_path_to_midi(sample[\"audio_path\"])\n",
    "    midi.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Example modal call"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import modal\n",
    "\n",
    "model_f_stem = modal.Cls.from_name(\"midi-transcription-dev\", \"MidiTranscriptionStub\")\n",
    "gen_id = \"45f04b32-8df3-4428-ab7a-47e561ff3847\"  # stone vocals\n",
    "input_s3_path = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.mp3\"\n",
    "\n",
    "output_s3_path = model_f_stem().transcribe.remote(gen_id)\n",
    "print(f\"Done with modal call for gen_id {gen_id}. See it at {output_s3_path}\")\n",
    "\n",
    "midi = Midi.from_s3(output_s3_path)\n",
    "midi.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}