Simplify the extensions implementation

modules/extensions.py

@@ -2,62 +2,58 @@ import extensions
 import modules.shared as shared
 import gradio as gr
 
-extension_state = {}
+state = {}
 available_extensions = []
 
 def load_extensions():
-    global extension_state
+    global state
-    for i,ext in enumerate(shared.args.extensions):
+    for i, name in enumerate(shared.args.extensions):
-        if ext in available_extensions:
+        if name in available_extensions:
-            print(f'Loading the extension "{ext}"... ', end='')
+            print(f'Loading the extension "{name}"... ', end='')
-            ext_string = f"extensions.{ext}.script"
+            import_string = f"extensions.{name}.script"
-            exec(f"import {ext_string}")
+            exec(f"import {import_string}")
-            extension_state[ext] = [True, i]
+            state[name] = [True, i]
             print(f'Ok.')
 
+def iterator():
+    for name in sorted(state, key=lambda x : state[x][1]):
+        if state[name][0] == True:
+            yield eval(f"extensions.{name}.script"), name
+
 def apply_extensions(text, typ):
-    for ext in sorted(extension_state, key=lambda x : extension_state[x][1]):
+    for extension, _ in iterator():
-        if extension_state[ext][0] == True:
+        if typ == "input" and hasattr(extension, "input_modifier"):
-            ext_string = f"extensions.{ext}.script"
+            text = extension.input_modifier(text)
-            if typ == "input" and hasattr(eval(ext_string), "input_modifier"):
+        elif typ == "output" and hasattr(extension, "output_modifier"):
-                text = eval(f"{ext_string}.input_modifier(text)")
+            text = extension.output_modifier(text)
-            elif typ == "output" and hasattr(eval(ext_string), "output_modifier"):
+        elif typ == "bot_prefix" and hasattr(extension, "bot_prefix_modifier"):
-                text = eval(f"{ext_string}.output_modifier(text)")
+            text = extension.bot_prefix_modifier(text)
-            elif typ == "bot_prefix" and hasattr(eval(ext_string), "bot_prefix_modifier"):
-                text = eval(f"{ext_string}.bot_prefix_modifier(text)")
     return text
 
-def update_extensions_parameters(*kwargs):
+def update_extensions_parameters(*args):
     i = 0
-    for ext in sorted(extension_state, key=lambda x : extension_state[x][1]):
+    for extension, _ in iterator():
-        if extension_state[ext][0] == True:
+        for param in extension.params:
-            params = eval(f"extensions.{ext}.script.params")
+            if len(args) >= i+1:
-            for param in params:
+                extension.params[param] = eval(f"args[{i}]")
-                if len(kwargs) >= i+1:
+                i += 1
-                    params[param] = eval(f"kwargs[{i}]")
-                    i += 1
-
-def get_params(name):
-    return eval(f"extensions.{name}.script.params")
 
 def create_extensions_block():
     extensions_ui_elements = []
     default_values = []
     if not (shared.args.chat or shared.args.cai_chat):
         gr.Markdown('## Extensions parameters')
-    for ext in sorted(extension_state, key=lambda x : extension_state[x][1]):
+    for extension, name in iterator():
-        if extension_state[ext][0] == True:
+        for param in extension.params:
-            params = get_params(ext)
+            _id = f"{name}-{param}"
-            for param in params:
+            default_value = shared.settings[_id] if _id in shared.settings else extension.params[param]
-                _id = f"{ext}-{param}"
+            default_values.append(default_value)
-                default_value = shared.settings[_id] if _id in shared.settings else params[param]
+            if type(extension.params[param]) == str:
-                default_values.append(default_value)
+                extensions_ui_elements.append(gr.Textbox(value=default_value, label=f"{name}-{param}"))
-                if type(params[param]) == str:
+            elif type(extension.params[param]) in [int, float]:
-                    extensions_ui_elements.append(gr.Textbox(value=default_value, label=f"{ext}-{param}"))
+                extensions_ui_elements.append(gr.Number(value=default_value, label=f"{name}-{param}"))
-                elif type(params[param]) in [int, float]:
+            elif type(extension.params[param]) == bool:
-                    extensions_ui_elements.append(gr.Number(value=default_value, label=f"{ext}-{param}"))
+                extensions_ui_elements.append(gr.Checkbox(value=default_value, label=f"{name}-{param}"))
-                elif type(params[param]) == bool:
-                    extensions_ui_elements.append(gr.Checkbox(value=default_value, label=f"{ext}-{param}"))
 
     update_extensions_parameters(*default_values)
     btn_extensions = gr.Button("Apply")