| """This module provides utility functions.""" |
| from scripts.Logger import Logger |
| import os |
| import re |
| import requests |
| from typing import List |
|
|
| import modules |
| from modules.sd_models import read_state_dict |
| from modules.sd_models_config import (find_checkpoint_config, config_default, config_sd2, config_sd2v, config_sd2_inpainting, |
| config_depth_model, config_unclip, config_unopenclip, config_inpainting, config_instruct_pix2pix, config_alt_diffusion) |
|
|
| import sys |
| sys.path.insert(0, os.path.join(os.path.dirname( |
| os.path.abspath(__file__)), "scripts")) |
|
|
|
|
| class Utils(): |
| """ |
| methods that are needed in different classes |
| """ |
|
|
| def __init__(self) -> None: |
| self.logger = Logger() |
| self.logger.debug = False |
| script_path = os.path.dirname( |
| os.path.dirname(os.path.abspath(__file__))) |
| self.held_md_file_name = os.path.join( |
| script_path, "HelpBatchCheckpointsPrompt.md") |
| self.held_md_url = f"https://raw.githubusercontent.com/h43lb1t0/BatchCheckpointPrompt/main/{self.held_md_file_name}.md" |
|
|
| def split_prompts(self, text: str) -> List[str]: |
| """Split the prompts by the ; and remove empty strings and newlines |
| |
| Args: |
| text (str): the input string |
| Returns: |
| List[str]: a list of prompts |
| """ |
| prompt_list = text.split(";") |
| return [prompt.replace('\n', '').strip( |
| ) for prompt in prompt_list if not prompt.isspace() and prompt != ''] |
|
|
|
|
| def remove_index_from_string(self, input: str) -> str: |
| """Remove the index from the string |
| |
| Args: |
| input (str): the input string |
| Returns: |
| str: the string without the index |
| """ |
| return re.sub(r"@index:\d+", "", input).strip() |
| |
| def remove_model_version_from_string(self, checkpoints_text: str) -> str: |
| """Remove the model version from the string |
| |
| Args: |
| input (str): the input string with all checkpoints |
| Returns: |
| str: the string without the model version |
| """ |
| patterns = [ |
| '@version:sd1', |
| '@version:sd2', |
| '@version:sd2v', |
| '@version:sd2-inpainting', |
| '@version:depth', |
| '@version:unclip', |
| '@version:unopenclip', |
| '@version:sd1-inpainting', |
| '@version:pix2pix', |
| '@version:alt' |
| ] |
| |
| |
| for pattern in patterns: |
| checkpoints_text = re.sub(pattern, '', checkpoints_text) |
|
|
| return checkpoints_text |
|
|
| def get_clean_checkpoint_path(self, checkpoint: str) -> str: |
| """Remove the checkpoint hash from the filename |
| |
| Args: |
| input (str): the input string with hash |
| Returns: |
| str: the string without the hash |
| """ |
| return re.sub(r' \[.*?\]', '', checkpoint).strip() |
|
|
| def getCheckpointListFromInput(self, checkpoints_text: str, clean: bool = True) -> List[str]: |
| """Get a list of checkpoints from the input string |
| |
| Args: |
| checkpoints_text (str): the input string with all checkpoints |
| clean (bool): remove the index and hash from the string |
| Returns: |
| List[str]: a list of checkpoints |
| """ |
| self.logger.debug_log(f"checkpoints: {checkpoints_text}") |
| checkpoints_text = self.remove_model_version_from_string(checkpoints_text) |
| if clean: |
| checkpoints_text = self.remove_index_from_string(checkpoints_text) |
| checkpoints_text = self.get_clean_checkpoint_path(checkpoints_text) |
| checkpoints = checkpoints_text.split(",") |
| checkpoints = [checkpoint.replace('\n', '').strip( |
| ) for checkpoint in checkpoints if checkpoints if not checkpoint.isspace() and checkpoint != ''] |
| return checkpoints |
|
|
| def get_help_md(self) -> str: |
| """Gets the help md file. |
| If the file is not localy found downloads it from the github repository |
| |
| Returns: |
| str: the help md file as a string |
| """ |
| md = "could not get help file. Check Github for more information" |
| if os.path.isfile(self.held_md_file_name): |
| with open(self.held_md_file_name) as f: |
| md = f.read() |
| else: |
| self.logger.debug_log("downloading help md") |
| result = requests.get(self.held_md_url) |
| if result.status_code == 200: |
| with open(self.held_md_file_name, "wb") as file: |
| file.write(result.content) |
| return self.get_help_md() |
| return md |
|
|
| def add_index_to_string(self, text: str, is_checkpoint: bool = True) -> str: |
| """Add the index to the string |
| |
| Args: |
| text (str): the input string |
| is_checkpoint (bool): if the string is a checkpoint lits or a prompt list |
| Returns: |
| str: the string with the index |
| """ |
| text_string = "" |
| if is_checkpoint: |
| checkpoint_List = self.getCheckpointListFromInput(text) |
| for i, checkpoint in enumerate(checkpoint_List): |
| text_string += f"{self.remove_index_from_string(checkpoint)} @index:{i},\n" |
| return text_string |
| else: |
| prompt_list = self.split_prompts(text) |
| for i, prompt in enumerate(prompt_list): |
| text_string += f"{self.remove_index_from_string(prompt)} @index:{i};\n\n" |
| return text_string |
|
|
| def add_model_version_to_string(self, checkpoints_text: str) -> str: |
| """Add the model version to the string. |
| EXPERIMENTAL! |
| |
| Args: |
| checkpoints_text (str): the input string with all checkpoints |
| Returns: |
| str: the string with the model version |
| """ |
| text_string = "" |
| checkpoints_not_cleaned = self.getCheckpointListFromInput( |
| checkpoints_text, clean=False) |
| checkpoints = self.getCheckpointListFromInput(checkpoints_text) |
| for i, checkpoint in enumerate(checkpoints): |
| info = modules.sd_models.get_closet_checkpoint_match(checkpoint) |
| state_dict = read_state_dict(info.filename) |
| version_string = find_checkpoint_config(state_dict, None) |
| if version_string == config_default: |
| version_string = "sd1" |
| elif version_string == config_sd2: |
| version_string = "sd2" |
| elif version_string == config_sd2v: |
| version_string = "sd2v" |
| elif version_string == config_sd2_inpainting: |
| version_string = "sd2-inpainting" |
| elif version_string == config_depth_model: |
| version_string = "depth" |
| elif version_string == config_unclip: |
| version_string = "unclip" |
| elif version_string == config_unopenclip: |
| version_string = "unopenclip" |
| elif version_string == config_inpainting: |
| version_string = "sd1-inpainting" |
| elif version_string == config_instruct_pix2pix: |
| version_string = "pix2pix" |
| elif version_string == config_alt_diffusion: |
| version_string = "alt" |
| checkpoint_partly_cleaned = checkpoints_not_cleaned[i].replace( |
| "\n", "").replace(",", "") |
| text_string += f"{checkpoint_partly_cleaned} @version:{version_string},\n\n" |
| return text_string |
|
|
| def remove_element_at_index(self, checkpoints: str, prompts: str, index: List[int]) -> List[str]: |
| """Remove the element at the given index from the string |
| |
| Args: |
| checkpoints (str): the input string with all checkpoints |
| prompts (str): the input string with all prompts |
| index (List[int]): the indices to remove |
| Returns: |
| List[str]: a list with the new checkpoints and prompts |
| """ |
|
|
| checkpoints_list = self.getCheckpointListFromInput(checkpoints) |
| prompts_list = self.split_prompts(prompts) |
| if (len(checkpoints_list) == len(prompts_list) or len(prompts_list) - len(index) <= 0 ): |
| if max(index) <= len(checkpoints_list) -1: |
| for i in index: |
| checkpoints_list.pop(i) |
| prompts_list.pop(i) |
| checkpoints = "" |
| for c in checkpoints_list: |
| checkpoints += f"{c}," |
| prompts = "" |
| for p in prompts_list: |
| prompts += f"{p};" |
| result = [self.add_index_to_string(checkpoints, True), self.add_index_to_string(prompts, False)] |
| self.logger.debug_log(f"result: {result}") |
| return result |
| else: |
| self.logger.debug_log("index is out of range") |
| return [checkpoints, prompts] |
| else: |
| self.logger.debug_log( |
| f"checkpoints and prompts are not the same length cp: {len(checkpoints_list)} p: {len(prompts_list)}") |
| return [checkpoints, prompts] |
