Skip to content

Commit

Permalink
Only execute callbacks if DP is enabled
Browse files Browse the repository at this point in the history
Ensure that callbacks (apart from settings and wildcards manager), only run when dynamic prompts is enabled. This builds on #648
  • Loading branch information
adieyal committed Nov 6, 2023
1 parent ddaca53 commit 5ef7740
Show file tree
Hide file tree
Showing 4 changed files with 74 additions and 26 deletions.
39 changes: 25 additions & 14 deletions sd_dynamic_prompts/callbacks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand All @@ -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,
)
Expand All @@ -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"]
Expand All @@ -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)
39 changes: 28 additions & 11 deletions sd_dynamic_prompts/dynamic_prompting.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}"
Expand Down Expand Up @@ -368,21 +385,21 @@ 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")
return p

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(
Expand Down
19 changes: 19 additions & 0 deletions sd_dynamic_prompts/script_protocol.py
Original file line number Diff line number Diff line change
@@ -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):
...
3 changes: 2 additions & 1 deletion sd_dynamic_prompts/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down

0 comments on commit 5ef7740

Please sign in to comment.