{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ePWjo4hLkSZh"
   },
   "source": [
    "# Orpheus MIRE Auto-Continuation Generator Notebook (ver. 1.0)\n",
    "\n",
    "***\n",
    "\n",
    "Powered by tegridy-tools: https://github.com/asigalov61/tegridy-tools\n",
    "\n",
    "***\n",
    "\n",
    "#### Project Los Angeles\n",
    "\n",
    "#### Tegridy Code 2026\n",
    "\n",
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "y1H5U8iiAIgD"
   },
   "source": [
    "# Setup Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "8Dt7FYceaCKF",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Install all dependencies (run only once per session)\n",
    "\n",
    "!git clone https://github.com/asigalov61/tegridy-tools\n",
    "!pip install tqdm\n",
    "!pip install ipywidgets\n",
    "\n",
    "!pip install einops\n",
    "!pip install einx\n",
    "!pip install scikit-learn\n",
    "!pip install torch-summary\n",
    "\n",
    "!pip install huggingface_hub"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "Lqp3urZyaDAp",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Import all needed modules\n",
    "\n",
    "print('=' * 70)\n",
    "print('Loading needed modules. Please wait...')\n",
    "\n",
    "import os\n",
    "\n",
    "os.environ[\"HF_XET_HIGH_PERFORMANCE\"] = \"1\"\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "import random\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "print('=' * 70)\n",
    "print('Loading TMIDIX module...')\n",
    "\n",
    "%cd ~/tegridy-tools/tegridy-tools/\n",
    "\n",
    "import TMIDIX\n",
    "\n",
    "%cd ~/tegridy-tools/tegridy-tools/X-Transformer/\n",
    "\n",
    "from x_transformer_2_3_1 import TransformerWrapper, Encoder, Decoder, AutoregressiveWrapper, top_p\n",
    "from x_transformer_2_3_1 import predict_masked_tokens_iter, analyze_generation_metrics\n",
    "\n",
    "%cd ~\n",
    "\n",
    "import torch\n",
    "from torch.amp import autocast\n",
    "\n",
    "from torchsummary import summary\n",
    "\n",
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "print('=' * 70)\n",
    "print('Done!')\n",
    "print('Enjoy! :)')\n",
    "print('=' * 70)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PcEkAnhyAIgL"
   },
   "source": [
    "# Download and Init Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download Orpheus MIRE Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_hub_download(repo_id='asigalov61/Orpheus-Music-Transformer',\n",
    "                filename='Orpheus_Music_Transformer_Masked_Encoder_Trained_Model_23000_steps_0.6548_loss_0.8132_acc.pth',\n",
    "                local_dir='./Models/',\n",
    "                )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Init Orpheus MIRE Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('=' * 70)\n",
    "print('Building model...')\n",
    "\n",
    "ENC_SEQ_LEN = 2048\n",
    "ENC_PAD_IDX = 18820\n",
    "DEVICE = 'cuda'\n",
    "full_path_to_trained_model = './Models/Orpheus_Music_Transformer_Masked_Encoder_Trained_Model_23000_steps_0.6548_loss_0.8132_acc.pth'\n",
    "\n",
    "enc_model = TransformerWrapper(\n",
    "    num_tokens = ENC_PAD_IDX+1,\n",
    "    max_seq_len = ENC_SEQ_LEN,\n",
    "    attn_layers = Encoder(dim = 2048,\n",
    "                          depth = 12,\n",
    "                          heads = 16,\n",
    "                          rotary_pos_emb = True,\n",
    "                          attn_flash = True\n",
    "                         )\n",
    "    )\n",
    "\n",
    "enc_model.load_state_dict(torch.load(full_path_to_trained_model))\n",
    "enc_model.cuda()\n",
    "enc_model.eval()\n",
    "\n",
    "summary(enc_model)\n",
    "\n",
    "print('=' * 70)\n",
    "print('Done!')\n",
    "print('=' * 70)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download Orpheus Large Base Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_hub_download(repo_id='asigalov61/Orpheus-Music-Transformer',\n",
    "                            filename='Orpheus_Music_Transformer_Large_Trained_Model_43860_steps_0.6682_loss_0.8054_acc.pth',\n",
    "                            local_dir='./Models/',\n",
    "                            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Init Orpheus Large Base Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SEQ_LEN = 8192\n",
    "PAD_IDX = 18819\n",
    "\n",
    "model = TransformerWrapper(\n",
    "    num_tokens = PAD_IDX+1,\n",
    "    max_seq_len = SEQ_LEN,\n",
    "    attn_layers = Decoder(dim = 2048,\n",
    "                          depth = 16,\n",
    "                          heads = 16,\n",
    "                          rotary_pos_emb = True,\n",
    "                          attn_flash = True\n",
    "                         )\n",
    "    )\n",
    "\n",
    "model = AutoregressiveWrapper(model, ignore_index = PAD_IDX, pad_value=PAD_IDX)\n",
    "\n",
    "print('=' * 70)\n",
    "print('Loading model checkpoint...')\n",
    "\n",
    "model_path = './Models/Orpheus_Music_Transformer_Large_Trained_Model_43860_steps_0.6682_loss_0.8054_acc.pth'\n",
    "\n",
    "model.load_state_dict(torch.load(model_path))\n",
    "\n",
    "print('=' * 70)\n",
    "\n",
    "model.cuda()\n",
    "model.eval()\n",
    "\n",
    "model = torch.compile(model)\n",
    "\n",
    "print('Done!')\n",
    "\n",
    "summary(model)\n",
    "\n",
    "dtype = torch.bfloat16\n",
    "\n",
    "ctx = autocast(device_type='cuda', dtype=dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load source MIDI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "midi_file = './tegridy-tools/tegridy-tools/seed-intro.mid'\n",
    "\n",
    "print('=' * 70)\n",
    "print('Loading MIDI File:', midi_file)\n",
    "print('=' * 70)\n",
    "\n",
    "raw_score = TMIDIX.midi2single_track_ms_score(midi_file)\n",
    "\n",
    "escore_notes = TMIDIX.advanced_score_processor(raw_score, return_enhanced_score_notes=True, apply_sustain=True)\n",
    "\n",
    "escore_notes = TMIDIX.augment_enhanced_score_notes(escore_notes[0], sort_drums_last=True)\n",
    "\n",
    "escore_notes = TMIDIX.remove_duplicate_pitches_from_escore_notes(escore_notes)\n",
    "\n",
    "escore_notes = TMIDIX.fix_escore_notes_durations(escore_notes, min_notes_gap=0)\n",
    "\n",
    "dscore = TMIDIX.delta_score_notes(escore_notes)\n",
    "\n",
    "dcscore = TMIDIX.chordify_score([d[1:] for d in dscore])\n",
    "\n",
    "melody_chords = [18816]\n",
    "\n",
    "#=======================================================\n",
    "# MAIN PROCESSING CYCLE\n",
    "#=======================================================\n",
    "\n",
    "for i, c in enumerate(dcscore):\n",
    "\n",
    "    # Delta start-times\n",
    "    delta_time = c[0][0]\n",
    "    melody_chords.append(delta_time)\n",
    "\n",
    "    for e in c:\n",
    "    \n",
    "        #=======================================================\n",
    "        \n",
    "        # Durations\n",
    "        dur = max(1, min(255, e[1]))\n",
    "\n",
    "        # Patches\n",
    "        pat = max(0, min(128, e[5]))\n",
    "        \n",
    "        # Pitches\n",
    "        ptc = max(1, min(127, e[3]))\n",
    "        \n",
    "        # Velocities\n",
    "        # Calculating octo-velocity\n",
    "        \n",
    "        vel = max(8, min(127, e[4]))\n",
    "        velocity = round(vel / 15)-1\n",
    "        \n",
    "        #=======================================================\n",
    "        # FINAL NOTE SEQ\n",
    "        #=======================================================\n",
    "        \n",
    "        # Writing final note\n",
    "        pat_ptc = (128 * pat) + ptc \n",
    "        dur_vel = (8 * dur) + velocity\n",
    "\n",
    "        melody_chords.extend([pat_ptc+256, dur_vel+16768]) # 18816\n",
    "\n",
    "print('Done!')\n",
    "print('=' * 70)\n",
    "print('Composition has', len(melody_chords), 'tokens')\n",
    "print('=' * 70)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_prime_tokens = 1024\n",
    "num_gen_tokens = 512\n",
    "num_gen_chunks = 6\n",
    "num_mem_tokens = 4096\n",
    "num_ref_toks = 768\n",
    "num_ref_iter = 25\n",
    "temperature = 0.95\n",
    "top_p_value = 0.99\n",
    "debug = False\n",
    "\n",
    "#===================================================================\n",
    "\n",
    "song = melody_chords[:num_prime_tokens]\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "for i in tqdm(range(num_gen_chunks)):\n",
    "\n",
    "    x = torch.LongTensor([song[-num_mem_tokens:]]).cuda()\n",
    "    \n",
    "    with ctx:\n",
    "        out, met = model.generate_metrics(x,\n",
    "                                          num_gen_tokens,\n",
    "                                          temperature=temperature,\n",
    "                                          filter_logits_fn=top_p,\n",
    "                                          filter_kwargs={'thres': top_p_value},\n",
    "                                          return_prime=False,\n",
    "                                          return_metrics=True,\n",
    "                                          verbose=False\n",
    "                                         )\n",
    "    \n",
    "    outs = out.tolist()\n",
    "\n",
    "    res = analyze_generation_metrics(met, torch.tensor(outs[0]), print_report=False)\n",
    "\n",
    "    ano_pos = [r['step'] for r in res['anomalies']]\n",
    "\n",
    "    # =============================================\n",
    "\n",
    "    mpos = [i+num_ref_toks for i in range(num_gen_tokens) if i in ano_pos]\n",
    "    \n",
    "    ref_seq = song[-num_ref_toks:] + outs[0]\n",
    "    tries = 0\n",
    "    max_tries = random.randint(5, 10)\n",
    "    \n",
    "    while tries < max_tries:\n",
    "    \n",
    "        results = predict_masked_tokens_iter(enc_model,\n",
    "                                             ref_seq,\n",
    "                                             mask_positions=mpos,\n",
    "                                             topk=1,\n",
    "                                             seq_len=ENC_SEQ_LEN,\n",
    "                                             mask_idx=ENC_PAD_IDX-1,\n",
    "                                             pad_idx=ENC_PAD_IDX,\n",
    "                                             vocab_size=ENC_PAD_IDX+1,\n",
    "                                             iterations=num_ref_iter,\n",
    "        )\n",
    "        ref_seq = results['predicted_ids']\n",
    "        \n",
    "        probs = [r['topk'][0][1] for r in  results['predictions']]\n",
    "\n",
    "        avg_prob = sum(probs) / len(probs)\n",
    "\n",
    "        if debug:\n",
    "            print(avg_prob)\n",
    "\n",
    "        tries += 1\n",
    "\n",
    "    if 18817 in ref_seq[num_ref_toks:]:\n",
    "        print('Outro token detected! Will stop generation soon!')\n",
    "\n",
    "    if 18818 not in ref_seq[num_ref_toks:]:\n",
    "        song.extend(ref_seq[num_ref_toks:])\n",
    "\n",
    "    else:\n",
    "        print('EOS token detected! Stopping generation...')\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert to MIDI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Sample INTs', song[:15])\n",
    "\n",
    "if len(song) != 0:\n",
    "\n",
    "    song_f = []\n",
    "    \n",
    "    time = 0\n",
    "    dur = 1\n",
    "    vel = 90\n",
    "    pitch = 60\n",
    "    channel = 0\n",
    "    patch = 0\n",
    "\n",
    "    patches = [-1] * 16\n",
    "\n",
    "    channels = [0] * 16\n",
    "    channels[9] = 1\n",
    "\n",
    "    for ss in song:\n",
    "\n",
    "        if 0 <= ss < 256:\n",
    "\n",
    "            time += ss * 16\n",
    "\n",
    "        if 256 <= ss < 16768:\n",
    "\n",
    "            patch = (ss-256) // 128\n",
    "\n",
    "            if patch < 128:\n",
    "\n",
    "                if patch not in patches:\n",
    "                  if 0 in channels:\n",
    "                      cha = channels.index(0)\n",
    "                      channels[cha] = 1\n",
    "                  else:\n",
    "                      cha = 15\n",
    "\n",
    "                  patches[cha] = patch\n",
    "                  channel = patches.index(patch)\n",
    "                else:\n",
    "                  channel = patches.index(patch)\n",
    "\n",
    "            if patch == 128:\n",
    "                channel = 9\n",
    "\n",
    "            pitch = (ss-256) % 128\n",
    "\n",
    "\n",
    "        if 16768 <= ss < 18816:\n",
    "\n",
    "            dur = ((ss-16768) // 8) * 16\n",
    "            vel = (((ss-16768) % 8)+1) * 15\n",
    "\n",
    "            song_f.append(['note', time, dur, channel, pitch, vel, patch])\n",
    "\n",
    "song_f = TMIDIX.remove_duplicate_pitches_from_escore_notes(song_f)\n",
    "\n",
    "song_f = TMIDIX.fix_escore_notes_durations(song_f, min_notes_gap=0)\n",
    "\n",
    "patches = [0 if x==-1 else x for x in patches]\n",
    "\n",
    "output_score, patches, overflow_patches = TMIDIX.patch_enhanced_score_notes(song_f)\n",
    "\n",
    "detailed_stats = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(output_score,\n",
    "                                                          output_signature = 'Orpheus Music Transformer',\n",
    "                                                          output_file_name = './Orpheus-Music-Transformer-Composition',\n",
    "                                                          track_name='Project Los Angeles',\n",
    "                                                          list_of_MIDI_patches=patches\n",
    "                                                          )\n",
    "\n",
    "print('Done!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Congrats! You did it! :)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}