From 9c4fd361d1626200e5c0ae5d1757eaf2c834502e Mon Sep 17 00:00:00 2001
From: Ashish Kumar Singh <ashishkmr472@gmail.com>
Date: Sat, 6 Jul 2024 02:38:01 +0530
Subject: [PATCH] feat: added more stuff to readme

---
 .env                       |   10 -
 .gitignore                 |    4 +-
 Diffusion flax linen.ipynb | 2223 +++++++++++++++++++++++++++++++++++-
 README.md                  |  143 ++-
 setup.py                   |    2 +
 5 files changed, 2359 insertions(+), 23 deletions(-)
 delete mode 100644 .env

diff --git a/.env b/.env
deleted file mode 100644
index 56ae081..0000000
--- a/.env
+++ /dev/null
@@ -1,10 +0,0 @@
-TF_USE_LEGACY_KERAS=False
-KERAS_BACKEND=tensorflow
-TF_GPU_THREAD_MODE=gpu_private
-TF_GPU_THREAD_COUNT=2
-JAX_ENABLE_X64=False
-XLA_FLAGS=--xla_gpu_enable_triton_softmax_fusion=true \
-            --xla_gpu_triton_gemm_any=True --xla_gpu_enable_async_collectives \
-            --xla_gpu_enable_latency_hiding_scheduler=true \ 
-            --xla_gpu_enable_highest_priority_async_stream=true \
-            --xla_gpu_disable_async_collectives=allreduce,allgather,reducescatter,collectivebroadcast,alltoall,collectivepermute
diff --git a/.gitignore b/.gitignore
index 18f83fe..b39bf05 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,4 +3,6 @@ __pycache__
 .ipynb_checkpoints
 checkpoints
 good models
-*(old).ipynb
\ No newline at end of file
+*(old).ipynb
+.env
+.env
diff --git a/Diffusion flax linen.ipynb b/Diffusion flax linen.ipynb
index 02e8c0b..0fc4f76 100644
--- a/Diffusion flax linen.ipynb	
+++ b/Diffusion flax linen.ipynb	
@@ -208,7 +208,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from jax.numpy import ndarray\n",
     "from flaxdiff.schedulers import CosineNoiseSchedule, NoiseScheduler, GeneralizedNoiseScheduler, KarrasVENoiseScheduler, EDMNoiseScheduler\n",
     "from flaxdiff.predictors import VPredictionTransform, EpsilonPredictionTransform, DiffusionPredictionTransform, DirectPredictionTransform, KarrasPredictionTransform"
    ]
@@ -1118,8 +1117,6 @@
     "IMAGE_SIZE = 64\n",
     "\n",
     "cosine_schedule = CosineNoiseSchedule(1000, beta_end=1)\n",
-    "# linear_schedule = LinearNoiseSchedule(1000)\n",
-    "# cosine_general_schedule = CosineGeneralNoiseScheduler(1)\n",
     "karas_ve_schedule = KarrasVENoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)\n",
     "edm_schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)"
    ]
