diff --git a/scripts/prompt_formatter.py b/scripts/prompt_formatter.py index 3da5a49..8d86ca7 100644 --- a/scripts/prompt_formatter.py +++ b/scripts/prompt_formatter.py @@ -1,17 +1,16 @@ """Enter format prompt.""" + import gradio as gr from modules import script_callbacks, scripts, shared -# from prompt_formatting_pipeline import pipeline from scripts import prompt_formatting_pipeline as pipeline +from scripts.prompt_formatting_definitions import UnderSpaceEnum -""" -Formatting settings -""" SPACE_COMMAS = True BRACKET2WEIGHT = True -SPACE2UNDERSCORE = False -IGNOREUNDERSCORES = True +# SPACE2UNDERSCORE = False +# IGNOREUNDERSCORES = True +PREFER_SPACING = UnderSpaceEnum.IGNORE ui_prompts = set() @@ -34,11 +33,9 @@ def format_prompt(*prompts: tuple[dict]): prompt = pipeline.remove_whitespace_excessive(prompt) # Replace Spaces and/or underscores, unless disabled - prompt = pipeline.space_to_underscore(prompt, opposite=IGNOREUNDERSCORES) + prompt = pipeline.space_to_underscore(prompt, mode=PREFER_SPACING) prompt = pipeline.align_brackets(prompt) - prompt = pipeline.space_and( - prompt - ) # for proper compositing alignment on colons + prompt = pipeline.space_and(prompt) # for proper compositing alignment on colon prompt = pipeline.space_bracekts(prompt) prompt = pipeline.align_colons(prompt) prompt = pipeline.align_commas(prompt, do_it=SPACE_COMMAS) @@ -76,6 +73,7 @@ def on_before_component(component: gr.component, **kwargs: dict): def on_ui_settings(): section = ("pformat", "Prompt Formatter") + shared.opts.add_option( "pformat_space_commas", shared.OptionInfo( @@ -96,36 +94,48 @@ def on_ui_settings(): section=section, ), ) + # shared.opts.add_option( + # "pfromat_space2underscore", + # shared.OptionInfo( + # False, + # "Convert spaces to underscores (default: underscore to spaces)", + # gr.Checkbox, + # {"interactive": True}, + # section=section, + # ), + # ) + # shared.opts.add_option( + # "pfromat_ignoreunderscores", + # shared.OptionInfo( + # True, + # "Do not convert either spaces or underscores (preserves DanBooru tag formatting)", + # gr.Checkbox, + # {"interactive": True}, + # section=section, + # ), + + shared.opts.add_option( - "pfromat_space2underscore", - shared.OptionInfo( - False, - "Convert spaces to underscores (default: underscore to spaces)", - gr.Checkbox, - {"interactive": True}, - section=section, - ), - ) - shared.opts.add_option( - "pfromat_ignoreunderscores", + 'pformat_preferspacing', shared.OptionInfo( - True, - "Do not convert either spaces or underscores (preserves DanBooru tag formatting)", - gr.Checkbox, - {"interactive": True}, - section=section, - ), + 'Space', + "Preference in formatter to use spaces or underscore or ignore.", + gr.Radio, + {"choices": ["Space", "Underscore", "Ignore"]}, + section=section + ) ) sync_settings() def sync_settings(): - global SPACE_COMMAS, BRACKET2WEIGHT, SPACE2UNDERSCORE, IGNOREUNDERSCORES + global SPACE_COMMAS, BRACKET2WEIGHT, SPACE2UNDERSCORE, IGNOREUNDERSCORES, PREFER_SPACING SPACE_COMMAS = shared.opts.pformat_space_commas BRACKET2WEIGHT = shared.opts.pfromat_bracket2weight - SPACE2UNDERSCORE = shared.opts.pfromat_space2underscore - IGNOREUNDERSCORES = shared.opts.pfromat_ignoreunderscores + # SPACE2UNDERSCORE = shared.opts.pfromat_space2underscore + # IGNOREUNDERSCORES = shared.opts.pfromat_ignoreunderscores + PREFER_SPACING = UnderSpaceEnum(shared.opts.pformat_preferspacing) script_callbacks.on_before_component(on_before_component) diff --git a/scripts/prompt_formatting_definitions.py b/scripts/prompt_formatting_definitions.py new file mode 100644 index 0000000..f961251 --- /dev/null +++ b/scripts/prompt_formatting_definitions.py @@ -0,0 +1,9 @@ +"""Definitions to be pulled by other scripts.""" + +from enum import Enum + + +class UnderSpaceEnum(Enum): + SPACE = 'Space' + UNDERSCORE = 'Underscore' + IGNORE = 'Ignore' diff --git a/scripts/prompt_formatting_pipeline.py b/scripts/prompt_formatting_pipeline.py index 22fc82b..b7fde35 100644 --- a/scripts/prompt_formatting_pipeline.py +++ b/scripts/prompt_formatting_pipeline.py @@ -4,6 +4,8 @@ import regex as re +from scripts.prompt_formatting_definitions import UnderSpaceEnum + brackets_opening = "([{<" brackets_closing = ")]}>" @@ -463,9 +465,7 @@ def calculate_weight(d: str, *, is_square_brackets: bool): return 1 / 1.1 ** int(d) if is_square_brackets else 1 * 1.1 ** int(d) - - -def space_to_underscore(prompt: str, *, opposite = True): +def space_to_underscore(prompt: str, mode: UnderSpaceEnum = UnderSpaceEnum.SPACE): """Replace space with underscore or vice versa. It's a but funky right now because it uses the tokenizer to chunk for sub. @@ -475,15 +475,18 @@ def space_to_underscore(prompt: str, *, opposite = True): This has been patched by requiring the match to be surrounded with a character, but I'm sure there's better solutions. It will work for now. """ - match = ( - r"(?)" - if opposite - else r"(?)" - ) - replace = "_" if opposite else " " + if mode == UnderSpaceEnum.IGNORE: + return prompt + + if mode == UnderSpaceEnum.SPACE: + match = r"(?)" + replace = " " + + elif mode == UnderSpaceEnum.UNDERSCORE: + match = r"(?)" + replace = "_" tokens: str = tokenize(prompt) - print(tokens) return ",".join(map(lambda t: re.sub(match, replace, t), tokens)) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 7255083..93f6c5e 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -7,6 +7,7 @@ import pytest from scripts import prompt_formatting_pipeline as pipeline +from scripts.prompt_formatting_definitions import UnderSpaceEnum def test_get_bracket_closing(): @@ -120,11 +121,10 @@ def test_bracket_to_weights(): assert pipeline.bracket_to_weights('((a), ((b)))') == '((a:1.10), (b:1.21):1.10)' def test_space_to_underscore(): - assert pipeline.space_to_underscore(', multiple subjects') == ', multiple_subjects' - assert pipeline.space_to_underscore('one two three') == 'one_two_three' - assert pipeline.space_to_underscore('this is a test') == 'this_is_a_test' - assert pipeline.space_to_underscore(', baz') == ', baz' + assert pipeline.space_to_underscore(', multiple subjects') == ', multiple subjects' + assert pipeline.space_to_underscore('one_two_three') == 'one two three' + assert pipeline.space_to_underscore('this is_a test') == 'this is a test' + assert pipeline.space_to_underscore(', baz', UnderSpaceEnum.IGNORE) == ', baz' - pipeline.BRACKET2WEIGHT = False - assert pipeline.space_to_underscore('some_var_name', opposite=pipeline.BRACKET2WEIGHT) == 'some var name' + assert pipeline.space_to_underscore('one two three', UnderSpaceEnum.UNDERSCORE) == 'one_two_three'