{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Decomposition\n",
    "\n",
    "This notebook calls modal workers to decompose a song into stems and midi."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.decomp import decomp, save_decomp_zip\n",
    "\n",
    "gen_id = \"a5e2198a-f352-4abb-9a24-7f81b143ded3\"\n",
    "\n",
    "# stems = decomp(gen_id)\n",
    "# save_decomp_zip(stems, name=gen_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sum([stem.midi for stem in stems]).show()\n",
    "# for stem in stems:\n",
    "#     print(stem.category)\n",
    "#     stem.midi.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Batched"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import concurrent.futures\n",
    "import os\n",
    "\n",
    "ids_to_decompose = {\n",
    "    \"stone\": \"a5e2198a-f352-4abb-9a24-7f81b143ded3\",\n",
    "    # \"sara\": \"a1936d21-d47a-44ad-8551-25feb86a58e9\",\n",
    "    # \"pop\": \"e3ccb02b-696e-4335-a674-9f7821b202f4\",\n",
    "    # \"chill\": \"63774ef1-ab73-472b-84dc-3c9d9a385339\",\n",
    "    # \"rainbow\": \"7aa56e05-3e52-4cd8-b7e4-cb009480b0e7\",\n",
    "    # \"surf\": \"90a879c4-1124-4c35-8f3a-723faf6130cc\",\n",
    "    # \"henry\": \"313c8579-456f-4218-828e-ca165ae2b432\",\n",
    "    # \"turkey\": \"0e72b9be-75e7-4d1c-bcb2-091a0d95f293\",\n",
    "}\n",
    "os.makedirs(\"decomps\", exist_ok=True)\n",
    "\n",
    "\n",
    "def process_decomp(item):\n",
    "    name, gen_id = item\n",
    "    stems = decomp(gen_id)\n",
    "    print(f\"Decomposed {name}, saving to decomps/{name}.zip\")\n",
    "    save_decomp_zip(stems, name=f\"decomps/{name}\")\n",
    "\n",
    "\n",
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    executor.map(process_decomp, ids_to_decompose.items())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}