diff --git a/notebooks/comfyui_colab.ipynb b/notebooks/comfyui_colab.ipynb
index d9726947..5108ec83 100644
--- a/notebooks/comfyui_colab.ipynb
+++ b/notebooks/comfyui_colab.ipynb
@@ -1,29 +1,13 @@
 {
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "provenance": []
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    },
-    "language_info": {
-      "name": "python"
-    },
-    "accelerator": "GPU",
-    "gpuClass": "standard"
-  },
   "cells": [
     {
       "cell_type": "markdown",
-      "source": [
-        "Git clone the repo and install the requirements. (ignore the pip errors about protobuf)"
-      ],
       "metadata": {
         "id": "aaaaaaaaaa"
-      }
+      },
+      "source": [
+        "Git clone the repo and install the requirements. (ignore the pip errors about protobuf)"
+      ]
     },
     {
       "cell_type": "code",
@@ -33,22 +17,55 @@
       },
       "outputs": [],
       "source": [
-        "!git clone https://github.com/comfyanonymous/ComfyUI\n",
-        "%cd ComfyUI\n",
-        "!pip install xformers -r requirements.txt"
+        "#@title Environment Setup\n",
+        "\n",
+        "from pathlib import Path\n",
+        "\n",
+        "OPTIONS = {}\n",
+        "\n",
+        "USE_GOOGLE_DRIVE = False  #@param {type:\"boolean\"}\n",
+        "UPDATE_COMFY_UI = True  #@param {type:\"boolean\"}\n",
+        "WORKSPACE = 'ComfyUI'\n",
+        "OPTIONS['USE_GOOGLE_DRIVE'] = USE_GOOGLE_DRIVE\n",
+        "OPTIONS['UPDATE_COMFY_UI'] = UPDATE_COMFY_UI\n",
+        "\n",
+        "if OPTIONS['USE_GOOGLE_DRIVE']:\n",
+        "    !echo \"Mounting Google Drive...\"\n",
+        "    %cd /\n",
+        "    \n",
+        "    from google.colab import drive\n",
+        "    drive.mount('/content/drive')\n",
+        "\n",
+        "    WORKSPACE = \"/content/drive/MyDrive/ComfyUI\"\n",
+        "    %cd /content/drive/MyDrive\n",
+        "\n",
+        "![ ! -d $WORKSPACE ] && echo -= Initial setup ComfyUI =- && git clone https://github.com/comfyanonymous/ComfyUI\n",
+        "%cd $WORKSPACE\n",
+        "\n",
+        "if OPTIONS['UPDATE_COMFY_UI']:\n",
+        "  !echo -= Updating ComfyUI =-\n",
+        "  !git pull\n",
+        "\n",
+        "!echo -= Install dependencies =-\n",
+        "!pip -q install xformers -r requirements.txt"
       ]
     },
     {
       "cell_type": "markdown",
-      "source": [
-        "Download some models/checkpoints/vae or custom comfyui nodes (uncomment the commands for the ones you want)"
-      ],
       "metadata": {
         "id": "cccccccccc"
-      }
+      },
+      "source": [
+        "Download some models/checkpoints/vae or custom comfyui nodes (uncomment the commands for the ones you want)"
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "dddddddddd"
+      },
+      "outputs": [],
       "source": [
         "# Checkpoints\n",
         "\n",
@@ -110,26 +127,26 @@
         "#!wget -c https://huggingface.co/sberbank-ai/Real-ESRGAN/resolve/main/RealESRGAN_x4.pth -P ./models/upscale_models/\n",
         "\n",
         "\n"
-      ],
-      "metadata": {
-        "id": "dddddddddd"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
+      "metadata": {
+        "id": "kkkkkkkkkkkkkk"
+      },
       "source": [
         "### Run ComfyUI with localtunnel (Recommended Way)\n",
         "\n",
         "\n"
-      ],
-      "metadata": {
-        "id": "kkkkkkkkkkkkkk"
-      }
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "jjjjjjjjjjjjj"
+      },
+      "outputs": [],
       "source": [
         "!npm install -g localtunnel\n",
         "\n",
@@ -154,15 +171,13 @@
         "threading.Thread(target=iframe_thread, daemon=True, args=(8188,)).start()\n",
         "\n",
         "!python main.py --dont-print-server"
-      ],
-      "metadata": {
-        "id": "jjjjjjjjjjjjj"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
+      "metadata": {
+        "id": "gggggggggg"
+      },
       "source": [
         "### Run ComfyUI with colab iframe (use only in case the previous way with localtunnel doesn't work)\n",
         "\n",
@@ -171,13 +186,15 @@
         "If you want to open it in another window use the link.\n",
         "\n",
         "Note that some UI features like live image previews won't work because the colab iframe blocks websockets."
-      ],
-      "metadata": {
-        "id": "gggggggggg"
-      }
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "hhhhhhhhhh"
+      },
+      "outputs": [],
       "source": [
         "import threading\n",
         "import time\n",
@@ -198,12 +215,23 @@
         "threading.Thread(target=iframe_thread, daemon=True, args=(8188,)).start()\n",
         "\n",
         "!python main.py --dont-print-server"
-      ],
-      "metadata": {
-        "id": "hhhhhhhhhh"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     }
-  ]
-}
\ No newline at end of file
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "provenance": []
+    },
+    "gpuClass": "standard",
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}