script.py 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051
  1. from pathlib import Path
  2. import gradio as gr
  3. from modules import shared
  4. from modules import ui as _ui
  5. params = {
  6. 'template': '%input%'
  7. }
  8. def get_available_templates():
  9. return ['None'] + sorted(set((k.stem for k in Path('extensions/prompt_template/templates').glob('*.txt'))), key=str.lower)
  10. def load_template(fname):
  11. if fname in ['None', '']:
  12. return '%input%'
  13. else:
  14. with open(Path(f'extensions/prompt_template/templates/{fname}.txt'), 'r', encoding='utf-8') as f:
  15. text = f.read()
  16. if text[-1] == '\n':
  17. text = text[:-1]
  18. return text
  19. def input_modifier(string):
  20. """
  21. This function is applied to your text inputs before
  22. they are fed into the model.
  23. """
  24. return params['template'].replace('%input%', string)
  25. def output_modifier(string):
  26. return f'\n{string}'
  27. def setup():
  28. shared.args.verbose = True
  29. def ui():
  30. # Gradio elements
  31. with gr.Row():
  32. with gr.Column():
  33. template = gr.Textbox(value=params['template'], info="%input% will be replaced with your user input.", label='Template')
  34. with gr.Column():
  35. with gr.Row():
  36. template_menu = gr.Dropdown(choices=get_available_templates(), value='None', label='Available templates')
  37. _ui.create_refresh_button(shared.gradio['model_menu'], lambda : None, lambda : {'choices': get_available_templates()}, 'refresh-button')
  38. template_menu.change(load_template, template_menu, template)
  39. template.change(lambda x: params.update({"template": x}), template, None)