1580158
/
sd-dynamic-prompts
mirror of https://github.com/adieyal/sd-dynamic-prompts.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Use orthodox method for pnginfo (#691)

pull/695/head
w-e-w 2023-12-11 23:23:26 +09:00 committed by GitHub
parent 2817d2ed51
commit 34ec7fdff9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 6 additions and 95 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
sd_dynamic_prompts/callbacks.py
View File

@ -9,11 +9,7 @@ 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,
strip_template_info,
)
from sd_dynamic_prompts.pnginfo_saver import 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
@ -21,27 +17,6 @@ from sd_dynamic_prompts.wildcards_tab import initialize as initialize_wildcards_
logger = logging.getLogger(__name__)
def register_pnginfo_saver(pnginfo_saver: PngInfoSaver) -> None:
def on_save(image_save_params: ImageSaveParams) -> None:
try:
if image_save_params.p:
png_info = image_save_params.pnginfo["parameters"]
image_prompts = PromptTemplates(
positive_template=image_save_params.p.prompt,
negative_template=image_save_params.p.negative_prompt,
)
updated_png_info = pnginfo_saver.update_pnginfo(
png_info,
image_prompts,
)
image_save_params.pnginfo["parameters"] = updated_png_info
except Exception:
logger.exception("Error save prompt file")
script_callbacks.on_before_image_saved(on_save)
def register_prompt_writer(prompt_writer: PromptWriter) -> None:
def on_save(image_save_params: ImageSaveParams) -> None:
image_name = Path(image_save_params.filename)

8
sd_dynamic_prompts/dynamic_prompting.py
View File

@ -30,7 +30,6 @@ from sd_dynamic_prompts.paths import (
get_magicprompt_models_txt_path,
get_wildcard_dir,
)
from sd_dynamic_prompts.pnginfo_saver import PngInfoSaver
from sd_dynamic_prompts.prompt_writer import PromptWriter
VERSION = __version__
@ -96,14 +95,12 @@ class Script(scripts.Script):
# When the Reload UI button in the settings tab is pressed, the script is loaded twice again
# Therefore we only register callbacks every second time the script is loaded
self._pnginfo_saver = PngInfoSaver()
self._prompt_writer = PromptWriter()
self._wildcard_manager = WildcardManager(get_wildcard_dir())
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()
callbacks.register_settings()
@ -374,7 +371,6 @@ class Script(scripts.Script):
ignore_whitespace = opts.dp_ignore_whitespace
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
@ -520,6 +516,10 @@ class Script(scripts.Script):
negative_prompts=all_negative_prompts,
)
if opts.dp_write_raw_template:
p.extra_generation_params["Template"] = original_prompt
p.extra_generation_params["Negative Template"] = original_negative_prompt
p.all_prompts = all_prompts
p.all_negative_prompts = all_negative_prompts
if no_image_generation:

34
sd_dynamic_prompts/pnginfo_saver.py
View File

@ -1,7 +1,6 @@
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import Any
logger = logging.getLogger(__name__)
@ -31,36 +30,3 @@ def strip_template_info(parameters: dict[str, Any]) -> None:
parameters["Negative prompt"] = (
parameters["Negative prompt"].split(split_by)[0].strip()
)
@dataclass
class PromptTemplates:
positive_template: str
negative_template: str
class PngInfoSaver:
def __init__(self):
self._enabled = True
@property
def enabled(self) -> bool:
return self._enabled
@enabled.setter
def enabled(self, enabled: bool) -> None:
self._enabled = enabled
def update_pnginfo(self, parameters: str, prompt_templates: PromptTemplates) -> str:
if not self._enabled:
return parameters
if prompt_templates.positive_template:
parameters += f"\n{TEMPLATE_LABEL}: {prompt_templates.positive_template}"
if prompt_templates.negative_template:
parameters += (
f"\n{NEGATIVE_TEMPLATE_LABEL}: {prompt_templates.negative_template}"
)
return parameters

32
tests/prompts/ui/test_pnginfo_saver.py
View File

@ -2,11 +2,7 @@ from typing import Any
import pytest
from sd_dynamic_prompts.pnginfo_saver import (
PngInfoSaver,
PromptTemplates,
strip_template_info,
)
from sd_dynamic_prompts.pnginfo_saver import strip_template_info
@pytest.fixture
@ -20,32 +16,6 @@ def basic_parameters():
}
def test_update_pnginfo() -> None:
png_info_saver = PngInfoSaver()
image_prompts = PromptTemplates("Template", "Negative Template")
parameters = "Parameters"
updated_parameters = png_info_saver.update_pnginfo(parameters, image_prompts)
assert (
updated_parameters
== f"Parameters\nTemplate: {image_prompts.positive_template}\nNegative Template: {image_prompts.negative_template}"
)
image_prompts.positive_template = ""
updated_parameters = png_info_saver.update_pnginfo(parameters, image_prompts)
assert (
updated_parameters
== f"Parameters\nNegative Template: {image_prompts.negative_template}"
)
image_prompts.positive_template = "Positive Template"
image_prompts.negative_template = ""
updated_parameters = png_info_saver.update_pnginfo(parameters, image_prompts)
assert (
updated_parameters == f"Parameters\nTemplate: {image_prompts.positive_template}"
)
BASIC_PARAMETERS = "Steps: 35, Sampler: Heun, CFG scale: 9, Seed: 77777, Size: 640x640, Model hash: d8691b4d16"