From e51ece21c04dc80477601c9e7b27c20f5867af05 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Fri, 24 Feb 2023 19:00:11 -0300
Subject: [PATCH] Add ui() function to extensions

---
 extensions/character_bias/script.py   |  9 ++++++++
 extensions/google_translate/script.py | 13 ++++++++++++
 extensions/silero_tts/script.py       | 10 +++++++++
 modules/extensions.py                 | 30 +++++++--------------------
 server.py                             |  3 +--
 5 files changed, 40 insertions(+), 25 deletions(-)

diff --git a/extensions/character_bias/script.py b/extensions/character_bias/script.py
index 9660a59a..aa949f29 100644
--- a/extensions/character_bias/script.py
+++ b/extensions/character_bias/script.py
@@ -1,3 +1,5 @@
+import gradio as gr
+
 params = {
     "bias string": " *I speak in an annoyingly cute way*",
 }
@@ -25,3 +27,10 @@ def bot_prefix_modifier(string):
     """
 
     return f'{string} {params["bias string"].strip()} '
+
+def ui():
+    # Gradio elements
+    string = gr.Textbox(value=params["bias string"], label='Character bias')
+
+    # Event functions to update the parameters in the backend
+    string.change(lambda x: params.update({"bias string": x}), string, None)
diff --git a/extensions/google_translate/script.py b/extensions/google_translate/script.py
index 064a7ec7..68bc54b2 100644
--- a/extensions/google_translate/script.py
+++ b/extensions/google_translate/script.py
@@ -1,9 +1,12 @@
+import gradio as gr
 from deep_translator import GoogleTranslator
 
 params = {
     "language string": "ja",
 }
 
+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
@@ -27,3 +30,13 @@ 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'])]
+
+    # Gradio elements
+    language = gr.Dropdown(value=language_name, choices=[k for k in language_codes], label='Language')
+
+    # Event functions to update the parameters in the backend
+    language.change(lambda x: params.update({"language string": language_codes[x]}), language, None)
diff --git a/extensions/silero_tts/script.py b/extensions/silero_tts/script.py
index 3328422b..6c0c4ef6 100644
--- a/extensions/silero_tts/script.py
+++ b/extensions/silero_tts/script.py
@@ -1,6 +1,7 @@
 import asyncio
 from pathlib import Path
 
+import gradio as gr
 import torch
 
 torch._C._jit_set_profiling_mode(False)
@@ -81,3 +82,12 @@ def bot_prefix_modifier(string):
     """
 
     return string
+
+def ui():
+    # Gradio elements
+    activate = gr.Checkbox(value=params['activate'], label='Activate TTS')
+    voice = gr.Dropdown(value=params['speaker'], choices=[f'en_{i}' for i in range(1, 118)], label='TTS voice')
+
+    # Event functions to update the parameters in the backend
+    activate.change(lambda x: params.update({"activate": x}), activate, None)
+    voice.change(lambda x: params.update({"speaker": x}), voice, None)
diff --git a/modules/extensions.py b/modules/extensions.py
index be829330..6af7c58a 100644
--- a/modules/extensions.py
+++ b/modules/extensions.py
@@ -32,31 +32,15 @@ def apply_extensions(text, typ):
             text = extension.bot_prefix_modifier(text)
     return text
 
-def update_extensions_parameters(*args):
-    i = 0
-    for extension, _ in iterator():
-        for param in extension.params:
-            if len(args) >= i+1:
-                extension.params[param] = eval(f"args[{i}]")
-                i += 1
-
 def create_extensions_block():
-    extensions_ui_elements = []
-    default_values = []
-    if not (shared.args.chat or shared.args.cai_chat):
-        gr.Markdown('## Extensions parameters')
+    # Updating the default values
     for extension, name in iterator():
         for param in extension.params:
             _id = f"{name}-{param}"
-            default_value = shared.settings[_id] if _id in shared.settings else extension.params[param]
-            default_values.append(default_value)
-            if type(extension.params[param]) == str:
-                extensions_ui_elements.append(gr.Textbox(value=default_value, label=f"{name}-{param}"))
-            elif type(extension.params[param]) in [int, float]:
-                extensions_ui_elements.append(gr.Number(value=default_value, label=f"{name}-{param}"))
-            elif type(extension.params[param]) == bool:
-                extensions_ui_elements.append(gr.Checkbox(value=default_value, label=f"{name}-{param}"))
+            if _id in shared.settings:
+                extension.params[param] = shared.settings[_id]
 
-    update_extensions_parameters(*default_values)
-    btn_extensions = gr.Button("Apply")
-    btn_extensions.click(update_extensions_parameters, [*extensions_ui_elements], [])
+    # Creating the extension ui elements
+    for extension, name in iterator():
+        if hasattr(extension, "ui"):
+            extension.ui()
diff --git a/server.py b/server.py
index 2df5f7e6..5db441b6 100644
--- a/server.py
+++ b/server.py
@@ -247,7 +247,7 @@ if shared.args.chat or shared.args.cai_chat:
                 with gr.Column():
                     shared.gradio['chat_prompt_size_slider'] = gr.Slider(minimum=shared.settings['chat_prompt_size_min'], maximum=shared.settings['chat_prompt_size_max'], step=1, label='Maximum prompt size in tokens', value=shared.settings['chat_prompt_size'])
 
-        create_settings_menus()
+            create_settings_menus()
         if shared.args.extensions is not None:
             with gr.Tab('Extensions'):
                 extensions_module.create_extensions_block()
@@ -351,7 +351,6 @@ else:
         shared.gradio['Stop'].click(None, None, None, cancels=gen_events)
 
 shared.gradio['interface'].queue()
-print([key for key in shared.gradio])
 if shared.args.listen:
     shared.gradio['interface'].launch(prevent_thread_lock=True, share=shared.args.share, server_name='0.0.0.0', server_port=shared.args.listen_port)
 else: