|
|
__author__ = "receyuki" |
|
|
__filename__ = "invokeai.py" |
|
|
__copyright__ = "Copyright 2023" |
|
|
__email__ = "receyuki@gmail.com" |
|
|
|
|
|
import json |
|
|
import re |
|
|
|
|
|
from ..format.base_format import BaseFormat |
|
|
from ..utility import remove_quotes |
|
|
|
|
|
|
|
|
class InvokeAI(BaseFormat): |
|
|
SETTING_KEY_INVOKEAI_METADATA = [ |
|
|
"", |
|
|
("scheduler", "refiner_scheduler"), |
|
|
"seed", |
|
|
("cfg_scale", "refiner_cfg_scale"), |
|
|
("steps", "refiner_steps"), |
|
|
"", |
|
|
] |
|
|
SETTING_KEY_METADATA = [ |
|
|
"model_weights", |
|
|
"sampler", |
|
|
"seed", |
|
|
"cfg_scale", |
|
|
"steps", |
|
|
"", |
|
|
] |
|
|
SETTING_KEY_DREAM = ["", "A", "S", "C", "s", ""] |
|
|
DREAM_MAPPING = { |
|
|
"Steps": "s", |
|
|
"Seed": "S", |
|
|
"Size": "", |
|
|
"CFG scale": "C", |
|
|
"Sampler": "A", |
|
|
} |
|
|
|
|
|
def __init__(self, info: dict = None, raw: str = ""): |
|
|
super().__init__(info, raw) |
|
|
if "invokeai_metadata" in self._info: |
|
|
self._invoke_invoke_metadata() |
|
|
elif "sd-metadata" in self._info: |
|
|
self._invoke_metadata() |
|
|
elif "Dream" in self._info: |
|
|
self._invoke_dream() |
|
|
|
|
|
def _invoke_invoke_metadata(self): |
|
|
data_json = json.loads(self._info.get("invokeai_metadata")) |
|
|
self._positive = data_json.pop("positive_prompt").strip() |
|
|
self._negative = data_json.pop("negative_prompt").strip() |
|
|
self._raw = "\n".join([self._positive, self._negative, str(data_json)]) |
|
|
self._setting = remove_quotes(str(data_json)).strip("{ }") |
|
|
self._width = str(data_json.get("width")) |
|
|
self._height = str(data_json.get("height")) |
|
|
|
|
|
has_refiner = True if data_json.get("refiner_model") else False |
|
|
|
|
|
for p, s in zip(super().PARAMETER_KEY, InvokeAI.SETTING_KEY_INVOKEAI_METADATA): |
|
|
match p: |
|
|
case "model": |
|
|
self._parameter[p] = remove_quotes( |
|
|
str( |
|
|
( |
|
|
data_json.get("model").get("model_name"), |
|
|
data_json.get("refiner_model").get("model_name"), |
|
|
) |
|
|
) |
|
|
if has_refiner |
|
|
else str(data_json.get("model").get("model_name")) |
|
|
) |
|
|
case "seed": |
|
|
self._parameter["seed"] = str(data_json.get("seed")) |
|
|
case "size": |
|
|
self._parameter["size"] = ( |
|
|
str(data_json.get("width")) + "x" + str(data_json.get("height")) |
|
|
) |
|
|
case _: |
|
|
self._parameter[p] = remove_quotes( |
|
|
str((data_json.get(s[0]), data_json.get(s[1]))) |
|
|
if has_refiner |
|
|
else str(data_json.get(s[0])) |
|
|
) |
|
|
|
|
|
def _invoke_metadata(self): |
|
|
data_json = json.loads(self._info.get("sd-metadata")) |
|
|
image = data_json.pop("image") |
|
|
prompt = ( |
|
|
image.get("prompt")[0].get("prompt") |
|
|
if isinstance(image.get("prompt"), list) |
|
|
else image.get("prompt") |
|
|
) |
|
|
|
|
|
self._positive, self._negative = self.split_prompt(prompt) |
|
|
|
|
|
raw_list = [ |
|
|
item |
|
|
for item in [ |
|
|
self._positive, |
|
|
self._negative, |
|
|
self._info.get("Dream"), |
|
|
self._info.get("sd-metadata"), |
|
|
] |
|
|
if item != "" |
|
|
] |
|
|
|
|
|
self._raw = "\n".join(raw_list).strip() |
|
|
|
|
|
image.pop("prompt") |
|
|
self._setting = remove_quotes( |
|
|
", ".join([str(data_json).strip("{ }"), str(image).strip("{ }")]) |
|
|
) |
|
|
|
|
|
self._width = str(image.get("width")) |
|
|
self._height = str(image.get("height")) |
|
|
|
|
|
for p, s in zip(super().PARAMETER_KEY, InvokeAI.SETTING_KEY_METADATA): |
|
|
match p: |
|
|
case "model": |
|
|
self._parameter["model"] = data_json.get(s) |
|
|
case "size": |
|
|
self._parameter["size"] = ( |
|
|
str(image.get("width")) + "x" + str(image.get("height")) |
|
|
) |
|
|
case _: |
|
|
self._parameter[p] = str(image.get(s)) |
|
|
|
|
|
def _invoke_dream(self): |
|
|
data = self._info.get("Dream") |
|
|
|
|
|
|
|
|
pattern = r'"(.*?)"\s*(.*?)$' |
|
|
prompt, setting = re.search(pattern, data).groups() |
|
|
self._positive, self._negative = self.split_prompt(prompt.strip('" ')) |
|
|
|
|
|
self._raw = "\n".join([self._positive, self.negative, self._info.get("Dream")]) |
|
|
|
|
|
|
|
|
pattern = r"-(\w+)\s+([\w.-]+)" |
|
|
setting_dict = dict(re.findall(pattern, setting)) |
|
|
setting_list = [] |
|
|
for key, value in InvokeAI.DREAM_MAPPING.items(): |
|
|
if key == "Size": |
|
|
setting_list.append( |
|
|
key + ": " + setting_dict.get("W") + "x" + setting_dict.get("H") |
|
|
) |
|
|
else: |
|
|
setting_list.append(key + ": " + setting_dict.get(value)) |
|
|
self._setting = ", ".join(setting_list) |
|
|
|
|
|
self._width = str(setting_dict.get("W")) |
|
|
self._height = str(setting_dict.get("H")) |
|
|
|
|
|
for p, s in zip(super().PARAMETER_KEY, InvokeAI.SETTING_KEY_DREAM): |
|
|
match p: |
|
|
case "model": |
|
|
self._parameter["model"] = "" |
|
|
case "size": |
|
|
self._parameter["size"] = ( |
|
|
str(setting_dict.get("W")) + "x" + str(setting_dict.get("H")) |
|
|
) |
|
|
case _: |
|
|
self._parameter[p] = setting_dict.get(s) |
|
|
|
|
|
@staticmethod |
|
|
def split_prompt(prompt: str): |
|
|
|
|
|
pattern = r"^(.*?)\[(.*?)\]$" |
|
|
match = re.match(pattern, prompt) |
|
|
|
|
|
if match: |
|
|
positive, negative = match.groups() |
|
|
positive = positive.strip() |
|
|
negative = negative.strip() |
|
|
else: |
|
|
positive = prompt.strip() |
|
|
negative = "" |
|
|
return positive, negative |
|
|
|