@@ -1441,7 +1438,2225 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "\t\tEpoch 1:   0%|                                                          | 0/511 [00:00<?, ?step/s]"
+      "\t\tEpoch 1: 600step [00:40, 14.90step/s, loss=0.1288]                                                \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 1 completed. Avg Loss: 0.10374092310667038, Time: 40.26s, Best Loss: 0.10139207541942596\n",
+      "\n",
+      "Epoch 2/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 2: 600step [00:25, 23.36step/s, loss=0.0980]                                                \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Saving model at epoch 1\n",
+      "\n",
+      "\tEpoch 2 completed. Avg Loss: 0.1007109209895134, Time: 25.69s, Best Loss: 0.1007109209895134\n",
+      "\n",
+      "Epoch 3/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 3: 600step [00:25, 23.51step/s, loss=0.0828]                                                \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 3 completed. Avg Loss: 0.10220622271299362, Time: 25.52s, Best Loss: 0.1007109209895134\n",
+      "\n",
+      "Epoch 4/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 4: 600step [00:25, 23.23step/s, loss=0.0966]                                                \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 4 completed. Avg Loss: 0.10254718363285065, Time: 25.83s, Best Loss: 0.1007109209895134\n",
+      "\n",
+      "Epoch 5/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 5: 600step [00:25, 23.64step/s, loss=0.0841]                                                \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 5 completed. Avg Loss: 0.10168876498937607, Time: 25.38s, Best Loss: 0.1007109209895134\n",
+      "\n",
+      "Epoch 6/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 6: 600step [00:25, 23.42step/s, loss=0.0965]                                                \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 6 completed. Avg Loss: 0.10296788066625595, Time: 25.62s, Best Loss: 0.1007109209895134\n",
+      "\n",
+      "Epoch 7/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 7: 600step [00:25, 23.37step/s, loss=0.1218]                                                \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 7 completed. Avg Loss: 0.10113122314214706, Time: 25.67s, Best Loss: 0.1007109209895134\n",
+      "\n",
+      "Epoch 8/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 8: 600step [00:25, 23.52step/s, loss=0.1103]                                                \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 8 completed. Avg Loss: 0.10133743286132812, Time: 25.51s, Best Loss: 0.1007109209895134\n",
+      "\n",
+      "Epoch 9/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 9: 600step [00:25, 23.56step/s, loss=0.1129]                                                \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 9 completed. Avg Loss: 0.10188990086317062, Time: 25.47s, Best Loss: 0.1007109209895134\n",
+      "\n",
+      "Epoch 10/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 10: 600step [00:25, 23.41step/s, loss=0.0797]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 10 completed. Avg Loss: 0.10143899917602539, Time: 25.64s, Best Loss: 0.1007109209895134\n",
+      "\n",
+      "Epoch 11/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 11: 600step [00:25, 23.31step/s, loss=0.1074]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 11 completed. Avg Loss: 0.10209371894598007, Time: 25.75s, Best Loss: 0.1007109209895134\n",
+      "\n",
+      "Epoch 12/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 12: 600step [00:25, 23.23step/s, loss=0.1169]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 12 completed. Avg Loss: 0.10186251252889633, Time: 25.83s, Best Loss: 0.1007109209895134\n",
+      "\n",
+      "Epoch 13/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 13: 600step [00:25, 23.59step/s, loss=0.1324]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 13 completed. Avg Loss: 0.10210812091827393, Time: 25.43s, Best Loss: 0.1007109209895134\n",
+      "\n",
+      "Epoch 14/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 14: 600step [00:25, 23.31step/s, loss=0.1314]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 14 completed. Avg Loss: 0.10178422182798386, Time: 25.74s, Best Loss: 0.1007109209895134\n",
+      "\n",
+      "Epoch 15/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 15: 600step [00:25, 23.60step/s, loss=0.0768]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Saving model at epoch 14\n",
+      "\n",
+      "\tEpoch 15 completed. Avg Loss: 0.10067716240882874, Time: 25.43s, Best Loss: 0.10067716240882874\n",
+      "\n",
+      "Epoch 16/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 16: 600step [00:25, 23.44step/s, loss=0.1046]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 16 completed. Avg Loss: 0.10220663249492645, Time: 25.60s, Best Loss: 0.10067716240882874\n",
+      "\n",
+      "Epoch 17/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 17: 600step [00:26, 22.99step/s, loss=0.1216]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 17 completed. Avg Loss: 0.10117293149232864, Time: 26.10s, Best Loss: 0.10067716240882874\n",
+      "\n",
+      "Epoch 18/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 18: 600step [00:25, 23.43step/s, loss=0.0915]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 18 completed. Avg Loss: 0.10148977488279343, Time: 25.61s, Best Loss: 0.10067716240882874\n",
+      "\n",
+      "Epoch 19/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 19: 600step [00:25, 23.44step/s, loss=0.0961]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 19 completed. Avg Loss: 0.10258100926876068, Time: 25.60s, Best Loss: 0.10067716240882874\n",
+      "\n",
+      "Epoch 20/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 20: 600step [00:25, 23.11step/s, loss=0.0953]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 20 completed. Avg Loss: 0.1025131568312645, Time: 25.96s, Best Loss: 0.10067716240882874\n",
+      "\n",
+      "Epoch 21/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 21: 600step [00:25, 23.63step/s, loss=0.1132]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Saving model at epoch 20\n",
+      "\n",
+      "\tEpoch 21 completed. Avg Loss: 0.10028696060180664, Time: 25.39s, Best Loss: 0.10028696060180664\n",
+      "\n",
+      "Epoch 22/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 22: 600step [00:25, 23.41step/s, loss=0.1074]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 22 completed. Avg Loss: 0.10159551352262497, Time: 25.63s, Best Loss: 0.10028696060180664\n",
+      "\n",
+      "Epoch 23/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 23: 600step [00:25, 23.40step/s, loss=0.0893]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 23 completed. Avg Loss: 0.10220605880022049, Time: 25.64s, Best Loss: 0.10028696060180664\n",
+      "\n",
+      "Epoch 24/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 24: 600step [00:25, 23.80step/s, loss=0.1231]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 24 completed. Avg Loss: 0.10277654230594635, Time: 25.22s, Best Loss: 0.10028696060180664\n",
+      "\n",
+      "Epoch 25/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 25: 600step [00:23, 25.04step/s, loss=0.0945]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 25 completed. Avg Loss: 0.10207752883434296, Time: 23.97s, Best Loss: 0.10028696060180664\n",
+      "\n",
+      "Epoch 26/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 26: 600step [00:23, 25.45step/s, loss=0.0953]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Saving model at epoch 25\n",
+      "\n",
+      "\tEpoch 26 completed. Avg Loss: 0.10001388192176819, Time: 23.58s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 27/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 27: 600step [00:24, 24.62step/s, loss=0.0978]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 27 completed. Avg Loss: 0.10173722356557846, Time: 24.37s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 28/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 28: 600step [00:23, 25.25step/s, loss=0.0861]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 28 completed. Avg Loss: 0.10213057696819305, Time: 23.76s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 29/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 29: 600step [00:23, 25.46step/s, loss=0.0795]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 29 completed. Avg Loss: 0.10150100290775299, Time: 23.57s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 30/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 30: 600step [00:23, 25.21step/s, loss=0.1007]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 30 completed. Avg Loss: 0.10075327754020691, Time: 23.80s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 31/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 31: 600step [00:23, 25.19step/s, loss=0.1069]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 31 completed. Avg Loss: 0.10015378147363663, Time: 23.82s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 32/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 32: 600step [00:23, 25.31step/s, loss=0.0937]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 32 completed. Avg Loss: 0.10194233804941177, Time: 23.71s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 33/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 33: 600step [00:24, 24.80step/s, loss=0.0885]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 33 completed. Avg Loss: 0.10142005234956741, Time: 24.20s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 34/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 34: 600step [00:23, 25.07step/s, loss=0.0863]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 34 completed. Avg Loss: 0.10025395452976227, Time: 23.93s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 35/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 35: 600step [00:23, 25.12step/s, loss=0.1099]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 35 completed. Avg Loss: 0.10105018317699432, Time: 23.89s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 36/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 36: 600step [00:23, 25.59step/s, loss=0.0984]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 36 completed. Avg Loss: 0.10205303877592087, Time: 23.45s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 37/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 37: 600step [00:23, 25.31step/s, loss=0.1219]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 37 completed. Avg Loss: 0.1012827530503273, Time: 23.71s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 38/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 38: 600step [00:23, 25.31step/s, loss=0.0837]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 38 completed. Avg Loss: 0.10118275880813599, Time: 23.71s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 39/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 39: 600step [00:23, 25.17step/s, loss=0.0878]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 39 completed. Avg Loss: 0.10167257487773895, Time: 23.84s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 40/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 40: 600step [00:23, 25.15step/s, loss=0.1019]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 40 completed. Avg Loss: 0.10078752785921097, Time: 23.86s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 41/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 41: 600step [00:23, 25.67step/s, loss=0.0949]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 41 completed. Avg Loss: 0.10077335685491562, Time: 23.37s, Best Loss: 0.10001388192176819\n",
+      "\n",
+      "Epoch 42/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 42: 600step [00:23, 25.20step/s, loss=0.1170]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Saving model at epoch 41\n",
+      "\n",
+      "\tEpoch 42 completed. Avg Loss: 0.09964023530483246, Time: 23.81s, Best Loss: 0.09964023530483246\n",
+      "\n",
+      "Epoch 43/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 43: 600step [00:23, 25.01step/s, loss=0.1037]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 43 completed. Avg Loss: 0.10168183594942093, Time: 23.99s, Best Loss: 0.09964023530483246\n",
+      "\n",
+      "Epoch 44/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 44: 600step [00:23, 25.04step/s, loss=0.0971]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 44 completed. Avg Loss: 0.10122662782669067, Time: 23.97s, Best Loss: 0.09964023530483246\n",
+      "\n",
+      "Epoch 45/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 45: 600step [00:24, 24.85step/s, loss=0.1060]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 45 completed. Avg Loss: 0.1012326180934906, Time: 24.15s, Best Loss: 0.09964023530483246\n",
+      "\n",
+      "Epoch 46/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 46: 600step [00:23, 25.31step/s, loss=0.0960]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 46 completed. Avg Loss: 0.10102282464504242, Time: 23.71s, Best Loss: 0.09964023530483246\n",
+      "\n",
+      "Epoch 47/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 47: 600step [00:23, 25.46step/s, loss=0.0957]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 47 completed. Avg Loss: 0.1019478291273117, Time: 23.56s, Best Loss: 0.09964023530483246\n",
+      "\n",
+      "Epoch 48/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 48: 600step [00:24, 24.91step/s, loss=0.0770]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Saving model at epoch 47\n",
+      "\n",
+      "\tEpoch 48 completed. Avg Loss: 0.09944416582584381, Time: 24.09s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 49/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 49: 600step [00:23, 25.43step/s, loss=0.1056]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 49 completed. Avg Loss: 0.10101470351219177, Time: 23.59s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 50/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 50: 600step [00:23, 25.32step/s, loss=0.1207]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 50 completed. Avg Loss: 0.09992846101522446, Time: 23.69s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 51/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 51: 600step [00:24, 24.81step/s, loss=0.0923]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 51 completed. Avg Loss: 0.10084119439125061, Time: 24.18s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 52/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 52: 600step [00:25, 23.68step/s, loss=0.0973]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 52 completed. Avg Loss: 0.09986752271652222, Time: 25.34s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 53/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 53: 600step [00:25, 23.38step/s, loss=0.0960]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 53 completed. Avg Loss: 0.10152383148670197, Time: 25.67s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 54/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 54: 600step [00:26, 22.97step/s, loss=0.1363]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 54 completed. Avg Loss: 0.10034225136041641, Time: 26.12s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 55/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 55: 600step [00:26, 22.77step/s, loss=0.1199]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 55 completed. Avg Loss: 0.10041052848100662, Time: 26.36s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 56/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 56: 600step [00:25, 23.75step/s, loss=0.1001]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 56 completed. Avg Loss: 0.10061133652925491, Time: 25.27s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 57/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 57: 600step [00:25, 23.39step/s, loss=0.1107]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 57 completed. Avg Loss: 0.10073479264974594, Time: 25.66s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 58/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 58: 600step [00:25, 23.56step/s, loss=0.1013]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 58 completed. Avg Loss: 0.10122673958539963, Time: 25.47s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 59/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 59: 600step [00:25, 23.48step/s, loss=0.0950]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 59 completed. Avg Loss: 0.10150661319494247, Time: 25.55s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 60/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 60: 600step [00:25, 23.51step/s, loss=0.0840]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 60 completed. Avg Loss: 0.10155317187309265, Time: 25.53s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 61/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 61: 600step [00:25, 23.10step/s, loss=0.1152]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 61 completed. Avg Loss: 0.10095465183258057, Time: 25.98s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 62/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 62: 600step [00:25, 23.50step/s, loss=0.0862]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 62 completed. Avg Loss: 0.10084925591945648, Time: 25.53s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 63/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 63: 600step [00:25, 23.20step/s, loss=0.0980]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 63 completed. Avg Loss: 0.10070839524269104, Time: 25.86s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 64/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 64: 600step [00:26, 23.04step/s, loss=0.1044]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 64 completed. Avg Loss: 0.10004094243049622, Time: 26.04s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 65/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 65: 600step [00:26, 22.75step/s, loss=0.1377]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 65 completed. Avg Loss: 0.1016031876206398, Time: 26.37s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 66/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 66: 600step [00:26, 22.34step/s, loss=0.0969]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 66 completed. Avg Loss: 0.10134512931108475, Time: 26.86s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 67/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 67: 600step [00:25, 23.48step/s, loss=0.0939]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 67 completed. Avg Loss: 0.10052471607923508, Time: 25.55s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 68/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 68: 600step [00:25, 23.50step/s, loss=0.1170]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 68 completed. Avg Loss: 0.10035669803619385, Time: 25.53s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 69/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 69: 600step [00:26, 22.89step/s, loss=0.0790]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 69 completed. Avg Loss: 0.1013682410120964, Time: 26.21s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 70/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 70: 600step [00:25, 23.60step/s, loss=0.0932]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 70 completed. Avg Loss: 0.1001950278878212, Time: 25.43s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 71/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 71: 600step [00:25, 23.72step/s, loss=0.1167]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 71 completed. Avg Loss: 0.1004190742969513, Time: 25.30s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 72/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 72: 600step [00:25, 23.30step/s, loss=0.0863]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 72 completed. Avg Loss: 0.10089786350727081, Time: 25.76s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 73/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 73: 600step [00:25, 23.93step/s, loss=0.1019]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 73 completed. Avg Loss: 0.10094337165355682, Time: 25.08s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 74/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 74: 600step [00:23, 25.17step/s, loss=0.0965]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 74 completed. Avg Loss: 0.10029532015323639, Time: 23.84s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 75/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 75: 600step [00:23, 25.36step/s, loss=0.1056]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 75 completed. Avg Loss: 0.10038549453020096, Time: 23.66s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 76/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 76: 600step [00:23, 25.29step/s, loss=0.1126]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 76 completed. Avg Loss: 0.1003800481557846, Time: 23.73s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 77/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 77: 600step [00:23, 25.67step/s, loss=0.0982]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 77 completed. Avg Loss: 0.0997966080904007, Time: 23.37s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 78/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 78: 600step [00:23, 25.22step/s, loss=0.0739]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 78 completed. Avg Loss: 0.09970474988222122, Time: 23.80s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 79/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 79: 600step [00:23, 25.46step/s, loss=0.1147]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 79 completed. Avg Loss: 0.09945365786552429, Time: 23.57s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 80/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 80: 600step [00:23, 25.64step/s, loss=0.0993]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 80 completed. Avg Loss: 0.10030060261487961, Time: 23.41s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 81/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 81: 600step [00:23, 25.50step/s, loss=0.0932]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 81 completed. Avg Loss: 0.0994800552725792, Time: 23.53s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 82/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 82: 600step [00:23, 25.25step/s, loss=0.0782]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 82 completed. Avg Loss: 0.10051275789737701, Time: 23.77s, Best Loss: 0.09944416582584381\n",
+      "\n",
+      "Epoch 83/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 83: 600step [00:23, 25.64step/s, loss=0.1201]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Saving model at epoch 82\n",
+      "\n",
+      "\tEpoch 83 completed. Avg Loss: 0.098627008497715, Time: 23.40s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 84/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 84: 600step [00:24, 24.85step/s, loss=0.1074]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 84 completed. Avg Loss: 0.10015368461608887, Time: 24.15s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 85/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 85: 600step [00:23, 25.10step/s, loss=0.0956]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 85 completed. Avg Loss: 0.10047619044780731, Time: 23.90s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 86/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 86: 600step [00:23, 25.51step/s, loss=0.1141]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 86 completed. Avg Loss: 0.10024626553058624, Time: 23.52s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 87/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 87: 600step [00:23, 25.22step/s, loss=0.0795]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 87 completed. Avg Loss: 0.10063982009887695, Time: 23.79s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 88/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 88: 600step [00:23, 25.66step/s, loss=0.1223]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 88 completed. Avg Loss: 0.10150663554668427, Time: 23.39s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 89/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 89: 600step [00:24, 24.95step/s, loss=0.0974]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 89 completed. Avg Loss: 0.1009196788072586, Time: 24.05s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 90/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 90: 600step [00:23, 25.18step/s, loss=0.1049]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 90 completed. Avg Loss: 0.10014201700687408, Time: 23.83s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 91/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 91: 600step [00:23, 25.54step/s, loss=0.1175]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 91 completed. Avg Loss: 0.10012254118919373, Time: 23.50s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 92/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 92: 600step [00:23, 25.50step/s, loss=0.0936]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 92 completed. Avg Loss: 0.0998297855257988, Time: 23.53s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 93/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 93: 600step [00:23, 25.69step/s, loss=0.1022]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 93 completed. Avg Loss: 0.10023824125528336, Time: 23.35s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 94/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 94: 600step [00:23, 25.63step/s, loss=0.0883]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 94 completed. Avg Loss: 0.1001194417476654, Time: 23.41s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 95/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 95: 600step [00:23, 25.42step/s, loss=0.0697]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 95 completed. Avg Loss: 0.09886610507965088, Time: 23.61s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 96/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 96: 600step [00:23, 25.41step/s, loss=0.1090]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 96 completed. Avg Loss: 0.09878651797771454, Time: 23.61s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 97/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 97: 600step [00:23, 25.39step/s, loss=0.1101]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 97 completed. Avg Loss: 0.09932038187980652, Time: 23.63s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 98/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 98: 600step [00:23, 25.56step/s, loss=0.0931]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 98 completed. Avg Loss: 0.1006971076130867, Time: 23.48s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 99/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 99: 600step [00:23, 25.62step/s, loss=0.1005]                                               \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 99 completed. Avg Loss: 0.10075321048498154, Time: 23.42s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 100/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 100: 600step [00:23, 25.65step/s, loss=0.0862]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 100 completed. Avg Loss: 0.09962140768766403, Time: 23.40s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 101/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 101: 600step [00:23, 25.26step/s, loss=0.1334]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 101 completed. Avg Loss: 0.09924189746379852, Time: 23.76s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 102/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 102: 600step [00:23, 25.10step/s, loss=0.1106]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 102 completed. Avg Loss: 0.10034596174955368, Time: 23.90s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 103/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 103: 600step [00:23, 25.62step/s, loss=0.1219]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 103 completed. Avg Loss: 0.09945647418498993, Time: 23.42s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 104/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 104: 600step [00:23, 25.56step/s, loss=0.0826]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 104 completed. Avg Loss: 0.10063814371824265, Time: 23.48s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 105/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 105: 600step [00:23, 25.65step/s, loss=0.1137]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 105 completed. Avg Loss: 0.09938663244247437, Time: 23.39s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 106/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 106: 600step [00:23, 25.50step/s, loss=0.0820]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 106 completed. Avg Loss: 0.09980179369449615, Time: 23.53s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 107/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 107: 600step [00:24, 24.69step/s, loss=0.1052]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 107 completed. Avg Loss: 0.10006403923034668, Time: 24.31s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 108/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 108: 600step [00:25, 23.53step/s, loss=0.1152]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 108 completed. Avg Loss: 0.09920285642147064, Time: 25.50s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 109/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 109: 600step [00:25, 23.55step/s, loss=0.1017]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 109 completed. Avg Loss: 0.09970049560070038, Time: 25.48s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 110/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 110: 600step [00:24, 24.06step/s, loss=0.1053]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 110 completed. Avg Loss: 0.09980332106351852, Time: 24.94s, Best Loss: 0.098627008497715\n",
+      "\n",
+      "Epoch 111/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 111: 600step [00:24, 24.78step/s, loss=0.1125]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Saving model at epoch 110\n",
+      "\n",
+      "\tEpoch 111 completed. Avg Loss: 0.09807940572500229, Time: 24.22s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 112/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 112: 600step [00:23, 25.81step/s, loss=0.0875]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 112 completed. Avg Loss: 0.09983931481838226, Time: 23.24s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 113/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 113: 600step [00:24, 24.75step/s, loss=0.1014]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 113 completed. Avg Loss: 0.0996389091014862, Time: 24.24s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 114/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 114: 600step [00:25, 23.82step/s, loss=0.1013]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 114 completed. Avg Loss: 0.10035794228315353, Time: 25.19s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 115/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 115: 600step [00:24, 24.63step/s, loss=0.1117]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 115 completed. Avg Loss: 0.10000094771385193, Time: 24.36s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 116/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 116: 600step [00:25, 23.83step/s, loss=0.1164]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 116 completed. Avg Loss: 0.09977006167173386, Time: 25.18s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 117/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 117: 600step [00:25, 23.88step/s, loss=0.1011]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 117 completed. Avg Loss: 0.09975694864988327, Time: 25.13s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 118/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 118: 600step [00:24, 24.09step/s, loss=0.0995]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 118 completed. Avg Loss: 0.10077869892120361, Time: 24.91s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 119/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 119: 600step [00:25, 23.72step/s, loss=0.1210]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 119 completed. Avg Loss: 0.09965817630290985, Time: 25.30s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 120/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 120: 600step [00:25, 23.89step/s, loss=0.1143]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 120 completed. Avg Loss: 0.10052836686372757, Time: 25.12s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 121/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 121: 600step [00:25, 23.96step/s, loss=0.1044]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 121 completed. Avg Loss: 0.09919138252735138, Time: 25.04s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 122/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 122: 600step [00:25, 23.60step/s, loss=0.0948]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 122 completed. Avg Loss: 0.10064461827278137, Time: 25.42s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 123/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 123: 600step [00:24, 24.26step/s, loss=0.1068]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 123 completed. Avg Loss: 0.09890001267194748, Time: 24.74s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 124/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 124: 600step [00:24, 24.41step/s, loss=0.1128]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 124 completed. Avg Loss: 0.09890596568584442, Time: 24.58s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 125/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 125: 600step [00:24, 24.41step/s, loss=0.1045]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 125 completed. Avg Loss: 0.09860695898532867, Time: 24.58s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 126/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 126: 600step [00:24, 24.52step/s, loss=0.1001]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 126 completed. Avg Loss: 0.09988345950841904, Time: 24.47s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 127/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 127: 600step [00:24, 24.39step/s, loss=0.1007]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 127 completed. Avg Loss: 0.0994444489479065, Time: 24.60s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 128/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 128: 600step [00:25, 23.65step/s, loss=0.0914]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 128 completed. Avg Loss: 0.1003241166472435, Time: 25.37s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 129/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 129: 600step [00:23, 25.18step/s, loss=0.0818]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 129 completed. Avg Loss: 0.09967852383852005, Time: 23.83s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 130/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 130: 600step [00:24, 24.79step/s, loss=0.0950]                                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\tEpoch 130 completed. Avg Loss: 0.10050974786281586, Time: 24.21s, Best Loss: 0.09807940572500229\n",
+      "\n",
+      "Epoch 131/2000\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\t\tEpoch 131:   0%|                                           | 0/511 [00:00<?, ?step/s, loss=0.1188]"
      ]
     }
    ],
