GPU

#@title <-- Tap this if you run on Mobile { display-mode: "form" }
#Taken from KoboldAI colab
%%html
<b>Press play on the audio player to keep the tab alive. (Uses only 13MB of data)</b>

<audio src="https://henk.tech/colabkobold/silence.m4a" controls>

#@markdown (RECOMMENDED) Generates an API key for you to use with the API
secure = False #@param {type:"boolean"}
#@markdown Allows to run SillyTavern Extras on CPU (use if you're out of daily GPU allowance)
use_cpu = False #@param {type:"boolean"}
#@markdown Allows to run Stable Diffusion pipeline on CPU (slow!)
use_sd_cpu = False #@param {type:"boolean"}
#@markdown ***
#@markdown Enables the WebSearch module
extras_enable_websearch = True #@param {type:"boolean"}
#@markdown ***
#@markdown Loads the image captioning module
extras_enable_caption = True #@param {type:"boolean"}
captioning_model = "Salesforce/blip-image-captioning-large" #@param [ "Salesforce/blip-image-captioning-large", "Salesforce/blip-image-captioning-base" ]
#@markdown * Salesforce/blip-image-captioning-large - good base model
#@markdown * Salesforce/blip-image-captioning-base - slightly faster but less accurate
#@markdown ***
#@markdown Loads the sentiment classification model
extras_enable_classify = True #@param {type:"boolean"}
classification_model = "nateraw/bert-base-uncased-emotion" #@param ["nateraw/bert-base-uncased-emotion", "joeddav/distilbert-base-uncased-go-emotions-student"]
#@markdown * nateraw/bert-base-uncased-emotion = 6 supported emotions

#@markdown * joeddav/distilbert-base-uncased-go-emotions-student = 28 supported emotions
#@markdown ***
#@markdown Loads the story summarization module
extras_enable_summarize = True #@param {type:"boolean"}
summarization_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" ]
#@markdown * slauw87/bart_summarisation - general purpose summarization model
#@markdown * Qiliang/bart-large-cnn-samsum-ChatGPT_v3 - summarization model optimized for chats
#@markdown * Qiliang/bart-large-cnn-samsum-ElectrifAi_v10 - nice results so far, but still being evaluated
#@markdown * distilbart-xsum-12-3 - faster, but pretty basic alternative
#@markdown ***
#@markdown Enables Silero text-to-speech module
extras_enable_silero_tts = True #@param {type:"boolean"}
#@markdown Enables Microsoft Edge text-to-speech module
extras_enable_edge_tts = True #@param {type:"boolean"}
#@markdown Enables RVC module
extras_enable_rvc = False #@param {type:"boolean"}
#@markdown ***
#@markdown Enables Whisper speech recognition module
extras_enable_whisper_stt = True #@param {type:"boolean"}
whisper_model = "base.en" #@param [ "tiny.en", "base.en", "small.en", "medium.en", "tiny", "base", "small", "medium", "large" ]
#@markdown There are five model sizes, four with English-only versions, offering speed and accuracy tradeoffs.
#@markdown The .en models for English-only applications tend to perform better, especially for the tiny.en and base.en models.
#@markdown ***
#@markdown Enables SD picture generation
extras_enable_sd = True #@param {type:"boolean"}
sd_model = "ckpt/anything-v4.5-vae-swapped" #@param [ "ckpt/anything-v4.5-vae-swapped", "hakurei/waifu-diffusion", "philz1337/clarity", "prompthero/openjourney", "ckpt/sd15", "stabilityai/stable-diffusion-2-1-base" ]
#@markdown * ckpt/anything-v4.5-vae-swapped - anime style model
#@markdown * hakurei/waifu-diffusion - anime style model
#@markdown * philz1337/clarity - realistic style model
#@markdown * prompthero/openjourney - midjourney style model
#@markdown * ckpt/sd15 - base SD 1.5
#@markdown * stabilityai/stable-diffusion-2-1-base - base SD 2.1
#@markdown ***
#@markdown Enables ChromaDB module
extras_enable_chromadb = True #@param {type:"boolean"}

import subprocess
import secrets

# ---
# SillyTavern extras
extras_url = '(disabled)'
params = []
if use_cpu:
    params.append('--cpu')
if use_sd_cpu:
    params.append('--sd-cpu')
if secure:
    params.append('--secure')
params.append('--share')
modules = []

if extras_enable_caption:
  modules.append('caption')
if extras_enable_summarize:
  modules.append('summarize')
if extras_enable_classify:
  modules.append('classify')
if extras_enable_sd:
  modules.append('sd')
if extras_enable_silero_tts:
  modules.append('silero-tts')
if extras_enable_edge_tts:
  modules.append('edge-tts')
if extras_enable_chromadb:
  modules.append('chromadb')
if extras_enable_whisper_stt:
  modules.append('whisper-stt')
  params.append(f'--stt-whisper-model-path={whisper_model}')
if extras_enable_rvc:
  modules.append('rvc')
  params.append('--max-content-length=2000')
  params.append('--rvc-save-file')


if extras_enable_websearch:
    print("Enabling WebSearch module")
    modules.append('websearch')
    !apt update
    !apt install -y chromium-chromedriver

params.append(f'--classification-model={classification_model}')
params.append(f'--summarization-model={summarization_model}')
params.append(f'--captioning-model={captioning_model}')
params.append(f'--sd-model={sd_model}')
params.append(f'--enable-modules={",".join(modules)}')


%cd /
!git clone https://github.com/SillyTavern/SillyTavern-extras
%cd /SillyTavern-extras
!git clone https://github.com/Cohee1207/tts_samples
!npm install -g localtunnel
%pip install -r requirements.txt
!wget https://github.com/cloudflare/cloudflared/releases/download/2023.5.0/cloudflared-linux-amd64 -O /tmp/cloudflared-linux-amd64
!chmod +x /tmp/cloudflared-linux-amd64

if extras_enable_rvc:
  print("Installing RVC requirements")
  %pip install -r requirements-rvc.txt

# Generate a random API key
api_key = secrets.token_hex(5)

# Write the API key to api_key.txt
with open('./api_key.txt', 'w') as f:
    f.write(api_key)
print(f"API Key generated: {api_key}")

cmd = f"python server.py {' '.join(params)}"
print(cmd)
extras_process = subprocess.Popen(
    cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd='/SillyTavern-extras', shell=True)
print('processId:', extras_process.pid)
while True:
    line = extras_process.stdout.readline().decode().strip()
    if line != None and line != '':
        print(line)