diff --git a/sd_dynamic_prompts/callbacks.py b/sd_dynamic_prompts/callbacks.py index 42f67076..b3bde713 100644 --- a/sd_dynamic_prompts/callbacks.py +++ b/sd_dynamic_prompts/callbacks.py @@ -9,7 +9,11 @@ from modules.generation_parameters_copypaste import parse_generation_parameters from modules.script_callbacks import ImageSaveParams -from sd_dynamic_prompts.pnginfo_saver import PngInfoSaver, PromptTemplates +from sd_dynamic_prompts.pnginfo_saver import ( + PngInfoSaver, + PromptTemplates, + strip_template_info, +) from sd_dynamic_prompts.prompt_writer import PromptWriter from sd_dynamic_prompts.settings import on_ui_settings from sd_dynamic_prompts.wildcards_tab import initialize as initialize_wildcards_tab @@ -47,17 +51,17 @@ def on_save(image_save_params: ImageSaveParams) -> None: script_callbacks.on_before_image_saved(on_save) -def register_on_infotext_pasted(pnginfo_saver: PngInfoSaver) -> None: +def register_on_infotext_pasted() -> None: def on_infotext_pasted(infotext: str, parameters: dict[str, Any]) -> None: new_parameters = {} if "Prompt" in parameters and "Template:" in parameters["Prompt"]: - parameters = pnginfo_saver.strip_template_info(parameters) + strip_template_info(parameters) new_parameters = parse_generation_parameters(parameters["Prompt"]) elif ( "Negative prompt" in parameters and "Template:" in parameters["Negative prompt"] ): - parameters = pnginfo_saver.strip_template_info(parameters) + strip_template_info(parameters) new_parameters = parse_generation_parameters(parameters["Negative prompt"]) new_parameters["Negative prompt"] = new_parameters["Prompt"] new_parameters["Prompt"] = parameters["Prompt"] diff --git a/sd_dynamic_prompts/dynamic_prompting.py b/sd_dynamic_prompts/dynamic_prompting.py index ae4c0128..622893f5 100644 --- a/sd_dynamic_prompts/dynamic_prompting.py +++ b/sd_dynamic_prompts/dynamic_prompting.py @@ -105,7 +105,7 @@ def __init__(self): callbacks.register_pnginfo_saver(self._pnginfo_saver) callbacks.register_prompt_writer(self._prompt_writer) - callbacks.register_on_infotext_pasted(self._pnginfo_saver) + callbacks.register_on_infotext_pasted() callbacks.register_settings() callbacks.register_wildcards_tab(self._wildcard_manager) diff --git a/sd_dynamic_prompts/pnginfo_saver.py b/sd_dynamic_prompts/pnginfo_saver.py index 16a03c15..1e8136da 100644 --- a/sd_dynamic_prompts/pnginfo_saver.py +++ b/sd_dynamic_prompts/pnginfo_saver.py @@ -10,6 +10,29 @@ NEGATIVE_TEMPLATE_LABEL = "Negative Template" +def strip_template_info(parameters: dict[str, Any]) -> None: + if "Prompt" in parameters and f"{TEMPLATE_LABEL}:" in parameters["Prompt"]: + parameters["Prompt"] = ( + parameters["Prompt"].split(f"{TEMPLATE_LABEL}:")[0].strip() + ) + elif "Negative prompt" in parameters: + split_by = None + if ( + f"\n{TEMPLATE_LABEL}:" in parameters["Negative prompt"] + and f"\n{NEGATIVE_TEMPLATE_LABEL}:" in parameters["Negative prompt"] + ): + split_by = f"{TEMPLATE_LABEL}" + elif f"\n{NEGATIVE_TEMPLATE_LABEL}:" in parameters["Negative prompt"]: + split_by = f"\n{NEGATIVE_TEMPLATE_LABEL}:" + elif f"\n{TEMPLATE_LABEL}:" in parameters["Negative prompt"]: + split_by = f"\n{TEMPLATE_LABEL}:" + + if split_by: + parameters["Negative prompt"] = ( + parameters["Negative prompt"].split(split_by)[0].strip() + ) + + @dataclass class PromptTemplates: positive_template: str @@ -41,27 +64,3 @@ def update_pnginfo(self, parameters: str, prompt_templates: PromptTemplates) -> ) return parameters - - def strip_template_info(self, parameters: dict[str, Any]) -> dict[str, Any]: - if "Prompt" in parameters and f"{TEMPLATE_LABEL}:" in parameters["Prompt"]: - parameters["Prompt"] = ( - parameters["Prompt"].split(f"{TEMPLATE_LABEL}:")[0].strip() - ) - elif "Negative prompt" in parameters: - split_by = None - if ( - f"\n{TEMPLATE_LABEL}:" in parameters["Negative prompt"] - and f"\n{NEGATIVE_TEMPLATE_LABEL}:" in parameters["Negative prompt"] - ): - split_by = f"{TEMPLATE_LABEL}" - elif f"\n{NEGATIVE_TEMPLATE_LABEL}:" in parameters["Negative prompt"]: - split_by = f"\n{NEGATIVE_TEMPLATE_LABEL}:" - elif f"\n{TEMPLATE_LABEL}:" in parameters["Negative prompt"]: - split_by = f"\n{TEMPLATE_LABEL}:" - - if split_by: - parameters["Negative prompt"] = ( - parameters["Negative prompt"].split(split_by)[0].strip() - ) - - return parameters diff --git a/tests/prompts/ui/test_pnginfo_saver.py b/tests/prompts/ui/test_pnginfo_saver.py index 073b0b6f..432187c0 100644 --- a/tests/prompts/ui/test_pnginfo_saver.py +++ b/tests/prompts/ui/test_pnginfo_saver.py @@ -2,7 +2,11 @@ import pytest -from sd_dynamic_prompts.pnginfo_saver import PngInfoSaver, PromptTemplates +from sd_dynamic_prompts.pnginfo_saver import ( + PngInfoSaver, + PromptTemplates, + strip_template_info, +) @pytest.fixture @@ -72,7 +76,6 @@ def test_remove_template_from_infotext( positive_prompt: str, negative_prompt: str, ) -> None: - png_info_saver = PngInfoSaver() if not negative_prompt: basic_parameters["Prompt"] = build_parameters(positive_prompt, negative_prompt) basic_parameters["Negative prompt"] = "" @@ -83,7 +86,7 @@ def test_remove_template_from_infotext( negative_prompt, ) - png_info_saver.strip_template_info(basic_parameters) + strip_template_info(basic_parameters) if negative_prompt: assert basic_parameters["Prompt"] == positive_prompt