| """This modules uses the JSON files created by the | |
| https://github.com/butaixianran/Stable-Diffusion-Webui-Civitai-Helper extension""" | |
| import sys | |
| import os | |
| sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "scripts")) | |
| from scripts.Logger import Logger | |
| from scripts.Utils import Utils | |
| import modules.shared as shared | |
| import modules.scripts as scripts | |
| import json | |
| class CivitaihelperPrompts(): | |
| """ | |
| This class is used to get the prompts from the JSON files created by the | |
| Stable-Diffusion-Webui-Civitai-Helper extension | |
| """ | |
| def __init__(self) -> None: | |
| self.model_path = self.get_custom_model_folder() | |
| self.utils = Utils() | |
| self.logger = Logger() | |
| self.logger.debug = False | |
| def get_custom_model_folder(self) -> str: | |
| """Get the custom model folder from the command line options or the default folder | |
| Returns: | |
| str: the model folder path | |
| """ | |
| if shared.cmd_opts.ckpt_dir and os.path.isdir(shared.cmd_opts.ckpt_dir): | |
| return shared.cmd_opts.ckpt_dir | |
| else: | |
| return os.path.join(scripts.basedir(), "models", "Stable-diffusion") | |
| def get_civitAi_prompt_from_model(self, path: str) -> str: | |
| """Get the prompt from the civitai.info file | |
| Args: | |
| path (str): the path to the model | |
| Returns: | |
| str: the prompt | |
| """ | |
| path = path.replace(".ckpt", ".civitai.info").replace( | |
| ".safetensors", ".civitai.info") | |
| path = self.utils.get_clean_checkpoint_path(path) | |
| path = os.path.join(self.model_path, path) | |
| self.logger.debug_log(f"{path} -> is file {os.path.isfile(path)}") | |
| if not os.path.exists(os.path.realpath(path)): | |
| return "{prompt};" | |
| model_info = None | |
| with open(os.path.realpath(path), 'r') as f: | |
| try: | |
| model_info = json.load(f) | |
| except Exception as e: | |
| self.logger.debug_log(f"Error loading civitai info: {e}", False) | |
| return "{prompt};" | |
| try: | |
| self.logger.debug_log(f"len: {len(model_info['images'])}") | |
| for i in range(0, len(model_info['images'])): | |
| try: | |
| info = model_info['images'][i]['meta']['prompt'] | |
| self.logger.debug_log(f"Prompt: {info}") | |
| if info: | |
| return f"{info};" | |
| except: | |
| pass | |
| return "{prompt};" | |
| except: | |
| return "{prompt};" | |
| def createCivitaiPromptString(self, checkpoints: str) -> str: | |
| """Create a string with all the prompts from the checkpoints | |
| Indices are added to the prompts | |
| Args: | |
| checkpoints (str): the checkpoints | |
| Returns: | |
| str: the prompt string with indices | |
| """ | |
| checkpoints_list = self.utils.getCheckpointListFromInput(checkpoints) | |
| prompts = "" | |
| prompts_with_info = "" | |
| for i, checkpoint in enumerate(checkpoints_list): | |
| prompts += self.get_civitAi_prompt_from_model(checkpoint) | |
| prompts_with_info += self.utils.add_index_to_string( | |
| prompts, is_checkpoint=False) | |
| self.logger.log_info("loaded all prompts") | |
| return prompts_with_info | |