236 lines
7.7 KiB
Python
236 lines
7.7 KiB
Python
import json
|
|
import re
|
|
|
|
DEBUG = False
|
|
|
|
|
|
def dprint(str):
|
|
if DEBUG:
|
|
print(f"debug:{str}")
|
|
|
|
|
|
class PromptInfoExtractor:
|
|
def __init__(self, prompt, config_filepath=None):
|
|
"""constructor
|
|
|
|
Args:
|
|
prompt (_type_): ComfyUI hidden object "prompt".
|
|
config_filepath (_type_, optional): Configration file path. Defaults to None.
|
|
"""
|
|
# self.load_data(prompt)
|
|
self._prompt = prompt
|
|
|
|
if DEBUG:
|
|
self._show_data()
|
|
|
|
if config_filepath:
|
|
self.load_config(config_filepath)
|
|
else:
|
|
self.config = {
|
|
"search_class_types": ["KSampler", "KSamplerAdvanced", "KSampler With Refiner (Fooocus)", "BNK_TiledKSampler"],
|
|
"output_format": "Steps: {steps}, Sampler: {sampler_name} {scheduler}, CFG scale: {cfg}, Seed: {seed}, Size: {width}x{height}, Model: {model_name}",
|
|
}
|
|
self.info = self.gather_info()
|
|
|
|
def load_data(self, json_filepath):
|
|
"""Load JSON data from the provided filepath."""
|
|
with open(json_filepath, "r") as file:
|
|
self._prompt = json.load(file)
|
|
|
|
def load_config(self, config_filepath):
|
|
"""Load configuration from the provided filepath."""
|
|
with open(config_filepath, "r") as config_file:
|
|
self.config = json.load(config_file)
|
|
|
|
def _show_data(self):
|
|
"""For debug. show information of ComfyUI hidden object "prompt" and "extra_onginfo"."""
|
|
dprint(f"prompt(hidden object){self._prompt}")
|
|
|
|
def gather_info(self):
|
|
ksampler_items = self.get_ksampler_items()
|
|
|
|
if not ksampler_items:
|
|
return None
|
|
|
|
key, ksampler_item = ksampler_items[0]
|
|
|
|
model_name = self.extract_model_name(ksampler_item)
|
|
latent_image_info = self.extract_latent_image_info(ksampler_item)
|
|
prompt_text = self.extract_prompt_info()
|
|
|
|
info_dict = {
|
|
"steps": ksampler_item["inputs"]["steps"],
|
|
"sampler_name": ksampler_item["inputs"]["sampler_name"],
|
|
"scheduler": ksampler_item["inputs"]["scheduler"],
|
|
"cfg": ksampler_item["inputs"]["cfg"],
|
|
"seed": ksampler_item["inputs"].get(
|
|
"seed", ksampler_item["inputs"].get("noise_seed", None)
|
|
),
|
|
"model_name": model_name,
|
|
"width": latent_image_info["inputs"]["width"],
|
|
"height": latent_image_info["inputs"]["height"],
|
|
}
|
|
if "prompt" in prompt_text:
|
|
info_dict["prompt"] = prompt_text["prompt"]
|
|
else:
|
|
info_dict["prompt"] = ""
|
|
|
|
if "negative" in prompt_text:
|
|
info_dict["negative"] = prompt_text["negative"]
|
|
else:
|
|
info_dict["negative"] = ""
|
|
|
|
return info_dict
|
|
|
|
def get_ksampler_items(self):
|
|
ksampler_items = [
|
|
(k, v)
|
|
for k, v in self._prompt.items()
|
|
if v["class_type"] in self.config["search_class_types"]
|
|
]
|
|
return sorted(ksampler_items, key=lambda x: int(x[0]))
|
|
|
|
def extract_model_name(self, item):
|
|
return (
|
|
self.get_ckpt_name(item["inputs"]["model"][0])
|
|
.replace("/", "_")
|
|
.replace("\\", "_")
|
|
)
|
|
|
|
def get_ckpt_name(self, node_number):
|
|
"""Recursively search for the 'ckpt_name' key starting from the specified node."""
|
|
node = self._prompt[node_number]
|
|
if "ckpt_name" in node["inputs"]:
|
|
return node["inputs"]["ckpt_name"]
|
|
if "model" in node["inputs"]:
|
|
return self.get_ckpt_name(node["inputs"]["model"][0])
|
|
return None
|
|
|
|
def extract_latent_image_info(self, item):
|
|
latent_image_node_number = item["inputs"]["latent_image"][0]
|
|
target_item = self._prompt[latent_image_node_number]
|
|
|
|
if target_item["class_type"] == "EmptyLatentImage":
|
|
return target_item
|
|
|
|
elif target_item["class_type"] == "SDXL Empty Latent Image":
|
|
resolusion_str = target_item["inputs"]["resolution"]
|
|
pattern = r"(\d+) x (\d+)"
|
|
match = re.search(pattern, resolusion_str)
|
|
|
|
latent_image_info = {
|
|
"inputs": {"width": int(match.group(1)), "height": int(match.group(2))}
|
|
}
|
|
return latent_image_info
|
|
|
|
else:
|
|
latent_image_info = {"inputs": {"width": 0, "height": 0}}
|
|
return latent_image_info
|
|
|
|
def extract_prompt_info(self):
|
|
positive_text = self.extract_text_by_key("positive")
|
|
negative_text = self.extract_text_by_key("negative")
|
|
|
|
info_dict = {}
|
|
if positive_text:
|
|
info_dict["prompt"] = positive_text
|
|
if negative_text:
|
|
info_dict["negative"] = negative_text
|
|
|
|
return info_dict
|
|
|
|
def extract_text_by_key(self, key):
|
|
"""Extract text by the given key, either 'positive' or 'negative'."""
|
|
ksampler_items = self.get_ksampler_items()
|
|
|
|
if not ksampler_items:
|
|
return None
|
|
|
|
ksampler_item = ksampler_items[0][1]
|
|
|
|
if key not in ksampler_item["inputs"]:
|
|
return None
|
|
|
|
target_node_number = ksampler_item["inputs"][key][0]
|
|
target_node = self._prompt.get(str(target_node_number), {})
|
|
extracted_text = self.extract_text_from_node_v2(target_node, key)
|
|
|
|
return extracted_text
|
|
|
|
def extract_text_from_node_v2(self, node, key):
|
|
# For CLIPTextEncoder
|
|
direct_text = node.get("inputs", {}).get("text")
|
|
if isinstance(direct_text, str):
|
|
return direct_text
|
|
|
|
# For SDXLCLIPTextEncoder
|
|
text_g = node.get("inputs", {}).get("text_g")
|
|
text_l = node.get("inputs", {}).get("text_l")
|
|
if text_g and text_l:
|
|
if text_g == text_l:
|
|
return text_g
|
|
else:
|
|
return f"text_g:{text_g} text_l:{text_l}"
|
|
elif text_g:
|
|
return text_g
|
|
elif text_l:
|
|
return text_l
|
|
|
|
# For SDXLPromptStyler
|
|
try:
|
|
node_number = node.get("inputs", {}).get("text")[0]
|
|
if key == "positive":
|
|
text_positive = self._prompt[node_number]["inputs"]["text_positive"]
|
|
if text_positive:
|
|
return text_positive
|
|
else:
|
|
text_negative = self._prompt[node_number]["inputs"]["text_negative"]
|
|
if text_negative:
|
|
return text_negative
|
|
except(Exception) as e:
|
|
return None
|
|
|
|
if DEBUG:
|
|
return "Prompt parse error"
|
|
else:
|
|
return None
|
|
|
|
def format_info(self, info_dict):
|
|
"""Format the gathered information based on the configuration."""
|
|
formatted_str = self.config["output_format"].format(**info_dict)
|
|
return formatted_str
|
|
|
|
def extract_and_format(self):
|
|
"""Extract and format the required information from the loaded JSON data."""
|
|
info = self.gather_info()
|
|
if not info:
|
|
return "No suitable data found."
|
|
return self.format_info(info)
|
|
|
|
def formatted_annotation(self):
|
|
annotation = ""
|
|
if len(self.info["prompt"]) > 0:
|
|
annotation += self.info["prompt"]
|
|
|
|
if len(self.info["negative"]) > 0:
|
|
if len(annotation) > 0:
|
|
annotation += "\n"
|
|
annotation += "Negative prompt: "
|
|
annotation += self.info["negative"]
|
|
|
|
if len(annotation) > 0:
|
|
annotation += "\n"
|
|
annotation += self.extract_and_format()
|
|
|
|
return annotation
|
|
|
|
def get_prompt_tags(self):
|
|
cleaned_string = re.sub(r":\d+\.\d+", "", self.info["prompt"])
|
|
items = cleaned_string.split(",")
|
|
|
|
return [
|
|
re.sub(r"[\(\)]", "", item).strip()
|
|
for item in items
|
|
if re.sub(r"[\(\)]", "", item).strip()
|
|
]
|