text-generation-webui/modules/extensions.py

63 lines
2.1 KiB
Python
Raw Normal View History

2023-03-15 19:11:16 +01:00
import gradio as gr
import extensions
2023-02-23 18:41:42 +01:00
import modules.shared as shared
2023-02-24 14:01:21 +01:00
state = {}
available_extensions = []
2023-03-19 14:22:24 +01:00
setup_called = False
2023-02-23 18:49:02 +01:00
def load_extensions():
2023-02-24 14:01:21 +01:00
global state
for i, name in enumerate(shared.args.extensions):
if name in available_extensions:
print(f'Loading the extension "{name}"... ', end='')
2023-03-16 03:29:56 +01:00
try:
exec(f"import extensions.{name}.script")
state[name] = [True, i]
print('Ok.')
except:
print('Fail.')
2023-02-23 18:49:02 +01:00
2023-02-24 16:41:27 +01:00
# This iterator returns the extensions in the order specified in the command-line
2023-02-24 14:01:21 +01:00
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
2023-02-24 16:41:27 +01:00
# Extension functions that map string -> string
def apply_extensions(text, typ):
2023-02-24 14:01:21 +01:00
for extension, _ in iterator():
if typ == "input" and hasattr(extension, "input_modifier"):
text = extension.input_modifier(text)
elif typ == "output" and hasattr(extension, "output_modifier"):
text = extension.output_modifier(text)
elif typ == "bot_prefix" and hasattr(extension, "bot_prefix_modifier"):
text = extension.bot_prefix_modifier(text)
return text
2023-02-23 18:55:21 +01:00
def create_extensions_block():
2023-03-19 14:25:49 +01:00
global setup_called
2023-02-24 23:00:11 +01:00
# Updating the default values
2023-02-24 14:01:21 +01:00
for extension, name in iterator():
if hasattr(extension, 'params'):
for param in extension.params:
_id = f"{name}-{param}"
if _id in shared.settings:
extension.params[param] = shared.settings[_id]
2023-02-23 18:55:21 +01:00
2023-03-19 14:22:24 +01:00
# Running setup function
if not setup_called:
for extension, name in iterator():
if hasattr(extension, "setup"):
extension.setup()
setup_called = True
2023-02-24 23:00:11 +01:00
# Creating the extension ui elements
2023-03-16 03:29:56 +01:00
if len(state) > 0:
with gr.Box(elem_id="extensions"):
gr.Markdown("Extensions")
for extension, name in iterator():
if hasattr(extension, "ui"):
extension.ui()