|
""""This module provides a function to get all""" |
|
from dataclasses import dataclass |
|
from typing import Union, List, Tuple |
|
import re |
|
import os |
|
from scripts.Logger import Logger |
|
from scripts.Utils import Utils |
|
|
|
import modules |
|
|
|
import modules.shared as shared |
|
|
|
@dataclass() |
|
class BatchParams: |
|
"""Dataclass to store the parameters for a batch |
|
|
|
Args: |
|
checkpoint (str): the checkpoint name |
|
prompt (str): the prompt |
|
hr_prompt (str): the hires. fix prompt |
|
neg_prompt (str): the negative prompt |
|
style (List[str]): the style (A1111 styles) |
|
batch_count (int, optional): the batch count. Defaults to -1. (don't overwrite the UI value) |
|
clip_skip (int, optional): the clip skip. Defaults to 1. |
|
width (int, optional): the width. Defaults to -1. (don't overwrite the UI value) |
|
height (int, optional): the height. Defaults to -1. (don't overwrite the UI value) |
|
""" |
|
checkpoint: str |
|
prompt: str |
|
hr_prompt: str |
|
neg_prompt: str |
|
style : List[str] |
|
batch_count: int = -1 |
|
clip_skip: int = 1 |
|
width: int = -1 |
|
height: int = -1 |
|
|
|
def __repr__(self) -> str: |
|
checkpointName: str = os.path.basename(self.checkpoint) |
|
return( f"BatchParams: {checkpointName},\n " |
|
f"prompt: {self.prompt},\n" |
|
f"style: {self.style},\n" |
|
f"neg_prompt: {self.neg_prompt},\n " |
|
f"batch_count: {self.batch_count},\n " |
|
f"clip_skip: {self.clip_skip}\n" |
|
f"size: {self.width}x{self.height}") |
|
|
|
logger = Logger() |
|
|
|
def get_all_batch_params(p: Union[modules.processing.StableDiffusionProcessingTxt2Img, modules.processing.StableDiffusionProcessingImg2Img], checkpoints_as_string: str, prompts_as_string: str) -> List[BatchParams]: |
|
"""Get all batch parameters from the input |
|
|
|
Args: |
|
p (Union[modules.processing.StableDiffusionProcessingTxt2Img, modules.processing.StableDiffusionProcessingImg2Img]): the processing object |
|
checkpoints_as_string (str): the checkpoints as string |
|
prompts_as_string (str): the prompts as string |
|
|
|
Returns: |
|
List[BatchParams]: the batch parameters |
|
""" |
|
|
|
def getRegexFromOpts(key: str, search_for_number: bool = True) -> Tuple[str, str]: |
|
"""Get the regex from the options. As the user can change the regex, |
|
it is checked if the regex is valid. |
|
|
|
Args: |
|
key (str): the key |
|
search_for_number (bool, optional): If true checks if the regex is valid. Defaults to True. |
|
|
|
Returns: |
|
Tuple[str, str]: the search pattern and the sub pattern |
|
""" |
|
sub_pattern = getattr(shared.opts, key) |
|
search_pattern = sub_pattern.replace("[", "([").replace("]", "])") |
|
|
|
if not re.search(r"\[0-9\]\+|\\d\+", sub_pattern) and search_for_number: |
|
raise RuntimeError(f'Can\'t find a number with the regex for {key}: "{sub_pattern}"') |
|
|
|
return search_pattern, sub_pattern |
|
|
|
utils = Utils() |
|
|
|
def get_batch_count_from_prompt(prompt: str) -> Tuple[int, str]: |
|
"""Extracts the batch count from the prompt if specified, else uses the default value |
|
|
|
Args: |
|
prompt (str): the prompt |
|
|
|
Returns: |
|
Tuple[int, str]: the batch count and the prompt |
|
""" |
|
search_pattern, sub_pattern = getRegexFromOpts("batchCountRegex") |
|
number_match = re.search(search_pattern, prompt) |
|
if number_match and number_match.group(1): |
|
|
|
number = int(number_match.group(1)) |
|
number = p.n_iter if number < 1 else number |
|
prompt = re.sub(sub_pattern, '', prompt) |
|
else: |
|
number = p.n_iter |
|
|
|
|
|
return number, prompt |
|
|
|
def get_clip_skip_from_prompt(prompt: str) -> Tuple[int, str]: |
|
"""Extracts the clip skip from the prompt if specified, else uses the default value |
|
|
|
Args: |
|
prompt (str): the prompt |
|
|
|
Returns: |
|
Tuple[int, str]: the clip skip and the prompt |
|
""" |
|
search_pattern, sub_pattern = getRegexFromOpts("clipSkipRegex") |
|
number_match = re.search(search_pattern, prompt) |
|
if number_match and number_match.group(1): |
|
|
|
number = int(number_match.group(1)) |
|
number = shared.opts.data["CLIP_stop_at_last_layers"] if number < 1 else number |
|
prompt = ( |
|
re.sub(sub_pattern, '', prompt)) |
|
else: |
|
number = shared.opts.data["CLIP_stop_at_last_layers"] |
|
|
|
|
|
return number, prompt |
|
|
|
def get_style_from_prompt(prompt: str) -> Tuple[List[str], str]: |
|
"""Extracts the style from the prompt if specified. |
|
|
|
Args: |
|
prompt (str): the prompt |
|
|
|
Returns: |
|
Tuple[List[str], str]: the styles and the prompt |
|
""" |
|
styles = [] |
|
search_pattern, sub_pattern = getRegexFromOpts("styleRegex", False) |
|
style_matches = re.findall(search_pattern, prompt) |
|
if style_matches: |
|
for i, stl in enumerate(style_matches): |
|
styles.append(stl) |
|
_, prompt_regex = getRegexFromOpts("promptRegex", False) |
|
replacement = prompt_regex if i == len(style_matches) - 1 else "" |
|
prompt = re.sub(sub_pattern, replacement, prompt, count=1) |
|
|
|
logger.debug_log(f"nr.: {i}, prompt: {prompt}", False) |
|
|
|
return styles, prompt |
|
|
|
def get_image_size_from_prompt(prompt: str) -> Tuple[int, int, str]: |
|
"""Extracts the image size from the prompt if specified, else uses the default value |
|
|
|
Args: |
|
prompt (str): the prompt |
|
|
|
Returns: |
|
Tuple[int, int, str]: the width, height and the prompt. |
|
If the width and height are not specified, -1 is returned. |
|
""" |
|
search_pattern, sub_pattern = getRegexFromOpts("widthHeightRegex", False) |
|
number_matches = re.search(search_pattern, prompt) |
|
if number_matches: |
|
try: |
|
width, height = map(int, number_matches.groups()) |
|
except ValueError: |
|
raise RuntimeError(f"Can't convert the image size to an integer: {number_matches[0]}") |
|
prompt = re.sub(sub_pattern, '', prompt) |
|
else: |
|
width, height = -1, -1 |
|
|
|
return width, height, prompt |
|
|
|
def split_postive_and_negative_postive_prompt(prompt: str) -> Tuple[str, str]: |
|
"""Splits the prompt into a positive and negative prompt. |
|
If a negative prompt is specified. |
|
|
|
Args: |
|
prompt (str): the prompt |
|
|
|
Returns: |
|
Tuple[str, str]: the positive and negative prompt |
|
""" |
|
pattern = getattr(shared.opts, "negPromptRegex") |
|
parts = re.split(pattern, prompt) |
|
if len(parts) > 1: |
|
neg_prompt = parts[1] |
|
else: |
|
neg_prompt = "" |
|
|
|
prompt = parts[0] |
|
|
|
return prompt, neg_prompt |
|
|
|
|
|
all_batch_params: List[BatchParams] = [] |
|
|
|
checkpoints: List[str] = utils.getCheckpointListFromInput(checkpoints_as_string) |
|
|
|
|
|
prompts: List[str] = utils.remove_index_from_string(prompts_as_string).split(";") |
|
prompts = [prompt.replace('\n', '').strip() for prompt in prompts if not prompt.isspace() and prompt != ''] |
|
|
|
if len(prompts) != len(checkpoints): |
|
logger.debug_log(f"len prompt: {len(prompts)}, len checkpoints{len(checkpoints)}") |
|
raise RuntimeError("amount of prompts don't match with amount of checkpoints") |
|
|
|
if len(prompts) == 0: |
|
raise RuntimeError("can't run without a checkpoint and prompt") |
|
|
|
|
|
for i in range(len(checkpoints)): |
|
|
|
info = modules.sd_models.get_closet_checkpoint_match(checkpoints[i]) |
|
if info is None: |
|
raise RuntimeError(f"Unknown checkpoint: {checkpoints[i]}") |
|
|
|
|
|
batch_count, prompts[i] = get_batch_count_from_prompt(prompts[i]) |
|
clip_skip, prompts[i] = get_clip_skip_from_prompt(prompts[i]) |
|
style, prompts[i] = get_style_from_prompt(prompts[i]) |
|
width, height, prompts[i] = get_image_size_from_prompt(prompts[i]) |
|
prompt_template, neg_prompt = split_postive_and_negative_postive_prompt(prompts[i]) |
|
|
|
|
|
_, prompt_regex = getRegexFromOpts("promptRegex", False) |
|
|
|
prompt = prompt_template.replace(prompt_regex, p.prompt) |
|
hr_prompt = prompt_template.replace(prompt_regex, p.hr_prompt) |
|
neg_prompt = p.negative_prompt + neg_prompt |
|
|
|
|
|
all_batch_params.append(BatchParams(checkpoints[i], prompt,hr_prompt, neg_prompt, style, batch_count, clip_skip, width, height)) |
|
|
|
logger.debug_log(f"batch_params: {all_batch_params[i]}", False) |
|
|
|
return all_batch_params |
