File size: 3,063 Bytes
c336648 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
from modules import script_callbacks
from lib.prompt_cleaning_hook.options import onUiSettings
from modules import scripts
from modules.processing import StableDiffusionProcessing
from src.prompt_cleaning_hook.cleanerModules import cleanerModules
from lib.prompt_cleaning_hook.options import getOption
from src.prompt_cleaning_hook.extension import extensionId, extensionTitle
from lib.prompt_cleaning_hook.logger import logger
def process(text: str, commaOnLinebreak: bool) -> dict[str, str | list[str]]:
changingModules: list[str] = []
for module in cleanerModules:
newText = module['process'](text)
if newText != text:
changingModules.append(module['id'])
text = newText
return {'text': text, 'changingModules': changingModules}
def logChange(result, original: str, isPositive: bool = True):
cleaned: str = result['text']
originalPromptSingleLine = original.replace('\n', '¶')
cleanedSingleLine = cleaned.replace('\n', '¶')
promptName = 'prompt' if isPositive else 'negative prompt'
modulesString = ' → '.join(result['changingModules'])
logMessage = f'Cleaned {promptName} from “{originalPromptSingleLine}” to “{cleanedSingleLine}” using modules {modulesString}'
logChanges = getOption('log_changes')
if logChanges:
logger.info(logMessage)
else:
logger.debug(logMessage)
class PromptCleaningHook(scripts.Script):
def __init__(self):
logger.debug('PromptCleaningHook initialized')
super().__init__()
def title(self):
return extensionTitle
def show(self, isImg2img):
return scripts.AlwaysVisible
def process(self, processing: StableDiffusionProcessing):
if not getOption('enabled'):
return
commaOnLinebreak = getOption('comma_on_linebreak')
saveMetadata = getOption('save_metadata')
changedCount = 0
originalPrompt = processing.prompt
originalNegativePrompt = processing.negative_prompt
for i in range(len(processing.all_prompts)):
cleanResult = process(processing.all_prompts[i], commaOnLinebreak)
if cleanResult['changingModules'].count == 0:
continue
logChange(cleanResult, processing.all_prompts[i])
changedCount += 1
processing.all_prompts[i] = cleanResult['text']
for i in range(len(processing.all_negative_prompts)):
cleanResult = process(processing.all_negative_prompts[i], commaOnLinebreak)
if cleanResult['changingModules'].count == 0:
continue
logChange(cleanResult, processing.all_negative_prompts[i], False)
changedCount += 1
processing.all_negative_prompts[i] = cleanResult['text']
if changedCount == 0:
return
logger.debug(f'Cleaned {changedCount}/{len(processing.all_prompts) + len(processing.all_negative_prompts)} prompts')
if saveMetadata:
processing.extra_generation_params.setdefault(f'{extensionId}_original_prompt', originalPrompt)
processing.extra_generation_params.setdefault(f'{extensionId}_original_negative_prompt', originalNegativePrompt)
script_callbacks.on_ui_settings(onUiSettings)
|