3v324v23's picture
lfs
1e3b872
__author__ = "receyuki"
__filename__ = "swarmui.py"
__copyright__ = "Copyright 2023"
__email__ = "receyuki@gmail.com"
import json
from ..format.base_format import BaseFormat
from ..utility import remove_quotes
class SwarmUI(BaseFormat):
SETTING_KEY = ["model", "", "seed", "cfgscale", "steps", ""]
def __init__(self, info: dict = None, raw: str = ""):
super().__init__(info, raw)
if not self._info:
self._info = json.loads(self._raw)
self._ss_format()
def _ss_format(self):
data_json = self._info.get("sui_image_params")
self._positive = data_json.get("prompt").strip()
self._negative = data_json.get("negativeprompt").strip()
self._raw = "\n".join([self._positive, self._negative, str(data_json)]).strip()
data_json.pop("prompt")
data_json.pop("negativeprompt")
self._setting = remove_quotes(str(data_json).strip("{ }"))
self._width = str(data_json.get("width"))
self._height = str(data_json.get("height"))
for p, s in zip(super().PARAMETER_KEY, SwarmUI.SETTING_KEY):
match p:
case "sampler":
comfyuisampler = data_json.get("comfyuisampler")
autowebuisampler = data_json.get("autowebuisampler")
self._parameter["sampler"] = str(
remove_quotes((comfyuisampler, autowebuisampler))
if comfyuisampler and autowebuisampler
else comfyuisampler or autowebuisampler
)
case "size":
self._parameter["size"] = (
str(data_json.get("width")) + "x" + str(data_json.get("height"))
)
case _:
self._parameter[p] = str(data_json.get(s))