diff --git a/README.md b/README.md
index f853706..3a4d44e 100644
--- a/README.md
+++ b/README.md
@@ -1,16 +1,143 @@
 # FlaxDiff
 
-## A versatile and easy to understand Diffusion library
+## A Versatile and Easy-to-Understand Diffusion Library
 
-In recent years, Diffusion and Score based multi step models have taken over the generative AI domain. However, the latest research in this field have been very math intensive. Its not very easy to understand how exactly some of the current state-of-the-art diffusion models work and generate such awesome images. Replicating the research in code can be at times quite daunting. 
+In recent years, diffusion and score-based multi-step models have revolutionized the generative AI domain. However, the latest research in this field has become highly math-intensive, making it challenging to understand how state-of-the-art diffusion models work and generate such impressive images. Replicating this research in code can be daunting.
 
-FlaxDiff a library of tools (schedulers, samplers, models, etc) designed and implemented in a very easy to understand way. The moto is understandability and readibility over performance. I started this project as a hobby to get familiar with Flax and Jax and learn about diffusion and the latest research in generative AI.
+FlaxDiff is a library of tools (schedulers, samplers, models, etc.) designed and implemented in an easy-to-understand way. The focus is on understandability and readability over performance. I started this project as a hobby to familiarize myself with Flax and Jax and to learn about diffusion and the latest research in generative AI.
 
