diff --git a/sd_dynamic_prompts/callbacks.py b/sd_dynamic_prompts/callbacks.py index 42f67076..56a7d8b6 100644 --- a/sd_dynamic_prompts/callbacks.py +++ b/sd_dynamic_prompts/callbacks.py @@ -4,21 +4,23 @@ from pathlib import Path from typing import Any -from dynamicprompts.wildcards import WildcardManager from modules import script_callbacks 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.prompt_writer import PromptWriter +from sd_dynamic_prompts.pnginfo_saver import PromptTemplates +from sd_dynamic_prompts.script_protocol import SDDPScript from sd_dynamic_prompts.settings import on_ui_settings from sd_dynamic_prompts.wildcards_tab import initialize as initialize_wildcards_tab logger = logging.getLogger(__name__) -def register_pnginfo_saver(pnginfo_saver: PngInfoSaver) -> None: +def register_pnginfo_saver(sddp: SDDPScript) -> None: def on_save(image_save_params: ImageSaveParams) -> None: + if not sddp.is_enabled: + return + try: if image_save_params.p: png_info = image_save_params.pnginfo["parameters"] @@ -27,7 +29,7 @@ def on_save(image_save_params: ImageSaveParams) -> None: negative_template=image_save_params.p.negative_prompt, ) - updated_png_info = pnginfo_saver.update_pnginfo( + updated_png_info = sddp.pnginfo_saver.update_pnginfo( png_info, image_prompts, ) @@ -38,26 +40,32 @@ def on_save(image_save_params: ImageSaveParams) -> None: script_callbacks.on_before_image_saved(on_save) -def register_prompt_writer(prompt_writer: PromptWriter) -> None: +def register_prompt_writer(sddp: SDDPScript) -> None: def on_save(image_save_params: ImageSaveParams) -> None: + if not sddp.is_enabled: + return + image_name = Path(image_save_params.filename) prompt_filename = image_name.with_suffix(".csv") - prompt_writer.write_prompts(prompt_filename) + sddp.prompt_writer.write_prompts(prompt_filename) script_callbacks.on_before_image_saved(on_save) -def register_on_infotext_pasted(pnginfo_saver: PngInfoSaver) -> None: +def register_on_infotext_pasted(sddp: SDDPScript) -> None: def on_infotext_pasted(infotext: str, parameters: dict[str, Any]) -> None: + if not sddp.is_enabled: + return + new_parameters = {} if "Prompt" in parameters and "Template:" in parameters["Prompt"]: - parameters = pnginfo_saver.strip_template_info(parameters) + parameters = sddp.pnginfo_saver.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) + parameters = sddp.pnginfo_saver.strip_template_info(parameters) new_parameters = parse_generation_parameters(parameters["Negative prompt"]) new_parameters["Negative prompt"] = new_parameters["Prompt"] new_parameters["Prompt"] = parameters["Prompt"] @@ -66,9 +74,12 @@ def on_infotext_pasted(infotext: str, parameters: dict[str, Any]) -> None: script_callbacks.on_infotext_pasted(on_infotext_pasted) -def register_settings(): - script_callbacks.on_ui_settings(on_ui_settings) +def register_settings(sddp: SDDPScript) -> None: + def on_ui_settings_callback(): + on_ui_settings(sddp) + + script_callbacks.on_ui_settings(on_ui_settings_callback) -def register_wildcards_tab(wildcard_manager: WildcardManager): - initialize_wildcards_tab(wildcard_manager) +def register_wildcards_tab(sddp: SDDPScript) -> None: + initialize_wildcards_tab(sddp.wildcard_manager) diff --git a/sd_dynamic_prompts/dynamic_prompting.py b/sd_dynamic_prompts/dynamic_prompting.py index 0ba449de..7e98a5da 100644 --- a/sd_dynamic_prompts/dynamic_prompting.py +++ b/sd_dynamic_prompts/dynamic_prompting.py @@ -99,15 +99,32 @@ def __init__(self): self._pnginfo_saver = PngInfoSaver() self._prompt_writer = PromptWriter() self._wildcard_manager = WildcardManager(get_wildcard_dir()) + self._is_enabled = False if loaded_count % 2 == 0: return - callbacks.register_pnginfo_saver(self._pnginfo_saver) - callbacks.register_prompt_writer(self._prompt_writer) - callbacks.register_on_infotext_pasted(self._pnginfo_saver) - callbacks.register_settings() - callbacks.register_wildcards_tab(self._wildcard_manager) + callbacks.register_pnginfo_saver(self) + callbacks.register_prompt_writer(self) + callbacks.register_on_infotext_pasted(self) + callbacks.register_settings(self) + callbacks.register_wildcards_tab(self) + + @property + def is_enabled(self): + return self._is_enabled + + @property + def pnginfo_saver(self): + return self._pnginfo_saver + + @property + def prompt_writer(self): + return self._prompt_writer + + @property + def wildcard_manager(self): + return self._wildcard_manager def title(self): return f"Dynamic Prompts v{VERSION}" @@ -368,8 +385,14 @@ def process( magic_model: str | None, magic_blocklist_regex: str | None, ): + self._is_enabled = is_enabled self._pnginfo_saver.enabled = opts.dp_write_raw_template self._prompt_writer.enabled = opts.dp_write_prompts_to_file + self._limit_jinja_prompts = opts.dp_limit_jinja_prompts + self._auto_purge_cache = opts.dp_auto_purge_cache + self._wildcard_manager.dedup_wildcards = not opts.dp_wildcard_manager_no_dedupe + self._wildcard_manager.sort_wildcards = not opts.dp_wildcard_manager_no_sort + self._wildcard_manager.shuffle_wildcards = opts.dp_wildcard_manager_shuffle if not is_enabled: logger.debug("Dynamic prompts disabled - exiting") @@ -377,12 +400,6 @@ def process( ignore_whitespace = opts.dp_ignore_whitespace - self._limit_jinja_prompts = opts.dp_limit_jinja_prompts - self._auto_purge_cache = opts.dp_auto_purge_cache - self._wildcard_manager.dedup_wildcards = not opts.dp_wildcard_manager_no_dedupe - self._wildcard_manager.sort_wildcards = not opts.dp_wildcard_manager_no_sort - self._wildcard_manager.shuffle_wildcards = opts.dp_wildcard_manager_shuffle - magicprompt_batch_size = opts.dp_magicprompt_batch_size parser_config = ParserConfig( diff --git a/sd_dynamic_prompts/script_protocol.py b/sd_dynamic_prompts/script_protocol.py new file mode 100644 index 00000000..5afeb18a --- /dev/null +++ b/sd_dynamic_prompts/script_protocol.py @@ -0,0 +1,19 @@ +from typing import Protocol + + +class SDDPScript(Protocol): + @property + def is_enabled(self): + ... + + @property + def pnginfo_saver(self): + ... + + @property + def prompt_writer(self): + ... + + @property + def wildcard_manager(self): + ... diff --git a/sd_dynamic_prompts/settings.py b/sd_dynamic_prompts/settings.py index 70c28481..8aeb7343 100644 --- a/sd_dynamic_prompts/settings.py +++ b/sd_dynamic_prompts/settings.py @@ -2,9 +2,10 @@ from modules import shared from sd_dynamic_prompts.helpers import load_magicprompt_models +from sd_dynamic_prompts.script_protocol import SDDPScript -def on_ui_settings(): +def on_ui_settings(script_protocol: SDDPScript): section = "dynamicprompts", "Dynamic Prompts" shared.opts.add_option( key="dp_ignore_whitespace",