{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Links**<br>\n",
        "Extensions API GitHub: https://github.com/Cohee1207/SillyTavern-extras/<br>\n",
        "SillyTavern community Discord (support and discussion): https://discord.gg/RZdyAEUPvj"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "lVftocpwCoYw"
      },
      "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",
        "#@markdown * Salesforce/blip-image-captioning-large - good base model\n",
        "#@markdown * Salesforce/blip-image-captioning-base - slightly faster but less accurate\n",
        "extras_enable_emotions = True #@param {type:\"boolean\"}\n",
        "#@markdown Loads the sentiment classification model\n",
        "Emotions_Model = \"nateraw/bert-base-uncased-emotion\" #@param [\"nateraw/bert-base-uncased-emotion\", \"joeddav/distilbert-base-uncased-go-emotions-student\"]\n",
        "#@markdown * nateraw/bert-base-uncased-emotion = 6 supported emotions<br>\n",
        "#@markdown * joeddav/distilbert-base-uncased-go-emotions-student = 28 supported emotions\n",
        "extras_enable_memory = True #@param {type:\"boolean\"}\n",
        "#@markdown Loads the story summarization module\n",
        "Memory_Model = \"slauw87/bart_summarisation\" #@param [ \"slauw87/bart_summarisation\", \"Qiliang/bart-large-cnn-samsum-ChatGPT_v3\", \"Qiliang/bart-large-cnn-samsum-ElectrifAi_v10\", \"distilbart-xsum-12-3\" ]\n",
        "#@markdown * slauw87/bart_summarisation - general purpose summarization model\n",
        "#@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\", \"philz1337/clarity\", \"ckpt/sd15\" ]\n",
        "#@markdown * ckpt/anything-v4.5-vae-swapped - anime style model\n",
        "#@markdown * philz1337/clarity - realistic style model\n",
        "#@markdown * ckpt/sd15 - base SD 1.5\n",
        "\n",
        "import subprocess\n",
        "\n",
        "# ---\n",
        "# SillyTavern extras\n",
        "extras_url = '(disabled)'\n",
        "params = []\n",
        "if use_cpu:\n",
        "    params.append('--cpu')\n",
        "params.append('--share')\n",
        "ExtrasModules = []\n",
        "\n",
        "if (extras_enable_captioning):\n",
        "  ExtrasModules.append('caption')\n",
        "if (extras_enable_memory):\n",
        "  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-model={SD_Model}')\n",
        "params.append(f'--enable-modules={\",\".join(ExtrasModules)}')\n",
        "\n",
        "\n",
        "%cd /\n",
        "!git clone https://github.com/Cohee1207/SillyTavern-extras\n",
        "%cd /SillyTavern-extras\n",
        "!npm install -g localtunnel\n",
        "!pip install -r requirements-complete.txt\n",
        "!pip install tensorflow==2.11\n",
        "\n",
        "\n",
        "cmd = f\"python server.py {' '.join(params)}\"\n",
        "print(cmd)\n",
        "extras_process = subprocess.Popen(\n",
        "    cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd='/SillyTavern-extras', shell=True)\n",
        "print('processId:', extras_process.pid)\n",
        "while True:\n",
        "    line = extras_process.stdout.readline().decode().strip()\n",
        "    if line != None and line != '':\n",
        "        print(line)\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "private_outputs": true,
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}