-I originally started this project in keras, as I have been intimately familiar with tensorlflow 2.0 in the past and used to love it, but as it turns out its pretty ancient now, and so I reimplmeneted everything in the newer framework Flax, powered by Jax due to its performance and ease of use. But the old notebooks and even old models with one of my first flax models is also provided.
+I initially started this project in Keras, being familiar with TensorFlow 2.0, but transitioned to Flax, powered by Jax, for its performance and ease of use. The old notebooks and models, including my first Flax models, are also provided.
 
-The `Diffusion flax linen.ipynb` notebook is my main workspace where I do all my latest experiements. Few checkpoints I feel good enough are uploaded to the `pretrained` folder along with a copy of my working notebook associated with that checkpoint. *You may need to copy that notebook over to the working root for it to work.*
+The `Diffusion_flax_linen.ipynb` notebook is my main workspace for experiments. Several checkpoints are uploaded to the `pretrained` folder along with a copy of the working notebook associated with each checkpoint. *You may need to copy the notebook to the working root for it to function properly.*
 
-## Disclaimer (and about me)
-I used to work as a Machine Learning Researcher at Hyperverge around 2019-2021 in the domain of computer vision, specifically facial anti spoofind and facial detection & recognition. Since switching to my latest job in 2021, I haven't been able to get the kind of R&D work to keep my skills and expertise sharpened. Therefore, I started this pet project to revisit and relearn the fundamentals and get familiar with the state of the art. Most of my full time work and experience now has been as a golang system engineer with some applied ML work (not R&D/desiging/training) sprinkled in. Therefore, the code may reflect the fact that I am a developer trying to tackle this vast maths heavy domain, and I may not end up doing a good job. Certain implementation may be off or outright wrong. Please forgive me for those mistakes, and do open up an issue to let me know.
+## Disclaimer (and About Me)
 
