Properly concatenate chat events

Simplify a function
SD Api Pics extension, v.1.1 (#596 )
2023-04-08 17:16:46 -03:00 · 2023-04-07 21:37:41 -03:00 · 2023-04-07 21:36:04 -03:00 · 2023-04-07 15:57:51 -03:00 · 2023-04-07 14:42:03 -03:00 · 2023-04-07 12:22:14 -03:00
38 changed files with 1073 additions and 450 deletions
--- a/.dockerignore
+++ b/.dockerignore
@@ -1,7 +1,6 @@
 .env
 Dockerfile
 /characters
-/extensions
 /loras
 /models
 /presets
--- a/19
+++ b/19
@@ -30,8 +30,7 @@ RUN apt-get update && \
    rm -rf /var/lib/apt/lists/*

 RUN --mount=type=cache,target=/root/.cache/pip pip3 install virtualenv
-
-COPY . /app/
+RUN mkdir /app

 WORKDIR /app

@@ -41,21 +40,29 @@ RUN test -n "${WEBUI_VERSION}" && git reset --hard ${WEBUI_VERSION} || echo "Usi
 RUN virtualenv /app/venv
 RUN . /app/venv/bin/activate && \
    pip3 install --upgrade pip setuptools && \
-    pip3 install torch torchvision torchaudio && \
-    pip3 install -r requirements.txt
+    pip3 install torch torchvision torchaudio

 COPY --from=builder /build /app/repositories/GPTQ-for-LLaMa
 RUN . /app/venv/bin/activate && \
    pip3 install /app/repositories/GPTQ-for-LLaMa/*.whl

-ENV CLI_ARGS=""
-
+COPY extensions/api/requirements.txt /app/extensions/api/requirements.txt
+COPY extensions/elevenlabs_tts/requirements.txt /app/extensions/elevenlabs_tts/requirements.txt
+COPY extensions/google_translate/requirements.txt /app/extensions/google_translate/requirements.txt
+COPY extensions/silero_tts/requirements.txt /app/extensions/silero_tts/requirements.txt
+COPY extensions/whisper_stt/requirements.txt /app/extensions/whisper_stt/requirements.txt
 RUN --mount=type=cache,target=/root/.cache/pip . /app/venv/bin/activate && cd extensions/api && pip3 install -r requirements.txt
 RUN --mount=type=cache,target=/root/.cache/pip . /app/venv/bin/activate && cd extensions/elevenlabs_tts && pip3 install -r requirements.txt
 RUN --mount=type=cache,target=/root/.cache/pip . /app/venv/bin/activate && cd extensions/google_translate && pip3 install -r requirements.txt
 RUN --mount=type=cache,target=/root/.cache/pip . /app/venv/bin/activate && cd extensions/silero_tts && pip3 install -r requirements.txt
 RUN --mount=type=cache,target=/root/.cache/pip . /app/venv/bin/activate && cd extensions/whisper_stt && pip3 install -r requirements.txt

+COPY requirements.txt /app/requirements.txt
+RUN . /app/venv/bin/activate && \
+    pip3 install -r requirements.txt
+
 RUN cp /app/venv/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda118.so /app/venv/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cpu.so

+COPY . /app/
+ENV CLI_ARGS=""
 CMD . /app/venv/bin/activate && python3 server.py ${CLI_ARGS}
--- a/README.md
+++ b/README.md
@@ -119,7 +119,7 @@ As an alternative to the recommended WSL method, you can install the web UI nati

 ```
 cp .env.example .env
-docker-compose up --build
+docker compose up --build
 ```

 Make sure to edit `.env.example` and set the appropriate CUDA version for your GPU.
@@ -192,14 +192,14 @@ Optionally, you can use the following command-line flags:
 #### Basic settings

 | Flag                                       | Description |
-|------------------|-------------|
-| `-h`, `--help`   | show this help message and exit |
+|--------------------------------------------|-------------|
+| `-h`, `--help`                             | Show this help message and exit. |
 | `--notebook`                               | Launch the web UI in notebook mode, where the output is written to the same text box as the input. |
 | `--chat`                                   | Launch the web UI in chat mode. |
 | `--model MODEL`                            | Name of the model to load by default. |
 | `--lora LORA`                              | Name of the LoRA to apply to the model by default. |
-|  `--model-dir MODEL_DIR`                    | Path to directory with all the models |
-|  `--lora-dir LORA_DIR`                      | Path to directory with all the loras |
+| `--model-dir MODEL_DIR`                    | Path to directory with all the models. |
+| `--lora-dir LORA_DIR`                      | Path to directory with all the loras. |
 | `--no-stream`                              | Don't stream the text output in real time. |
 | `--settings SETTINGS_FILE`                 | Load the default interface settings from this json file. See `settings-template.json` for an example. If you create a file called `settings.json`, this file will be loaded by default without the need to use the `--settings` flag. |
 | `--extensions EXTENSIONS [EXTENSIONS ...]` | The list of extensions to load. If you want to load more than one extension, write the names separated by spaces. |
@@ -208,7 +208,7 @@ Optionally, you can use the following command-line flags:
 #### Accelerate/transformers

 | Flag                                        | Description |
-|------------------|-------------|
+|---------------------------------------------|-------------|
 | `--cpu`                                     | Use the CPU to generate text. |
 | `--auto-devices`                            | Automatically split the model across the available GPU(s) and CPU. |
 |  `--gpu-memory GPU_MEMORY [GPU_MEMORY ...]` | Maxmimum GPU memory in GiB to be allocated per GPU. Example: `--gpu-memory 10` for a single GPU, `--gpu-memory 10 5` for two GPUs. You can also set values in MiB like `--gpu-memory 3500MiB`. |
@@ -222,13 +222,13 @@ Optionally, you can use the following command-line flags:
 #### llama.cpp

 | Flag        | Description |
-|------------------|-------------|
+|-------------|-------------|
 | `--threads` | Number of threads to use in llama.cpp. |

 #### GPTQ

 | Flag                      | Description |
-|------------------|-------------|
+|---------------------------|-------------|
 | `--wbits WBITS`           | GPTQ: Load a pre-quantized model with specified precision in bits. 2, 3, 4 and 8 are supported. |
 | `--model_type MODEL_TYPE` | GPTQ: Model type of pre-quantized model. Currently LLaMA, OPT, and GPT-J are supported. |
 | `--groupsize GROUPSIZE`   | GPTQ: Group size. |
@@ -246,7 +246,7 @@ Optionally, you can use the following command-line flags:
 #### DeepSpeed

 | Flag                                  | Description |
-|------------------|-------------|
+|---------------------------------------|-------------|
 | `--deepspeed`                         | Enable the use of DeepSpeed ZeRO-3 for inference via the Transformers integration. |
 | `--nvme-offload-dir NVME_OFFLOAD_DIR` | DeepSpeed: Directory to use for ZeRO-3 NVME offloading. |
 | `--local_rank LOCAL_RANK`             | DeepSpeed: Optional argument for distributed setups. |
@@ -254,14 +254,14 @@ Optionally, you can use the following command-line flags:
 #### RWKV

 | Flag                            | Description |
-|------------------|-------------|
+|---------------------------------|-------------|
 | `--rwkv-strategy RWKV_STRATEGY` | RWKV: The strategy to use while loading the model. Examples: "cpu fp32", "cuda fp16", "cuda fp16i8". |
 | `--rwkv-cuda-on`                | RWKV: Compile the CUDA kernel for better performance. |

 #### Gradio

 | Flag                                  | Description |
-|------------------|-------------|
+|---------------------------------------|-------------|
 | `--listen`                            | Make the web UI reachable from your local network. |
 | `--listen-port LISTEN_PORT`           | The listening port that the server will use. |
 | `--share`                             | Create a public URL. This is useful for running the web UI on Google Colab or similar. |
--- a/api-example-stream.py
+++ b/api-example-stream.py
@@ -17,6 +17,7 @@ def random_hash():
    letters = string.ascii_lowercase + string.digits
    return ''.join(random.choice(letters) for i in range(9))

+
 async def run(context):
    server = "127.0.0.1"
    params = {
@@ -69,6 +70,7 @@ async def run(context):

 prompt = "What I would like to say is the following: "

+
 async def get_result():
    async for response in run(prompt):
        # Print intermediate steps
--- a/convert-to-flexgen.py
+++ b/convert-to-flexgen.py
@@ -17,6 +17,7 @@ parser = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpForma
 parser.add_argument('MODEL', type=str, default=None, nargs='?', help="Path to the input model.")
 args = parser.parse_args()

+
 def disable_torch_init():
    """
    Disable the redundant torch default initialization to accelerate model creation.
@@ -31,12 +32,14 @@ def disable_torch_init():
    torch_layer_norm_init_backup = torch.nn.LayerNorm.reset_parameters
    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)

+
 def restore_torch_init():
    """Rollback the change made by disable_torch_init."""
    import torch
    setattr(torch.nn.Linear, "reset_parameters", torch_linear_init_backup)
    setattr(torch.nn.LayerNorm, "reset_parameters", torch_layer_norm_init_backup)

+
 if __name__ == '__main__':
    path = Path(args.MODEL)
    model_name = path.name
--- a/css/chat.css
+++ b/css/chat.css
@@ -36,3 +36,8 @@ div.svelte-362y77>*, div.svelte-362y77>.form>* {
 .wrap.svelte-6roggh.svelte-6roggh {
  max-height: 92.5%;
 }
+
+/* This is for the microphone button in the whisper extension */
+.sm.svelte-1ipelgc {
+  width: 100%;
+}
--- a/download-model.py
+++ b/download-model.py
@@ -29,6 +29,7 @@ parser.add_argument('--clean', action='store_true', help='Does not resume the pr
 parser.add_argument('--check', action='store_true', help='Validates the checksums of model files.')
 args = parser.parse_args()

+
 def get_file(url, output_folder):
    filename = Path(url.rsplit('/', 1)[1])
    output_path = output_folder / filename
@@ -54,6 +55,7 @@ def get_file(url, output_folder):
                t.update(len(data))
                f.write(data)

+
 def sanitize_branch_name(branch_name):
    pattern = re.compile(r"^[a-zA-Z0-9._-]+$")
    if pattern.match(branch_name):
@@ -61,6 +63,7 @@ def sanitize_branch_name(branch_name):
    else:
        raise ValueError("Invalid branch name. Only alphanumeric characters, period, underscore and dash are allowed.")

+
 def select_model_from_default_options():
    models = {
        "OPT 6.7B": ("facebook", "opt-6.7b", "main"),
@@ -106,6 +109,7 @@ EleutherAI/pythia-1.4b-deduped

    return model, branch

+
 def get_download_links_from_huggingface(model, branch):
    base = "https://huggingface.co"
    page = f"/api/models/{model}/tree/{branch}?cursor="
@@ -172,9 +176,11 @@ def get_download_links_from_huggingface(model, branch):

    return links, sha256, is_lora

+
 def download_files(file_list, output_folder, num_threads=8):
    thread_map(lambda url: get_file(url, output_folder), file_list, max_workers=num_threads, disable=True)

+
 if __name__ == '__main__':
    model = args.MODEL
    branch = args.branch
--- a/extensions/api/script.py
+++ b/extensions/api/script.py
@@ -9,6 +9,7 @@ params = {
    'port': 5000,
 }

+
 class Handler(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path == '/api/v1/model':
@@ -32,7 +33,7 @@ class Handler(BaseHTTPRequestHandler):
            self.end_headers()

            prompt = body['prompt']
-            prompt_lines = [l.strip() for l in prompt.split('\n')]
+            prompt_lines = [k.strip() for k in prompt.split('\n')]

            max_context = body.get('max_context_length', 2048)

@@ -95,5 +96,6 @@ def run_server():
        print(f'Starting KoboldAI compatible api at http://{server_addr[0]}:{server_addr[1]}/api')
    server.serve_forever()

+
 def setup():
    Thread(target=run_server, daemon=True).start()
--- a/extensions/character_bias/script.py
+++ b/extensions/character_bias/script.py
@@ -5,6 +5,7 @@ params = {
    "bias string": " *I am so happy*",
 }

+
 def input_modifier(string):
    """
    This function is applied to your text inputs before
@@ -13,6 +14,7 @@ def input_modifier(string):

    return string

+
 def output_modifier(string):
    """
    This function is applied to the model outputs.
@@ -20,6 +22,7 @@ def output_modifier(string):

    return string

+
 def bot_prefix_modifier(string):
    """
    This function is only applied in chat mode. It modifies
@@ -27,11 +30,12 @@ def bot_prefix_modifier(string):
    behavior.
    """

-    if params['activate'] == True:
+    if params['activate']:
        return f'{string} {params["bias string"].strip()} '
    else:
        return string

+
 def ui():
    # Gradio elements
    activate = gr.Checkbox(value=params['activate'], label='Activate character bias')
--- a/extensions/elevenlabs_tts/script.py
+++ b/extensions/elevenlabs_tts/script.py
@@ -2,10 +2,11 @@ import re
 from pathlib import Path

 import gradio as gr
-import modules.shared as shared
 from elevenlabslib import ElevenLabsUser
 from elevenlabslib.helpers import save_bytes_to_path

+import modules.shared as shared
+
 params = {
    'activate': True,
    'api_key': '12345',
@@ -22,6 +23,8 @@ if not shared.args.no_stream:
    raise ValueError

 # Check if the API is valid and refresh the UI accordingly.
+
+
 def check_valid_api():

    global user, user_info, params
@@ -29,7 +32,7 @@ def check_valid_api():
    user = ElevenLabsUser(params['api_key'])
    user_info = user._get_subscription_data()
    print('checking api')
-    if params['activate'] == False:
+    if not params['activate']:
        return gr.update(value='Disconnected')
    elif user_info is None:
        print('Incorrect API Key')
@@ -39,6 +42,8 @@ def check_valid_api():
        return gr.update(value='Connected')

 # Once the API is verified, get the available voices and update the dropdown list
+
+
 def refresh_voices():

    global user, user_info
@@ -51,11 +56,13 @@ def refresh_voices():
    else:
        return

+
 def remove_surrounded_chars(string):
    # this expression matches to 'as few symbols as possible (0 upwards) between any asterisks' OR
    # 'as few symbols as possible (0 upwards) between an asterisk and the end of the string'
    return re.sub('\*[^\*]*?(\*|$)', '', string)

+
 def input_modifier(string):
    """
    This function is applied to your text inputs before
@@ -64,6 +71,7 @@ def input_modifier(string):

    return string

+
 def output_modifier(string):
    """
    This function is applied to the model outputs.
@@ -71,9 +79,9 @@ def output_modifier(string):

    global params, wav_idx, user, user_info

-    if params['activate'] == False:
+    if not params['activate']:
        return string
-    elif user_info == None:
+    elif user_info is None:
        return string

    string = remove_surrounded_chars(string)
@@ -94,6 +102,7 @@ def output_modifier(string):
    wav_idx += 1
    return string

+
 def ui():

    # Gradio elements
--- a/extensions/google_translate/script.py
+++ b/extensions/google_translate/script.py
@@ -7,6 +7,7 @@ params = {

 language_codes = {'Afrikaans': 'af', 'Albanian': 'sq', 'Amharic': 'am', 'Arabic': 'ar', 'Armenian': 'hy', 'Azerbaijani': 'az', 'Basque': 'eu', 'Belarusian': 'be', 'Bengali': 'bn', 'Bosnian': 'bs', 'Bulgarian': 'bg', 'Catalan': 'ca', 'Cebuano': 'ceb', 'Chinese (Simplified)': 'zh-CN', 'Chinese (Traditional)': 'zh-TW', 'Corsican': 'co', 'Croatian': 'hr', 'Czech': 'cs', 'Danish': 'da', 'Dutch': 'nl', 'English': 'en', 'Esperanto': 'eo', 'Estonian': 'et', 'Finnish': 'fi', 'French': 'fr', 'Frisian': 'fy', 'Galician': 'gl', 'Georgian': 'ka', 'German': 'de', 'Greek': 'el', 'Gujarati': 'gu', 'Haitian Creole': 'ht', 'Hausa': 'ha', 'Hawaiian': 'haw', 'Hebrew': 'iw', 'Hindi': 'hi', 'Hmong': 'hmn', 'Hungarian': 'hu', 'Icelandic': 'is', 'Igbo': 'ig', 'Indonesian': 'id', 'Irish': 'ga', 'Italian': 'it', 'Japanese': 'ja', 'Javanese': 'jw', 'Kannada': 'kn', 'Kazakh': 'kk', 'Khmer': 'km', 'Korean': 'ko', 'Kurdish': 'ku', 'Kyrgyz': 'ky', 'Lao': 'lo', 'Latin': 'la', 'Latvian': 'lv', 'Lithuanian': 'lt', 'Luxembourgish': 'lb', 'Macedonian': 'mk', 'Malagasy': 'mg', 'Malay': 'ms', 'Malayalam': 'ml', 'Maltese': 'mt', 'Maori': 'mi', 'Marathi': 'mr', 'Mongolian': 'mn', 'Myanmar (Burmese)': 'my', 'Nepali': 'ne', 'Norwegian': 'no', 'Nyanja (Chichewa)': 'ny', 'Pashto': 'ps', 'Persian': 'fa', 'Polish': 'pl', 'Portuguese (Portugal, Brazil)': 'pt', 'Punjabi': 'pa', 'Romanian': 'ro', 'Russian': 'ru', 'Samoan': 'sm', 'Scots Gaelic': 'gd', 'Serbian': 'sr', 'Sesotho': 'st', 'Shona': 'sn', 'Sindhi': 'sd', 'Sinhala (Sinhalese)': 'si', 'Slovak': 'sk', 'Slovenian': 'sl', 'Somali': 'so', 'Spanish': 'es', 'Sundanese': 'su', 'Swahili': 'sw', 'Swedish': 'sv', 'Tagalog (Filipino)': 'tl', 'Tajik': 'tg', 'Tamil': 'ta', 'Telugu': 'te', 'Thai': 'th', 'Turkish': 'tr', 'Ukrainian': 'uk', 'Urdu': 'ur', 'Uzbek': 'uz', 'Vietnamese': 'vi', 'Welsh': 'cy', 'Xhosa': 'xh', 'Yiddish': 'yi', 'Yoruba': 'yo', 'Zulu': 'zu'}

+
 def input_modifier(string):
    """
    This function is applied to your text inputs before
@@ -15,6 +16,7 @@ def input_modifier(string):

    return GoogleTranslator(source=params['language string'], target='en').translate(string)

+
 def output_modifier(string):
    """
    This function is applied to the model outputs.
@@ -22,6 +24,7 @@ def output_modifier(string):

    return GoogleTranslator(source='en', target=params['language string']).translate(string)

+
 def bot_prefix_modifier(string):
    """
    This function is only applied in chat mode. It modifies
@@ -31,6 +34,7 @@ def bot_prefix_modifier(string):

    return string

+
 def ui():
    # Finding the language name from the language code to use as the default value
    language_name = list(language_codes.keys())[list(language_codes.values()).index(params['language string'])]
--- a/extensions/llama_prompts/script.py
+++ b/extensions/llama_prompts/script.py
@@ -1,15 +1,18 @@
 import gradio as gr
-import modules.shared as shared
 import pandas as pd

+import modules.shared as shared
+
 df = pd.read_csv("https://raw.githubusercontent.com/devbrones/llama-prompts/main/prompts/prompts.csv")

+
 def get_prompt_by_name(name):
    if name == 'None':
        return ''
    else:
        return df[df['Prompt name'] == name].iloc[0]['Prompt'].replace('\\n', '\n')

+
 def ui():
    if not shared.is_chat():
        choices = ['None'] + list(df['Prompt name'])
--- a/extensions/sd_api_pictures/README.MD
+++ b/extensions/sd_api_pictures/README.MD
@@ -0,0 +1,78 @@
+## Description:
+TL;DR: Lets the bot answer you with a picture!  
+
+Stable Diffusion API pictures for TextGen, v.1.1.0  
+An extension to [oobabooga's textgen-webui](https://github.com/oobabooga/text-generation-webui) allowing you to receive pics generated by [Automatic1111's SD-WebUI API](https://github.com/AUTOMATIC1111/stable-diffusion-webui)
+
+<details>
+<summary>Interface overview</summary>
+
+![Interface](https://raw.githubusercontent.com/Brawlence/texgen-webui-SD_api_pics/main/illust/Interface.jpg)
+
+</details>
+
+Load it in the `--chat` mode with `--extension sd_api_pictures` alongside `send_pictures` (it's not really required, but completes the picture, *pun intended*).  
+
+The image generation is triggered either:  
+- manually through the 'Force the picture response' button while in `Manual` or `Immersive/Interactive` modes OR  
+- automatically in `Immersive/Interactive` mode if the words `'send|main|message|me'` are followed by `'image|pic|picture|photo|snap|snapshot|selfie|meme'` in the user's prompt  
+- always on in Picturebook/Adventure mode (if not currently suppressed by 'Suppress the picture response')  
+
+## Prerequisites
+
+One needs an available instance of Automatic1111's webui running with an `--api` flag. Ain't tested with a notebook / cloud hosted one but should be possible.   
+To run it locally in parallel on the same machine, specify custom `--listen-port` for either Auto1111's or ooba's webUIs.  
+
+## Features:
+- API detection (press enter in the API box)  
+- VRAM management (model shuffling)  
+- Three different operation modes (manual, interactive, always-on)  
+- persistent settings via settings.json
+
+The model input is modified only in the interactive mode; other two are unaffected. The output pic description is presented differently for Picture-book / Adventure mode.  
+
+Connection check (insert the Auto1111's address and press Enter):  
+![API-check](https://raw.githubusercontent.com/Brawlence/texgen-webui-SD_api_pics/main/illust/API-check.gif) 
+
+### Persistents settings
+
+Create or modify the `settings.json` in the `text-generation-webui` root directory to override the defaults
+present in script.py, ex:
+
+```json
+{
+    "sd_api_pictures-manage_VRAM": 1,
+    "sd_api_pictures-save_img": 1,
+    "sd_api_pictures-prompt_prefix": "(Masterpiece:1.1), detailed, intricate, colorful, (solo:1.1)",
+    "sd_api_pictures-sampler_name": "DPM++ 2M Karras"
+}
+```
+
+will automatically set the `Manage VRAM` & `Keep original images` checkboxes and change the texts in `Prompt Prefix` and `Sampler name` on load.
+
+---
+
+## Demonstrations:
+
+Those are examples of the version 1.0.0, but the core functionality is still the same
+
+<details>
+<summary>Conversation 1</summary>
+
+![EXA1](https://user-images.githubusercontent.com/42910943/224866564-939a3bcb-e7cf-4ac0-a33f-b3047b55054d.jpg)
+![EXA2](https://user-images.githubusercontent.com/42910943/224866566-38394054-1320-45cf-9515-afa76d9d7745.jpg)
+![EXA3](https://user-images.githubusercontent.com/42910943/224866568-10ea47b7-0bac-4269-9ec9-22c387a13b59.jpg)
+![EXA4](https://user-images.githubusercontent.com/42910943/224866569-326121ad-1ea1-4874-9f6b-4bca7930a263.jpg)
+
+
+</details>
+
+<details>
+<summary>Conversation 2</summary>
+
+![Hist1](https://user-images.githubusercontent.com/42910943/224865517-c6966b58-bc4d-4353-aab9-6eb97778d7bf.jpg)
+![Hist2](https://user-images.githubusercontent.com/42910943/224865527-b2fe7c2e-0da5-4c2e-b705-42e233b07084.jpg)
+![Hist3](https://user-images.githubusercontent.com/42910943/224865535-a38d94e7-8975-4a46-a655-1ae1de41f85d.jpg)
+
+</details>
+
--- a/extensions/sd_api_pictures/script.py
+++ b/extensions/sd_api_pictures/script.py
@@ -1,93 +1,151 @@
 import base64
 import io
 import re
+import time
+from datetime import date
 from pathlib import Path

 import gradio as gr
-import modules.chat as chat
 import modules.shared as shared
 import requests
 import torch
+from modules.models import reload_model, unload_model
 from PIL import Image

 torch._C._jit_set_profiling_mode(False)

 # parameters which can be customized in settings.json of webui
 params = {
-    'enable_SD_api': False,
    'address': 'http://127.0.0.1:7860',
+    'mode': 0,  # modes of operation: 0 (Manual only), 1 (Immersive/Interactive - looks for words to trigger), 2 (Picturebook Adventure - Always on)
+    'manage_VRAM': False,
    'save_img': False,
-    'SD_model': 'NeverEndingDream', # not really used right now
-    'prompt_prefix': '(Masterpiece:1.1), (solo:1.3), detailed, intricate, colorful',
+    'SD_model': 'NeverEndingDream',  # not used right now
+    'prompt_prefix': '(Masterpiece:1.1), detailed, intricate, colorful',
    'negative_prompt': '(worst quality, low quality:1.3)',
-    'side_length': 512,
-    'restore_faces': False
+    'width': 512,
+    'height': 512,
+    'restore_faces': False,
+    'seed': -1,
+    'sampler_name': 'DDIM',
+    'steps': 32,
+    'cfg_scale': 7
 }

+
+def give_VRAM_priority(actor):
+    global shared, params
+
+    if actor == 'SD':
+        unload_model()
+        print("Requesting Auto1111 to re-load last checkpoint used...")
+        response = requests.post(url=f'{params["address"]}/sdapi/v1/reload-checkpoint', json='')
+        response.raise_for_status()
+
+    elif actor == 'LLM':
+        print("Requesting Auto1111 to vacate VRAM...")
+        response = requests.post(url=f'{params["address"]}/sdapi/v1/unload-checkpoint', json='')
+        response.raise_for_status()
+        reload_model()
+
+    elif actor == 'set':
+        print("VRAM mangement activated -- requesting Auto1111 to vacate VRAM...")
+        response = requests.post(url=f'{params["address"]}/sdapi/v1/unload-checkpoint', json='')
+        response.raise_for_status()
+
+    elif actor == 'reset':
+        print("VRAM mangement deactivated -- requesting Auto1111 to reload checkpoint")
+        response = requests.post(url=f'{params["address"]}/sdapi/v1/reload-checkpoint', json='')
+        response.raise_for_status()
+
+    else:
+        raise RuntimeError(f'Managing VRAM: "{actor}" is not a known state!')
+
+    response.raise_for_status()
+    del response
+
+
+if params['manage_VRAM']:
+    give_VRAM_priority('set')
+
+samplers = ['DDIM', 'DPM++ 2M Karras']  # TODO: get the availible samplers with http://{address}}/sdapi/v1/samplers
 SD_models = ['NeverEndingDream']  # TODO: get with http://{address}}/sdapi/v1/sd-models and allow user to select

 streaming_state = shared.args.no_stream  # remember if chat streaming was enabled
 picture_response = False  # specifies if the next model response should appear as a picture
-pic_id = 0
+

 def remove_surrounded_chars(string):
    # this expression matches to 'as few symbols as possible (0 upwards) between any asterisks' OR
    # 'as few symbols as possible (0 upwards) between an asterisk and the end of the string'
    return re.sub('\*[^\*]*?(\*|$)', '', string)

-# I don't even need input_hijack for this as visible text will be commited to history as the unmodified string
+
+def triggers_are_in(string):
+    string = remove_surrounded_chars(string)
+    # regex searches for send|main|message|me (at the end of the word) followed by
+    # a whole word of image|pic|picture|photo|snap|snapshot|selfie|meme(s),
+    # (?aims) are regex parser flags
+    return bool(re.search('(?aims)(send|mail|message|me)\\b.+?\\b(image|pic(ture)?|photo|snap(shot)?|selfie|meme)s?\\b', string))
+
+
 def input_modifier(string):
    """
    This function is applied to your text inputs before
    they are fed into the model.
    """
-    global params, picture_response
-    if not params['enable_SD_api']:
+
+    global params
+
+    if not params['mode'] == 1:  # if not in immersive/interactive mode, do nothing
        return string

-    commands = ['send', 'mail', 'me']
-    mediums = ['image', 'pic', 'picture', 'photo']
-    subjects = ['yourself', 'own']
-    lowstr = string.lower()
-
-    # TODO: refactor out to separate handler and also replace detection with a regexp
-    if any(command in lowstr for command in commands) and any(case in lowstr for case in mediums): # trigger the generation if a command signature and a medium signature is found
-        picture_response = True
-        shared.args.no_stream = True                                                               # Disable streaming cause otherwise the SD-generated picture would return as a dud
-        shared.processing_message = "*Is sending a picture...*"
-        string = "Please provide a detailed description of your surroundings, how you look and the situation you're in and what you are doing right now"
-        if any(target in lowstr for target in subjects):                                           # the focus of the image should be on the sending character
-            string = "Please provide a detailed and vivid description of how you look and what you are wearing"
+    if triggers_are_in(string):  # if we're in it, check for trigger words
+        toggle_generation(True)
+        string = string.lower()
+        if "of" in string:
+            subject = string.split('of', 1)[1]  # subdivide the string once by the first 'of' instance and get what's coming after it
+            string = "Please provide a detailed and vivid description of " + subject
+        else:
+            string = "Please provide a detailed description of your appearance, your surroundings and what you are doing right now"

    return string

 # Get and save the Stable Diffusion-generated picture
 def get_SD_pictures(description):

-    global params, pic_id
+    global params
+
+    if params['manage_VRAM']:
+        give_VRAM_priority('SD')

    payload = {
        "prompt": params['prompt_prefix'] + description,
-        "seed": -1,
-        "sampler_name": "DPM++ 2M Karras",
-        "steps": 32,
-        "cfg_scale": 7,
-        "width": params['side_length'],
-        "height": params['side_length'],
+        "seed": params['seed'],
+        "sampler_name": params['sampler_name'],
+        "steps": params['steps'],
+        "cfg_scale": params['cfg_scale'],
+        "width": params['width'],
+        "height": params['height'],
        "restore_faces": params['restore_faces'],
        "negative_prompt": params['negative_prompt']
    }

+    print(f'Prompting the image generator via the API on {params["address"]}...')
    response = requests.post(url=f'{params["address"]}/sdapi/v1/txt2img', json=payload)
+    response.raise_for_status()
    r = response.json()

    visible_result = ""
    for img_str in r['images']:
        image = Image.open(io.BytesIO(base64.b64decode(img_str.split(",", 1)[0])))
        if params['save_img']:
-            output_file = Path(f'extensions/sd_api_pictures/outputs/{pic_id:06d}.png')
+            variadic = f'{date.today().strftime("%Y_%m_%d")}/{shared.character}_{int(time.time())}'
+            output_file = Path(f'extensions/sd_api_pictures/outputs/{variadic}.png')
+            output_file.parent.mkdir(parents=True, exist_ok=True)
            image.save(output_file.as_posix())
-            pic_id += 1
+            visible_result = visible_result + f'<img src="/file/extensions/sd_api_pictures/outputs/{variadic}.png" alt="{description}" style="max-width: unset; max-height: unset;">\n'
+        else:
            # lower the resolution of received images for the chat, otherwise the log size gets out of control quickly with all the base64 values in visible history
            image.thumbnail((300, 300))
            buffered = io.BytesIO()
@@ -97,6 +155,9 @@ def get_SD_pictures(description):
            img_str = "data:image/jpeg;base64," + base64.b64encode(image_bytes).decode()
            visible_result = visible_result + f'<img src="{img_str}" alt="{description}">\n'

+    if params['manage_VRAM']:
+        give_VRAM_priority('LLM')
+
    return visible_result

 # TODO: how do I make the UI history ignore the resulting pictures (I don't want HTML to appear in history)
@@ -105,7 +166,8 @@ def output_modifier(string):
    """
    This function is applied to the model outputs.
    """
-    global pic_id, picture_response, streaming_state
+
+    global picture_response, params

    if not picture_response:
        return string
@@ -118,17 +180,19 @@ def output_modifier(string):

    if string == '':
        string = 'no viable description in reply, try regenerating'
+        return string

-    # I can't for the love of all that's holy get the name from shared.gradio['name1'], so for now it will be like this
-    text = f'*Description: "{string}"*'
+    text = ""
+    if (params['mode'] < 2):
+        toggle_generation(False)
+        text = f'*Sends a picture which portrays: “{string}”*'
+    else:
+        text = string

-    image = get_SD_pictures(string)
+    string = get_SD_pictures(string) + "\n" + text

-    picture_response = False
+    return string

-    shared.processing_message = "*Is typing...*"
-    shared.args.no_stream = streaming_state
-    return image + "\n" + text

 def bot_prefix_modifier(string):
    """
@@ -139,41 +203,92 @@ def bot_prefix_modifier(string):

    return string

-def force_pic():
-    global picture_response
-    picture_response = True
+
+def toggle_generation(*args):
+    global picture_response, shared, streaming_state
+
+    if not args:
+        picture_response = not picture_response
+    else:
+        picture_response = args[0]
+
+    shared.args.no_stream = True if picture_response else streaming_state  # Disable streaming cause otherwise the SD-generated picture would return as a dud
+    shared.processing_message = "*Is sending a picture...*" if picture_response else "*Is typing...*"
+
+
+def filter_address(address):
+    address = address.strip()
+    # address = re.sub('http(s)?:\/\/|\/$','',address) # remove starting http:// OR https:// OR trailing slash
+    address = re.sub('\/$', '', address)  # remove trailing /s
+    if not address.startswith('http'):
+        address = 'http://' + address
+    return address
+
+
+def SD_api_address_update(address):
+
+    global params
+
+    msg = "✔️ SD API is found on:"
+    address = filter_address(address)
+    params.update({"address": address})
+    try:
+        response = requests.get(url=f'{params["address"]}/sdapi/v1/sd-models')
+        response.raise_for_status()
+        # r = response.json()
+    except:
+        msg = "❌ No SD API endpoint on:"
+
+    return gr.Textbox.update(label=msg)
+

 def ui():

    # Gradio elements
-    with gr.Accordion("Stable Diffusion api integration", open=True):
+    # gr.Markdown('### Stable Diffusion API Pictures') # Currently the name of extension is shown as the title
+    with gr.Accordion("Parameters", open=True):
        with gr.Row():
-            with gr.Column():
-                enable = gr.Checkbox(value=params['enable_SD_api'], label='Activate SD Api integration')
-                save_img = gr.Checkbox(value=params['save_img'], label='Keep original received images in the outputs subdir')
-            with gr.Column():
-                address = gr.Textbox(placeholder=params['address'], value=params['address'], label='Stable Diffusion host address')
+            address = gr.Textbox(placeholder=params['address'], value=params['address'], label='Auto1111\'s WebUI address')
+            mode = gr.Dropdown(["Manual", "Immersive/Interactive", "Picturebook/Adventure"], value="Manual", label="Mode of operation", type="index")
+            with gr.Column(scale=1, min_width=300):
+                manage_VRAM = gr.Checkbox(value=params['manage_VRAM'], label='Manage VRAM')
+                save_img = gr.Checkbox(value=params['save_img'], label='Keep original images and use them in chat')

-        with gr.Row():
-            force_btn = gr.Button("Force the next response to be a picture")
-            generate_now_btn = gr.Button("Generate an image response to the input")
+            force_pic = gr.Button("Force the picture response")
+            suppr_pic = gr.Button("Suppress the picture response")

        with gr.Accordion("Generation parameters", open=False):
            prompt_prefix = gr.Textbox(placeholder=params['prompt_prefix'], value=params['prompt_prefix'], label='Prompt Prefix (best used to describe the look of the character)')
            with gr.Row():
+                with gr.Column():
                    negative_prompt = gr.Textbox(placeholder=params['negative_prompt'], value=params['negative_prompt'], label='Negative Prompt')
-                dimensions = gr.Slider(256,702,value=params['side_length'],step=64,label='Image dimensions')
-                # model = gr.Dropdown(value=SD_models[0], choices=SD_models, label='Model')
+                    sampler_name = gr.Textbox(placeholder=params['sampler_name'], value=params['sampler_name'], label='Sampler')
+                with gr.Column():
+                    width = gr.Slider(256, 768, value=params['width'], step=64, label='Width')
+                    height = gr.Slider(256, 768, value=params['height'], step=64, label='Height')
+            with gr.Row():
+                steps = gr.Number(label="Steps:", value=params['steps'])
+                seed = gr.Number(label="Seed:", value=params['seed'])
+                cfg_scale = gr.Number(label="CFG Scale:", value=params['cfg_scale'])

    # Event functions to update the parameters in the backend
-    enable.change(lambda x: params.update({"enable_SD_api": x}), enable, None)
+    address.change(lambda x: params.update({"address": filter_address(x)}), address, None)
+    mode.select(lambda x: params.update({"mode": x}), mode, None)
+    mode.select(lambda x: toggle_generation(x > 1), inputs=mode, outputs=None)
+    manage_VRAM.change(lambda x: params.update({"manage_VRAM": x}), manage_VRAM, None)
+    manage_VRAM.change(lambda x: give_VRAM_priority('set' if x else 'reset'), inputs=manage_VRAM, outputs=None)
    save_img.change(lambda x: params.update({"save_img": x}), save_img, None)
-    address.change(lambda x: params.update({"address": x}), address, None)
+
+    address.submit(fn=SD_api_address_update, inputs=address, outputs=address)
    prompt_prefix.change(lambda x: params.update({"prompt_prefix": x}), prompt_prefix, None)
    negative_prompt.change(lambda x: params.update({"negative_prompt": x}), negative_prompt, None)
-    dimensions.change(lambda x: params.update({"side_length": x}), dimensions, None)
-    # model.change(lambda x: params.update({"SD_model": x}), model, None)
+    width.change(lambda x: params.update({"width": x}), width, None)
+    height.change(lambda x: params.update({"height": x}), height, None)

-    force_btn.click(force_pic)
-    generate_now_btn.click(force_pic)
-    generate_now_btn.click(chat.cai_chatbot_wrapper, shared.input_params, shared.gradio['display'], show_progress=shared.args.no_stream)
+    sampler_name.change(lambda x: params.update({"sampler_name": x}), sampler_name, None)
+    steps.change(lambda x: params.update({"steps": x}), steps, None)
+    seed.change(lambda x: params.update({"seed": x}), seed, None)
+    cfg_scale.change(lambda x: params.update({"cfg_scale": x}), cfg_scale, None)
+
+    force_pic.click(lambda x: toggle_generation(True), inputs=force_pic, outputs=None)
+    suppr_pic.click(lambda x: toggle_generation(False), inputs=suppr_pic, outputs=None)
--- a/extensions/send_pictures/script.py
+++ b/extensions/send_pictures/script.py
@@ -17,11 +17,13 @@ input_hijack = {
 processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
 model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base", torch_dtype=torch.float32).to("cpu")

+
 def caption_image(raw_image):
    inputs = processor(raw_image.convert('RGB'), return_tensors="pt").to("cpu", torch.float32)
    out = model.generate(**inputs, max_new_tokens=100)
    return processor.decode(out[0], skip_special_tokens=True)

+
 def generate_chat_picture(picture, name1, name2):
    text = f'*{name1} sends {name2} a picture that contains the following: "{caption_image(picture)}"*'
    # lower the resolution of sent images for the chat, otherwise the log size gets out of control quickly with all the base64 values in visible history
@@ -32,6 +34,7 @@ def generate_chat_picture(picture, name1, name2):
    visible_text = f'<img src="data:image/jpeg;base64,{img_str}" alt="{text}">'
    return text, visible_text

+
 def ui():
    picture_select = gr.Image(label='Send a picture', type='pil')

--- a/extensions/silero_tts/requirements.txt
+++ b/extensions/silero_tts/requirements.txt
@@ -1,6 +1,5 @@
 ipython
+num2words
 omegaconf
 pydub
 PyYAML
-torch
-torchaudio
--- a/extensions/silero_tts/script.py
+++ b/extensions/silero_tts/script.py
@@ -1,14 +1,16 @@
-import re
 import time
 from pathlib import Path

 import gradio as gr
-import modules.chat as chat
-import modules.shared as shared
 import torch

+from extensions.silero_tts import tts_preprocessor
+from modules import chat, shared
+from modules.html_generator import chat_html_wrapper
+
 torch._C._jit_set_profiling_mode(False)

+
 params = {
    'activate': True,
    'speaker': 'en_56',
@@ -20,6 +22,7 @@ params = {
    'autoplay': True,
    'voice_pitch': 'medium',
    'voice_speed': 'medium',
+    'local_cache_path': ''  # User can override the default cache path to something other via settings.json
 }

 current_params = params.copy()
@@ -37,26 +40,31 @@ table = str.maketrans({
    '"': "&quot;",
 })

+
 def xmlesc(txt):
    return txt.translate(table)

+
 def load_model():
+    torch_cache_path = torch.hub.get_dir() if params['local_cache_path'] == '' else params['local_cache_path']
+    model_path = torch_cache_path + "/snakers4_silero-models_master/src/silero/model/" + params['model_id'] + ".pt"
+    if Path(model_path).is_file():
+        print(f'\nUsing Silero TTS cached checkpoint found at {torch_cache_path}')
+        model, example_text = torch.hub.load(repo_or_dir=torch_cache_path + '/snakers4_silero-models_master/', model='silero_tts', language=params['language'], speaker=params['model_id'], source='local', path=model_path, force_reload=True)
+    else:
+        print(f'\nSilero TTS cache not found at {torch_cache_path}. Attempting to download...')
        model, example_text = torch.hub.load(repo_or_dir='snakers4/silero-models', model='silero_tts', language=params['language'], speaker=params['model_id'])
    model.to(params['device'])
    return model
-model = load_model()

-def remove_surrounded_chars(string):
-    # this expression matches to 'as few symbols as possible (0 upwards) between any asterisks' OR
-    # 'as few symbols as possible (0 upwards) between an asterisk and the end of the string'
-    return re.sub('\*[^\*]*?(\*|$)','',string)

-def remove_tts_from_history(name1, name2):
+def remove_tts_from_history(name1, name2, mode):
    for i, entry in enumerate(shared.history['internal']):
        shared.history['visible'][i] = [shared.history['visible'][i][0], entry[1]]
-    return chat.generate_chat_output(shared.history['visible'], name1, name2, shared.character)
+    return chat_html_wrapper(shared.history['visible'], name1, name2, mode)

-def toggle_text_in_history(name1, name2):
+
+def toggle_text_in_history(name1, name2, mode):
    for i, entry in enumerate(shared.history['visible']):
        visible_reply = entry[1]
        if visible_reply.startswith('<audio'):
@@ -65,7 +73,8 @@ def toggle_text_in_history(name1, name2):
                shared.history['visible'][i] = [shared.history['visible'][i][0], f"{visible_reply.split('</audio>')[0]}</audio>\n\n{reply}"]
            else:
                shared.history['visible'][i] = [shared.history['visible'][i][0], f"{visible_reply.split('</audio>')[0]}</audio>"]
-    return chat.generate_chat_output(shared.history['visible'], name1, name2, shared.character)
+    return chat_html_wrapper(shared.history['visible'], name1, name2, mode)
+

 def input_modifier(string):
    """
@@ -81,6 +90,7 @@ def input_modifier(string):
    shared.args.no_stream = True  # Disable streaming cause otherwise the audio output will stutter and begin anew every time the message is being updated
    return string

+
 def output_modifier(string):
    """
    This function is applied to the model outputs.
@@ -94,15 +104,11 @@ def output_modifier(string):
            current_params = params.copy()
            break

-    if params['activate'] == False:
+    if not params['activate']:
        return string

    original_string = string
-    string = remove_surrounded_chars(string)
-    string = string.replace('"', '')
-    string = string.replace('“', '')
-    string = string.replace('\n', ' ')
-    string = string.strip()
+    string = tts_preprocessor.preprocess(string)

    if string == '':
        string = '*Empty reply, try regenerating*'
@@ -121,6 +127,7 @@ def output_modifier(string):
    shared.args.no_stream = streaming_state  # restore the streaming option to the previous value
    return string

+
 def bot_prefix_modifier(string):
    """
    This function is only applied in chat mode. It modifies
@@ -130,17 +137,25 @@ def bot_prefix_modifier(string):

    return string

+
+def setup():
+    global model
+    model = load_model()
+
+
 def ui():
    # Gradio elements
    with gr.Accordion("Silero TTS"):
        with gr.Row():
            activate = gr.Checkbox(value=params['activate'], label='Activate TTS')
            autoplay = gr.Checkbox(value=params['autoplay'], label='Play TTS automatically')
+
        show_text = gr.Checkbox(value=params['show_text'], label='Show message text under audio player')
        voice = gr.Dropdown(value=params['speaker'], choices=voices_by_gender, label='TTS voice')
        with gr.Row():
            v_pitch = gr.Dropdown(value=params['voice_pitch'], choices=voice_pitches, label='Voice pitch')
            v_speed = gr.Dropdown(value=params['voice_speed'], choices=voice_speeds, label='Voice speed')
+
        with gr.Row():
            convert = gr.Button('Permanently replace audios with the message texts')
            convert_cancel = gr.Button('Cancel', visible=False)
@@ -150,13 +165,13 @@ def ui():
    convert_arr = [convert_confirm, convert, convert_cancel]
    convert.click(lambda: [gr.update(visible=True), gr.update(visible=False), gr.update(visible=True)], None, convert_arr)
    convert_confirm.click(lambda: [gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)], None, convert_arr)
-    convert_confirm.click(remove_tts_from_history, [shared.gradio['name1'], shared.gradio['name2']], shared.gradio['display'])
+    convert_confirm.click(remove_tts_from_history, [shared.gradio[k] for k in ['name1', 'name2', 'Chat mode']], shared.gradio['display'])
    convert_confirm.click(lambda: chat.save_history(timestamp=False), [], [], show_progress=False)
    convert_cancel.click(lambda: [gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)], None, convert_arr)

    # Toggle message text in history
    show_text.change(lambda x: params.update({"show_text": x}), show_text, None)
-    show_text.change(toggle_text_in_history, [shared.gradio['name1'], shared.gradio['name2']], shared.gradio['display'])
+    show_text.change(toggle_text_in_history, [shared.gradio[k] for k in ['name1', 'name2', 'Chat mode']], shared.gradio['display'])
    show_text.change(lambda: chat.save_history(timestamp=False), [], [], show_progress=False)

    # Event functions to update the parameters in the backend
--- a/extensions/silero_tts/test_tts.py
+++ b/extensions/silero_tts/test_tts.py
@@ -0,0 +1,81 @@
+import time
+from pathlib import Path
+
+import torch
+import tts_preprocessor
+
+torch._C._jit_set_profiling_mode(False)
+
+
+params = {
+    'activate': True,
+    'speaker': 'en_49',
+    'language': 'en',
+    'model_id': 'v3_en',
+    'sample_rate': 48000,
+    'device': 'cpu',
+    'show_text': True,
+    'autoplay': True,
+    'voice_pitch': 'medium',
+    'voice_speed': 'medium',
+}
+
+current_params = params.copy()
+voices_by_gender = ['en_99', 'en_45', 'en_18', 'en_117', 'en_49', 'en_51', 'en_68', 'en_0', 'en_26', 'en_56', 'en_74', 'en_5', 'en_38', 'en_53', 'en_21', 'en_37', 'en_107', 'en_10', 'en_82', 'en_16', 'en_41', 'en_12', 'en_67', 'en_61', 'en_14', 'en_11', 'en_39', 'en_52', 'en_24', 'en_97', 'en_28', 'en_72', 'en_94', 'en_36', 'en_4', 'en_43', 'en_88', 'en_25', 'en_65', 'en_6', 'en_44', 'en_75', 'en_91', 'en_60', 'en_109', 'en_85', 'en_101', 'en_108', 'en_50', 'en_96', 'en_64', 'en_92', 'en_76', 'en_33', 'en_116', 'en_48', 'en_98', 'en_86', 'en_62', 'en_54', 'en_95', 'en_55', 'en_111', 'en_3', 'en_83', 'en_8', 'en_47', 'en_59', 'en_1', 'en_2', 'en_7', 'en_9', 'en_13', 'en_15', 'en_17', 'en_19', 'en_20', 'en_22', 'en_23', 'en_27', 'en_29', 'en_30', 'en_31', 'en_32', 'en_34', 'en_35', 'en_40', 'en_42', 'en_46', 'en_57', 'en_58', 'en_63', 'en_66', 'en_69', 'en_70', 'en_71', 'en_73', 'en_77', 'en_78', 'en_79', 'en_80', 'en_81', 'en_84', 'en_87', 'en_89', 'en_90', 'en_93', 'en_100', 'en_102', 'en_103', 'en_104', 'en_105', 'en_106', 'en_110', 'en_112', 'en_113', 'en_114', 'en_115']
+voice_pitches = ['x-low', 'low', 'medium', 'high', 'x-high']
+voice_speeds = ['x-slow', 'slow', 'medium', 'fast', 'x-fast']
+
+# Used for making text xml compatible, needed for voice pitch and speed control
+table = str.maketrans({
+    "<": "&lt;",
+    ">": "&gt;",
+    "&": "&amp;",
+    "'": "&apos;",
+    '"': "&quot;",
+})
+
+
+def xmlesc(txt):
+    return txt.translate(table)
+
+
+def load_model():
+    model, example_text = torch.hub.load(repo_or_dir='snakers4/silero-models', model='silero_tts', language=params['language'], speaker=params['model_id'])
+    model.to(params['device'])
+    return model
+
+
+model = load_model()
+
+
+def output_modifier(string):
+    """
+    This function is applied to the model outputs.
+    """
+
+    global model, current_params
+
+    original_string = string
+    string = tts_preprocessor.preprocess(string)
+    processed_string = string
+
+    if string == '':
+        string = '*Empty reply, try regenerating*'
+    else:
+        output_file = Path(f'extensions/silero_tts/outputs/test_{int(time.time())}.wav')
+        prosody = '<prosody rate="{}" pitch="{}">'.format(params['voice_speed'], params['voice_pitch'])
+        silero_input = f'<speak>{prosody}{xmlesc(string)}</prosody></speak>'
+        model.save_wav(ssml_text=silero_input, speaker=params['speaker'], sample_rate=int(params['sample_rate']), audio_path=str(output_file))
+
+        autoplay = 'autoplay' if params['autoplay'] else ''
+        string = f'<audio src="file/{output_file.as_posix()}" controls {autoplay}></audio>'
+
+        if params['show_text']:
+            string += f'\n\n{original_string}\n\nProcessed:\n{processed_string}'
+
+    print(string)
+
+
+if __name__ == '__main__':
+    import sys
+    output_modifier(sys.argv[1])
--- a/extensions/silero_tts/tts_preprocessor.py
+++ b/extensions/silero_tts/tts_preprocessor.py
@@ -0,0 +1,194 @@
+import re
+
+from num2words import num2words
+
+punctuation = r'[\s,.?!/)\'\]>]'
+alphabet_map = {
+    "A": " Ei ",
+    "B": " Bee ",
+    "C": " See ",
+    "D": " Dee ",
+    "E": " Eee ",
+    "F": " Eff ",
+    "G": " Jee ",
+    "H": " Eich ",
+    "I": " Eye ",
+    "J": " Jay ",
+    "K": " Kay ",
+    "L": " El ",
+    "M": " Emm ",
+    "N": " Enn ",
+    "O": " Ohh ",
+    "P": " Pee ",
+    "Q": " Queue ",
+    "R": " Are ",
+    "S": " Ess ",
+    "T": " Tee ",
+    "U": " You ",
+    "V": " Vee ",
+    "W": " Double You ",
+    "X": " Ex ",
+    "Y": " Why ",
+    "Z": " Zed "  # Zed is weird, as I (da3dsoul) am American, but most of the voice models sound British, so it matches
+}
+
+
+def preprocess(string):
+    # the order for some of these matter
+    # For example, you need to remove the commas in numbers before expanding them
+    string = remove_surrounded_chars(string)
+    string = string.replace('"', '')
+    string = string.replace('\u201D', '').replace('\u201C', '')  # right and left quote
+    string = string.replace('\u201F', '')  # italic looking quote
+    string = string.replace('\n', ' ')
+    string = convert_num_locale(string)
+    string = replace_negative(string)
+    string = replace_roman(string)
+    string = hyphen_range_to(string)
+    string = num_to_words(string)
+
+    # TODO Try to use a ML predictor to expand abbreviations. It's hard, dependent on context, and whether to actually
+    # try to say the abbreviation or spell it out as I've done below is not agreed upon
+
+    # For now, expand abbreviations to pronunciations
+    # replace_abbreviations adds a lot of unnecessary whitespace to ensure separation
+    string = replace_abbreviations(string)
+    string = replace_lowercase_abbreviations(string)
+
+    # cleanup whitespaces
+    # remove whitespace before punctuation
+    string = re.sub(rf'\s+({punctuation})', r'\1', string)
+    string = string.strip()
+    # compact whitespace
+    string = ' '.join(string.split())
+
+    return string
+
+
+def remove_surrounded_chars(string):
+    # this expression matches to 'as few symbols as possible (0 upwards) between any asterisks' OR
+    # 'as few symbols as possible (0 upwards) between an asterisk and the end of the string'
+    return re.sub(r'\*[^*]*?(\*|$)', '', string)
+
+
+def convert_num_locale(text):
+    # This detects locale and converts it to American without comma separators
+    pattern = re.compile(r'(?:\s|^)\d{1,3}(?:\.\d{3})+(,\d+)(?:\s|$)')
+    result = text
+    while True:
+        match = pattern.search(result)
+        if match is None:
+            break
+
+        start = match.start()
+        end = match.end()
+        result = result[0:start] + result[start:end].replace('.', '').replace(',', '.') + result[end:len(result)]
+
+    # removes comma separators from existing American numbers
+    pattern = re.compile(r'(\d),(\d)')
+    result = pattern.sub(r'\1\2', result)
+
+    return result
+
+
+def replace_negative(string):
+    # handles situations like -5. -5 would become negative 5, which would then be expanded to negative five
+    return re.sub(rf'(\s)(-)(\d+)({punctuation})', r'\1negative \3\4', string)
+
+
+def replace_roman(string):
+    # find a string of roman numerals.
+    # Only 2 or more, to avoid capturing I and single character abbreviations, like names
+    pattern = re.compile(rf'\s[IVXLCDM]{{2,}}{punctuation}')
+    result = string
+    while True:
+        match = pattern.search(result)
+        if match is None:
+            break
+
+        start = match.start()
+        end = match.end()
+        result = result[0:start + 1] + str(roman_to_int(result[start + 1:end - 1])) + result[end - 1:len(result)]
+
+    return result
+
+
+def roman_to_int(s):
+    rom_val = {'I': 1, 'V': 5, 'X': 10, 'L': 50, 'C': 100, 'D': 500, 'M': 1000}
+    int_val = 0
+    for i in range(len(s)):
+        if i > 0 and rom_val[s[i]] > rom_val[s[i - 1]]:
+            int_val += rom_val[s[i]] - 2 * rom_val[s[i - 1]]
+        else:
+            int_val += rom_val[s[i]]
+    return int_val
+
+
+def hyphen_range_to(text):
+    pattern = re.compile(r'(\d+)[-–](\d+)')
+    result = pattern.sub(lambda x: x.group(1) + ' to ' + x.group(2), text)
+    return result
+
+
+def num_to_words(text):
+    # 1000 or 10.23
+    pattern = re.compile(r'\d+\.\d+|\d+')
+    result = pattern.sub(lambda x: num2words(float(x.group())), text)
+    return result
+
+
+def replace_abbreviations(string):
+    # abbreviations 1 to 4 characters long. It will get things like A and I, but those are pronounced with their letter
+    pattern = re.compile(rf'(^|[\s(.\'\[<])([A-Z]{{1,4}})({punctuation}|$)')
+    result = string
+    while True:
+        match = pattern.search(result)
+        if match is None:
+            break
+
+        start = match.start()
+        end = match.end()
+        result = result[0:start] + replace_abbreviation(result[start:end]) + result[end:len(result)]
+
+    return result
+
+
+def replace_lowercase_abbreviations(string):
+    # abbreviations 1 to 4 characters long, separated by dots i.e. e.g.
+    pattern = re.compile(rf'(^|[\s(.\'\[<])(([a-z]\.){{1,4}})({punctuation}|$)')
+    result = string
+    while True:
+        match = pattern.search(result)
+        if match is None:
+            break
+
+        start = match.start()
+        end = match.end()
+        result = result[0:start] + replace_abbreviation(result[start:end].upper()) + result[end:len(result)]
+
+    return result
+
+
+def replace_abbreviation(string):
+    result = ""
+    for char in string:
+        result += match_mapping(char)
+
+    return result
+
+
+def match_mapping(char):
+    for mapping in alphabet_map.keys():
+        if char == mapping:
+            return alphabet_map[char]
+
+    return char
+
+
+def __main__(args):
+    print(preprocess(args[1]))
+
+
+if __name__ == "__main__":
+    import sys
+    __main__(sys.argv)
--- a/extensions/whisper_stt/script.py
+++ b/extensions/whisper_stt/script.py
@@ -1,5 +1,6 @@
 import gradio as gr
 import speech_recognition as sr
+from modules import shared

 input_hijack = {
    'state': False,
@@ -7,7 +8,7 @@ input_hijack = {
 }


-def do_stt(audio, text_state=""):
+def do_stt(audio):
    transcription = ""
    r = sr.Recognizer()

@@ -21,34 +22,23 @@ def do_stt(audio, text_state=""):
    except sr.RequestError as e:
        print("Could not request results from Whisper", e)

-    input_hijack.update({"state": True, "value": [transcription, transcription]})
-
-    text_state += transcription + " "
-    return text_state, text_state
+    return transcription


-def update_hijack(val):
-    input_hijack.update({"state": True, "value": [val, val]})
-    return val
-
-
-def auto_transcribe(audio, audio_auto, text_state=""):
+def auto_transcribe(audio, auto_submit):
    if audio is None:
        return "", ""
-    if audio_auto:
-        return do_stt(audio, text_state)
-    return "", ""
+
+    transcription = do_stt(audio)
+    if auto_submit:
+        input_hijack.update({"state": True, "value": [transcription, transcription]})
+
+    return transcription, None


 def ui():
-    tr_state = gr.State(value="")
-    output_transcription = gr.Textbox(label="STT-Input",
-                                      placeholder="Speech Preview. Click \"Generate\" to send",
-                                      interactive=True)
-    output_transcription.change(fn=update_hijack, inputs=[output_transcription], outputs=[tr_state])
-    audio_auto = gr.Checkbox(label="Auto-Transcribe", value=True)
    with gr.Row():
        audio = gr.Audio(source="microphone")
-        audio.change(fn=auto_transcribe, inputs=[audio, audio_auto, tr_state], outputs=[output_transcription, tr_state])
-        transcribe_button = gr.Button(value="Transcribe")
-        transcribe_button.click(do_stt, inputs=[audio, tr_state], outputs=[output_transcription, tr_state])
+        auto_submit = gr.Checkbox(label='Submit the transcribed audio automatically', value=True)
+    audio.change(fn=auto_transcribe, inputs=[audio, auto_submit], outputs=[shared.gradio['textbox'], audio])
+    audio.change(None, auto_submit, None, _js="(check) => {if (check) { document.getElementById('Generate').click() }}")
--- a/modules/GPTQ_loader.py
+++ b/modules/GPTQ_loader.py
@@ -17,9 +17,11 @@ from quant import make_quant


 def _load_quant(model, checkpoint, wbits, groupsize=-1, faster_kernel=False, exclude_layers=['lm_head'], kernel_switch_threshold=128):
-    config = AutoConfig.from_pretrained(model)
+
    def noop(*args, **kwargs):
        pass
+
+    config = AutoConfig.from_pretrained(model)
    torch.nn.init.kaiming_uniform_ = noop
    torch.nn.init.uniform_ = noop
    torch.nn.init.normal_ = noop
@@ -64,6 +66,7 @@ def _load_quant(model, checkpoint, wbits, groupsize=-1, faster_kernel=False, exc

    return model

+
 def load_quantized(model_name):
    if not shared.args.model_type:
        # Try to determine model type from model name
--- a/modules/LoRA.py
+++ b/modules/LoRA.py
@@ -4,15 +4,9 @@ import torch
 from peft import PeftModel

 import modules.shared as shared
-from modules.models import load_model
-from modules.text_generation import clear_torch_cache
+from modules.models import reload_model


-def reload_model():
-    shared.model = shared.tokenizer = None
-    clear_torch_cache()
-    shared.model, shared.tokenizer = load_model(shared.model_name)
-
 def add_lora_to_model(lora_name):

    # If a LoRA had been previously loaded, or if we want
--- a/modules/RWKV.py
+++ b/modules/RWKV.py
@@ -54,6 +54,7 @@ class RWKVModel:
                reply += token
                yield reply

+
 class RWKVTokenizer:
    def __init__(self):
        pass
--- a/modules/api.py
+++ b/modules/api.py
@@ -28,6 +28,7 @@ def generate_reply_wrapper(string):
    for i in generate_reply(params[0], generate_params):
        yield i

+
 def create_apis():
    t1 = gr.Textbox(visible=False)
    t2 = gr.Textbox(visible=False)
--- a/modules/callbacks.py
+++ b/modules/callbacks.py
@@ -30,6 +30,7 @@ class _SentinelTokenStoppingCriteria(transformers.StoppingCriteria):
                        return True
        return False

+
 class Stream(transformers.StoppingCriteria):
    def __init__(self, callback_func=None):
        self.callback_func = callback_func
@@ -39,6 +40,7 @@ class Stream(transformers.StoppingCriteria):
            self.callback_func(input_ids[0])
        return False

+
 class Iteratorize:

    """
@@ -96,6 +98,7 @@ class Iteratorize:
        self.stop_now = True
        clear_torch_cache()

+
 def clear_torch_cache():
    gc.collect()
    if not shared.args.cpu:
--- a/modules/chat.py
+++ b/modules/chat.py
@@ -12,7 +12,7 @@ from PIL import Image
 import modules.extensions as extensions_module
 import modules.shared as shared
 from modules.extensions import apply_extensions
-from modules.html_generator import (fix_newlines, chat_html_wrapper,
+from modules.html_generator import (chat_html_wrapper, fix_newlines,
                                    make_thumbnail)
 from modules.text_generation import (encode, generate_reply,
                                     get_max_prompt_length)
@@ -23,7 +23,6 @@ def generate_chat_prompt(user_input, max_new_tokens, name1, name2, context, chat
    end_of_turn = kwargs['end_of_turn'] if 'end_of_turn' in kwargs else ''
    impersonate = kwargs['impersonate'] if 'impersonate' in kwargs else False
    also_return_rows = kwargs['also_return_rows'] if 'also_return_rows' in kwargs else False
-
    rows = [f"{context.strip()}\n"]

    # Finding the maximum prompt size
@@ -68,6 +67,7 @@ def generate_chat_prompt(user_input, max_new_tokens, name1, name2, context, chat
    else:
        return prompt

+
 def extract_message_from_reply(reply, name1, name2, stop_at_newline):
    next_character_found = False

@@ -98,22 +98,25 @@ def extract_message_from_reply(reply, name1, name2, stop_at_newline):
    reply = fix_newlines(reply)
    return reply, next_character_found

+
 def chatbot_wrapper(text, generate_state, name1, name2, context, mode, end_of_turn, regenerate=False):
    if mode == 'instruct':
        stopping_strings = [f"\n{name1}", f"\n{name2}"]
    else:
        stopping_strings = [f"\n{name1}:", f"\n{name2}:"]

-    eos_token = '\n' if generate_state['stop_at_newline'] else None
+    # Defining some variables
+    cumulative_reply = ''
+    just_started = True
    name1_original = name1
+    visible_text = custom_generate_chat_prompt = None
+    eos_token = '\n' if generate_state['stop_at_newline'] else None
    if 'pygmalion' in shared.model_name.lower():
        name1 = "You"

    # Check if any extension wants to hijack this function call
-    visible_text = None
-    custom_generate_chat_prompt = None
    for extension, _ in extensions_module.iterator():
-        if hasattr(extension, 'input_hijack') and extension.input_hijack['state'] == True:
+        if hasattr(extension, 'input_hijack') and extension.input_hijack['state']:
            extension.input_hijack['state'] = False
            text, visible_text = extension.input_hijack['value']
        if custom_generate_chat_prompt is None and hasattr(extension, 'custom_generate_chat_prompt'):
@@ -123,6 +126,7 @@ def chatbot_wrapper(text, generate_state, name1, name2, context, mode, end_of_tu
        visible_text = text
    text = apply_extensions(text, "input")

+    # Generating the prompt
    kwargs = {'end_of_turn': end_of_turn, 'is_instruct': mode == 'instruct'}
    if custom_generate_chat_prompt is None:
        prompt = generate_chat_prompt(text, generate_state['max_new_tokens'], name1, name2, context, generate_state['chat_prompt_size'], **kwargs)
@@ -134,8 +138,6 @@ def chatbot_wrapper(text, generate_state, name1, name2, context, mode, end_of_tu
        yield shared.history['visible'] + [[visible_text, shared.processing_message]]

    # Generate
-    cumulative_reply = ''
-    just_started = True
    for i in range(generate_state['chat_generation_attempts']):
        reply = None
        for reply in generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", generate_state, eos_token=eos_token, stopping_strings=stopping_strings):
@@ -167,12 +169,15 @@ def chatbot_wrapper(text, generate_state, name1, name2, context, mode, end_of_tu

    yield shared.history['visible']

+
 def impersonate_wrapper(text, generate_state, name1, name2, context, mode, end_of_turn):
    if mode == 'instruct':
        stopping_strings = [f"\n{name1}", f"\n{name2}"]
    else:
        stopping_strings = [f"\n{name1}:", f"\n{name2}:"]

+    # Defining some variables
+    cumulative_reply = ''
    eos_token = '\n' if generate_state['stop_at_newline'] else None
    if 'pygmalion' in shared.model_name.lower():
        name1 = "You"
@@ -182,7 +187,6 @@ def impersonate_wrapper(text, generate_state, name1, name2, context, mode, end_o
    # Yield *Is typing...*
    yield shared.processing_message

-    cumulative_reply = ''
    for i in range(generate_state['chat_generation_attempts']):
        reply = None
        for reply in generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", generate_state, eos_token=eos_token, stopping_strings=stopping_strings):
@@ -197,12 +201,14 @@ def impersonate_wrapper(text, generate_state, name1, name2, context, mode, end_o

    yield reply

+
 def cai_chatbot_wrapper(text, generate_state, name1, name2, context, mode, end_of_turn):
    for history in chatbot_wrapper(text, generate_state, name1, name2, context, mode, end_of_turn):
        yield chat_html_wrapper(history, name1, name2, mode)

+
 def regenerate_wrapper(text, generate_state, name1, name2, context, mode, end_of_turn):
-    if (shared.character != 'None' and len(shared.history['visible']) == 1) or len(shared.history['internal']) == 0:
+    if (len(shared.history['visible']) == 1 and not shared.history['visible'][0][0]) or len(shared.history['internal']) == 0:
        yield chat_html_wrapper(shared.history['visible'], name1, name2, mode)
    else:
        last_visible = shared.history['visible'].pop()
@@ -213,6 +219,7 @@ def regenerate_wrapper(text, generate_state, name1, name2, context, mode, end_of
            shared.history['visible'][-1] = [last_visible[0], history[-1][1]]
            yield chat_html_wrapper(shared.history['visible'], name1, name2, mode)

+
 def remove_last_message(name1, name2, mode):
    if len(shared.history['visible']) > 0 and shared.history['internal'][-1][0] != '<|BEGIN-VISIBLE-CHAT|>':
        last = shared.history['visible'].pop()
@@ -222,12 +229,14 @@ def remove_last_message(name1, name2, mode):

    return chat_html_wrapper(shared.history['visible'], name1, name2, mode), last[0]

+
 def send_last_reply_to_input():
    if len(shared.history['internal']) > 0:
        return shared.history['internal'][-1][1]
    else:
        return ''

+
 def replace_last_reply(text, name1, name2, mode):
    if len(shared.history['visible']) > 0:
        shared.history['visible'][-1][1] = text
@@ -235,9 +244,11 @@ def replace_last_reply(text, name1, name2, mode):

    return chat_html_wrapper(shared.history['visible'], name1, name2, mode)

+
 def clear_html():
    return chat_html_wrapper([], "", "")

+
 def clear_chat_log(name1, name2, greeting, mode):
    shared.history['visible'] = []
    shared.history['internal'] = []
@@ -248,12 +259,14 @@ def clear_chat_log(name1, name2, greeting, mode):

    return chat_html_wrapper(shared.history['visible'], name1, name2, mode)

+
 def redraw_html(name1, name2, mode):
    return chat_html_wrapper(shared.history['visible'], name1, name2, mode)

+
 def tokenize_dialogue(dialogue, name1, name2, mode):
    history = []
-
+    messages = []
    dialogue = re.sub('<START>', '', dialogue)
    dialogue = re.sub('<start>', '', dialogue)
    dialogue = re.sub('(\n|^)[Aa]non:', '\\1You:', dialogue)
@@ -262,7 +275,6 @@ def tokenize_dialogue(dialogue, name1, name2, mode):
    if len(idx) == 0:
        return history

-    messages = []
    for i in range(len(idx) - 1):
        messages.append(dialogue[idx[i]:idx[i + 1]].strip())
    messages.append(dialogue[idx[-1]:].strip())
@@ -288,6 +300,7 @@ def tokenize_dialogue(dialogue, name1, name2, mode):

    return history

+
 def save_history(timestamp=True):
    if timestamp:
        fname = f"{shared.character}_{datetime.now().strftime('%Y%m%d-%H%M%S')}.json"
@@ -299,6 +312,7 @@ def save_history(timestamp=True):
        f.write(json.dumps({'data': shared.history['internal'], 'data_visible': shared.history['visible']}, indent=2))
    return Path(f'logs/{fname}')

+
 def load_history(file, name1, name2):
    file = file.decode('utf-8')
    try:
@@ -323,10 +337,12 @@ def load_history(file, name1, name2):
        shared.history['internal'] = tokenize_dialogue(file, name1, name2)
        shared.history['visible'] = copy.deepcopy(shared.history['internal'])

+
 def replace_character_names(text, name1, name2):
    text = text.replace('{{user}}', name1).replace('{{char}}', name2)
    return text.replace('<USER>', name1).replace('<BOT>', name2)

+
 def build_pygmalion_style_context(data):
    context = ""
    if 'char_persona' in data and data['char_persona'] != '':
@@ -336,6 +352,7 @@ def build_pygmalion_style_context(data):
    context = f"{context.strip()}\n<START>\n"
    return context

+
 def generate_pfp_cache(character):
    cache_folder = Path("cache")
    if not cache_folder.exists():
@@ -348,6 +365,7 @@ def generate_pfp_cache(character):
            return img
    return None

+
 def load_character(character, name1, name2, mode):
    shared.character = character
    shared.history['internal'] = []
@@ -404,9 +422,11 @@ def load_character(character, name1, name2, mode):

    return name1, name2, picture, greeting, context, end_of_turn, chat_html_wrapper(shared.history['visible'], name1, name2, mode, reset_cache=True)

+
 def load_default_history(name1, name2):
    load_character("None", name1, name2, "chat")

+
 def upload_character(json_file, img, tavern=False):
    json_file = json_file if type(json_file) == str else json_file.decode('utf-8')
    data = json.loads(json_file)
@@ -425,6 +445,7 @@ def upload_character(json_file, img, tavern=False):
    print(f'New character saved to "characters/{outfile_name}.json".')
    return outfile_name

+
 def upload_tavern_character(img, name1, name2):
    _img = Image.open(io.BytesIO(img))
    _img.getexif()
@@ -433,12 +454,13 @@ def upload_tavern_character(img, name1, name2):
    _json = {"char_name": _json['name'], "char_persona": _json['description'], "char_greeting": _json["first_mes"], "example_dialogue": _json['mes_example'], "world_scenario": _json['scenario']}
    return upload_character(json.dumps(_json), img, tavern=True)

+
 def upload_your_profile_picture(img, name1, name2, mode):
    cache_folder = Path("cache")
    if not cache_folder.exists():
        cache_folder.mkdir()

-    if img == None:
+    if img is None:
        if Path("cache/pfp_me.png").exists():
            Path("cache/pfp_me.png").unlink()
    else:
--- a/modules/extensions.py
+++ b/modules/extensions.py
@@ -9,25 +9,32 @@ state = {}
 available_extensions = []
 setup_called = set()

+
 def load_extensions():
-    global state
+    global state, setup_called
    for i, name in enumerate(shared.args.extensions):
        if name in available_extensions:
            print(f'Loading the extension "{name}"... ', end='')
            try:
                exec(f"import extensions.{name}.script")
+                extension = eval(f"extensions.{name}.script")
+                if extension not in setup_called and hasattr(extension, "setup"):
+                    setup_called.add(extension)
+                    extension.setup()
                state[name] = [True, i]
                print('Ok.')
            except:
                print('Fail.')
                traceback.print_exc()

+
 # This iterator returns the extensions in the order specified in the command-line
 def iterator():
    for name in sorted(state, key=lambda x: state[x][1]):
-        if state[name][0] == True:
+        if state[name][0]:
            yield eval(f"extensions.{name}.script"), name

+
 # Extension functions that map string -> string
 def apply_extensions(text, typ):
    for extension, _ in iterator():
@@ -39,6 +46,7 @@ def apply_extensions(text, typ):
            text = extension.bot_prefix_modifier(text)
    return text

+
 def create_extensions_block():
    global setup_called

@@ -51,14 +59,9 @@ def create_extensions_block():
                    extension.params[param] = shared.settings[_id]

    should_display_ui = False
-
-    # Running setup function
    for extension, name in iterator():
        if hasattr(extension, "ui"):
            should_display_ui = True
-        if extension not in setup_called and hasattr(extension, "setup"):
-            setup_called.add(extension)
-            extension.setup()

    # Creating the extension ui elements
    if should_display_ui:
--- a/modules/html_generator.py
+++ b/modules/html_generator.py
@@ -24,6 +24,7 @@ with open(Path(__file__).resolve().parent / '../css/html_cai_style.css', 'r') as
 with open(Path(__file__).resolve().parent / '../css/html_instruct_style.css', 'r') as f:
    instruct_css = f.read()

+
 def fix_newlines(string):
    string = string.replace('\n', '\n\n')
    string = re.sub(r"\n{3,}", "\n\n", string)
@@ -31,6 +32,8 @@ def fix_newlines(string):
    return string

 # This could probably be generalized and improved
+
+
 def convert_to_markdown(string):
    string = string.replace('\\begin{code}', '```')
    string = string.replace('\\end{code}', '```')
@@ -40,11 +43,13 @@ def convert_to_markdown(string):
    string = fix_newlines(string)
    return markdown.markdown(string, extensions=['fenced_code'])

+
 def generate_basic_html(string):
    string = convert_to_markdown(string)
    string = f'<style>{readable_css}</style><div class="container">{string}</div>'
    return string

+
 def process_post(post, c):
    t = post.split('\n')
    number = t[0].split(' ')[1]
@@ -59,6 +64,7 @@ def process_post(post, c):
    src = f'<span class="name">Anonymous </span> <span class="number">No.{number}</span>\n{src}'
    return src

+
 def generate_4chan_html(f):
    posts = []
    post = ''
@@ -98,6 +104,7 @@ def generate_4chan_html(f):

    return output

+
 def make_thumbnail(image):
    image = image.resize((350, round(image.size[1] / image.size[0] * 350)), Image.Resampling.LANCZOS)
    if image.size[1] > 470:
@@ -105,6 +112,7 @@ def make_thumbnail(image):

    return image

+
 def get_image_cache(path):
    cache_folder = Path("cache")
    if not cache_folder.exists():
@@ -119,6 +127,7 @@ def get_image_cache(path):

    return image_cache[path][1]

+
 def generate_instruct_html(history):
    output = f'<style>{instruct_css}</style><div class="chat" id="chat">'
    for i, _row in enumerate(history[::-1]):
@@ -151,6 +160,7 @@ def generate_instruct_html(history):

    return output

+
 def generate_cai_chat_html(history, name1, name2, reset_cache=False):
    output = f'<style>{cai_css}</style><div class="chat" id="chat">'

@@ -200,9 +210,11 @@ def generate_cai_chat_html(history, name1, name2, reset_cache=False):
    output += "</div>"
    return output

+
 def generate_chat_html(history, name1, name2):
    return generate_cai_chat_html(history, name1, name2)

+
 def chat_html_wrapper(history, name1, name2, mode, reset_cache=False):
    if mode == "cai-chat":
        return generate_cai_chat_html(history, name1, name2, reset_cache)
--- a/modules/llamacpp_model_alternative.py
+++ b/modules/llamacpp_model_alternative.py
@@ -6,8 +6,6 @@ Documentation:
 https://abetlen.github.io/llama-cpp-python/
 '''

-import multiprocessing
-
 from llama_cpp import Llama

 from modules import shared
--- a/modules/models.py
+++ b/modules/models.py
@@ -1,3 +1,4 @@
+import gc
 import json
 import os
 import re
@@ -16,11 +17,10 @@ import modules.shared as shared

 transformers.logging.set_verbosity_error()

-local_rank = None
-
 if shared.args.flexgen:
    from flexgen.flex_opt import CompressionConfig, ExecutionEnv, OptLM, Policy

+local_rank = None
 if shared.args.deepspeed:
    import deepspeed
    from transformers.deepspeed import (HfDeepSpeedConfig,
@@ -181,6 +181,23 @@ def load_model(model_name):
    print(f"Loaded the model in {(time.time()-t0):.2f} seconds.")
    return model, tokenizer

+
+def clear_torch_cache():
+    gc.collect()
+    if not shared.args.cpu:
+        torch.cuda.empty_cache()
+
+
+def unload_model():
+    shared.model = shared.tokenizer = None
+    clear_torch_cache()
+
+
+def reload_model():
+    unload_model()
+    shared.model, shared.tokenizer = load_model(shared.model_name)
+
+
 def load_soft_prompt(name):
    if name == 'None':
        shared.soft_prompt = False
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -61,6 +61,7 @@ settings = {
    }
 }

+
 def str2bool(v):
    if isinstance(v, bool):
        return v
@@ -71,6 +72,7 @@ def str2bool(v):
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

+
 parser = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=54))

 # Basic settings
@@ -145,5 +147,6 @@ if args.cai_chat:
    print("Warning: --cai-chat is deprecated. Use --chat instead.")
    args.chat = True

+
 def is_chat():
    return args.chat
--- a/modules/text_generation.py
+++ b/modules/text_generation.py
@@ -1,4 +1,3 @@
-import gc
 import re
 import time
 import traceback
@@ -12,7 +11,7 @@ from modules.callbacks import (Iteratorize, Stream,
                               _SentinelTokenStoppingCriteria)
 from modules.extensions import apply_extensions
 from modules.html_generator import generate_4chan_html, generate_basic_html
-from modules.models import local_rank
+from modules.models import clear_torch_cache, local_rank


 def get_max_prompt_length(tokens):
@@ -21,6 +20,7 @@ def get_max_prompt_length(tokens):
        max_length -= shared.soft_prompt_tensor.shape[1]
    return max_length

+
 def encode(prompt, tokens_to_generate=0, add_special_tokens=True):
    if any((shared.is_RWKV, shared.is_llamacpp)):
        input_ids = shared.tokenizer.encode(str(prompt))
@@ -44,6 +44,7 @@ def encode(prompt, tokens_to_generate=0, add_special_tokens=True):
        else:
            return input_ids.cuda()

+
 def decode(output_ids):
    # Open Assistant relies on special tokens like <|endoftext|>
    if re.match('.*(oasst|galactica)-*', shared.model_name.lower()):
@@ -53,6 +54,7 @@ def decode(output_ids):
        reply = reply.replace(r'<|endoftext|>', '')
        return reply

+
 def generate_softprompt_input_tensors(input_ids):
    inputs_embeds = shared.model.transformer.wte(input_ids)
    inputs_embeds = torch.cat((shared.soft_prompt_tensor, inputs_embeds), dim=1)
@@ -61,6 +63,8 @@ def generate_softprompt_input_tensors(input_ids):
    return inputs_embeds, filler_input_ids

 # Removes empty replies from gpt4chan outputs
+
+
 def fix_gpt4chan(s):
    for i in range(10):
        s = re.sub("--- [0-9]*\n>>[0-9]*\n---", "---", s)
@@ -69,6 +73,8 @@ def fix_gpt4chan(s):
    return s

 # Fix the LaTeX equations in galactica
+
+
 def fix_galactica(s):
    s = s.replace(r'\[', r'$')
    s = s.replace(r'\]', r'$')
@@ -79,6 +85,7 @@ def fix_galactica(s):
    s = re.sub(r"\n{3,}", "\n\n", s)
    return s

+
 def formatted_outputs(reply, model_name):
    if not shared.is_chat():
        if 'galactica' in model_name.lower():
@@ -92,10 +99,6 @@ def formatted_outputs(reply, model_name):
    else:
        return reply

-def clear_torch_cache():
-    gc.collect()
-    if not shared.args.cpu:
-        torch.cuda.empty_cache()

 def set_manual_seed(seed):
    if seed != -1:
@@ -103,9 +106,11 @@ def set_manual_seed(seed):
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)

+
 def stop_everything_event():
    shared.stop_everything = True

+
 def generate_reply(question, generate_state, eos_token=None, stopping_strings=[]):
    clear_torch_cache()
    set_manual_seed(generate_state['seed'])
@@ -115,22 +120,22 @@ def generate_reply(question, generate_state, eos_token=None, stopping_strings=[]

    original_question = question
    if not shared.is_chat():
-        question = apply_extensions(question, "input")
+        question = apply_extensions(question, 'input')
    if shared.args.verbose:
-        print(f"\n\n{question}\n--------------------\n")
+        print(f'\n\n{question}\n--------------------\n')

    # These models are not part of Hugging Face, so we handle them
    # separately and terminate the function call earlier
    if any((shared.is_RWKV, shared.is_llamacpp)):
        for k in ['temperature', 'top_p', 'top_k', 'repetition_penalty']:
            generate_params[k] = generate_state[k]
-        generate_params["token_count"] = generate_state["max_new_tokens"]
+        generate_params['token_count'] = generate_state['max_new_tokens']
        try:
            if shared.args.no_stream:
                reply = shared.model.generate(context=question, **generate_params)
                output = original_question + reply
                if not shared.is_chat():
-                    reply = original_question + apply_extensions(reply, "output")
+                    reply = original_question + apply_extensions(reply, 'output')
                yield formatted_outputs(reply, shared.model_name)
            else:
                if not shared.is_chat():
@@ -141,7 +146,7 @@ def generate_reply(question, generate_state, eos_token=None, stopping_strings=[]
                for reply in shared.model.generate_with_streaming(context=question, **generate_params):
                    output = original_question + reply
                    if not shared.is_chat():
-                        reply = original_question + apply_extensions(reply, "output")
+                        reply = original_question + apply_extensions(reply, 'output')
                    yield formatted_outputs(reply, shared.model_name)

        except Exception:
@@ -150,7 +155,7 @@ def generate_reply(question, generate_state, eos_token=None, stopping_strings=[]
            t1 = time.time()
            original_tokens = len(encode(original_question)[0])
            new_tokens = len(encode(output)[0]) - original_tokens
-            print(f"Output generated in {(t1-t0):.2f} seconds ({new_tokens/(t1-t0):.2f} tokens/s, {new_tokens} tokens, context {original_tokens})")
+            print(f'Output generated in {(t1-t0):.2f} seconds ({new_tokens/(t1-t0):.2f} tokens/s, {new_tokens} tokens, context {original_tokens})')
            return

    input_ids = encode(question, generate_state['max_new_tokens'])
@@ -166,31 +171,30 @@ def generate_reply(question, generate_state, eos_token=None, stopping_strings=[]
        t = [encode(string, 0, add_special_tokens=False) for string in stopping_strings]
        stopping_criteria_list.append(_SentinelTokenStoppingCriteria(sentinel_token_ids=t, starting_idx=len(input_ids[0])))

-    generate_params["max_new_tokens"] = generate_state['max_new_tokens']
    if not shared.args.flexgen:
-        for k in ["do_sample", "temperature", "top_p", "typical_p", "repetition_penalty", "encoder_repetition_penalty", "top_k", "min_length", "no_repeat_ngram_size", "num_beams", "penalty_alpha", "length_penalty", "early_stopping"]:
+        for k in ['max_new_tokens', 'do_sample', 'temperature', 'top_p', 'typical_p', 'repetition_penalty', 'encoder_repetition_penalty', 'top_k', 'min_length', 'no_repeat_ngram_size', 'num_beams', 'penalty_alpha', 'length_penalty', 'early_stopping']:
            generate_params[k] = generate_state[k]
-        generate_params["eos_token_id"] = eos_token_ids
-        generate_params["stopping_criteria"] = stopping_criteria_list
+        generate_params['eos_token_id'] = eos_token_ids
+        generate_params['stopping_criteria'] = stopping_criteria_list
        if shared.args.no_stream:
-            generate_params["min_length"] = 0
+            generate_params['min_length'] = 0
    else:
-        for k in ["do_sample", "temperature"]:
+        for k in ['max_new_tokens', 'do_sample', 'temperature']:
            generate_params[k] = generate_state[k]
-        generate_params["stop"] = generate_state["eos_token_ids"][-1]
+        generate_params['stop'] = generate_state['eos_token_ids'][-1]
        if not shared.args.no_stream:
-            generate_params["max_new_tokens"] = 8
+            generate_params['max_new_tokens'] = 8

    if shared.args.no_cache:
-        generate_params.update({"use_cache": False})
+        generate_params.update({'use_cache': False})
    if shared.args.deepspeed:
-        generate_params.update({"synced_gpus": True})
+        generate_params.update({'synced_gpus': True})
    if shared.soft_prompt:
        inputs_embeds, filler_input_ids = generate_softprompt_input_tensors(input_ids)
-        generate_params.update({"inputs_embeds": inputs_embeds})
-        generate_params.update({"inputs": filler_input_ids})
+        generate_params.update({'inputs_embeds': inputs_embeds})
+        generate_params.update({'inputs': filler_input_ids})
    else:
-        generate_params.update({"inputs": input_ids})
+        generate_params.update({'inputs': input_ids})

    try:
        # Generate the entire reply at once.
@@ -205,7 +209,7 @@ def generate_reply(question, generate_state, eos_token=None, stopping_strings=[]
            new_tokens = len(output) - len(input_ids[0])
            reply = decode(output[-new_tokens:])
            if not shared.is_chat():
-                reply = original_question + apply_extensions(reply, "output")
+                reply = original_question + apply_extensions(reply, 'output')

            yield formatted_outputs(reply, shared.model_name)

@@ -232,7 +236,7 @@ def generate_reply(question, generate_state, eos_token=None, stopping_strings=[]
                    new_tokens = len(output) - len(input_ids[0])
                    reply = decode(output[-new_tokens:])
                    if not shared.is_chat():
-                        reply = original_question + apply_extensions(reply, "output")
+                        reply = original_question + apply_extensions(reply, 'output')

                    if output[-1] in eos_token_ids:
                        break
@@ -250,7 +254,7 @@ def generate_reply(question, generate_state, eos_token=None, stopping_strings=[]
                new_tokens = len(output) - len(original_input_ids[0])
                reply = decode(output[-new_tokens:])
                if not shared.is_chat():
-                    reply = original_question + apply_extensions(reply, "output")
+                    reply = original_question + apply_extensions(reply, 'output')

                if np.count_nonzero(np.isin(input_ids[0], eos_token_ids)) < np.count_nonzero(np.isin(output, eos_token_ids)):
                    break
@@ -259,10 +263,10 @@ def generate_reply(question, generate_state, eos_token=None, stopping_strings=[]
                input_ids = np.reshape(output, (1, output.shape[0]))
                if shared.soft_prompt:
                    inputs_embeds, filler_input_ids = generate_softprompt_input_tensors(input_ids)
-                    generate_params.update({"inputs_embeds": inputs_embeds})
-                    generate_params.update({"inputs": filler_input_ids})
+                    generate_params.update({'inputs_embeds': inputs_embeds})
+                    generate_params.update({'inputs': filler_input_ids})
                else:
-                    generate_params.update({"inputs": input_ids})
+                    generate_params.update({'inputs': input_ids})

            yield formatted_outputs(reply, shared.model_name)

@@ -272,5 +276,5 @@ def generate_reply(question, generate_state, eos_token=None, stopping_strings=[]
        t1 = time.time()
        original_tokens = len(original_input_ids[0])
        new_tokens = len(output) - original_tokens
-        print(f"Output generated in {(t1-t0):.2f} seconds ({new_tokens/(t1-t0):.2f} tokens/s, {new_tokens} tokens, context {original_tokens})")
+        print(f'Output generated in {(t1-t0):.2f} seconds ({new_tokens/(t1-t0):.2f} tokens/s, {new_tokens} tokens, context {original_tokens})')
        return
--- a/modules/training.py
+++ b/modules/training.py
@@ -19,9 +19,11 @@ CURRENT_STEPS = 0
 MAX_STEPS = 0
 CURRENT_GRADIENT_ACCUM = 1

+
 def get_dataset(path: str, ext: str):
    return ['None'] + sorted(set([k.stem for k in Path(path).glob(f'*.{ext}') if k.stem != 'put-trainer-datasets-here']), key=str.lower)

+
 def create_train_interface():
    with gr.Tab('Train LoRA', elem_id='lora-train-tab'):
        lora_name = gr.Textbox(label="Name", info="The name of your new LoRA file")
@@ -67,10 +69,12 @@ def create_train_interface():
                                      cutoff_len, dataset, eval_dataset, format, raw_text_file, overlap_len, newline_favor_len], [output])
        stop_button.click(do_interrupt, [], [], cancels=[], queue=False)

+
 def do_interrupt():
    global WANT_INTERRUPT
    WANT_INTERRUPT = True

+
 class Callbacks(transformers.TrainerCallback):
    def on_step_begin(self, args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs):
        global CURRENT_STEPS, MAX_STEPS
@@ -79,6 +83,7 @@ class Callbacks(transformers.TrainerCallback):
        if WANT_INTERRUPT:
            control.should_epoch_stop = True
            control.should_training_stop = True
+
    def on_substep_end(self, args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs):
        global CURRENT_STEPS
        CURRENT_STEPS += 1
@@ -86,6 +91,7 @@ class Callbacks(transformers.TrainerCallback):
            control.should_epoch_stop = True
            control.should_training_stop = True

+
 def clean_path(base_path: str, path: str):
    """"Strips unusual symbols and forcibly builds a path as relative to the intended directory."""
    # TODO: Probably could do with a security audit to guarantee there's no ways this can be bypassed to target an unwanted path.
@@ -95,6 +101,7 @@ def clean_path(base_path: str, path: str):
        return path
    return f'{Path(base_path).absolute()}/{path}'

+
 def do_train(lora_name: str, micro_batch_size: int, batch_size: int, epochs: int, learning_rate: str, lora_rank: int, lora_alpha: int, lora_dropout: float,
             cutoff_len: int, dataset: str, eval_dataset: str, format: str, raw_text_file: str, overlap_len: int, newline_favor_len: int):
    global WANT_INTERRUPT, CURRENT_STEPS, MAX_STEPS, CURRENT_GRADIENT_ACCUM
@@ -145,7 +152,7 @@ def do_train(lora_name: str, micro_batch_size: int, batch_size: int, epochs: int
    # == Prep the dataset, format, etc ==
    if raw_text_file not in ['None', '']:
        print("Loading raw text file dataset...")
-        with open(clean_path('training/datasets', f'{raw_text_file}.txt'), 'r') as file:
+        with open(clean_path('training/datasets', f'{raw_text_file}.txt'), 'r', encoding='utf-8') as file:
            raw_text = file.read()
        tokens = shared.tokenizer.encode(raw_text)
        del raw_text  # Note: could be a gig for a large dataset, so delete redundant data as we go to be safe on RAM
@@ -302,10 +309,12 @@ def do_train(lora_name: str, micro_batch_size: int, batch_size: int, epochs: int
        print("Training complete!")
        yield f"Done! LoRA saved to `{lora_name}`"

+
 def split_chunks(arr, step):
    for i in range(0, len(arr), step):
        yield arr[i:i + step]

+
 def cut_chunk_for_newline(chunk: str, max_length: int):
    if '\n' not in chunk:
        return chunk
@@ -319,6 +328,7 @@ def cut_chunk_for_newline(chunk: str, max_length: int):
        chunk = chunk[:last_newline]
    return chunk

+
 def format_time(seconds: float):
    if seconds < 120:
        return f"`{seconds:.0f}` seconds"
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -13,6 +13,7 @@ with open(Path(__file__).resolve().parent / '../css/main.js', 'r') as f:
 with open(Path(__file__).resolve().parent / '../css/chat.js', 'r') as f:
    chat_js = f.read()

+
 class ToolButton(gr.Button, gr.components.FormComponent):
    """Small button with single emoji as text, fits inside gradio forms"""

@@ -22,6 +23,7 @@ class ToolButton(gr.Button, gr.components.FormComponent):
    def get_block_name(self):
        return "button"

+
 def create_refresh_button(refresh_component, refresh_method, refreshed_args, elem_id):
    def refresh():
        refresh_method()
--- a/server.py
+++ b/server.py
@@ -15,12 +15,11 @@ import gradio as gr
 from PIL import Image

 import modules.extensions as extensions_module
-from modules import chat, shared, training, ui, api
+from modules import api, chat, shared, training, ui
 from modules.html_generator import chat_html_wrapper
 from modules.LoRA import add_lora_to_model
-from modules.models import load_model, load_soft_prompt
-from modules.text_generation import (clear_torch_cache, generate_reply,
-                                     stop_everything_event)
+from modules.models import load_model, load_soft_prompt, unload_model
+from modules.text_generation import generate_reply, stop_everything_event

 # Loading custom settings
 settings_file = None
@@ -34,15 +33,18 @@ if settings_file is not None:
    for item in new_settings:
        shared.settings[item] = new_settings[item]

+
 def get_available_models():
    if shared.args.flexgen:
        return sorted([re.sub('-np$', '', item.name) for item in list(Path(f'{shared.args.model_dir}/').glob('*')) if item.name.endswith('-np')], key=str.lower)
    else:
        return sorted([re.sub('.pth$', '', item.name) for item in list(Path(f'{shared.args.model_dir}/').glob('*')) if not item.name.endswith(('.txt', '-np', '.pt', '.json'))], key=str.lower)

+
 def get_available_presets():
    return sorted(set((k.stem for k in Path('presets').glob('*.txt'))), key=str.lower)

+
 def get_available_prompts():
    prompts = []
    prompts += sorted(set((k.stem for k in Path('prompts').glob('[0-9]*.txt'))), key=str.lower, reverse=True)
@@ -50,10 +52,12 @@ def get_available_prompts():
    prompts += ['None']
    return prompts

+
 def get_available_characters():
    paths = (x for x in Path('characters').iterdir() if x.suffix in ('.json', '.yaml', '.yml'))
    return ['None'] + sorted(set((k.stem for k in paths if k.stem != "instruction-following")), key=str.lower)

+
 def get_available_instruction_templates():
    path = "characters/instruction-following"
    paths = []
@@ -61,18 +65,18 @@ def get_available_instruction_templates():
        paths = (x for x in Path(path).iterdir() if x.suffix in ('.json', '.yaml', '.yml'))
    return ['None'] + sorted(set((k.stem for k in paths)), key=str.lower)

+
 def get_available_extensions():
    return sorted(set(map(lambda x: x.parts[1], Path('extensions').glob('*/script.py'))), key=str.lower)

+
 def get_available_softprompts():
    return ['None'] + sorted(set((k.stem for k in Path('softprompts').glob('*.zip'))), key=str.lower)

+
 def get_available_loras():
    return ['None'] + sorted([item.name for item in list(Path(shared.args.lora_dir).glob('*')) if not item.name.endswith(('.txt', '-np', '.pt', '.json'))], key=str.lower)

-def unload_model():
-    shared.model = shared.tokenizer = None
-    clear_torch_cache()

 def load_model_wrapper(selected_model):
    if selected_model != shared.model_name:
@@ -84,10 +88,12 @@ def load_model_wrapper(selected_model):

    return selected_model

+
 def load_lora_wrapper(selected_lora):
    add_lora_to_model(selected_lora)
    return selected_lora

+
 def load_preset_values(preset_menu, state, return_dict=False):
    generate_params = {
        'do_sample': True,
@@ -118,6 +124,7 @@ def load_preset_values(preset_menu, state, return_dict=False):
        state.update(generate_params)
        return state, *[generate_params[k] for k in ['do_sample', 'temperature', 'top_p', 'typical_p', 'repetition_penalty', 'encoder_repetition_penalty', 'top_k', 'min_length', 'no_repeat_ngram_size', 'num_beams', 'penalty_alpha', 'length_penalty', 'early_stopping']]

+
 def upload_soft_prompt(file):
    with zipfile.ZipFile(io.BytesIO(file)) as zf:
        zf.extract('meta.json')
@@ -130,12 +137,14 @@ def upload_soft_prompt(file):

    return name

+
 def save_prompt(text):
    fname = f"{datetime.now().strftime('%Y-%m-%d-%H%M%S')}.txt"
    with open(Path(f'prompts/{fname}'), 'w', encoding='utf-8') as f:
        f.write(text)
    return f"Saved to prompts/{fname}"

+
 def load_prompt(fname):
    if fname in ['None', '']:
        return ''
@@ -146,6 +155,7 @@ def load_prompt(fname):
                text = text[:-1]
            return text

+
 def create_prompt_menus():
    with gr.Row():
        with gr.Column():
@@ -161,6 +171,7 @@ def create_prompt_menus():
    shared.gradio['prompt_menu'].change(load_prompt, [shared.gradio['prompt_menu']], [shared.gradio['textbox']], show_progress=False)
    shared.gradio['save_prompt'].click(save_prompt, [shared.gradio['textbox']], [shared.gradio['status']], show_progress=False)

+
 def create_model_menus():
    with gr.Row():
        with gr.Column():
@@ -175,6 +186,7 @@ def create_model_menus():
    shared.gradio['model_menu'].change(load_model_wrapper, shared.gradio['model_menu'], shared.gradio['model_menu'], show_progress=True)
    shared.gradio['lora_menu'].change(load_lora_wrapper, shared.gradio['lora_menu'], shared.gradio['lora_menu'], show_progress=True)

+
 def create_settings_menus(default_preset):
    generate_params = load_preset_values(default_preset if not shared.args.flexgen else 'Naive', {}, return_dict=True)
    for k in ['max_new_tokens', 'seed', 'stop_at_newline', 'chat_prompt_size', 'chat_generation_attempts']:
@@ -209,7 +221,6 @@ def create_settings_menus(default_preset):
            with gr.Box():
                gr.Markdown('Contrastive search')
                shared.gradio['penalty_alpha'] = gr.Slider(0, 5, value=generate_params['penalty_alpha'], label='penalty_alpha')
-
            with gr.Box():
                gr.Markdown('Beam search (uses a lot of VRAM)')
                with gr.Row():
@@ -219,7 +230,6 @@ def create_settings_menus(default_preset):
                        shared.gradio['length_penalty'] = gr.Slider(-5, 5, value=generate_params['length_penalty'], label='length_penalty')
                shared.gradio['early_stopping'] = gr.Checkbox(value=generate_params['early_stopping'], label='early_stopping')

-
    with gr.Accordion('Soft prompt', open=False):
        with gr.Row():
            shared.gradio['softprompts_menu'] = gr.Dropdown(choices=available_softprompts, value='None', label='Soft prompt')
@@ -233,6 +243,7 @@ def create_settings_menus(default_preset):
    shared.gradio['softprompts_menu'].change(load_soft_prompt, shared.gradio['softprompts_menu'], shared.gradio['softprompts_menu'], show_progress=True)
    shared.gradio['upload_softprompt'].upload(upload_soft_prompt, shared.gradio['upload_softprompt'], shared.gradio['softprompts_menu'])

+
 def set_interface_arguments(interface_mode, extensions, bool_active):
    modes = ["default", "notebook", "chat", "cai_chat"]
    cmd_list = vars(shared.args)
@@ -251,6 +262,7 @@ def set_interface_arguments(interface_mode, extensions, bool_active):

    shared.need_restart = True

+
 available_models = get_available_models()
 available_presets = get_available_presets()
 available_characters = get_available_characters()
@@ -299,8 +311,8 @@ else:
    default_text = load_prompt(shared.settings['prompts'][next((k for k in shared.settings['prompts'] if re.match(k.lower(), shared.model_name.lower())), 'default')])
 title = 'Text generation web UI'

-def create_interface():

+def create_interface():
    gen_events = []
    if shared.args.extensions is not None and len(shared.args.extensions) > 0:
        extensions_module.load_extensions()
@@ -312,7 +324,7 @@ def create_interface():
                shared.gradio['display'] = gr.HTML(value=chat_html_wrapper(shared.history['visible'], shared.settings['name1'], shared.settings['name2'], 'cai-chat'))
                shared.gradio['textbox'] = gr.Textbox(label='Input')
                with gr.Row():
-                    shared.gradio['Generate'] = gr.Button('Generate')
+                    shared.gradio['Generate'] = gr.Button('Generate', elem_id='Generate')
                    shared.gradio['Stop'] = gr.Button('Stop', elem_id="stop")
                with gr.Row():
                    shared.gradio['Impersonate'] = gr.Button('Impersonate')
@@ -382,55 +394,71 @@ def create_interface():
                create_settings_menus(default_preset)

            shared.input_params = [shared.gradio[k] for k in ['Chat input', 'generate_state', 'name1', 'name2', 'context', 'Chat mode', 'end_of_turn']]
+            clear_arr = [shared.gradio[k] for k in ['Clear history-confirm', 'Clear history', 'Clear history-cancel']]
+            reload_inputs = [shared.gradio[k] for k in ['name1', 'name2', 'Chat mode']]

            def set_chat_input(textbox):
                return textbox, ""

-            gen_events.append(shared.gradio['Generate'].click(set_chat_input, shared.gradio['textbox'], [shared.gradio['Chat input'], shared.gradio['textbox']], show_progress=False))
-            gen_events.append(shared.gradio['Generate'].click(chat.cai_chatbot_wrapper, shared.input_params, shared.gradio['display'], show_progress=shared.args.no_stream))
-            gen_events.append(shared.gradio['textbox'].submit(set_chat_input, shared.gradio['textbox'], [shared.gradio['Chat input'], shared.gradio['textbox']], show_progress=False))
-            gen_events.append(shared.gradio['textbox'].submit(chat.cai_chatbot_wrapper, shared.input_params, shared.gradio['display'], show_progress=shared.args.no_stream))
-            gen_events.append(shared.gradio['Regenerate'].click(chat.regenerate_wrapper, shared.input_params, shared.gradio['display'], show_progress=shared.args.no_stream))
+            gen_events.append(shared.gradio['Generate'].click(
+                set_chat_input, shared.gradio['textbox'], [shared.gradio['Chat input'], shared.gradio['textbox']], show_progress=False).then(
+                chat.cai_chatbot_wrapper, shared.input_params, shared.gradio['display'], show_progress=shared.args.no_stream).then(
+                lambda x: '', shared.gradio['textbox'], shared.gradio['textbox'], show_progress=False).then(
+                lambda: chat.save_history(timestamp=False), [], [], show_progress=False)
+            )
+
+            gen_events.append(shared.gradio['textbox'].submit(
+                set_chat_input, shared.gradio['textbox'], [shared.gradio['Chat input'], shared.gradio['textbox']], show_progress=False).then(
+                chat.cai_chatbot_wrapper, shared.input_params, shared.gradio['display'], show_progress=shared.args.no_stream).then(
+                lambda x: '', shared.gradio['textbox'], shared.gradio['textbox'], show_progress=False).then(
+                lambda: chat.save_history(timestamp=False), [], [], show_progress=False)
+            )
+
+            gen_events.append(shared.gradio['Regenerate'].click(
+                chat.regenerate_wrapper, shared.input_params, shared.gradio['display'], show_progress=shared.args.no_stream).then(
+                lambda x: '', shared.gradio['textbox'], shared.gradio['textbox'], show_progress=False).then(
+                lambda: chat.save_history(timestamp=False), [], [], show_progress=False)
+            )
+
+            shared.gradio['Replace last reply'].click(
+                chat.replace_last_reply, [shared.gradio[k] for k in ['textbox', 'name1', 'name2', 'Chat mode']], shared.gradio['display'], show_progress=shared.args.no_stream).then(
+                lambda x: '', shared.gradio['textbox'], shared.gradio['textbox'], show_progress=False).then(
+                lambda: chat.save_history(timestamp=False), [], [], show_progress=False)
+
+            shared.gradio['Clear history-confirm'].click(
+                lambda: [gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)], None, clear_arr).then(
+                chat.clear_chat_log, [shared.gradio[k] for k in ['name1', 'name2', 'greeting', 'Chat mode']], shared.gradio['display']).then(
+                lambda: chat.save_history(timestamp=False), [], [], show_progress=False)
+
+            shared.gradio['Stop'].click(
+                stop_everything_event, [], [], queue=False, cancels=gen_events if shared.args.no_stream else None).then(
+                chat.redraw_html, reload_inputs, [shared.gradio['display']])
+
+            shared.gradio['Chat mode'].change(
+                lambda x: gr.update(visible=x == 'instruct'), shared.gradio['Chat mode'], shared.gradio['Instruction templates']).then(
+                chat.redraw_html, reload_inputs, [shared.gradio['display']])
+
+            shared.gradio['Instruction templates'].change(
+                lambda character, name1, name2, mode: chat.load_character(character, name1, name2, mode), [shared.gradio[k] for k in ['Instruction templates', 'name1', 'name2', 'Chat mode']], [shared.gradio[k] for k in ['name1', 'name2', 'character_picture', 'greeting', 'context', 'end_of_turn', 'display']]).then(
+                chat.redraw_html, reload_inputs, [shared.gradio['display']])
+
+            shared.gradio['upload_chat_history'].upload(
+                chat.load_history, [shared.gradio[k] for k in ['upload_chat_history', 'name1', 'name2']], []).then(
+                chat.redraw_html, reload_inputs, [shared.gradio['display']])
+
            gen_events.append(shared.gradio['Impersonate'].click(chat.impersonate_wrapper, shared.input_params, shared.gradio['textbox'], show_progress=shared.args.no_stream))
-            shared.gradio['Stop'].click(stop_everything_event, [], [], queue=False, cancels=gen_events if shared.args.no_stream else None)
-
            shared.gradio['Copy last reply'].click(chat.send_last_reply_to_input, [], shared.gradio['textbox'], show_progress=shared.args.no_stream)
-            shared.gradio['Replace last reply'].click(chat.replace_last_reply, [shared.gradio[k] for k in ['textbox', 'name1', 'name2', 'Chat mode']], shared.gradio['display'], show_progress=shared.args.no_stream)
-
-            # Clear history with confirmation
-            clear_arr = [shared.gradio[k] for k in ['Clear history-confirm', 'Clear history', 'Clear history-cancel']]
            shared.gradio['Clear history'].click(lambda: [gr.update(visible=True), gr.update(visible=False), gr.update(visible=True)], None, clear_arr)
-            shared.gradio['Clear history-confirm'].click(lambda :[gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)], None, clear_arr)
-            shared.gradio['Clear history-confirm'].click(chat.clear_chat_log, [shared.gradio[k] for k in ['name1', 'name2', 'greeting', 'Chat mode']], shared.gradio['display'])
            shared.gradio['Clear history-cancel'].click(lambda: [gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)], None, clear_arr)
-            shared.gradio['Chat mode'].change(lambda x : gr.update(visible= x=='instruct'), shared.gradio['Chat mode'], shared.gradio['Instruction templates'])
-
            shared.gradio['Remove last'].click(chat.remove_last_message, [shared.gradio[k] for k in ['name1', 'name2', 'Chat mode']], [shared.gradio['display'], shared.gradio['textbox']], show_progress=False)
            shared.gradio['download_button'].click(chat.save_history, inputs=[], outputs=[shared.gradio['download']])
            shared.gradio['Upload character'].click(chat.upload_character, [shared.gradio['upload_json'], shared.gradio['upload_img_bot']], [shared.gradio['character_menu']])
-
-            # Clearing stuff and saving the history
-            for i in ['Generate', 'Regenerate', 'Replace last reply']:
-                shared.gradio[i].click(lambda x: '', shared.gradio['textbox'], shared.gradio['textbox'], show_progress=False)
-                shared.gradio[i].click(lambda : chat.save_history(timestamp=False), [], [], show_progress=False)
-            shared.gradio['Clear history-confirm'].click(lambda : chat.save_history(timestamp=False), [], [], show_progress=False)
-            shared.gradio['textbox'].submit(lambda x: '', shared.gradio['textbox'], shared.gradio['textbox'], show_progress=False)
-            shared.gradio['textbox'].submit(lambda : chat.save_history(timestamp=False), [], [], show_progress=False)
-
            shared.gradio['character_menu'].change(chat.load_character, [shared.gradio[k] for k in ['character_menu', 'name1', 'name2', 'Chat mode']], [shared.gradio[k] for k in ['name1', 'name2', 'character_picture', 'greeting', 'context', 'end_of_turn', 'display']])
-            shared.gradio['Instruction templates'].change(lambda character, name1, name2, mode: chat.load_character(character, name1, name2, mode), [shared.gradio[k] for k in ['Instruction templates', 'name1', 'name2', 'Chat mode']], [shared.gradio[k] for k in ['name1', 'name2', 'character_picture', 'greeting', 'context', 'end_of_turn', 'display']])
-            shared.gradio['upload_chat_history'].upload(chat.load_history, [shared.gradio[k] for k in ['upload_chat_history', 'name1', 'name2']], [])
            shared.gradio['upload_img_tavern'].upload(chat.upload_tavern_character, [shared.gradio['upload_img_tavern'], shared.gradio['name1'], shared.gradio['name2']], [shared.gradio['character_menu']])
            shared.gradio['your_picture'].change(chat.upload_your_profile_picture, [shared.gradio[k] for k in ['your_picture', 'name1', 'name2', 'Chat mode']], shared.gradio['display'])

-            reload_inputs = [shared.gradio[k] for k in ['name1', 'name2', 'Chat mode']]
-            shared.gradio['upload_chat_history'].upload(chat.redraw_html, reload_inputs, [shared.gradio['display']])
-            shared.gradio['Stop'].click(chat.redraw_html, reload_inputs, [shared.gradio['display']])
-            shared.gradio['Instruction templates'].change(chat.redraw_html, reload_inputs, [shared.gradio['display']])
-            shared.gradio['Chat mode'].change(chat.redraw_html, reload_inputs, [shared.gradio['display']])
-
            shared.gradio['interface'].load(None, None, None, _js=f"() => {{{ui.main_js+ui.chat_js}}}")
-            shared.gradio['interface'].load(lambda : chat.load_default_history(shared.settings['name1'], shared.settings['name2']), None, None)
+            shared.gradio['interface'].load(chat.load_default_history, [shared.gradio[k] for k in ['name1', 'name2']], None)
            shared.gradio['interface'].load(chat.redraw_html, reload_inputs, [shared.gradio['display']], show_progress=True)

        elif shared.args.notebook:
@@ -523,10 +551,12 @@ def create_interface():
            shared.gradio['interface_modes_menu'] = gr.Dropdown(choices=modes, value=current_mode, label="Mode")
            shared.gradio['extensions_menu'] = gr.CheckboxGroup(choices=get_available_extensions(), value=shared.args.extensions, label="Available extensions")
            shared.gradio['bool_menu'] = gr.CheckboxGroup(choices=bool_list, value=bool_active, label="Boolean command-line flags")
-            shared.gradio['reset_interface'] = gr.Button("Apply and restart the interface", type="primary")
+            shared.gradio['reset_interface'] = gr.Button("Apply and restart the interface")

-            shared.gradio['reset_interface'].click(set_interface_arguments, [shared.gradio[k] for k in ['interface_modes_menu', 'extensions_menu', 'bool_menu']], None)
-            shared.gradio['reset_interface'].click(lambda : None, None, None, _js='() => {document.body.innerHTML=\'<h1 style="font-family:monospace;margin-top:20%;color:lightgray;text-align:center;">Reloading...</h1>\'; setTimeout(function(){location.reload()},2500); return []}')
+            # Reset interface event
+            shared.gradio['reset_interface'].click(
+                set_interface_arguments, [shared.gradio[k] for k in ['interface_modes_menu', 'extensions_menu', 'bool_menu']], None).then(
+                lambda: None, None, None, _js='() => {document.body.innerHTML=\'<h1 style="font-family:monospace;margin-top:20%;color:lightgray;text-align:center;">Reloading...</h1>\'; setTimeout(function(){location.reload()},2500); return []}')

        if shared.args.extensions is not None:
            extensions_module.create_extensions_block()
@@ -562,6 +592,7 @@ def create_interface():
    else:
        shared.gradio['interface'].launch(prevent_thread_lock=True, share=shared.args.share, server_port=shared.args.listen_port, inbrowser=shared.args.auto_launch, auth=auth)

+
 create_interface()

 while True:
Author	SHA1	Message	Date
oobabooga	25652f0994	Properly concatenate chat events	2023-04-08 17:16:46 -03:00
oobabooga	0b458bf82d	Simplify a function	2023-04-07 21:37:41 -03:00
Φφ	ffd102e5c0	SD Api Pics extension, v.1.1 (#596 )	2023-04-07 21:36:04 -03:00
oobabooga	5543a5089d	Auto-submit the whisper extension transcription	2023-04-07 15:57:51 -03:00
oobabooga	1dc464dcb0	Sort imports	2023-04-07 14:42:03 -03:00
oobabooga	962e33dc10	Change button style	2023-04-07 12:22:14 -03:00
oobabooga	42ea6a3fc0	Change the timing for setup() calls	2023-04-07 12:20:57 -03:00
Φφ	e563b015d8	Silero TTS offline cache (#628 )	2023-04-07 12:15:57 -03:00
oobabooga	1c413ed593	Remove torch from silero	2023-04-07 11:51:50 -03:00
da3dsoul	3f922d4bfb	Extract the Preprocessing for Silero into a file and Improve it (#757 )	2023-04-07 11:46:29 -03:00
Maya	744bf7cbf2	Get rid of type parameter warning (#883 ) Fix annoying `The 'type' parameter has been deprecated. Use the Number component instead` warning	2023-04-07 11:17:16 -03:00
oobabooga	768354239b	Change training file encoding	2023-04-07 11:15:52 -03:00
oobabooga	6762e62a40	Simplifications	2023-04-07 11:14:32 -03:00
oobabooga	a453d4e9c4	Reorganize some chat functions	2023-04-07 11:07:03 -03:00
MarlinMr	ec979cd9c4	Use updated docker compose (#877 )	2023-04-07 10:48:47 -03:00
MarlinMr	2c0018d946	Cosmetic change of README.md (#878 )	2023-04-07 10:47:10 -03:00
Maya	8fa182cfa7	Fix regeneration of first message in instruct mode (#881 )	2023-04-07 10:45:42 -03:00
Alastair D'Silva	862aad637b	Tweak COPY order in Dockerfile (#863 )	2023-04-07 00:56:44 -03:00
oobabooga	46c4654226	More PEP8 stuff	2023-04-07 00:52:02 -03:00
oobabooga	ea6e77df72	Make the code more like PEP8 for readability (#862 )	2023-04-07 00:15:45 -03:00