62 lines
2.1 KiB
Python
62 lines
2.1 KiB
Python
from PIL import Image
|
|
|
|
import piexif
|
|
import piexif.helper
|
|
from scripts.iib.tool import parse_generation_parameters, replace_punctuation
|
|
from scripts.iib.parsers.model import ImageGenerationInfo, ImageGenerationParams
|
|
from PIL.ExifTags import TAGS
|
|
import json
|
|
|
|
|
|
class StableSwarmUIParser:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def get_exif_data(clz, image: Image) -> str:
|
|
items = image.info or {}
|
|
|
|
if "exif" in items:
|
|
exif = piexif.load(items["exif"])
|
|
exif_bytes = (
|
|
(exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b"")
|
|
)
|
|
|
|
unicode_start = exif_bytes.find(b"UNICODE")
|
|
if unicode_start == -1:
|
|
raise ValueError("'UNICODE' markup isn't found")
|
|
|
|
unicode_data = exif_bytes[unicode_start + len("UNICODE") + 1 :]
|
|
geninfo = unicode_data.decode("utf-16")
|
|
return geninfo
|
|
|
|
@classmethod
|
|
def parse(clz, img: Image, file_path):
|
|
if not clz.test(img, file_path):
|
|
raise Exception("The input image does not match the current parser.")
|
|
exif_data = json.loads(clz.get_exif_data(img))["sui_image_params"]
|
|
prompt = exif_data.pop("prompt")
|
|
negativeprompt = exif_data.pop("negativeprompt")
|
|
steps = exif_data.pop("steps")
|
|
meta_kv = [f"Steps: {steps}", "Source Identifier: StableSwarmUI"]
|
|
for key, value in exif_data.items():
|
|
value = replace_punctuation(str(value))
|
|
meta_kv.append(f"{key}: {value}")
|
|
meta = ", ".join(meta_kv)
|
|
info = "\n".join([prompt, f"Negative prompt: {negativeprompt}", meta])
|
|
params = parse_generation_parameters(info)
|
|
return ImageGenerationInfo(
|
|
info,
|
|
ImageGenerationParams(
|
|
meta=params["meta"], pos_prompt=params["pos_prompt"], extra=params
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def test(clz, img: Image, file_path: str):
|
|
try:
|
|
exif = clz.get_exif_data(img)
|
|
return exif.find("sui_image_params") != -1
|
|
except Exception as e:
|
|
return False
|