accompany_rhythm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import os\n",
    "import sys\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "\n",
    "code_model_dir = './model'\n",
    "code_utils_dir = os.path.join(code_model_dir, 'utils')\n",
    "sys.path.extend([code_model_dir, code_utils_dir])\n",
    "\n",
    "MODEL_FP = 'model/pretrained/LakhNES/model.pt'\n",
    "VOCAB_FP = 'model/pretrained/LakhNES/vocab.txt'\n",
    "USE_CUDA = True\n",
    "\n",
    "device = torch.device(\"cuda\" if USE_CUDA else \"cpu\")\n",
    "\n",
    "# Load the best saved model.\n",
    "with open(MODEL_FP, 'rb') as f:\n",
    "    model = torch.load(f)\n",
    "model.backward_compatible()\n",
    "model = model.to(device)\n",
    "\n",
    "# Make sure model uses vanilla softmax.\n",
    "if model.sample_softmax > 0:\n",
    "  raise NotImplementedError()\n",
    "if model.crit.n_clusters != 0:\n",
    "  raise NotImplementedError()\n",
    "\n",
    "# Load the vocab.\n",
    "idx2sym = ['<S>']\n",
    "ins_to_midi_pitches = defaultdict(list)\n",
    "with open(VOCAB_FP, 'r') as f:\n",
    "  for line in f:\n",
    "    evt = line.strip().split(',')[-1]\n",
    "    if 'NOTEON' in evt:\n",
    "        toks = evt.split('_')\n",
    "        ins_to_midi_pitches[toks[0]].append(int(toks[2]))\n",
    "    idx2sym.append(evt)\n",
    "sym2idx = {s:i for i, s in enumerate(idx2sym)}\n",
    "wait_amts = set([int(s[3:]) for s in idx2sym if s[:2] == 'WT'])\n",
    "print(len(idx2sym))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TX1_PATH = 'data/nesmdb_tx1/test/*.tx1.txt'\n",
    "\n",
    "import glob as glob\n",
    "import os\n",
    "\n",
    "TX1_FPS = sorted(glob.glob(TX1_PATH))\n",
    "\n",
    "def evts_to_ids(evts):\n",
    "    ids = []\n",
    "    for s in evts:\n",
    "        if s in sym2idx:\n",
    "            ids.append(sym2idx[s])\n",
    "        else:\n",
    "            assert s[:2] == 'WT'\n",
    "            wait_amt = int(s.split('_')[1])\n",
    "            closest = min(wait_amts, key=lambda x:abs(x - wait_amt))\n",
    "            ids.append(sym2idx['WT_{}'.format(closest)])\n",
    "    return ids\n",
    "\n",
    "def ids_to_evts(ids):\n",
    "    return [idx2sym[i] for i in ids]\n",
    "\n",
    "def evts_to_rhythm(evts):\n",
    "    rhythm = []\n",
    "    for e in evts:\n",
    "        t = e.split('_')\n",
    "        if t[0] == 'WT' or t[0] == 'NO' or t[0] == '<S>':\n",
    "            rhythm.append(e)\n",
    "    return rhythm\n",
    "\n",
    "fn_to_evts = {}\n",
    "for fp in TX1_FPS:\n",
    "    fn = os.path.split(fp)[1].split('.')[0]\n",
    "    with open(fp, 'r') as f:\n",
    "        evts = f.read().strip().splitlines()\n",
    "    \n",
    "    if len(evts) > 0 and evts[0][:2] == 'WT':\n",
    "        evts = evts[1:]\n",
    "    if len(evts) > 0 and evts[-1][:2] == 'WT':\n",
    "        evts = evts[:-1]\n",
    "        \n",
    "    evts = ['<S>'] + evts + ['<S>']\n",
    "    fn_to_evts[fn] = evts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tempfile\n",
    "import xmlrpc.client\n",
    "from scipy.io.wavfile import read as wavread, write as wavwrite\n",
    "from IPython.display import display, Audio\n",
    "\n",
    "s = xmlrpc.client.ServerProxy('http://localhost:1337')\n",
    "\n",
    "def tx1_to_wav(tx1):\n",
    "    tf = tempfile.NamedTemporaryFile()\n",
    "    wf = tempfile.NamedTemporaryFile()\n",
    "    \n",
    "    with open(tf.name, 'w') as f:\n",
    "        f.write('\\n'.join(tx1))\n",
    "    \n",
    "    s.tx1_to_wav(tf.name, wf.name)\n",
    "    fs, wav = wavread(wf.name)\n",
    "    \n",
    "    tf.close()\n",
    "    wf.close()\n",
    "    \n",
    "    return wav\n",
    "\n",
    "from IPython.display import display, Audio\n",
    "\n",
    "def paprev(tx1_ids, fn, displaywav=True):\n",
    "    if len(tx1_ids) > 0 and tx1_ids[0] == 0:\n",
    "        tx1_ids = tx1_ids[1:]\n",
    "    if len(tx1_ids) > 0 and tx1_ids[-1] == 0:\n",
    "        tx1_ids = tx1_ids[:-1]\n",
    "    \n",
    "    tx1 = [idx2sym[i] for i in tx1_ids]\n",
    "    \n",
    "    with open(fn, 'w') as f:\n",
    "        f.write('\\n'.join(tx1))\n",
    "    \n",
    "    wav = tx1_to_wav(tx1)\n",
    "    \n",
    "    wavfp = fn.replace('.tx1.txt', '.wav')\n",
    "    \n",
    "    wav = np.copy(wav)\n",
    "    wav *= 32767.\n",
    "    wav = np.clip(wav, -32768., 32767.)\n",
    "    wav = wav.astype(np.int16)\n",
    "    wavwrite(wavfp, 44100, wav)\n",
    "\n",
    "    if displaywav:\n",
    "        display(Audio(wav, rate=44100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "out_dir = 'generated/rhythms'\n",
    "if not os.path.isdir(out_dir):\n",
    "    os.makedirs(out_dir)\n",
    "\n",
    "fn = '214_MadoolanoTsubasa_04_05MainTowerCastle'\n",
    "genlen = 512\n",
    "\n",
    "rhythm_evts = fn_to_evts[fn]\n",
    "assert len(rhythm_evts) > genlen\n",
    "\n",
    "rhythm_evts = rhythm_evts[:genlen + 1]\n",
    "# Quantize\n",
    "rhythm_evts = ids_to_evts(evts_to_ids(rhythm_evts))\n",
    "\n",
    "paprev(evts_to_ids(rhythm_evts), '{}/{}_full.tx1.txt'.format(out_dir, fn))\n",
    "paprev(evts_to_ids(evts_to_rhythm(rhythm_evts)), '{}/{}_rhythm.tx1.txt'.format(out_dir, fn))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "temp = 0.95\n",
    "topk = 32\n",
    "memlen = 512\n",
    "\n",
    "from utils import TxlSimpleSampler\n",
    "import uuid\n",
    "\n",
    "sampler = TxlSimpleSampler(model, device, mem_len=memlen)\n",
    "\n",
    "inp = 0\n",
    "nll = 0.\n",
    "rhythm = []\n",
    "for i in range(genlen):\n",
    "    tarevt = rhythm_evts[i + 1]\n",
    "    tar = sym2idx[tarevt]\n",
    "    toks = tarevt.split('_')\n",
    "    if 'NOTEOFF' in tarevt or toks[0] == 'NO' or toks[0] == 'WT':\n",
    "        #print(tarevt)\n",
    "        sampler.sample_next_token_updating_mem(inp, exclude_eos=False)\n",
    "    else:\n",
    "        _, probs = sampler.sample_next_token_updating_mem(inp, exclude_eos=False)\n",
    "        _probs = probs.cpu().numpy()\n",
    "        \n",
    "        ins = toks[0]\n",
    "        \n",
    "        ins_pitches = ins_to_midi_pitches[ins]\n",
    "        \n",
    "        samplefrom = [sym2idx['{}_NOTEON_{}'.format(ins, p)] for p in ins_pitches]\n",
    "        \n",
    "        _probmask = np.zeros_like(_probs)\n",
    "        _probmask[samplefrom] = 1.\n",
    "        _probs *= _probmask\n",
    "        \n",
    "        if topk is not None:\n",
    "            ind = np.argpartition(_probs, -topk)[-topk:]\n",
    "            _probmask = np.zeros_like(_probs)\n",
    "            _probmask[samplefrom] = 1.\n",
    "            _probs *= _probmask\n",
    "        \n",
    "        _probs /= np.sum(_probs)\n",
    "        tar = np.random.choice(range(len(idx2sym)), p=_probs)\n",
    "        assert tar in samplefrom\n",
    "    rhythm.append(tar)\n",
    "\n",
    "    inp = tar\n",
    "\n",
    "paprev(rhythm, '{}/{}_jamrhythm_{}.tx1.txt'.format(out_dir, fn, str(uuid.uuid4().hex)[:4]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jupyter3_Python_3",
   "language": "python",
   "name": "jupyter3_python_3"
  },
  "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}