| 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 |
|
|
