196 lines
7.8 KiB
Python
196 lines
7.8 KiB
Python
from dataclasses import replace
|
|
import os
|
|
import logging
|
|
from typing import NamedTuple, Optional
|
|
import unittest
|
|
import datetime
|
|
|
|
from ppp_classes import IFWILDCARDS_CHOICES, ONWARNING_CHOICES, PPPStateOptions
|
|
from ppp_enmappings import PPPExtraNetworkMappings # pylint: disable=import-error
|
|
from ppp_wildcards import PPPWildcards # pylint: disable=import-error
|
|
from ppp import PromptPostProcessor # pylint: disable=import-error
|
|
from ppp_logging import DEBUG_LEVEL, PromptPostProcessorLogFactory # pylint: disable=import-error
|
|
|
|
|
|
class PromptPair(NamedTuple):
|
|
prompt: str = ""
|
|
negative_prompt: str = ""
|
|
|
|
|
|
class TestPromptPostProcessorBase(unittest.TestCase):
|
|
"""
|
|
A test case class for testing the PromptPostProcessor class.
|
|
"""
|
|
|
|
def setUp(self, enable_file_logging=False):
|
|
"""
|
|
Set up the test case by initializing the necessary objects and configurations.
|
|
|
|
Args:
|
|
enable_file_logging (bool): Whether to enable logging to a file. Defaults to True.
|
|
"""
|
|
self.enable_file_logging = enable_file_logging
|
|
test_name = self.id().split(".")[-1] # Extract the test method name
|
|
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
|
|
if self.enable_file_logging:
|
|
log_filename = f"tests/logs/{test_name}_{timestamp}.log"
|
|
else:
|
|
log_filename = None # Disable file logging
|
|
|
|
self.lf = PromptPostProcessorLogFactory(log_filename)
|
|
self.ppp_logger = self.lf.log
|
|
self.ppp_logger.setLevel(logging.DEBUG)
|
|
self.grammar_content = None
|
|
self.interrupted = False
|
|
self.defopts = PPPStateOptions(
|
|
debug_level=DEBUG_LEVEL.full,
|
|
gen_onwarning=ONWARNING_CHOICES.stop,
|
|
wil_process_wildcards=True,
|
|
wil_ifwildcards=IFWILDCARDS_CHOICES.ignore,
|
|
wil_choice_separator=", ",
|
|
wil_keep_choices_order=False,
|
|
stn_separator=", ",
|
|
stn_ignore_repeats=True,
|
|
cup_do_cleanup=True,
|
|
cup_cleanup_variables=True,
|
|
cup_emptyconstructs=True,
|
|
cup_extraseparators=True,
|
|
cup_extraseparators2=True,
|
|
cup_extraseparators_include_eol=False,
|
|
cup_extraspaces=True,
|
|
cup_breaks=True,
|
|
cup_breaks_eol=False,
|
|
cup_ands=True,
|
|
cup_ands_eol=False,
|
|
cup_extranetworktags=True,
|
|
cup_mergeattention=True,
|
|
rem_removeextranetworktags=False,
|
|
)
|
|
self.def_env_info = {
|
|
"app": "tests",
|
|
"ppp_config": None,
|
|
"model_class": "SDXL",
|
|
"property_base": {"is_sdxl": True},
|
|
"models_path": "./webui/models",
|
|
"model_filename": "./webui/models/Stable-diffusion/testmodel.safetensors",
|
|
}
|
|
self.interrupted = False
|
|
self.wildcards_obj = PPPWildcards(self.lf.log)
|
|
self.extranetwork_maps_obj = PPPExtraNetworkMappings(self.lf.log)
|
|
self.wildcards_obj.refresh_wildcards(
|
|
DEBUG_LEVEL.full,
|
|
[
|
|
os.path.abspath(os.path.join(os.path.dirname(__file__), "wildcards")),
|
|
os.path.abspath(os.path.join(os.path.dirname(__file__), "wildcards2")),
|
|
],
|
|
"""
|
|
yaml_input:
|
|
wildcardI:
|
|
- choice1
|
|
- choice2
|
|
- choice3
|
|
""",
|
|
)
|
|
self.extranetwork_maps_obj.refresh_extranetwork_mappings(
|
|
DEBUG_LEVEL.full,
|
|
[os.path.abspath(os.path.join(os.path.dirname(__file__), "enmappings"))],
|
|
"""
|
|
""",
|
|
)
|
|
grammar_filename = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../grammar.lark")
|
|
with open(grammar_filename, "r", encoding="utf-8") as file:
|
|
self.grammar_content = file.read()
|
|
|
|
def interrupt(self):
|
|
self.interrupted = True
|
|
|
|
def process(
|
|
self,
|
|
input_prompts: PromptPair,
|
|
expected_output_prompts: Optional[PromptPair | list[PromptPair]] = None,
|
|
seed: int = 1,
|
|
ppp: Optional[str | PromptPostProcessor] = None,
|
|
interrupted: bool = False,
|
|
variables: dict[str, str] | None = None,
|
|
):
|
|
"""
|
|
Process the prompt and compare the results with the expected prompts.
|
|
|
|
Args:
|
|
input_prompts (PromptPair): The input prompts.
|
|
expected_output_prompts (PromptPair | list[PromptPair], optional): The expected prompts.
|
|
seed (int, optional): The seed value. Defaults to 1.
|
|
ppp (Optional[str | PromptPostProcessor], optional): The PromptPostProcessor instance or type. Defaults to None.
|
|
interrupted (bool, optional): The interrupted flag. Defaults to False.
|
|
variables (dict[str,str]|None, optional): Output variables to check. Defaults to None.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
if isinstance(ppp, str):
|
|
if ppp == "nocup":
|
|
the_obj = PromptPostProcessor(
|
|
self.ppp_logger,
|
|
self.def_env_info,
|
|
replace(
|
|
self.defopts,
|
|
cup_do_cleanup=False,
|
|
cup_cleanup_variables=False,
|
|
cup_emptyconstructs=False,
|
|
cup_extraseparators=False,
|
|
cup_extraseparators2=False,
|
|
cup_extraseparators_include_eol=False,
|
|
cup_extraspaces=False,
|
|
cup_breaks=False,
|
|
cup_breaks_eol=False,
|
|
cup_ands=False,
|
|
cup_ands_eol=False,
|
|
cup_extranetworktags=False,
|
|
cup_mergeattention=False,
|
|
),
|
|
self.grammar_content,
|
|
self.interrupt,
|
|
self.wildcards_obj,
|
|
self.extranetwork_maps_obj,
|
|
)
|
|
else:
|
|
the_obj = ppp
|
|
if not the_obj:
|
|
the_obj = PromptPostProcessor(
|
|
self.ppp_logger,
|
|
self.def_env_info,
|
|
self.defopts,
|
|
self.grammar_content,
|
|
self.interrupt,
|
|
self.wildcards_obj,
|
|
self.extranetwork_maps_obj,
|
|
)
|
|
out = (
|
|
[PromptPair("", "")]
|
|
if expected_output_prompts is None
|
|
else expected_output_prompts if isinstance(expected_output_prompts, list) else [expected_output_prompts]
|
|
)
|
|
for eo in out:
|
|
result_prompt, result_negative_prompt, output_variables = the_obj.process_prompt(
|
|
input_prompts.prompt,
|
|
input_prompts.negative_prompt,
|
|
seed,
|
|
)
|
|
self.assertEqual(self.interrupted, interrupted, "Interrupted flag is incorrect")
|
|
if not self.interrupted:
|
|
if expected_output_prompts is not None:
|
|
self.assertEqual(result_prompt, eo.prompt, "Incorrect prompt")
|
|
self.assertEqual(result_negative_prompt, eo.negative_prompt, "Incorrect negative prompt")
|
|
if variables is not None:
|
|
for var_name, var_value in variables.items():
|
|
self.assertIn(
|
|
var_name, output_variables, f"Variable '{var_name}' not found in output variables"
|
|
)
|
|
self.assertEqual(
|
|
output_variables[var_name],
|
|
var_value,
|
|
f"Variable '{var_name}' has incorrect value",
|
|
)
|
|
seed += 1
|