comments.py 1.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849
  1. from modules import scripts, shared, script_callbacks
  2. import re
  3. def strip_comments(text):
  4. text = re.sub('(^|\n)#[^\n]*(\n|$)', '\n', text) # while line comment
  5. text = re.sub('#[^\n]*(\n|$)', '\n', text) # in the middle of the line comment
  6. return text
  7. class ScriptStripComments(scripts.Script):
  8. def title(self):
  9. return "Comments"
  10. def show(self, is_img2img):
  11. return scripts.AlwaysVisible
  12. def process(self, p, *args):
  13. if not shared.opts.enable_prompt_comments:
  14. return
  15. p.all_prompts = [strip_comments(x) for x in p.all_prompts]
  16. p.all_negative_prompts = [strip_comments(x) for x in p.all_negative_prompts]
  17. p.main_prompt = strip_comments(p.main_prompt)
  18. p.main_negative_prompt = strip_comments(p.main_negative_prompt)
  19. if getattr(p, 'enable_hr', False):
  20. p.all_hr_prompts = [strip_comments(x) for x in p.all_hr_prompts]
  21. p.all_hr_negative_prompts = [strip_comments(x) for x in p.all_hr_negative_prompts]
  22. p.hr_prompt = strip_comments(p.hr_prompt)
  23. p.hr_negative_prompt = strip_comments(p.hr_negative_prompt)
  24. def before_token_counter(params: script_callbacks.BeforeTokenCounterParams):
  25. if not shared.opts.enable_prompt_comments:
  26. return
  27. params.prompt = strip_comments(params.prompt)
  28. script_callbacks.on_before_token_counter(before_token_counter)
  29. shared.options_templates.update(shared.options_section(('sd', "Stable Diffusion", "sd"), {
  30. "enable_prompt_comments": shared.OptionInfo(True, "Enable comments").info("Use # anywhere in the prompt to hide the text between # and the end of the line from the generation."),
  31. }))