File size: 3,337 Bytes
bb7f1f4 |
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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
"""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
|