Simplify the extensions implementation

This commit is contained in:
oobabooga 2023-02-24 10:01:21 -03:00
parent e26118eba9
commit fe5057f932

View File

@ -2,62 +2,58 @@ import extensions
import modules.shared as shared import modules.shared as shared
import gradio as gr import gradio as gr
extension_state = {} state = {}
available_extensions = [] available_extensions = []
def load_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.') 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): 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 return text
def update_extensions_parameters(*kwargs): def update_extensions_parameters(*args):
i = 0 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:
params[param] = eval(f"kwargs[{i}]")
i += 1 i += 1
def get_params(name):
return eval(f"extensions.{name}.script.params")
def create_extensions_block(): def create_extensions_block():
extensions_ui_elements = [] extensions_ui_elements = []
default_values = [] default_values = []
if not (shared.args.chat or shared.args.cai_chat): if not (shared.args.chat or shared.args.cai_chat):
gr.Markdown('## Extensions parameters') 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_value = shared.settings[_id] if _id in shared.settings else params[param]
default_values.append(default_value) default_values.append(default_value)
if type(params[param]) == str: if type(extension.params[param]) == str:
extensions_ui_elements.append(gr.Textbox(value=default_value, label=f"{ext}-{param}")) extensions_ui_elements.append(gr.Textbox(value=default_value, label=f"{name}-{param}"))
elif type(params[param]) in [int, float]: elif type(extension.params[param]) in [int, float]:
extensions_ui_elements.append(gr.Number(value=default_value, label=f"{ext}-{param}")) extensions_ui_elements.append(gr.Number(value=default_value, label=f"{name}-{param}"))
elif type(params[param]) == bool: elif type(extension.params[param]) == bool:
extensions_ui_elements.append(gr.Checkbox(value=default_value, label=f"{ext}-{param}")) extensions_ui_elements.append(gr.Checkbox(value=default_value, label=f"{name}-{param}"))
update_extensions_parameters(*default_values) update_extensions_parameters(*default_values)
btn_extensions = gr.Button("Apply") btn_extensions = gr.Button("Apply")