59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
from PIL import Image
|
|
|
|
from scripts.iib.tool import (
|
|
comfyui_exif_data_to_str,
|
|
is_img_created_by_comfyui,
|
|
is_img_created_by_comfyui_with_webui_gen_info,
|
|
get_comfyui_exif_data,
|
|
parse_generation_parameters,
|
|
read_sd_webui_gen_info_from_image,
|
|
)
|
|
from scripts.iib.parsers.model import ImageGenerationInfo, ImageGenerationParams
|
|
from scripts.iib.logger import logger
|
|
|
|
|
|
class ComfyUIParser:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def parse(clz, img, file_path):
|
|
info = ""
|
|
params = None
|
|
if not clz.test(img, file_path):
|
|
raise Exception("The input image does not match the current parser.")
|
|
width, height = img.size
|
|
try:
|
|
if is_img_created_by_comfyui_with_webui_gen_info(img):
|
|
info = read_sd_webui_gen_info_from_image(img, file_path)
|
|
info += ", Source Identifier: ComfyUI"
|
|
params = parse_generation_parameters(info)
|
|
else:
|
|
params = get_comfyui_exif_data(img)
|
|
info = comfyui_exif_data_to_str(params)
|
|
except Exception as e:
|
|
logger.error("parse comfyui image failed. prompt:", exc_info=e)
|
|
logger.error(img.info.get("prompt"))
|
|
return ImageGenerationInfo(
|
|
params=ImageGenerationParams(
|
|
meta={"final_width": width, "final_height": height}
|
|
)
|
|
)
|
|
return ImageGenerationInfo(
|
|
info,
|
|
ImageGenerationParams(
|
|
meta=params["meta"] | {"final_width": width, "final_height": height},
|
|
pos_prompt=params["pos_prompt"],
|
|
extra=params,
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def test(clz, img: Image, file_path: str) -> bool:
|
|
try:
|
|
return is_img_created_by_comfyui(
|
|
img
|
|
) or is_img_created_by_comfyui_with_webui_gen_info(img)
|
|
except Exception:
|
|
return False
|