{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "<a href=\"https://colab.research.google.com/github/OliBomby/Mapperatorinator/blob/main/colab/mai_mod_inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "# Beatmap Modding with MaiMod\n",
    "\n",
    "This notebook is an interactive demo of an AI-driven osu! Beatmap Modding Tool created by OliBomby. This model is capable of finding various faults and inconsistencies in beatmaps which other automated modding tools can not detect. Run this tool on your beatmaps to get suggestions on how to improve them.\n",
    "\n",
    "### Instructions for running:\n",
    "\n",
    "* Make sure to use a GPU runtime, click:  __Runtime >> Change Runtime Type >> GPU__\n",
    "* __Execute each cell in order__. Press ▶️ on the left of each cell to execute the cell.\n",
    "* __Setup Environment__: run the first cell to clone the repository and install the required dependencies. You only need to run this cell once per session.\n",
    "* __Upload Audio__: choose the beatmap song .mp3 or .ogg file from your computer. You can find these files in stable by using File > Open Song Folder, or in lazer by using File > Edit Externally.\n",
    "* __Upload Beatmap__: choose the beatmap .osu file from your computer.\n",
    "* __Generate Suggestions__ to generate suggestions for your uploaded beatmap.\n"
   ],
   "id": "3c19902455e25588"
  },
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "#@title Setup Environment { display-mode: \"form\" }\n",
    "#@markdown Run this cell to clone the repository and install the required dependencies. You only need to run this cell once per session.\n",
    "\n",
    "!git clone https://github.com/OliBomby/Mapperatorinator.git\n",
    "%cd Mapperatorinator\n",
    "\n",
    "!pip install transformers==4.53.3\n",
    "!pip install hydra-core\n",
    "!pip install slider git+https://github.com/OliBomby/slider.git@gedagedigedagedaoh\n",
    "\n",
    "import os\n",
    "from google.colab import files\n",
    "from mai_mod import main\n",
    "from hydra import compose, initialize_config_dir\n",
    "\n",
    "input_audio = \"\"\n",
    "input_beatmap = \"\""
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "#@title Upload Audio { display-mode: \"form\" }\n",
    "#@markdown Run this cell to upload the song of the beatmap that you want to mod. Please upload a .mp3 or .ogg file. You can find these files in stable by using File > Open Song Folder, or in lazer by using File > Edit Externally.\n",
    "\n",
    "def upload_audio():\n",
    "    data = list(files.upload().keys())\n",
    "    if len(data) > 1:\n",
    "        print('Multiple files uploaded; using only one.')\n",
    "    file = data[0]\n",
    "    if not file.endswith('.mp3') and not file.endswith('.ogg'):\n",
    "        print('Invalid file format. Please upload a .mp3 or .ogg file.')\n",
    "        return \"\"\n",
    "    return data[0]\n",
    "\n",
    "input_audio = upload_audio()"
   ],
   "id": "624a60c5777279e7",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "#@title Upload Beatmap { display-mode: \"form\" }\n",
    "#@markdown Run this cell to upload the beatmap **.osu** file of the beatmap that you want to mod. You can find these files in stable by using File > Open Song Folder, or in lazer by using File > Edit Externally.\n",
    "\n",
    "def upload_beatmap():\n",
    "    data = list(files.upload().keys())\n",
    "    if len(data) > 1:\n",
    "        print('Multiple files uploaded; using only one.')\n",
    "    file = data[0]\n",
    "    if not file.endswith('.osu'):\n",
    "        print('Invalid file format. Please upload a .osu file.\\nIn stable you can find the .osu file in the song folder (File > Open Song Folder).\\nIn lazer you can find the .osu file by using File > Edit Externally.')\n",
    "        return \"\"\n",
    "    return file\n",
    "\n",
    "input_beatmap = upload_beatmap()"
   ],
   "id": "63884394491f6664",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "#@title Generate Suggestions { display-mode: \"form\" }\n",
    "#@markdown Run this cell to generate suggestions for your uploaded beatmap. The suggestions will be printed in the output.\n",
    "\n",
    "# Validate stuff\n",
    "assert os.path.exists(input_beatmap), \"Please upload a beatmap.\"\n",
    "assert os.path.exists(input_audio), \"Please upload an audio file.\"\n",
    "    \n",
    "# Create config\n",
    "config = \"mai_mod\"\n",
    "with initialize_config_dir(version_base=\"1.1\", config_dir=\"/content/Mapperatorinator/configs\"):\n",
    "    conf = compose(config_name=config)\n",
    "\n",
    "# Do inference\n",
    "conf.audio_path = input_audio\n",
    "conf.beatmap_path = input_beatmap\n",
    "conf.precision = \"fp32\"  # For some reason AMP causes OOM in Colab\n",
    "\n",
    "main(conf)"
   ],
   "id": "166eb3e5f9398554",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "accelerator": "GPU",
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}