{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### This notebook is to generate samples from the conditions saved in test_slices.pkl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "import torch\n",
    "from model import init_ldm_model\n",
    "from model.model_sdf import Diffpro_SDF\n",
    "from model.sampler_sdf import SDFSampler\n",
    "\n",
    "import numpy as np\n",
    "import pickle\n",
    "from train.train_params import params_chord_cond, params_chord_lsh_cond\n",
    "from data.prepare_training_pianoroll.convert_to_midi import extend_piano_roll, piano_roll_to_midi, save_midi\n",
    "from generation_utils.fine_grained_control import *\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "WITH_MELODY = True # whether to use melody as condition\n",
    "MODEL_PATH = 'results/default_model/default_model/chkpts/weights.pt' # path of the model\n",
    "CONDITION_DATA_PATH = 'data/train_test_slices/test_slices.pkl' # path of the condition data\n",
    "SAVE_CHORD_IN_MIDI = False # whether to save the chords in midi file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/music/chord_trainer_clean/model/model_sdf.py:30: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  trained_leaner = torch.load(chkpt_fpath, map_location=device)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1767, 8, 64, 64)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/music/chord_trainer_clean/model/sampler_sdf.py:97: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
      "  self.autocast = torch.cuda.amp.autocast(enabled=is_autocast)\n"
     ]
    }
   ],
   "source": [
    "# load the model\n",
    "if WITH_MELODY:\n",
    "    ldm_model = init_ldm_model(params_chord_lsh_cond, debug_mode=False)\n",
    "else:\n",
    "    ldm_model = init_ldm_model(params_chord_cond, debug_mode=False)\n",
    "model = Diffpro_SDF.load_trained(ldm_model, MODEL_PATH).to(device)\n",
    "sampler = SDFSampler(model.ldm, 64, 64, is_autocast=False, device=device, debug_mode=False)\n",
    "\n",
    "\n",
    "# Load condition dataset\n",
    "with open(CONDITION_DATA_PATH, 'rb') as f:\n",
    "    test_data = pickle.load(f)\n",
    "\n",
    "print(test_data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[10, 2, 64, 64]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"overflow-x: scroll;\">Sample...</pre>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p_sample\n",
      "p_sample\n",
      "p_sample\n",
      "p_sample\n",
      "p_sample\n",
      "p_sample\n",
      "p_sample\n",
      "p_sample\n",
      "p_sample\n",
      "p_sample\n"
     ]
    }
   ],
   "source": [
    "background_cond = test_data[0:10,2:,:,:].copy() # use the sample with index 0-9 for generation\n",
    "background_cond = torch.Tensor(background_cond).to(device)\n",
    "\n",
    "output_x = sampler.generate(background_cond=background_cond, batch_size=background_cond.shape[0], \n",
    "                            same_noise_all_measure=False, X0EditFunc=X0EditFunc, use_classifier_free_guidance=True,\n",
    "                            use_lsh=True, device=device)\n",
    "output_x = torch.clamp(output_x, min=0, max=1)\n",
    "output_x = output_x.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save into midi files\n",
    "for i in range(output_x.shape[0]):\n",
    "    full_roll = extend_piano_roll(output_x[i]) # accompaniment roll\n",
    "    full_lsh_roll = None # melody roll\n",
    "    if background_cond.shape[1]>=6:\n",
    "        if background_cond[:,4:6,:,:].min()>=0:\n",
    "            full_lsh_roll = extend_piano_roll(background_cond[i,4:6,:,:].cpu().numpy())\n",
    "\n",
    "    midi_file = piano_roll_to_midi(full_roll, None, full_lsh_roll, bpm=80)\n",
    "    filename = f\"generated_samples/sample_{i}.mid\"\n",
    "    save_midi(midi_file, filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "music_sampling",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
