notebooks/train_diffusion_human.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os\n",
    "import sys\n",
    "\n",
    "# change path to parent directory for paths\n",
    "sys.path.append(os.path.dirname(os.getcwd()))\n",
    "os.chdir(os.path.dirname(os.getcwd()))\n",
    "\n",
    "import accelerate\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "import torch\n",
    "import yaml\n",
    "from diffusers.optimization import get_scheduler\n",
    "from omegaconf import OmegaConf\n",
    "from tqdm.auto import tqdm\n",
    "from einops import rearrange\n",
    "\n",
    "from ldns.networks import AutoEncoder, CountWrapper\n",
    "from ldns.utils.plotting_utils import *\n",
    "from ldns.losses import latent_regularizer\n",
    "from ldns.networks import Denoiser\n",
    "from diffusers.training_utils import EMAModel\n",
    "from diffusers.schedulers import DDPMScheduler\n",
    "\n",
    "import lovely_tensors as lt\n",
    "lt.monkey_patch()\n",
    "\n",
    "matplotlib.rc_file('matplotlibrc')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## load config and model path\n",
    "\n",
    "cfg_ae = OmegaConf.load(\"conf/autoencoder-human.yaml\")\n",
    "\n",
    "\n",
    "cfg_yaml = \"\"\"\n",
    "denoiser_model:\n",
    "  C_in: 32\n",
    "  C: 384\n",
    "  kernel: s4\n",
    "  num_blocks: 6\n",
    "  bidirectional: True\n",
    "  num_train_timesteps: 1000\n",
    "training:\n",
    "  lr: 0.001\n",
    "  num_epochs: 2000\n",
    "  num_warmup_epochs: 100\n",
    "  batch_size: 256\n",
    "  random_seed: 42\n",
    "  precision: \"no\"\n",
    "exp_name: diffusion_human\n",
    "\"\"\"\n",
    "cfg = OmegaConf.create(yaml.safe_load(cfg_yaml))\n",
    "cfg.dataset = cfg_ae.dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data path: data/human/competitionData\n",
      "Train: 8417, Val: 213, Test: 640\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from ldns.data.human import get_human_dataloaders\n",
    "\n",
    "# set seed\n",
    "torch.manual_seed(cfg.training.random_seed)\n",
    "np.random.seed(cfg.training.random_seed)\n",
    "\n",
    "train_dataloader, val_dataloader, test_dataloader = get_human_dataloaders(\n",
    "        cfg_ae.dataset.datapath, batch_size=cfg_ae.training.batch_size\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ae_model = AutoEncoder(\n",
    "    C_in=cfg_ae.model.C_in,\n",
    "    C=cfg_ae.model.C,\n",
    "    C_latent=cfg_ae.model.C_latent,\n",
    "    L=cfg_ae.dataset.max_seqlen,\n",
    "    num_blocks=cfg_ae.model.num_blocks,\n",
    "    num_blocks_decoder=cfg_ae.model.get(\"num_blocks_decoder\", cfg_ae.model.num_blocks),\n",
    "    num_lin_per_mlp=cfg_ae.model.get(\"num_lin_per_mlp\", 2),  # default 2\n",
    "    bidirectional=cfg_ae.model.get(\"bidirectional\", False),\n",
    ")\n",
    "\n",
    "ae_model = CountWrapper(ae_model)\n",
    "ae_model.load_state_dict(torch.load(f\"exp/{cfg_ae.exp_name}/model.pt\"))\n",
    "\n",
    "accelerator = accelerate.Accelerator(\n",
    "    mixed_precision=\"no\",\n",
    ")\n",
    "\n",
    "ae_model = accelerator.prepare(ae_model)\n",
    "print(cfg_ae.exp_name)\n",
    "\n",
    "(\n",
    "    train_dataloader,\n",
    "    val_dataloader,\n",
    "    test_dataloader,\n",
    ") = accelerator.prepare(\n",
    "    train_dataloader,\n",
    "    val_dataloader,\n",
    "    test_dataloader,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ldns.data.human import LatentHumanDataset\n",
    "\n",
    "# create training dataset\n",
    "latent_dataset_train = LatentHumanDataset(train_dataloader, ae_model, clip=False)\n",
    "\n",
    "# create validation and test datasets using training set statistics\n",
    "latent_dataset_val = LatentHumanDataset(\n",
    "    val_dataloader,\n",
    "    ae_model,\n",
    "    latent_means=latent_dataset_train.latent_means,\n",
    "    latent_stds=latent_dataset_train.latent_stds,\n",
    "    clip=False,\n",
    ")\n",
    "latent_dataset_test = LatentHumanDataset(\n",
    "    test_dataloader,\n",
    "    ae_model,\n",
    "    latent_means=latent_dataset_train.latent_means,\n",
    "    latent_stds=latent_dataset_train.latent_stds,\n",
    "    clip=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'signal': tensor[128, 512] n=65536 (0.2Mb) x∈[0., 6.000] μ=0.421 σ=0.747,\n",
       " 'latent': tensor[48, 512] n=24576 (96Kb) x∈[-2.907, 2.980] μ=-0.154 σ=0.891,\n",
       " 'mask': tensor[1, 512] 2Kb x∈[0., 1.000] μ=0.648 σ=0.478}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'signal': tensor[128, 512] n=65536 (0.2Mb) x∈[0., 9.000] μ=0.354 σ=0.717,\n",
       " 'latent': tensor[48, 512] n=24576 (96Kb) x∈[-2.887, 3.879] μ=0.049 σ=0.971,\n",
       " 'mask': tensor[1, 512] 2Kb x∈[0., 1.000] μ=0.582 σ=0.494}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(latent_dataset_train[0])\n",
    "display(latent_dataset_train[1])\n",
    "\n",
    "element = latent_dataset_train[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fd7c539e250>]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAdcAAAFgCAYAAADpSzMMAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/H5lhTAAAACXBIWXMAABcSAAAXEgFnn9JSAABQTklEQVR4nO3deXhTVfoH8G/Spk33hbYUKNDSQtkXFQREdivKMuIgDOoIOgriMoo4DI6goI6CLP50FGXAUUFAkDKAoogylEVkESwoIFCgQKFQlrZJl6RNc39/lKa52ZOmuVm+n+fJ85Bz701OL8l9c8495z0yQRAEEBERkdvIpa4AERGRv2FwJSIicjMGVyIiIjdjcCUiInIzBlciIiI3Y3AlIiJyMwZXIiIiN2NwJSIicjMGVyIiIjdjcCUiInIzBlciIiI3C5a6Ap4mCAJKS0sNz2NiYiCTySSsERER+ZuAC66lpaWIi4szPC8uLkZsbKx0FSIiIr/DbmEiIiI3Y3AlIiJyMwZXIiIiN2NwJSIicjMGVyIiIjdjcCUiInIzBlciIiI3Y3AlIiJyM8mCa0FBAWJjY5GTk2N339WrV6NTp04ICwtDhw4d8NlnnzV+BYmIiFwkSXC9cOECsrKyRGkIrcnOzsZDDz2ErKwsbNiwAQMHDsTEiRPxxRdfeKCmREREzpMJgiB46s30ej2WL1+OF198EYIg4MaNG9i+fTsGDhxo9ZjMzEx0794da9asMZSNGzcOhw4dwqlTp5yuQ0lJCdMfEhFRo/JobuEjR47gySefxFNPPYWhQ4di+PDhNvfPz8/HyZMnMWfOHFH5mDFjsHbtWpw6dQpt27a1erxGo4FGoxGVOdJaJgfo9cCN04BOK3VNiMiS4FAgPh2Qc2iNFDwaXFu1aoW8vDykpKQ4dK/1+PHjAIB27dqJyjMyMgAAJ06csBlc586daxaYyU0+vx84s13qWhCRLW0GAY9skLoWAcmjP2ni4+ORkpLi8P51rczo6GhReVRUFABApVK5r3LkuPLrDKxEvuDMdqD8mtS1CEheveScXq+3uV3O7g5p1Jh0BYdEAuCauETeQQCqyuqf1lRJV5UA5tXBNSYmBgCgVqtF5XUt1rrt1syYMQPPP/+8qKy0tBSpqaluq2NA0teIn884D8iDpKkLEYnpa4DX4sXPyeO8OrhmZmYCAPLy8tCjRw9DeV5eHgCgQ4cONo9XKpVQKpWNV8FAJZj0KMjYg0DkNUy/j6bfV/IIr74qZmRkIC0tDevWrROVZ2dno23btmyBSkUw+SUsY5cwkdcw/T6afl/JI7yq5apSqXDs2DGkp6cjMTERAPDKK6/g0UcfRZMmTTBq1Chs3LgRa9euZRIJKRlPjZaxO5jI68iC6oOq51IZkBGvarkeOnQIffr0webNmw1lEydOxEcffYTvv/8e9913H3bs2IHly5dj3LhxEtY0wBl3M7FLmMj7GH8v2S0sCY9maPIGzNDkBkW/A4tvr/13sBKYeUXa+hCR2BtNAd3NBDpP7QOS2ktbnwDEZgc5jy1XIu/GlqvkeGUk5zG4Enk3BlfJ8cpIzjMefcgBTUTex/h7ydHCkmBwJeeJWq6chkPkdYy/l2y5SoLBlZxnnJaSmZmIvI/x99JOGllqHAyu5DzecyXybrznKjleGcl5DK5E3o3BVXK8MpLzOKCJyLtxQJPkGFzJeWy5Enk3tlwlxysjOc94CSuuqUvkfYy/l1xyThK8MpLz2HIl8m5suUqOV0ZyHoMrkXdjcJUcr4zkPFFw5YAmIq8jGtDE4CoFBldyHluuRN6NLVfJ8cpIzhMNaGLLlcjriDI0cUCTFBhcyXlsuRJ5N7ZcJccrIzmPifuJvBsT90uOwZWcxwxNRN6NGZokx+BKzmO3MJF3Y7ew5HhlJOcJQv2/OaCJyPsYfy+Nv6/kMQyu5Dzj0YdsuRJ5HxnTH0qNV0ZyHruFibwbu4UlxysjOY8Dmoi8Gwc0SY7BlZzHqThE3o1TcSTH4ErOM/6yckATkfeRM7ew1DweXLdu3YqePXsiPDwcaWlpWLBgAQQbo9ny8vIgk8nMHp07d/ZgrUmEA5qIvBsHNEku2JNvtnfvXowYMQLjxo3D66+/jt27d2P69OnQ6XSYMWOGxWNyc3MBANu2bUN4eLih3Pjf5GEc0ETk3TigSXIeDa6vvvoqevTogRUrVgAAhg0bhurqarz55pt47rnnEBYWZnZMbm4uUlJSMHjwYE9WlWzhknNE3o1LzknOY80OrVaLnJwcjB49WlQ+ZswYqNVq7N692+Jxubm56N69u0vvqdFoUFJSInqUlpa69FpkhC1XIu/GlqvkPHZlPHPmDKqqqtCuXTtReUZGBgDgxIkTFo/Lzc2FWq1G3759oVQqkZycjBkzZqC6utrue86dOxdxcXGiR2pqaoP/loDHAU1E3o0DmiTnsW7huhZjdHS0qDwqKgoAoFKpzI65du0aLl68CJ1Oh7fffhutW7fGtm3bMG/ePFy4cAErV65s/IqTOdGAJk7FIfI6xt9LDmiShMeCq15v+9eTXG7eiI6IiMDWrVvRtm1bQ4tzwIABCA0NxcyZMzFz5kx06NChMapLtrBbmMi7sVtYch67MsbExAAA1Gq1qLyuxVq33VhYWBjuuusus67c4cOHAwAOHz5s8z1nzJiB4uJi0SM/P9/Fv4AMOKCJyLtxQJPkPNZyTU9PR1BQEPLy8kTldc8ttUBPnTqF//3vfxg3bhxiY2MN5ZWVlQCAxMREm++pVCqhVCobWHMyI3CeK5FXE7Vc2S0sBY9dGZVKJfr374/169eLkkZkZ2cjJiYGvXr1MjumsLAQTz75JL788ktR+Zo1axAdHY1bb7210etNFnDJOSLvxiXnJOfRea4zZ87E0KFDMXbsWDz22GPYs2cP5s+fj7lz5yI8PBwqlQrHjh1Deno6EhMT0a9fPwwZMgTTpk1DZWUlOnbsiM2bN+O9997DokWLRK1Z8iBmaCLybszQJDmPXhkHDx6M7OxsnDhxAvfddx9WrlyJ+fPnY/r06QCAQ4cOoU+fPti8eXNt5eRyrF+/Hk888QTeeecdjBgxAlu3bsW///1vPP/8856sOhnjgCYi78YBTZLzaMsVAEaPHm2WSKLOwIEDzfIMR0dHY+HChVi4cKEnqkeOYHAl8m4MrpLjlZGcxwFNRN6NA5okxysjOY8Zmoi8GzM0SY7BlZzHAU1E3o0DmiTHKyM5z/i+OIMrkfcRdQtzKo4UeGUk5zFDE5F3Y4YmyTG4kvM4oInIu3FAk+R4ZSTniQY08SNE5HXknIojNV4ZyXmc50rk3TjPVXK8MpLzOFqYyLtxtLDkeGUk53FAE5F344AmyTG4kvPYLUzk3dgtLDleGcl5zNBE5N2YoUlyDK7kPLZcibwbW66S45WRnCca0CSTrh5EZJnx95IDmiTB4ErO44AmIu/GAU2SY3Al5zFDE5F3Y4YmyfHKSM7jgCYi78YBTZJjcCXncVUcIu/GVXEkxysjOY8Zmoi8GzM0SY5XRnIeBzQReTcOaJIcgys5T+BUHCKvZvy95IAmSTC4kvM4oInIu3FAk+QYXMl5zNBE5N2YoUlyvDKS8zigici7cUCT5Dx+Zdy6dSt69uyJ8PBwpKWlYcGCBRDsDBVfvXo1OnXqhLCwMHTo0AGfffaZh2pLFomm4rBbuCH2nrmOv687gn1nrktdFfInogFNnIojBY8G171792LEiBFo37491q9fj4ceegjTp0/HvHnzrB6TnZ2Nhx56CFlZWdiwYQMGDhyIiRMn4osvvvBgzUmE3cJuodXV4E//3os1P1/AuH/vRZWO3XfkJuwWllywJ9/s1VdfRY8ePbBixQoAwLBhw1BdXY0333wTzz33HMLCwsyO+cc//oEHHngA77zzDgDg7rvvxo0bNzBr1iz86U9/8mT1qY7x6EMOaHJZkUoren5FpUHL+HCJakN+RTSgid3CUvBYcNVqtcjJycGcOXNE5WPGjMHbb7+N3bt346677hJty8/Px8mTJy0es3btWpw6dQpt27a1+p4ajQYajUZUVlpa2sC/pO7FS4Fze9zzWr6m4kb9v9lyJfI+xt/LiuvAiW+lq4uUWvUBwmIleWuPBdczZ86gqqoK7dq1E5VnZGQAAE6cOGEWXI8fPw4ANo+xFVznzp1rFpjdpjgfWM2WM+e5uo+e98bIXYy/l9fzAvdaNWkHENZdkrf2WLOjrsUYHR0tKo+KigIAqFQqtxxDHhbeROoa+KzqGr3N50QuC0+QugYBz2MtV73e9oVDLjeP864c4zHBSiCpo3TvLzkZ0GYg0Lqf1BXxWdU14paqlgOayF1a3wH0fho4kwMggHtEgpXSvbWn3igmJgYAoFarReV1rc+67Q09xtiMGTPw/PPPi8pKS0uRmprqcL2tSswEnvqp4a9DAcu85RrAF0Fyr6BgYNibUtcioHksuKanpyMoKAh5eXmi8rrnHTp0MDsmMzPTsE+PHj0cOsaYUqmEUindLxciW6pMgiun4hD5D4/1qyqVSvTv3x/r168XJY3Izs5GTEwMevXqZXZMRkYG0tLSsG7dOlF5dnY22rZt654WKJFEqnUMrkT+yqPzXGfOnImhQ4di7NixeOyxx7Bnzx7Mnz8fc+fORXh4OFQqFY4dO4b09HQkJiYCAF555RU8+uijaNKkCUaNGoWNGzdi7dq1TCJBPqmyqgYFxRXISIo06wauquF8RCJ/4dERQYMHD0Z2djZOnDiB++67DytXrsT8+fMxffp0AMChQ4fQp08fbN682XDMxIkT8dFHH+H777/Hfffdhx07dmD58uUYN26cJ6tO1GCVVTUYtCAHd72zE/O/O4FqvWnLlfdcifyFTLCX2NfPlJSUIC4uzvC8uLgYsbGx0lWIAkb2wQJM+/Kw4fmSP9+KySsOGp7/8ZYULHigK2ScO0zk85heh8hDKqp0ouf518pFz7MPFWDvmRsgIt/H4ErkIUEm87IPnS8222d69mGzMiLyPQyuRB5Spq0WPf/u6BWzfThimMg/MLgSeYhao7O7jyKIX0kif8BvMlEj0lTXT69xJLiGMLgS+QV+k4kagSAIeHLFQXSZ/R0+zDmN62ValGnZciUKFB5NIkEUKA6eK8aWo5cBAPO2/I55W3536DhFMKfhEPkD/kwmagQXSypdOo4tVyL/wG8yUSOQu5gIQiHlMopE5Db8JhM1giC57eAabGU7u4WJ/APvuRK50ZbfLmNxTh6OFJTa3K9j82iL+wSz5UrkF/hNJnITXY0eT35+0G5gBYDMplEWy00XUCci38TgSuQm+dfL7e90U2YygyuRP2NwJXKT44Vqh/dtZ6XlyvSHRP6BwZXITY4Xqhzet0lkiMVyLYMrkV9gcCVyk9NXyxzet0lEqMVydgsT+QeOFiZyka5GjzU/X8C240Uo0+qw/6xja7FmJEUiPsJyy7WKwZXILzC4EjlJV6PHnz/ej5/OXHfp+LfHdEVIsOVOI+N7rr8WlGLT4YsY0bU5urWMdem9yLfcKK9CXLgCMheTkJD3YHAlctKmw5dcCqx/HZyBwR2aoruNQFldIwCoDeCTVvyMwlINNh2+hN1/H8zUiH5u3pbf8WHOafTLSMCKv/RigPVx/LYSOcmZgUvGnhvazmZgBepbroWlGhSWagAAV1Ra9H97O0orq20dSj7uw5zTAIDdedfw87liiWtDDcXgSuSkiFDnO3weuDXFLCXi53+5HcnRSjSNrh/cVHfPtaKqRrRvYakGbzu4sg75HtOBbFfVWolqQu7C4ErkpIgQ54PrrJEdzcr6tU3ATy8NxpeT+xrKqnR67Dx5FSUVVWb7r9x33un3Je9VWFqJLb8VoqJKZ/Zjih3Cvo/3XF2UV1SG5774RepqSCoyNBhP3NkGQzs2lboqHlUjCE7t//dh7RGtVFjcJpPJzAY3PfKf/S7XjXxDkVqDkf/6EdfKtOiXkYD5D3QVbW/o7daKKh22/34VZdrAvpVwV8dkqyPzGxuDq4s01TU4esm1e2/+JK+oDHvaDcbWo1egCJIhq2My5HZWhPF1pq0Me3R2ptcogvz7fJG5N74+jmtltV2/u/Ou4fWvj4m2V9U49wPO1MT/HMD+fMemhvmzTs1jJAuu7BamBrleXoWVe8/j2dW/4MnPD+Hb3y5LXaVGp6m2H1w7NosGAMhlwAO3tbS5r7VpOeS/frkgHrD0za/i701llc7p16zRC/gx7xrOXC1jYPUCHm25vvvuu/jXv/6FixcvokOHDnjjjTdw77332jxm2bJleOKJJ8zKn376abz//vuNVVW7WsaH46OHb5Hs/aVUVSPgr6vru8RfM/rV/fSqQxjedbgU1fKYCgcufO+N74Evf76AXmnxSI5R2tzXmeAqCAKnaPiBSju9H872jgDAm98cx8e7z5qV901v4vRr+QtXBh+6i8feedGiRZg+fTpeffVV3Hbbbfj4448xatQo5OTkoF+/flaPy83NRWZmJj799FNReXJyciPX2LaYMAWGdW4maR2koteLg2ugqayyn0UpPTECL93bwaHXUzixhmtpZTViw6Xp5iL3aYzgaimwRoYGY9UTvZ1+LWo4jwTXyspKvP7665g2bRpmzZoFABg2bBj69u2LOXPm4Pvvv7d6bG5uLnr27InevfkB8RZyuQwhQfKATdVXWW2/5epM69KZe9RX1VoGVx+2Mfci/vW/PJTbCZ72gq+jYsIsD6SjxueR4Lpv3z6UlJRg9OjRhjKZTIb7778fL730EiorKxEWFmZ2nCAIOHLkiOg4Z2g0Gmg0GlFZaan9hazJvtBg68FVrxf8clDTjpNXsTH3InacuOr2176tdZxDiQNcadGQ9PaeuY6NuZewer9j06nc9f/M4Codj4ykOH78OACgXbt2ovKMjAzU1NTg9OnTFo87ffo01Go1Dhw4gMzMTCgUCmRmZmL58uUOve/cuXMRFxcneqSmpjbob6FaoQrrH526UZD+pEyrw4T/7Mf6Qxfttjpc8eHDt+L2tHi7+1U6MJiKvIumugZTPj/ocGAFHOsdcQSDq3Qa3HItLy/HihUrrG5v3ry5obUYHR0t2hYVVbtgtEpleUpLbm4uAODs2bNYuHAhFAoFli9fjgkTJkCr1Voc6ESeERocZHXb5l8L8egdaR6sTeM7ZmXa1XND2uLfO880OOglRoXi6UEZ2HfW9hxXBlffk1dUhuIK5+abXlVrnRq8VqO3PHUnNpzBVSoNDq7FxcWYMmWK1e0DBgxAVlaWzdeQWxnQ0b9/f3z11VcYNGgQIiIiAAB33303ioqK8Morr+Dxxx/nyEmJ2Gq5frYn3++Ca7mVEcJdU2Lw00uDMeJfu1FQXNmg9wgLsf6DpY6G3cI+59z1CqeP+eF4Ee5bvAfZT/ZBsJUFG1b8lI///JiPh25vhT/1amVxH7ZcpdPgbuGUlBQIgmD1kZOTg5iYGACAWq0WHVvXYq3bbiopKQkjRowwBNY6w4cPx+XLl3HlyhWbdZsxYwaKi4tFj/z8fBf/UjJmq+V6qURjdZuvKtdaDq5hiiDEhofgb3dnGspuax3n0nuEKewHV7Zcfcui70/i6VWHXDr28IUSq/PGK6p0mLXxKM5eK8cbm4/jcqnlH3YMrtLxyICmzMzaC09eXh569uxpKM/Ly0NISAjatGlj8bhdu3bhzJkzmDBhgqi8srISQUFBiI+3fY9KqVRCqbQ9x5Bco7TRcq2q0aNKp/er5AglVrr16lqbI7s2x/FCNc5dL8f0Ye1deg+lA8H153PFuFRSiVHdWqBVk3CX3oc84/CFEry37VSDXuPZ1b/g9rR4JEWLr2OXS8U/YC9Y6TWJZnCVjEeufn379kVERATWrVtnKBMEAevXr8eAAQMQGhpq8bjt27dj4sSJOHnypKFMr9dj3bp16Nu3L0JCOCVBKqF2AqcjiRZ8yY1y80T6ABB+M4m/XC7DjHva48OHb0VaQoTFfe0x7RZumxSJO9smiMpW7TuPBVtPYsrKgy69B3nOV4cv2d3HkdR8c0xSIwK1yxAau1Fm+fPZRKLUf+Sh4BoeHo4XX3wR8+fPx6xZs/Dtt99i7NixOHjwIF599VXDfgUFBdi7dy+02toPzuTJk5GcnIyRI0dizZo1+PrrrzFixAj89ttvmD9/vieqTlbYa2Ut2HoC566Xe6g2jc9acHWkK9dRpq8VE6bAir/cjpHdmpvte/SSCipNYCdl93bbfi+yu09smAJL/nwrutlY53fzkUKzsiK1uOV6RW35Vszg9kl260CNw2P9dq+88gpee+01fPbZZ7j//vtx5swZbNq0CXfccYdhn2XLlqFPnz4oLKz9MDVt2hS7du1C165d8de//hVjx45FeXk5tm3bhttvv91TVScL7LVcP997Ho/8Zz/0VkYx+prrVoJrfKT7WgamwbWuJRtuJYBfbOAAKmo8muoa5Bv9uMye0sfifkpFEO7ulIyNT9+Bff8Y4vDrm3YLF6nMp79N7Jtq1p1MnuOx9IdyuRwzZ87EzJkzre4ze/ZszJ49W1SWkZGBL7/8spFrR86yNaCpzrnrFSiv0iHKynJrvuRGufnFKzRYjggHRvg6yvQHS12wtTaKuKC4Eh2aRVvcRtI6c7UcdSsThgTJ0S0l1uJ+xmMXmjoQCPOKylCjF3BZZRJcLbRc7eW0psbFJefIJaaBYMrAdHyx/7zZfL4qnX+kSLxu4Z5WQmSoW6eCmWa1qguq1rrgC4qdn+JBnpF3tczw77SECKvTaUx/OHVoFo3jheZzqgVBwI951/Hwx/tqjzP5TFhqucZxjquk/Gc4J3mU6aT1x/ulWfzlXd3AdSm9haV7rk3c2CVsSbghuFr+mjZ0Xi01ns1H6gczZSRFAqjtpjWlNOkBevdP3S2+Xs9//mAIrID5lCxL91zjmINaUgyu5BKtSV7hKKXC4vJO/tByFQQBxRXmwbWxF2Gua7FaGzR1qYTB1Rv9flmF747Wz8G/5ea85+eGtDXb17RXol3TKIyyMIDtmpXRwHVMRw8DQBxHCkuKwZVcYho0Q4LlCLLQRVpV4/tJD1QancUWeJMIy1PIGmJQZqLh34/0SQVg/Z5rY+Q4poY7dK7E8O+EyFCM79USQG2we+O+zqJ9LXX5h7twH9/Sj1i2XKXFe67kEktf5goLycardL7fLWxtGk5CI3QLv/aHzvh0Tz5uax1nmC9r7Z4rUyF6p1NF9ZnohnZIMsyFBswDp6Uuf3clX+E9V2kxuJJLLAZXrfnF3h/WfLU0UhhonG7hlvHhmDWio6jMWrewpR8zJL28ovrBTHX3W+uY/l8mRJr3foRYGfzkLKY+lBa7hcklloKmpTUoq/0guFoaKQw4NnXCHazNKXbXgtrkXqeu1AfXtk2jRNtMF7ofc2uK2fG2FsVw1KN3pFodoUyewZYrueSJO9Nw8Obi3v0yalP0WVo5xh8GNFnrFvbUPMLgIMvTfRhcvY+uRi+ag5pqkv/5ttQ4dG8Zi9wLJZg9siNaxpvnhw4Jcn3u9McTbkOPVnGNPtiO7GNwJZdkdUzGc0Pa4sKNCky9qx2A2l/hn/yYL9rPH7qFrWVnauah4NorrQmilMFQa3TISIo0dDtyhRzvU1opnudtGuQUQXJkT+mLGr1g9d6qM/dcnxqYjns6N8PSXWcwqH0ihnRo6nylqVEwuJJL5HKZIajW+evgtjh6UYX9+TcMZf7QcrWWw9dT3cKRocHY9sIA/H5ZjfiIEIz4124AlrvhSVrGSVSC5TJEWpieFiSXIUhuPfmIM8F1ysB0RCkVeG98D+cqSo2OnfLkNnERIVj7ZB90S6lfn9cfgmuZxvLAIUeWiHOXpGgl+rdLFA1S0er0fpO72V+UGM2Hjg0PcSmDl7Xg2q6peHDUuNta+kVqUX/F4EpupzAaSOEPA5qsLZQuBdM5r+wa9i7GLddYF6fCmA5ga58chdNv3ovvnu8vKjdNl0nehcGV3M74l7dULdfTV8vw8e6zbsliVGYhuPZuE9/g13WF6TxJBlfvYtxydXWeqWlwXTS2O4LkMrNWMAcDezfecyW3k7Lluv33Iuw6dQ3/+fEsAGBT7kVsfKZfg15TbdQtnBQVit5tmuBvd2c26DVdZZqLliOGvUuJqOXq2ohd03mu1qbmBMsZXb0Zgyu5nXHLVXuz5frWt8fx7a+XMfWuthjdw3xunzsUllZi0oqfRakKDxeUQlNd06D7o8ZTjGaP6oR7uzRrUD0bQi6XQamQQ1Nde17ZcvUuxW5ouZrOTzVuyd7aOs4wBW5cz5YuvT55Bn/6kNsZ//KuqtHj4LliLNlxBudvVGDqmsON9r6r9p23mAPY2mhfRxkPaLI0+tPTjLP8cMSwdympbHjL1ZTx2slvj+mKcbe1xNt/7Mq1fL2c9FcK8jvGLddqnYDtvxd55H3VVkb1qiqrkRTl+rSZMqO0jpZW/vG08JBgw8AZdgsDP52+jiK1BsM6J4sCkRSul9WnynR1QJMp427h9MRIzBvT1S2vS42LLVdyO3HLtQanjRaOBmqz2DQGa8G1tLJho33LtPWtkSil9MHVONl7ZYDnF/7lfDHGL92L577Ixce7z0pdHfx2sX6h87QmEW55TXflGibP4v8auZ0iuH5UY3WNgDNXy0XbNY00glhtpftXVelat3CNXsDne88Z7m8C3tNyrVNZ5ftTnRpi1sbfDP9+e8sJCWsCXFVrcdFodHr3VrEuvY4giG9tWMstTd5N+isF+R3j3KhqjQ5nr5kE1+qaRrl3ab3l6lpwXbD1BD7MOS0q84p7rkbTcYxb1YHIWoIPKbyx+Zjh302jQ9EsJsyl12lt0uJ1JREFSY8/icjtjFuuq/efN8svrGmkEa5XyywvDefKgKYrKo1ZYAW8I7gmG6Vd/GD7aY/d0/YG569X4KFlezHl84NQa6ptphF0lxoHsmAVFFdgY+4lw/NBmUkuv19mchQm9k1FbLgCc0Z1cvl1SFrSXynI74TauUdk3M3qTkVGq5EYK62oRrlWh+0ninBLqzg0j7Xfoth96ppZWZgiyCMXc3vaJNa3bM7fqMCjnx7AJxN7YlB71y/ovmLal7k4kF87FSU1IUI0p7rOj3nXsO/sDYzr2RItHPi/tkZXo8cj/9mPXy+W4q37u2BE1+ZW9/0xr/7zEiyX4e/D2rv8vkDtlK9XR3Zkq9WHseVKbmfpgmesMVqu1TV6qGx0C09bexjPrPoFI/6126H3P3FFbVbmqVVw7ElLMB8o85fPDkhQE8+rC6wA8GHOafx+Wfz/dLGkEo9+egDvbTuFkf/ajXd/OIV8k9sSjtqYewl7Tl+HWqPDM6t+sbnv7rzrhn+P79UKcW5Y8o2B1bcxuJLb2VvVQ6tzf3A1zoxjSqWpxpajlwHUrs265bfLdl/veKHKrMxb5hW2SYg0K2P+/lrvfH/SkHLzRnkV3vnhJJ78/KBLr3X0kvlnwJqz1+pHxN+WGufS+5F/kSS4fvXVVw7/KtPpdJg1axZatmyJ8PBw3Hnnndi3b18j15Aawl5wbYxuYeOcrqaul4m3ORLcTVtEAPBYvzTnK9YIUhPMF9gGLOdADjTrDhaYlf1+We1Sb0l0mPiuma3za5zMI5or1RAkCK45OTl48MEHHd7/hRdewKJFizB9+nSsWbMGwcHBGDp0KPLy8hqxltQQpt3CTw9KR0ZSfWvripV7ow1RbKPleqG4QvRcbueHXUWVDlfV9YOjUpuEY/qwTNza2jtaJFFKhShLU51z113r/vQVDfnx4MqgNtNFEi6XWv/caoyCq+nKRRSYPBZc1Wo1Xn75ZQwdOhQhIY7dj7hw4QI+/PBDLFiwAM8++yxGjhyJLVu2ID4+HvPmzWvkGpOrEqNCRc/v7dJMlPjghbWHcc3KyF5X3Si33nI9f0McXAU7XajGgVUmA354YQCeGpjRoPq5W5NI8++QpdSP/kIQBPz5Y9d7rK6XVeGD7Xn49MezDi8mYdrDYutHoXGOZ0s/fCjweGy08Mcff4ylS5figw8+QGFhIebMmWP3mG3btkGn02H06NGGstDQUIwYMQIbNmywe7xGo4FGI/5ClJaWOl13cs6gzCTc2yUZ+dcq8MzgDHRqHmO2msv7/8vDbDdOM7DVLWx6kbTXAioyCq5NIkLMEql7g7jwEBQUi5fTa6zMV97gRnkVfjlf4vLxb2w+hh9vDjqKiwjBH7q3sHuM8YINAFBoo+Vq3C1s2uKlwOSxq8bIkSORn5+PyZMnO3zM8ePHERUVheTkZFF5RkYGLl26hLKyMitH1po7dy7i4uJEj9TUVFeqT04ICZZj8UO34pvn7jSsIGO6Ks1FN6yzWmdj7kXMWP+rw/vbW/y8SFUfXBMiQ23sKR1FkHnXtul8Yn9Srm3YILgfjUbzfrH/gkPHVJi8Z+HNz2xlVQ2KVBpDJiW9XjCs/gSYf9YpMDW45VpeXo4VK1ZY3d68eXOMGjUK6enpTr92aWkpoqPNR2hGRUUBAFQqFSIjzUdOkvdRmqxJaSk4uKK6Ro8Xv3RupR17Lder6voWSlK0d0y/MWVpLU+dH3cLN3RlI2P2Uliu2HsO2QcLkHuhRFT+3bHL6JvRBFM+P4QitRYdmkXj87/0MgumvOdKgBuCa3FxMaZMmWJ1+4ABAzBq1CiXXluvt/1LXM7Fgn2GaZabIDf935VpdE7fa3SmWzgpyjtbrsEWfpzo7HxffJk7g6utXv4rKg1mbfjN4rbfLqrwxw9/Mjw/XqjCv3edwRN3thHtx25hAtzQLZySkgJBEKw+cnJyXH7tmJgYqNXmUyJUKpVhuy0zZsxAcXGx6JGfn+9yfch15SZLo7mp4QqNC3NmV+47j9mbjpolSK/jC8HVUqKOhnadejNVA1c2MmZtTnRekRq3v7nNqdf68ucCs5zWpuMLKDB5ddMvMzMTKpUKV69eFZXn5eWhdevWCAuzndpMqVQiNjZW9LAXkKlxVJgMDjENtq5ydT3TT/fkW00mYbwmp7fec+3cwvx2ybOrf8HRS/45YM+dLVdrCzk8vybX6de6UV4lSjiiVMgh94IUmSQ9rw6ud911FwBg3bp1hjKtVouvv/4aWVlZUlWLXGA6OMTVZeBMNSQhxabDlyyWGwf+SC9Yv9WSKQMzLObNnbTctWxE3s7W56VrinM/mK29lvFarKYGZiZa3XbBaKoXp+FQHa8KrgUFBdi7dy+02tqWQ+vWrTFhwgRMnToVixYtwtdff4177rkHJSUlmD59usS1JWeYTmuwlgfYWZUmmXde/4Pj03tM15mtY9zKjgjxzuAaGRqMHX8biPREcZ5hd47ClsK+M9cxI/sI9p65Liq3tpzg8K7NsOmZfujTponD71Hiwg+7x+5Iw/CbI99NvfXt74Z/M7hSHa8KrsuWLUOfPn1QWFhoKFuyZAmmTJmCefPmYdy4cdDpdPj++++RkeFdk/rJNtO1Ld3Xcq0Pri3jw9CxueP5f09cUeObXwvNyn1lzmJwkNzlNUO9UZVOj4mfHMAXBy5g4if7RWkqrXUL193ftDTAy5qKqhrctWgHvvm1EHoHkzKHhwQhJc7+uVZ68eeFPEuS4Dp79myLg0nqyo3nooaGhuKdd97BlStXUF5ejp07d+L222/3YG3JHd64r7PoubvuoRnfcw1TBCHcyZbmKxvNR4Yad2F7c3AFnAsq3u6KSmPoidBU63GppH5KlLUBTXVTvIKdvM95qqgMT608hKdXHXJof6XCseCqbaTlFMn3eFXLlfxX5xYxWPGXXobnZVqdw60GW4xHC4cpgpxezPxamXlmJ+NuYWeDtadZmu/qKwRBwMkrasMqNsaJGGqf1//fqq21XG92w9qa2nV7Wjz+1LMlLMXfb3+77NCgOJkMaOFAcG2MvNnkm3z3m0k+p0uL+oEngmB91KYzjC+MoYogl1qapr0oxvdxw0O9u+XqrmQcUpi+7giy3tmJ0Yt/RJVObxZA6+6zCoKAc9crLL2EoeU6pIP1heLXTO6DuX/singra6wa55K2JD0xAh2So9Ekwv7IcR3X/qObGFzJY6KVClEwuF7e8OT9xgn7wxRBdrPvWGK8ok6VTi9KSuHt3cL2Fqb3VoIg4Muby8MdvaTClwcv4KFl4sT8ak01qmv0eHDpPouL1wP191zH3tYS93ZJRlpCBJ4eVJ8N7oMHbzH821qvxtUyy63NjKRIfPJoT2x4+g7I5TJ0bB6N1k0sL/dHZMq7+7zIr8jlMjSJCMXlm11nV9VVyLDe4LDrUkml2UjNUJO1ZEOC5HZz7l4qqTS0aky7CMMV3v0V8dV7rqb/Jy//1/zet6pShy/2n8dPJiOHjdV3C8uw+KFbDeXDuzSHTq9H15RYQ1mUlXVWjXNJG4tWBmNQZv0HVBEkx6onemPXyauIUios3q+9pVWsWRkFJu++cpDfSYgKMQTXhi479/aW30XPlQo5ZCZrtT7cuzWul2tx7JIKp4osL/RQWKpB55td1hXV4oEz3p4nVmExx7DeK1fyMWY679kSlaYay386Z3OfkGDLf6elUePWW66WP4eWXrtFbBj+1KuVxcXXm8co8dofOpuVU2BicCWPMs541NDgerxQ3FVoKRAKEPDun3oAAFJnbLb4OpdL6+eGGqcQVATJrF68vYWllmt5VQ1iwry73qbzni1Ra3Q21+kF4FQ2pCgrCUGstVxtdbkrFUGIVgYb5mv3aBWL7Cf7MjsTGXj3N5D8TqJRcLU3kMSYpalbptmTLC311am5/ew9N8rr77lWiua4ev9vT0sBwN7CBN6gwoERuqrKaosjfI0FyRwPZtaybVn7HJreYjC1aGx3w7+n392egZVEvP/qQX4lIcr5luue09cwbe1hpCVE4D8TexqCaKFJNqK6C+37D/bAP9b/im4tYzG6h/1FsYuNFlovF03D8e4uYcDyaGF769V6A0fqqNJUm03PMeVMPIu2cs/VWrewvcFiQzs2xbZpAxAkkyE1IcLmvhR4GFzJo8Tdwra7/Oo88vF+6PQCCks1mL7uCN4b3wPvfH8Sl0rFozzrpm6M6Noc93RuhiAHr7zGXY+VPpKdqY6le6uXSirRrmmUBLVxnGMtV50DwdWJlquVe67WWvqOjMROT+R60mQZu4XJo4yTzZ+4bHl6hSnjuYObDl/Csl1n8O62U2b7Gc+bdTSwAuKWa4WvdQtb+DtdWd3F0xxpuW7+tdCQYMIaJ2KrYU6so3Xx1WlO5B346SGPMp6qcLGk0qVE829sPm6xXK11LSmFccvV97qFzb/CJRXVZovTextHWq6OcKblajqS3F5deAuVGoLBlTwqKVopmoj/c/4Nt732s4PbunTcJaMAb7yggLXRpd7E2pSbuhSOmuoanListrowvFQcGS1sTbeWsQBq5zAP7dDU4eOsBWJfuEdNvofBlTyuk9EcxIJi9yyR1je9CW5Pi7e5T5SVe27FFdVYe+ACAKDIaORoYpR3LpRuzFr6w4qqGmh1NbjrnR24+/92Ys5XxzxcM9scmedqzbvjuuOpgelY/pdeiAm3PEjJEmuN3HKtDjo7iUaInMXgSh5nvEaqpcn4rnh6UIbVbr86H/25PoPPzOEdRNsW5+QBAIqMEq8nRindUrfGZO3ecrlWh5wTV3HhRu2Pl0/35Nu9f+lJ1lqu3W+2Sm1JTYjA9GHt0duJNVwBoFW85dSFFdU1FgdOeVdbn3wNgyt5nHGyB0dWJHGEI/dH+6Y3wVfP9MOXT/bBn/u0Fm3Lv14BQRBELdckH2i5WuvtLdfWmCVgOF6o8kCNHGPtPueIrpYXJHeHYZ2SDV3KE/umGsqtLSLhZT3p5GMYXMnjRMHVTS1XRxL2y2QydEmJQc/UeIQGB+HPvcUB9tD5Euw5XZ/H1ieCq5Xy8irz7EaHzhc3foUcZHyfs01iBLq3jMXk/m2cuofqLLlchg1P9cXBmUMx4572om32MkEROcv7R2yQ3wlTON5ydfReWJiF7Ez2vPaHTljz8wVDd+kzJonYk6K9v1vY2kClIrUWu05dFZU5OvXJE4xbrqO6NcfzQ9sBEHfLNwaZTIYmkaEQBAFyGVA3qPo/u8826vtS4GFwJY8TBVc7LVd7SQTquLLUnEwmQ4vYMJy9Vg6gNoG/MV9oueqtBNe/rv7FrMzRpB2eYNwNa5zcwd5CCf8c7Z7E+DKZDBEhwVDfbEGv/+WiW16XqA67hcnjnOkWdjS4ujontXms9dapL4wWdmY66w03rJ/rLsapL43Ps63EHeue7IMHe7VyWx3s/SATOKSJGoAtV/I4pQPdwhdLKvHU5wfNUhxaYy/JujVpCRH4Mc98vdBJ/dv4RIYeZwbdeNN9xatWpjwFyWtXIrI0svm2VNtTrZwVHur9SULId3n/1YP8jnEr03gqjqa6BgfPFaO6Ro//+/4kDheUOrRyzqwRHe1Ow7HGeNSosUn927j0ep7mTOvqupcEV71eELVck0ymPHkqM1aED6S3JN/FTxd5nKV7roIg4JH/7Mf+s85lbJLLgL/0S3O5LhlJUXikT2uzRbl9ITsT4FzLVa3RoUqnl3yN2tLKalTX1FfctPs9XBGEEriWytIZEXZarmlNuNINuY4tV/I44+BaN2r06CWV04EVcO6eozXGGaOA2gTvocG+0WXYy05WKlPGixRIxXiJt5BgOaJNfsh46k5nWoJ4RZtH+rTGE3emQS4DWsaH4fE7faP3gryTb/w8J7+itNAt/IuEczA7NBMH18hQx1PqSa1najyeHJCO3XlXESSX4/CFEpv7P7v6F9zVoSkevzPN5a70hhLdb40MNavHlUaejlPnqYHp+ObXQpRWViMxKhQz7mmP8JBgTOqfjrhwhdW8zUSOYHAlj7M0z/UXO0EBqE3UXtUIOWDbmKzJWdmApPJSqE2I0B5Ldpy2G1z3n72B/WdvoH2zKNzZNtEj9TP168VSw78tjcj21II+LePD8fWz/bAx9yIGtU8yjFT2hVHi5P0k+Wn21VdfOfyredmyZZDJZGaPZ555ppFrSY0l3GQqjiAIOH213O5xaQkRaBbj/sQOpotol7spJaOnmf5IsOX/fjBfD9cTavSCKGFDn3Tn8gO7W8v4cDwzuC06NY+RtB7kfzzecs3JycGDDz7o8P65ubnIzMzEp59+KipPTk52c83IU4xbrnoBqKrRi5Z6s0apkOOtP96C+xfvaczq+awB7RIRHxHi0JSbmDBpur4LSytF+ZsfvSPVbB/jzEl1erl5Gg5RY/NYcFWr1Zg7dy7mzZuHmBjHfyXm5uaiZ8+e6N27dyPWjjxJaTLV4ugllUNTbkKDg3BLq7jGqpbPCwmWo1tKDLafuGp331iJgutlo3nLCZEhZtNwAOCdcd3x3Be5hucZSZF4e0xXT1SPyG08Flw//vhjLF26FB988AEKCwsxZ84cu8cIgoAjR45g9OjRLr2nRqOBRiMeHFFaWmplb/IU0zzAjrZEQxWNdxdj8UO34KmVtbmF54zq1Gjv09hsZTgyduZaOQRBcPugptKKajyx4meoKqvx3vgeaNc0SrTdOClIspUu/j90b4EOzaIRrVSgaXTt/U+pBl8Rucpj91xHjhyJ/Px8TJ482eFjTp8+DbVajQMHDiAzMxMKhQKZmZlYvny5Q8fPnTsXcXFxokdqaqqLfwG5iyJIjmAr65Da0pjTY+7pnIz3xvfAvD92wXg3ptjzNHu5eevkXihplPuub2w+hv1nb+D3y2pMW3vYbPvl0krDv5Ojw6y+TrumUUiOURrGWBD5mga3XMvLy7FixQqr25s3b45Ro0YhPT3d6dfOzc0FAJw9exYLFy6EQqHA8uXLMWHCBGi1WjzxxBOuVpsklhQV6nBqwzp1A6H+b1x3PL8mFwDw1v1d3FIfmUyGUd2au+W1pGRtdaDwkCCzNVTf3XYKU+9q59b3//JggeHfxqOC6xgvjtAYg9OIvEWDg2txcTGmTJlidfuAAQMwatQol167f//++OqrrzBo0CBERNRmS7n77rtRVFSEV155BY8//jh/1fqolPhwp4NrbHjtfcIRXZtBpxegq9Hj/ltSGqN6PstS6sAerWJxZ0YC3vtfngQ1Eisssd8tTOQPGtwtnJKSAkEQrD5ycnJcfu2kpCSMGDHCEFjrDB8+HJcvX8aVK1dsHj9jxgwUFxeLHvn5+S7Xh9ynVXy408fUDcIJDpJjzK0p+FOvVpKn8vM2lrqF54/phikDMyzuX+OpSaU3FamNgqsPrJdL5CqvvjLt2rULn332mVl5ZWUlgoKCEB9ve3i+UqlEbGys6OHMSGVqPC3jnA+uMeEhjVAT/2LaLTx7ZEdkJEUiLCQIcy10oZd4OB2iSlOfoEOq6UBEnuDVwXX79u2YOHEiTp48aSjT6/VYt24d+vbti5AQXmx9Vct464NZrOHF2D7TbuG4iPrviKVWraeXoSszCq6+sjgCkSu8KrgWFBRg79690Gpr5zxOnjwZycnJGDlyJNasWYOvv/4aI0aMwG+//Yb58+dLXFtqiNZNXO8WJuuUJi1X43VuTbcBnl+GTq2pTxYSyeBKfsyrguuyZcvQp08fFBYWAgCaNm2KXbt2oWvXrvjrX/+KsWPHory8HNu2bcPtt98ucW2pIdok2E/VF2HS0ooJZ3C1x3Seq/E9aUsjiT3Vci2tqMaqfedFqSWjlfz/JP8lSXCdPXs2BAsLUdaVG89FzcjIwJdffokrV66goqICO3bsQL9+/TxYW2oMxt2V1sSa3GNly9U+027hkKD655a6ha+X2c+M5Y66TM8+jH/891fRdnYLkz/zqpYrkTHTYMB7rvaZdv3ab7m6b1HyuuUD69QF13KtDt8dNR/ZHxHK4Er+i8GVJGNvOo5pEqdoBle7TFuu4nuu5l/3cjcur6fSiAN1XaA/ZGGt3jBFEBRcL5X8GD/dJJkPH77F5vZbW4uT9FsakENipq39EDsDmircGFzVGvFr1QXPA/nmwZWDmcjfMbiSZDo1j8FHNgJsx2bReOv+LuiXkYBPH+3pwZr5LtOuX0WQzOo2AGYpERuiXCsOrtU3F7Y/f918rV7ebyV/x084Scq0NTX2thT8fK4YqU0iMLJbc8SGh/h0In1PMz2fN+ObxW0AUKF1X3DVVOtFz+uyP11RmQ+acmXhBiJfwuBKkjIdNP7m6C4I5r04l5l2Cxvfgw0PCUL75Cj8flltKKuodl9wrTR5reqauuBqnkP65JUyt70vkTfiVYwk1Se9CeJvTsvplRbPwNpAkaHBGHNr7WIGd3dqipZGg8ZkMhmWTbgNHZtFG8oq3XjPtdKki/lamRbT1x3GmWvm3cKpLiQRIfIlbLmSpJSKIKyd3Bu7T13DvV2aSV0dv7DggW74292ZSIoKNduWEheOZwdnYMrNheHdec9VqzN/rbU/F1jYE3hmcFu3vS+RN2JwJcllJEUhIylK6mr4laY2Vpwx7jp2Z3A1bbmaCpLLsHVqf1Roa9AlhQtokH9jcCUKMMYpEt05Fcc0iYSpxMhQpCfaT3tJ5A94g4sowIQ3VsvVZLSwqR6tYt32XkTejsGVKMAYB1d7XbmOOnutHPO2/G5zn6cHWV6wncgfMbgSBRjjbmGdXkCVznaL0xF/zz5ic/tb93dB5xa8z0qBg8GVKMCYzoV1x33X/WdvWN02eUAbJgKhgMPgShRgTJP7u/O+qyWW0i4S+TsGV6IAowiSI8QoWQeDK5H7MbgSBSDxXFf3TcexhKsZUSBicCUKQJFGC5WXaRsWXPV6weZ2tlwpEDG4EgUg4yXfyjQNC66lldU2t4daWKSdyN/xU08UgIxbrqaLnDvrenmVze1suVIgYnAlCkCRSvd1C9+wF1xDGFwp8DC4EgWgKKXC8O+GBtfC0kqb2zmgiQIRgytRADLuFlZpbN8ztaVMq7ObnYndwhSIGFyJApC7BjR9sD0PGqOE/V1TYqA0GcBk+pwoEHjsU69SqfDiiy8iPT0dERER6NKlCxYvXgy93nZeU51Oh1mzZqFly5YIDw/HnXfeiX379nmo1kT+yV1TcT7MOS16/vDtrZEQKV6knd3CFIg8ElwFQcC4cePw6aef4oUXXsCmTZswcuRIPPvss/jnP/9p89gXXngBixYtwvTp07FmzRoEBwdj6NChyMvL80TVifySccs158RVt71uqybhqDGZ98puYQpEHlks/ZdffsGWLVuwdu1aPPDAAwCAIUOGoLi4GG+//TZmzpwJmUxmdtyFCxfw4Ycf4r333sOUKVMAAFlZWWjXrh3mzZuHpUuXeqL6RH7HuOVaWlmNNQfOY1xP55Prd24Rjd8uqgzPOzSLRpFaK9rHeBUeokDhsW7hSZMmYciQIaKy9u3bo6ysDEVFRRaP2bZtG3Q6HUaPHm0oCw0NxYgRI/DNN9/YfU+NRoOSkhLRo7S0tGF/CJEfMG65AsDfs3+Frsb5ped0NfWt1Il9UxETpjBvuXIqDgUgjwTXW265BUuWLEF8fLyofMOGDUhMTERiYqLF444fP46oqCgkJyeLyjMyMnDp0iWUlZXZfN+5c+ciLi5O9EhNTW3Q30LkD0KDzQPerlPXnH6dyur6pP+D2ycBABIiQwxlwzolmx1DFAga3F9TXl6OFStWWN3evHlzjBo1yqz83XffRU5ODhYuXAi53HKMLy0tRXR0tFl5VFQUgNpBUpGRkS7WnChwZSSZf29yThRh0M0A6SiNUXCta6FOvasdZm74DU2jlHhjdOeGVZTIRzU4uBYXFxvuh1oyYMAAs+D6/vvvY+rUqRg7diymTp1q9Vh7I4mtBWUisq1lfDjmjOqEVzcdNZTtOOn8wCbjaTjKm63hh25vjaEdmiI2XGGxhUwUCBocnVJSUiAIgtVHTk6OYV+9Xo9p06bh2Wefxfjx47Fy5UqLA5nqxMTEQK1Wm5WrVCrDdltmzJiB4uJi0SM/P9+lv5PI30zom4pd0wcZnudfr0Clk2u7GncLG89nbRqtZGClgOaxYXxVVVUYP3481q9fj2nTpmH+/Pk2AysAZGZmQqVS4erVq6L7snl5eWjdujXCwsJsHq9UKqFUKt1SfyJ/1DRa/P24Xq5FSki4Q8fq9QKqdEYtV065ITLwWL/qxIkT8d///hfvvPMOFixYYDewAsBdd90FAFi3bp2hTKvV4uuvv0ZWVlaj1ZUoUIQEyxETVp9n+HqZ7ST8xrQ68W0bBleieh5puW7cuBGrV6/GqFGj0Lt3b+zdu1e0vUePHggNDUVBQQEKCgoMz1u3bo0JEyZg6tSpqKysRLt27bBo0SKUlJRg+vTpnqg6kd9rEhliWJP1ernWzt71jLuEAaY5JDLmkeCanZ0NANi0aRM2bdpktv3s2bNITU3FsmXLMGfOHMNzAFiyZAni4uIwb948lJWV4dZbb8X333+PjIwMT1SdyO8lRITizNVyAMA1J1quGrPgypYrUR2ZIAiC/d38R0lJCeLi4gzPi4uLERsbK12FiCQ25fOD+Pa3ywCAvw9rjykD0x067szVMgxeuAMAECyXIe/NexutjkS+hv04RAGuiVHSh+tlWhSpNYZuYlvEI4XZaiUyxqSfRAGuSUT9KjYbD1/CZz/l1/776X7o2Nw8iUsd0RxXBlciEbZciQKcccv1qlqL6hoB1TUCnlj+s83jVu07b/g3BzMRifEbQRTgjKfiGLtYUmn1mHKtDtmHCgzPq11I+k/kzxhciQJctNJycAWsB81rZeIpO1dUjk/hIQoEDK5EAS46zPrQi7wiyytPlWl1jVUdIr/A4EoU4Gy1XC/cqLBYXq4Vz3F94NYUt9aJyNcxuBIFOGv3XAGgsFRjsbzcpOX6/F3t3FonIl/H4EoU4KJtBNdLpZYHNRl3C3dNiUGLWNuLaBAFGgZXogAXGixHSJDlS8GlEvst1/AQznElMsXgShTgZDKZ1UFNhVam45QbrfsaGcpcNESmGFyJyOqgJkfuuUYwuBKZYXAlIkRZue+qspJjmMGVyDYGVyJCtNJygDRds7WO8YAmdgsTmWNwJSIkRoVaLNfpBYtZmkQt1xAGVyJTDK5EhORopdVtpouiA0CZURKJiFCOFiYyxeBKREiOsR5cLXUNl7NbmMgmBlciQlKU9eCqrbbQLVzFAU1EtjC4EpHTLddSo1HEtjI8EQUqBlcisnnP9ZMfz6JIJZ7vahxcYxlcicwwuBIREiJDEBJs+XKwev8FPPbZAQiCAADQ6wVRcLWV+J8oUDG4EhGCg+RY+EA3dG8Zi9f/0AltEiJE23+7qIL65iAmtVaHm3EWABAbzuBKZIojEYgIADCyW3OM7NYcALBq/wWz7SXl1YhWKlBaIc7aFGVjPViiQMWWKxGZCVOYXxoKimsXThcNZlIGI0gu81i9iHyFx4KrSqXCiy++iPT0dERERKBLly5YvHgx9HrzYf7Gli1bBplMZvZ45plnPFRzosATZmEZuQeX7cPFkkqUVFYZymLYJUxkkUe6hQVBwLhx43DgwAHMmTMH7du3x7Zt2/Dss8/i+vXrmDVrltVjc3NzkZmZiU8//VRUnpyc3Mi1JgpcymDLWZfmb/kde8/cMDyPDQvxVJWIfIpHgusvv/yCLVu2YO3atXjggQcAAEOGDEFxcTHefvttzJw5EzKZ5a6l3Nxc9OzZE7179/ZEVYkIgNLKAugbci+JnhtnaiKieh7rFp40aRKGDBkiKmvfvj3KyspQVFRk8RhBEHDkyBF0797dpffUaDQoKSkRPUpLS116LaJAYq3laspW8gmiQOaR4HrLLbdgyZIliI+PF5Vv2LABiYmJSExMtHjc6dOnoVarceDAAWRmZkKhUCAzMxPLly936H3nzp2LuLg40SM1NbWhfw6R37M259XUsM68PUNkSYO7hcvLy7FixQqr25s3b45Ro0aZlb/77rvIycnBwoULIZdb/iLn5uYCAM6ePYuFCxdCoVBg+fLlmDBhArRaLZ544omGVp+ILDh2yX4PT0iwHH/u3doDtSHyPQ0OrsXFxZgyZYrV7QMGDDALru+//z6mTp2KsWPHYurUqVaP7d+/P7766isMGjQIERG1k9rvvvtuFBUV4ZVXXsHjjz9u9V4tEbnujowEHC6wHWDH92zJ7x+RFQ3uFk5JSYEgCFYfOTk5hn31ej2mTZuGZ599FuPHj8fKlSttfjmTkpIwYsQIQ2CtM3z4cFy+fBlXrlyxWbcZM2aguLhY9MjPz2/In0sUEMbe1tJmvmEASLKznSiQeSxDU1VVFcaPH4/169dj2rRpmD9/vt1fvbt27cKZM2cwYcIEUXllZSWCgoLM7uGaUiqVUCp5ASByVmpCBHZOH4QguQxaXQ2+O3oZf8/+FVW6+nnpCZGchkNkjcdGC0+cOBH//e9/8c4772DBggUOdSdt374dEydOxMmTJw1ler0e69atQ9++fRESwi83UWMJCZYjSC5DeEgwRvdIwd6XxKP9g6yMlSAiD7VcN27ciNWrV2PUqFHo3bs39u7dK9reo0cPhIaGoqCgAAUFBYbnkydPxocffoiRI0fitddeQ0REBBYvXozffvsNO3fu9ETVieim+IgQ9GgVi1/OlyBYLkO/jASpq0TktWSCYLy+ReN45JFHbI4oPnv2LFJTUzF79mzMmTPH8BwA8vLy8NJLL2Hnzp1Qq9Xo2bMn/vnPf6Jfv34u1aWkpARxcXGG58XFxYiNjXXptYgCzbnr5Vi57zx6t4nH4PZNpa4OkdfySHD1JgyuRETU2HjThIiIyM0YXImIiNyMwZWIiMjNGFyJiIjcjMGViIjIzRhciYiI3IzBlYiIyM0YXImIiNzMY4n7vYVpzozSUvvrVhIRUWCLiYlxaonFgAuuKpVK9LwuzSIREZE1zmbzY7cwERGRmzG4EhERuVnAJe7X6/W4cOGC4Xl0dLRT/eh1SktLRV3K+fn5iImJcUcVfRrPi3U8N5bxvFjG82KZVOeF91ztkMvlaN26tdtfNyYmhqvrWMDzYh3PjWU8L5bxvFjmreeF3cJERERuxuBKRETkZgyuREREbhZw91zdRalU4tVXXxU9J54XW3huLON5sYznxTJfOS8BN1qYiIiosbFbmIiIyM0YXImIiNyMwZWIiMjNGFyJiIjcjMGViIjIzRhciYiI3IzBlYiIyM0YXF20detW9OzZE+Hh4UhLS8OCBQsQKFOGCwoKEBsbi5ycHFF5Xl4eRo4cidjYWCQkJGDKlClmi9OXlZXh6aefRnJyMiIjI3HvvffixIkTHqy9e+n1enz00Ufo2rUrIiMj0aZNG0ydOlX0dwfieQFqz82CBQvQtm1bhIWFoVu3bli5cqVon59//hkDBw5EZGQkmjdvjn/84x+oqqoS7XPlyhU89NBDaNKkCWJiYjB+/HgUFhZ68k9pNPfff79ohRcgcD8vGo0GCoUCMplM9IiMjDTs41OfF4Gc9tNPPwkKhUJ4+OGHhW+//VZ4+eWXBZlMJrz11ltSV63RnT9/XujQoYMAQNi+fbuhvLi4WGjZsqXQs2dPYcOGDcK///1vITY2Vrj77rtFx48cOVJITEwUPvnkEyE7O1vo2rWr0KxZM+HGjRse/kvc46233hKCgoKEGTNmCN9//73wwQcfCPHx8cLQoUMFvV4fsOdFEATh5ZdfFhQKhfDWW28JP/zwg/DCCy8IAIRVq1YJgiAIp0+fFmJiYoRhw4YJmzdvFhYsWCCEhoYKkydPNrxGdXW10KNHDyE9PV1Yu3atsHLlSqFly5ZC586dhaqqKqn+NLdYsWKFAEBo3bq1oSyQPy8HDhwQAAiff/658NNPPxke+/fvFwTB9z4vDK4uyMrKEnr16iUqmz59uhAVFSVUVFRIVKvGVVNTI3zyySdCkyZNhPj4eLPg+uabbwrh4eHC1atXDWXffPONAEDYvXu3IAiCsGfPHgGA8M033xj2KSoqEiIiIoQ33njDY3+Lu9TU1AixsbHCU089JSr/4osvBADCgQMHAvK8CIIglJeXCxEREcKLL74oKh8wYIDQu3dvQRAEYdKkSUJKSoqg1WoN2xcvXizI5XLh3LlzgiAIwqpVqwQAwtGjRw37HD16VJDJZMLnn3/ugb+kcVy8eFGIi4sTUlJSRME1UD8vgiAIS5cuFYKDgwWNRmNxu699Xtgt7CStVoucnByMHj1aVD5mzBio1Wrs3r1bopo1riNHjuDJJ5/EI488ghUrVpht/+6773DnnXciISHBUJaVlYWoqCh88803hn0iIiKQlZVl2CcxMREDBgww7ONLVCoV/vznP+PBBx8Ulbdv3x4AcPr06YA8LwAQGhqKPXv2YNq0aaLykJAQaDQaALV/9/DhwxESEmLYPmbMGOj1enz33XeGfTIzM9GxY0fDPh07dkSHDh189twAwOOPP46srCwMGTJEVB6onxcAyM3NRfv27REaGmpxu699XhhcnXTmzBlUVVWhXbt2ovKMjAwA8Pn7Hta0atUKeXl5WLRoEcLDw822Hz9+3OycBAUFIS0tzXBOjh8/jjZt2iAoKEi0X0ZGhk+et9jYWLz33nu44447ROUbNmwAAHTq1CkgzwtQ+zd27doVycnJEAQBV65cwdy5c/HDDz/gqaeeQmVlJc6dO2d2bhITExEdHS06N6b7AL59bpYtW4aDBw/i/fffN9sWqJ8XoDa4BgcHIysrCxEREYiPj8fkyZOhVqt98vPC4Oqk0tJSAEB0dLSoPCoqCgDMBh74i/j4eKSkpFjdXlpaanZOgNrzUndOHNnH1+3btw9z587FyJEj0blzZ54XAF988QWSk5Px0ksv4d5778XDDz9s9XsE+Pe5OXfuHF544QUsXrxY1DqtE6ifF0EQcOTIEeTl5eEPf/gDvv32W7z88stYvXo17r33XhQXFwPwrc8Ll5xzkl6vt7ldLg/M3yu2zkvdOXFkH1/2448/YsSIEUhLS8Mnn3wCgOcFAHr16oUdO3bgyJEjmDVrFoYNG4ZVq1bZPMYfz40gCHjsscdw77334o9//KPFfQL18yIIAjZt2oTExER06tQJANC/f38kJyfj4Ycfxvbt220e743nhsHVSTExMQAAtVotKq/7VVS3PdDExMSYnROg9ry0aNHCsM+VK1cs7uPr523NmjWYOHEi2rVrhy1btqBJkyYAeF4AID09Henp6ejfvz+io6MxYcIEnD59GoD59wgQ/922zp+vnZsPPvgAR44cwa+//gqdTgcAhul7Op0Ocrk8YD8vcrkcAwcONCsfPnw4AODs2bMAfOvz4ps/cySUnp6OoKAg5OXlicrrnnfo0EGKakkuMzPT7JzU1NTg7NmzhnOSmZmJs2fPmv26zMvL8+nztmDBAowfPx59+vTBzp070axZM8O2QD0vV69exfLly1FUVCQqv+WWWwAAly5dQosWLczOTVFREdRqtejcmO4D+Oa5WbduHa5du4ZmzZpBoVBAoVBg+fLlOHfuHBQKBV577bWA/bxcunQJS5cuxfnz50XllZWVAIBmzZr53OeFwdVJSqUS/fv3x/r160VJI7KzsxETE4NevXpJWDvpZGVlYceOHbh69aqhbOvWrSgrKzOMaszKyoJarTaM7ANqL8I7d+4UjXz0JUuWLMHf/vY3jB07Flu2bDH7dRyo56WyshITJkzAxx9/LCrfunUrAKBr167IysrC119/Da1Wa9ienZ2NoKAgDB48GEDtuTl+/DiOHTtm2OfYsWM4fvy4z52bJUuW4MCBA6LHiBEj0KxZMxw4cACTJk0K2M+LTqfDpEmTsGTJElH5mjVrEBQUhDvvvNP3Pi8enfjjJ7Zt2ybIZDJhzJgxwjfffCPMnDlTkMlkwrx586Sumkds377dbJ5rUVGRkJCQIHTr1k1Yv369sHTpUiEuLk645557RMcOHDhQiIuLE5YuXSqsX79e6Nq1q9CiRQufnPxeWFgohIWFCampqcKuXbtEE99/+uknoaioKCDPS53HHntMUCqVwvz584UffvhBePXVV4XQ0FDhL3/5iyAIgnD8+HFBqVQKgwYNEr766ith4cKFQmhoqDBlyhTDa2g0GiEzM1No1aqVsGrVKmHVqlVCq1athC5dugjV1dVS/WluM2HCBNE810D+vDz66KOCQqEQXn/9deGHH34QZs+eLYSEhAjPPfecIAi+93lhcHXR+vXrhS5dugghISFCWlqasGDBAqmr5DGWgqsgCMKvv/4qDBkyRAgLCxOSkpKESZMmCSqVSrTPjRs3hIkTJwqxsbFCdHS0cM899wi///67B2vvPh9//LEAwOrjk08+EQQh8M5LHa1WK7zxxhtC27ZthZCQECE9PV2YN2+eUFNTY9hn586dwu233y6EhoYKLVq0EGbMmGGWSef8+fPC6NGjhcjISCEuLk4YN26ccOnSJU//OY3CNLgKQuB+XjQajfD6668L7dq1E0JDQ4X09HRh7ty5Pvt5kQlCgCTEJSIi8hDecyUiInIzBlciIiI3Y3AlIiJyMwZXIiIiN2NwJSIicjMGVyIiIjdjcCUiInIzBlciIiI3Y3AlIiJyMwZXIiIiN2NwJSIicjMGVyIiIjdjcCUiInIzBlciIiI3Y3AlIiJyMwZXIiIiN/t/eRsTM85MWr4AAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 508.5x381 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(element['latent'][0])\n",
    "plt.plot(element['mask'][0])\n",
    "\n",
    "# mask and (padded) latents visualized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_latent_dataloader = torch.utils.data.DataLoader(\n",
    "    latent_dataset_train,\n",
    "    batch_size=cfg.training.batch_size,\n",
    "    shuffle=True,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    ")\n",
    "\n",
    "val_latent_dataloader = torch.utils.data.DataLoader(\n",
    "    latent_dataset_val,\n",
    "    batch_size=cfg.training.batch_size,\n",
    "    shuffle=False,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    ")\n",
    "\n",
    "num_batches = len(train_latent_dataloader)\n",
    "\n",
    "# check if signal length is power of 2\n",
    "if cfg.dataset.max_seqlen & (cfg.dataset.max_seqlen - 1) != 0:\n",
    "    cfg.training.precision = \"no\"  # torch.fft doesnt support half if L!=2^x\n",
    "\n",
    "# prepare the denoiser model and dataset\n",
    "(\n",
    "    train_latent_dataloader,\n",
    "    val_latent_dataloader,\n",
    ") = accelerator.prepare(\n",
    "    train_latent_dataloader,\n",
    "    val_latent_dataloader,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "## initialize denoiser\n",
    "\n",
    "denoiser = Denoiser(\n",
    "    C_in=cfg.denoiser_model.C_in + 1, # 1 for mask (length of required latent)\n",
    "    C=cfg.denoiser_model.C,\n",
    "    L=cfg.dataset.max_seqlen,\n",
    "    num_blocks=cfg.denoiser_model.num_blocks,\n",
    "    bidirectional=cfg.denoiser_model.get(\"bidirectional\", True),\n",
    ")\n",
    "\n",
    "# initial values may be way off so better to scale down the output layer\n",
    "denoiser.conv_out.weight.data = denoiser.conv_out.weight.data * 0.1\n",
    "denoiser.conv_out.bias.data = denoiser.conv_out.bias.data * 0.1\n",
    "\n",
    "    \n",
    "scheduler = DDPMScheduler(\n",
    "    num_train_timesteps=cfg.denoiser_model.num_train_timesteps,\n",
    "    clip_sample=False,\n",
    "    beta_schedule=\"linear\",\n",
    ")\n",
    "\n",
    "\n",
    "optimizer = torch.optim.AdamW(\n",
    "    denoiser.parameters(), lr=cfg.training.lr\n",
    ")  # default wd=0.01 for now\n",
    "\n",
    "\n",
    "\n",
    "num_batches = len(train_latent_dataloader)\n",
    "lr_scheduler = get_scheduler(\n",
    "    name=\"cosine\",\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=num_batches * cfg.training.num_warmup_epochs,  # warmup for 10% of epochs\n",
    "    num_training_steps=num_batches * cfg.training.num_epochs * 1.3,  # total number of steps\n",
    ")\n",
    "\n",
    "# prepare the denoiser model and dataset\n",
    "(\n",
    "    denoiser,\n",
    "    optimizer,\n",
    "    lr_scheduler,\n",
    ") = accelerator.prepare(\n",
    "    denoiser,\n",
    "    optimizer,\n",
    "    lr_scheduler,\n",
    ")\n",
    "\n",
    "ema_model = EMAModel(denoiser, min_value=0.9999)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.688417 M parameters\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       " Denoiser(\n",
       "   (conv_in): Conv1d(33, 384, kernel_size=(1,), stride=(1,))\n",
       "   (blocks): ModuleList(\n",
       "     (0-5): 6 x DenoiserBlock(\n",
       "       (time_mixer): FFTConv(\n",
       "         (activation): Identity()\n",
       "         (kernel): SSMKernelDPLR()\n",
       "         (drop): Identity()\n",
       "         (drop_kernel): Identity()\n",
       "       )\n",
       "       (channel_mixer): Sequential(\n",
       "         (0): Rearrange('b c l -> b l c')\n",
       "         (1): Linear(in_features=384, out_features=768, bias=True)\n",
       "         (2): GELU(approximate='none')\n",
       "         (3): Linear(in_features=768, out_features=384, bias=True)\n",
       "         (4): Rearrange('b l c -> b c l')\n",
       "       )\n",
       "       (norm1): InstanceNorm1d(384, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "       (norm2): InstanceNorm1d(384, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "       (act1): GELU(approximate='none')\n",
       "       (act2): Identity()\n",
       "       (ada_ln): Sequential(\n",
       "         (0): GELU(approximate='none')\n",
       "         (1): Linear(in_features=96, out_features=2304, bias=True)\n",
       "       )\n",
       "     )\n",
       "   )\n",
       "   (conv_out): Conv1d(384, 33, kernel_size=(1,), stride=(1,))\n",
       "   (t_emb): TimestepEmbedder(\n",
       "     (mlp): Sequential(\n",
       "       (0): Linear(in_features=256, out_features=96, bias=True)\n",
       "       (1): SiLU()\n",
       "       (2): Linear(in_features=96, out_features=96, bias=True)\n",
       "     )\n",
       "   )\n",
       " ))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "from ldns.utils import count_parameters\n",
    "\n",
    "print(count_parameters(denoiser)/1e6, \"M parameters\")\n",
    "\n",
    "\n",
    "denoiser\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_spikes_with_mask(ema_denoiser, scheduler, ae, cfg, lengths=None, batch_size=1, device=\"cuda\"):\n",
    "    \"\"\"Sample spike trains from the diffusion model with variable length masks.\n",
    "    \n",
    "    Uses DDPM sampling with EMA-averaged model to generate latents, which are then \n",
    "    decoded to rates and spikes. \n",
    "    Args:\n",
    "        ema_denoiser: EMA-averaged denoiser model\n",
    "        scheduler: DDPM scheduler for sampling\n",
    "        ae: Trained autoencoder for decoding latents to rates\n",
    "        cfg: Config object with model parameters\n",
    "        lengths: Optional list/tensor of sequence lengths for masks\n",
    "        batch_size: Number of samples to generate\n",
    "        device: Device to run sampling on\n",
    "        \n",
    "    Returns:\n",
    "        dict containing:\n",
    "            rates: Generated firing rates\n",
    "            spikes: Sampled spike trains\n",
    "            latents: Generated latents\n",
    "            masks: Generated masks\n",
    "            mask_lengths: Lengths used for masks\n",
    "    \"\"\"\n",
    "    # start with random noise matching model input shape\n",
    "    z_t = torch.randn(\n",
    "        (batch_size, cfg.denoiser_model.C_in, cfg.dataset.max_seqlen)\n",
    "    ).to(device)\n",
    "\n",
    "    # generate lengths for masks if not provided\n",
    "    if lengths is None:\n",
    "        lengths = torch.linspace(100, 512, batch_size).long().to(device)\n",
    "    else:\n",
    "        if isinstance(lengths, int):\n",
    "            lengths = torch.tensor([lengths] * batch_size).to(device)\n",
    "        elif isinstance(lengths, list):\n",
    "            lengths = torch.tensor(lengths).long().to(device)\n",
    "    \n",
    "    # create masks with 1s in center and 0s for padding\n",
    "    masks = torch.zeros(batch_size, cfg.dataset.max_seqlen).to(device)\n",
    "    for i, l in enumerate(lengths):\n",
    "        padding_left = (cfg.dataset.max_seqlen - l) // 2\n",
    "        padding_right = cfg.dataset.max_seqlen - l - padding_left\n",
    "        masks[i, padding_left:padding_left + l] = 1.0\n",
    "\n",
    "    masks = masks.unsqueeze(1)\n",
    "\n",
    "    # get EMA model and prepare for inference\n",
    "    ema_denoiser_avg = ema_denoiser.averaged_model\n",
    "    ema_denoiser_avg.eval()\n",
    "    scheduler.set_timesteps(cfg.denoiser_model.num_train_timesteps)\n",
    "\n",
    "    # iteratively denoise using DDPM\n",
    "    for t in tqdm(scheduler.timesteps, desc=\"Sampling DDPM (different masks)\"):\n",
    "        with torch.no_grad():\n",
    "            model_output = ema_denoiser_avg(\n",
    "                torch.cat([z_t, masks], dim=1), torch.tensor([t] * batch_size).to(device).long()\n",
    "            )[:, :-1]\n",
    "        z_t = scheduler.step(model_output, t, z_t, return_dict=False)[0]\n",
    "\n",
    "    # scale latents back to original range\n",
    "    z_t = z_t * latent_dataset_train.latent_stds.to(z_t.device) + latent_dataset_train.latent_means.to(z_t.device)\n",
    "\n",
    "    # decode latents to rates and sample spikes\n",
    "    with torch.no_grad():\n",
    "        rates = ae.decode(z_t).cpu()\n",
    "    \n",
    "    spikes = torch.poisson(rates)\n",
    "\n",
    "    return {\n",
    "        \"rates\": rates,\n",
    "        \"spikes\": spikes,\n",
    "        \"latents\": z_t.cpu(),\n",
    "        \"masks\": masks.cpu(),\n",
    "        \"mask_lengths\": lengths,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_real_vs_sampled_rates_and_spikes(\n",
    "    real_rates,\n",
    "    sampled_rates,\n",
    "    real_spikes,\n",
    "    sampled_spikes,\n",
    "    real_masks,\n",
    "    sampled_masks,\n",
    "    batch_idx=0,\n",
    "):\n",
    "    \"\"\"Plot real and sampled rates and spikes side by side for comparison.\n",
    "    \n",
    "    Args:\n",
    "        real_rates (torch.Tensor): Ground truth firing rates [batch, channels, length]\n",
    "        sampled_rates (torch.Tensor): Model generated firing rates [batch, channels, length] \n",
    "        real_spikes (torch.Tensor): Ground truth spike trains [batch, channels, length]\n",
    "        sampled_spikes (torch.Tensor): Model generated spike trains [batch, channels, length]\n",
    "        real_masks (torch.Tensor): Masks for real data [batch, 1, length]\n",
    "        sampled_masks (torch.Tensor): Masks for sampled data [batch, 1, length]\n",
    "        batch_idx (int): Which batch element to plot. Defaults to 0.\n",
    "    \"\"\"\n",
    "    B, C, L = real_rates.shape\n",
    "\n",
    "    # create 2x2 subplot with appropriate size\n",
    "    fig, axs = plt.subplots(2, 2, figsize=cm2inch(12, 8), dpi=300)\n",
    "\n",
    "    # select single batch element to plot\n",
    "    real_rates = real_rates[batch_idx]\n",
    "    sampled_rates = sampled_rates[batch_idx]\n",
    "    real_spikes = real_spikes[batch_idx]\n",
    "    sampled_spikes = sampled_spikes[batch_idx]\n",
    "    real_masks = real_masks[batch_idx]\n",
    "    sampled_masks = sampled_masks[batch_idx]\n",
    "\n",
    "    # get indices where mask is 1 (non-padding)\n",
    "    real_mask_idx_with_1 = torch.arange(real_masks[0].nonzero().flatten().numel())\n",
    "    sampled_mask_idx_with_1 = sampled_masks[0].nonzero().flatten()\n",
    "\n",
    "    # plot real rates with colorbar\n",
    "    im = axs[0, 0].imshow(real_rates[:, real_mask_idx_with_1], cmap=\"viridis\", alpha=1.0, aspect=\"auto\")\n",
    "    axs[0, 0].set_title(\"Real rates\")\n",
    "    fig.colorbar(im, ax=axs[0, 0], orientation=\"vertical\", fraction=0.046, pad=0.04)\n",
    "\n",
    "    # plot sampled rates with colorbar\n",
    "    im = axs[0, 1].imshow(\n",
    "        sampled_rates[:, sampled_mask_idx_with_1], cmap=\"viridis\", alpha=1.0, aspect=\"auto\"\n",
    "    )\n",
    "    axs[0, 1].set_title(\"Sampled rates\") \n",
    "    fig.colorbar(im, ax=axs[0, 1], orientation=\"vertical\", fraction=0.046, pad=0.04)\n",
    "\n",
    "    # plot real spikes with colorbar\n",
    "    im = axs[1, 0].imshow(real_spikes[:, real_mask_idx_with_1], cmap=\"Greys\", alpha=1.0, aspect=\"auto\")\n",
    "    axs[1, 0].set_title(\"Real spikes\")\n",
    "    fig.colorbar(im, ax=axs[1, 0], orientation=\"vertical\", fraction=0.046, pad=0.04)\n",
    "\n",
    "    # plot sampled spikes with colorbar\n",
    "    im = axs[1, 1].imshow(\n",
    "        sampled_spikes[:, sampled_mask_idx_with_1], cmap=\"Greys\", alpha=1.0, aspect=\"auto\"\n",
    "    )\n",
    "    axs[1, 1].set_title(\"Sampled spikes\")\n",
    "    fig.colorbar(im, ax=axs[1, 1], orientation=\"vertical\", fraction=0.046, pad=0.04)\n",
    "\n",
    "    # print shapes for debugging\n",
    "    print(\n",
    "        real_rates[:, real_mask_idx_with_1].shape,\n",
    "        sampled_rates[:, sampled_mask_idx_with_1].shape,\n",
    "        real_spikes[:, real_mask_idx_with_1].shape,\n",
    "        sampled_spikes[:, sampled_mask_idx_with_1].shape,\n",
    "    )\n",
    "\n",
    "    # remove y-ticks from right plots for cleaner visualization\n",
    "    for i, ax in enumerate(axs.flatten()):\n",
    "        if i% 2 != 0:\n",
    "            ax.set_yticks([])\n",
    "\n",
    "    fig.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define loss function - smooth L1 loss with small beta for faster convergence\n",
    "loss_fn = torch.nn.SmoothL1Loss(beta=0.04, reduction=\"none\")\n",
    "\n",
    "# training loop\n",
    "pbar = tqdm(range(0, cfg.training.num_epochs), desc=\"epochs\")\n",
    "for epoch in pbar:\n",
    "    for i, batch in enumerate(train_latent_dataloader):\n",
    "        denoiser.train()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # get latent vectors and masks from batch\n",
    "        z = batch[\"latent\"] \n",
    "        z_mask = batch[\"mask\"]\n",
    "\n",
    "        # sample random timesteps\n",
    "        t = torch.randint(\n",
    "            0, cfg.denoiser_model.num_train_timesteps, (z.shape[0],), device=\"cpu\"\n",
    "        ).long()\n",
    "\n",
    "        # add noise to latents according to timestep\n",
    "        noise = torch.randn_like(z)\n",
    "        noisy_z = scheduler.add_noise(z, noise, t)\n",
    "\n",
    "        # predict noise and remove mask dimension\n",
    "        noise_pred = denoiser(torch.cat([noisy_z, z_mask], dim=1), t)\n",
    "        noise_pred = noise_pred[:,:-1] # remove mask conditioning dim\n",
    "\n",
    "        # compute masked loss\n",
    "        loss = loss_fn(noise, noise_pred)\n",
    "        loss = loss * z_mask  # zero out loss for padding\n",
    "        loss = loss.mean()\n",
    "\n",
    "        # optimization step\n",
    "        accelerator.backward(loss)\n",
    "        accelerator.clip_grad_norm_(denoiser.parameters(), 1.0)\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "\n",
    "        # update progress bar\n",
    "        if i % 10 == 0:\n",
    "            pbar.set_postfix(\n",
    "                {\n",
    "                    \"loss\": loss.item(),\n",
    "                    \"lr\": lr_scheduler.get_last_lr()[0],\n",
    "                    \"epoch\": epoch,\n",
    "                }\n",
    "            )\n",
    "\n",
    "        # update EMA model\n",
    "        ema_model.step(denoiser)\n",
    "\n",
    "    # generate and plot samples every 50 epochs\n",
    "    if (epoch) % 50 == 0:\n",
    "        denoiser.eval()\n",
    "\n",
    "        # get validation batch and generate samples\n",
    "        val_batch = next(iter(val_latent_dataloader))\n",
    "        val_batch['mask'][:2,0].sum(-1)\n",
    "        ret_dict = sample_spikes_with_mask(ema_model,\n",
    "                                           scheduler, \n",
    "                                           ae_model, \n",
    "                                           cfg, \n",
    "                                           lengths=[int(i.item()) for i in val_batch['mask'][:2,0].sum(-1)], \n",
    "                                           batch_size=2, \n",
    "                                           device=(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        )\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            val_batch_rates = ae_model(val_batch[\"signal\"])[0].cpu()\n",
    "\n",
    "        # plot real vs generated samples for first two items in batch\n",
    "        plot_real_vs_sampled_rates_and_spikes(\n",
    "            val_batch_rates,\n",
    "            ret_dict[\"rates\"],\n",
    "            val_batch[\"signal\"].cpu(),\n",
    "            ret_dict[\"spikes\"],\n",
    "            val_batch[\"mask\"].cpu(),\n",
    "            ret_dict[\"masks\"],\n",
    "            batch_idx=0,\n",
    "        )\n",
    "        plot_real_vs_sampled_rates_and_spikes(\n",
    "            val_batch_rates,\n",
    "            ret_dict[\"rates\"],\n",
    "            val_batch[\"signal\"].cpu(),\n",
    "            ret_dict[\"spikes\"],\n",
    "            val_batch[\"mask\"].cpu(),\n",
    "            ret_dict[\"masks\"],\n",
    "            batch_idx=1,\n",
    "        )\n",
    "\n",
    "pbar.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save model\n",
    "os.makedirs(f\"exp/{cfg.exp_name}\", exist_ok=True)\n",
    "torch.save(ema_model.averaged_model.state_dict(), f\"exp/{cfg.exp_name}/model.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We evaluate the model in `notebooks/plotting_diffusion_human.ipynb`"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "timeseries",
   "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.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}