diff --git a/extensions/prompt_template/script.py b/extensions/prompt_template/script.py new file mode 100644 index 0000000..7218159 --- /dev/null +++ b/extensions/prompt_template/script.py @@ -0,0 +1,51 @@ +from pathlib import Path + +import gradio as gr + +from modules import shared +from modules import ui as _ui + +params = { + 'template': '%input%' +} + +def get_available_templates(): + return ['None'] + sorted(set((k.stem for k in Path('extensions/prompt_template/templates').glob('*.txt'))), key=str.lower) + +def load_template(fname): + if fname in ['None', '']: + return '%input%' + else: + with open(Path(f'extensions/prompt_template/templates/{fname}.txt'), 'r', encoding='utf-8') as f: + text = f.read() + if text[-1] == '\n': + text = text[:-1] + return text + +def input_modifier(string): + """ + This function is applied to your text inputs before + they are fed into the model. + """ + + return params['template'].replace('%input%', string) + +def output_modifier(string): + return f'\n{string}' + +def setup(): + shared.args.verbose = True + +def ui(): + # Gradio elements + + with gr.Row(): + with gr.Column(): + template = gr.Textbox(value=params['template'], info="%input% will be replaced with your user input.", label='Template') + with gr.Column(): + with gr.Row(): + template_menu = gr.Dropdown(choices=get_available_templates(), value='None', label='Available templates') + _ui.create_refresh_button(shared.gradio['model_menu'], lambda : None, lambda : {'choices': get_available_templates()}, 'refresh-button') + + template_menu.change(load_template, template_menu, template) + template.change(lambda x: params.update({"template": x}), template, None)