From ead53164a86f371051355e3cc87b32feaac5e040 Mon Sep 17 00:00:00 2001
From: SillyLossy <titan0599@gmail.com>
Date: Fri, 12 May 2023 00:46:15 +0300
Subject: [PATCH] Add SD and TTS to colab.

---
 colab/GPU.ipynb | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/colab/GPU.ipynb b/colab/GPU.ipynb
index 74bbee704..c31548bf9 100644
--- a/colab/GPU.ipynb
+++ b/colab/GPU.ipynb
@@ -20,6 +20,7 @@
       "outputs": [],
       "source": [
         "#@markdown Enables hosting of extensions backend for SillyTavern Extras\n",
+        "use_cpu = False #@param {type:\"boolean\"}\n",
         "extras_enable_captioning = True #@param {type:\"boolean\"}\n",
         "#@markdown Loads the image captioning module\n",
         "Captions_Model = \"Salesforce/blip-image-captioning-large\" #@param [ \"Salesforce/blip-image-captioning-large\", \"Salesforce/blip-image-captioning-base\" ]\n",
@@ -37,6 +38,14 @@
         "#@markdown * Qiliang/bart-large-cnn-samsum-ChatGPT_v3 - summarization model optimized for chats\n",
         "#@markdown * Qiliang/bart-large-cnn-samsum-ElectrifAi_v10 - nice results so far, but still being evaluated\n",
         "#@markdown * distilbart-xsum-12-3 - faster, but pretty basic alternative\n",
+        "extras_enable_tts = True #@param {type:\"boolean\"}\n",
+        "#@markdown Enables Silero text-to-speech module\n",
+        "extras_enable_sd = True #@param {type:\"boolean\"}\n",
+        "#@markdown Enables SD picture generation\n",
+        "SD_Model = \"ckpt/anything-v4.5-vae-swapped\" #@param [ \"ckpt/anything-v4.5-vae-swapped\", \"ckpt/majicmix-realistic\", \"ckpt/sd15\" ]\n",
+        "#@markdown * ckpt/anything-v4.5-vae-swapped - anime style model\n",
+        "#@markdown * ckpt/majicmix-realistic - realistic style model\n",
+        "#@markdown * ckpt/sd15 - base SD 1.5\n",
         "\n",
         "import subprocess\n",
         "\n",
@@ -44,7 +53,8 @@
         "# SillyTavern extras\n",
         "extras_url = '(disabled)'\n",
         "params = []\n",
-        "params.append('--cpu')\n",
+        "if use_cpu:\n",
+        "    params.append('--cpu')\n",
         "params.append('--share')\n",
         "ExtrasModules = []\n",
         "\n",
@@ -54,10 +64,16 @@
         "  ExtrasModules.append('summarize')\n",
         "if (extras_enable_emotions):\n",
         "  ExtrasModules.append('classify')\n",
+        "if (extras_enable_sd):\n",
+        "  ExtrasModules.append('sd')\n",
+        "if (extras_enable_tts):\n",
+        "  ExtrasModules.append('tts')\n",
         "\n",
         "params.append(f'--classification-model={Emotions_Model}')\n",
         "params.append(f'--summarization-model={Memory_Model}')\n",
         "params.append(f'--captioning-model={Captions_Model}')\n",
+        "params.append(f'--sd-local')\n",
+        "params.append(f'--sd-model={SD_Model}')\n",
         "params.append(f'--enable-modules={\",\".join(ExtrasModules)}')\n",
         "\n",
         "\n",
@@ -65,7 +81,7 @@
         "!git clone https://github.com/Cohee1207/SillyTavern-extras\n",
         "%cd /SillyTavern-extras\n",
         "!npm install -g localtunnel\n",
-        "!pip install -r requirements.txt\n",
+        "!pip install -r requirements-complete.txt\n",
         "!pip install tensorflow==2.11\n",
         "\n",
         "\n",