-## 
\ No newline at end of file
+I worked as a Machine Learning Researcher at Hyperverge from 2019-2021, focusing on computer vision, specifically facial anti-spoofing and facial detection & recognition. Since switching to my current job in 2021, I haven't engaged in as much R&D work, leading me to start this pet project to revisit and relearn the fundamentals and get familiar with the state-of-the-art. My current role involves primarily Golang system engineering with some applied ML work just sprinkled in. Therefore, the code may reflect my learning journey. Please forgive any mistakes and do open an issue to let me know.
+
+## Features
+
+### Schedulers
+Implemented in `flaxdiff.schedulers`:
+- **LinearNoiseSchedule** (`flaxdiff.schedulers.LinearNoiseSchedule`): A beta-parameterized discrete scheduler.
+- **CosineNoiseSchedule** (`flaxdiff.schedulers.CosineNoiseSchedule`): A beta-parameterized discrete scheduler.
+- **ExpNoiseSchedule** (`flaxdiff.schedulers.ExpNoiseSchedule`): A beta-parameterized discrete scheduler.
+- **CosineContinuousNoiseScheduler** (`flaxdiff.schedulers.CosineContinuousNoiseScheduler`): A continuous scheduler.
+- **CosineGeneralNoiseScheduler** (`flaxdiff.schedulers.CosineGeneralNoiseScheduler`): A continuous sigma parameterized cosine scheduler.
+- **KarrasVENoiseScheduler** (`flaxdiff.schedulers.KarrasVENoiseScheduler`): A sigma-parameterized continuous scheduler proposed by Karras et al. 2022, best suited for inference.
+- **EDMNoiseScheduler** (`flaxdiff.schedulers.EDMNoiseScheduler`): A sigma-parameterized continuous scheduler based on the Exponential Diffusion Model (EDM), best suited for training with the KarrasKarrasVENoiseScheduler.
+
+### Model Predictors
+Implemented in `flaxdiff.predictors`:
+- **EpsilonPredictor** (`flaxdiff.predictors.EpsilonPredictor`): Predicts the noise in the data.
+- **X0Predictor** (`flaxdiff.predictors.X0Predictor`): Predicts the original data from the noisy data.
+- **VPredictor** (`flaxdiff.predictors.VPredictor`): Predicts a linear combination of the data and noise, commonly used in the EDM.
+- **KarrasEDMPredictor** (`flaxdiff.predictors.KarrasEDMPredictor`): A generalized predictor for the EDM, integrating various parameterizations.
+
+### Samplers
+Implemented in `flaxdiff.samplers`:
+- **DDPMSampler** (`flaxdiff.samplers.DDPMSampler`): Implements the Denoising Diffusion Probabilistic Model (DDPM) sampling process.
+- **DDIMSampler** (`flaxdiff.samplers.DDIMSampler`): Implements the Denoising Diffusion Implicit Model (DDIM) sampling process.
+- **EulerSampler** (`flaxdiff.samplers.EulerSampler`): An ODE solver sampler using Euler's method.
+- **HeunSampler** (`flaxdiff.samplers.HeunSampler`): An ODE solver sampler using Heun's method.
+- **RK4Sampler** (`flaxdiff.samplers.RK4Sampler`): An ODE solver sampler using the Runge-Kutta method.
+- **MultiStepDPM** (`flaxdiff.samplers.MultiStepDPM`): Implements a multi-step sampling method inspired by the Multistep DPM solver as presented here: [tonyduan/diffusion](https://github.com/tonyduan/diffusion/blob/fcc0ed829baf29e1493b460b073e735a848c08ea/src/samplers.py#L44)) 
+
+### Training
+Implemented in `flaxdiff.trainer`:
+- **DiffusionTrainer** (`flaxdiff.trainer.DiffusionTrainer`): A class designed to facilitate the training of diffusion models. It manages the training loop, loss calculation, and model updates.
+
+### Models
+Implemented in `flaxdiff.models`:
+- **UNet** (`flaxdiff.models.simple_unet.SimpleUNet`): A sample UNET architecture for diffusion models.
+- **Layers**: A library of layers including upsampling (`flaxdiff.models.simple_unet.Upsample`), downsampling (`flaxdiff.models.simple_unet.Downsample`), Time embeddings (`flaxdiff.models.simple_unet.FouriedEmbedding`), attention (`flaxdiff.models.simple_unet.AttentionBlock`), and residual blocks (`flaxdiff.models.simple_unet.ResidualBlock`).
+
+## Installation
+
+To install FlaxDiff, you need to have Python 3.10 or higher. Install the required dependencies using:
+
+```bash
+pip install -r requirements.txt
+```
+
+The models were trained and tested with jax==0.4.28 and flax==0.8.4. However, when I updated to the latest jax==0.4.30 and flax==0.8.5, 
+the models stopped training. There seems to have been some major change breaking the training dynamics and therefore I would recommend
+sticking to the versions mentioned in the requirements.txt
+
+## Getting Started
+
+### Training Example
+
+Here is a simplified example to get you started with training a diffusion model using FlaxDiff:
+
+```python
+from flaxdiff.schedulers import EDMNoiseScheduler
+from flaxdiff.predictors import KarrasPredictionTransform
+from flaxdiff.models.simple_unet import SimpleUNet as UNet
+from flaxdiff.trainer import DiffusionTrainer
+import jax
+import optax
+from datetime import datetime
+
+BATCH_SIZE = 16
+IMAGE_SIZE = 64
+
+# Define noise scheduler
+edm_schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)
+
+# Define model
+unet = UNet(emb_features=256, 
+            feature_depths=[64, 128, 256, 512],
+            attention_configs=[{"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}],
+            num_res_blocks=2,
+            num_middle_res_blocks=1)
+
+# Load dataset
+data, datalen = get_dataset("oxford_flowers102", batch_size=BATCH_SIZE, image_scale=IMAGE_SIZE)
+batches = datalen // BATCH_SIZE
+
+# Define optimizer
+solver = optax.adam(2e-4)
+
+# Create trainer
+trainer = DiffusionTrainer(unet, optimizer=solver, 
+                           noise_schedule=edm_schedule,
+                           rngs=jax.random.PRNGKey(4), 
+                           name="Diffusion_SDE_VE_" + datetime.now().strftime("%Y-%m-%d_%H:%M:%S"),
+                           model_output_transform=KarrasPredictionTransform(sigma_data=edm_schedule.sigma_data))
+
+# Train the model
+final_state = trainer.fit(data, batches, epochs=2000)
+```
+
+### Inference Example
+
+Here is a simplified example for generating images using a trained model:
+
+```python
+from flaxdiff.samplers import DiffusionSampler
+
+class EulerSampler(DiffusionSampler):
+    def take_next_step(self, current_samples, reconstructed_samples, pred_noise, current_step, state, next_step=None):
+        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step)
+        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step)
+        dt = next_sigma - current_sigma
+        x_0_coeff = (current_alpha * next_sigma - next_alpha * current_sigma) / dt
+        dx = (current_samples - x_0_coeff * reconstructed_samples) / current_sigma
+        next_samples = current_samples + dx * dt
+        return next_samples, state
+
+# Create sampler
+sampler = EulerSampler(trainer.model, trainer.state.ema_params, edm_schedule, model_output_transform=trainer.model_output_transform)
+
+# Generate images
+samples = sampler.generate_images(num_images=64, diffusion_steps=100, start_step=1000, end_step=0)
+plotImages(samples, dpi=300)
+```
+
+## Contribution
+
+Feel free to contribute by opening issues or submitting pull requests. Let's make FlaxDiff better together!
+
+## License
+
+This project is licensed under the MIT License.
diff --git a/setup.py b/setup.py
index c54d7f8..e76291e 100644
--- a/setup.py
+++ b/setup.py
@@ -11,6 +11,8 @@
     packages=find_packages(),
     version='0.1.0',
     description='A versatile and easy to understand Diffusion library',
+    long_description=open('README.md').read(),
+    long_description_content_type='text/markdown',
     author='Ashish Kumar Singh',
     author_email='ashishkmr472@gmail.com',
     install_requires=required_packages,