png_info_helper.py

import os 
import re
import json
from json import JSONDecodeError
from PIL import Image, PngImagePlugin
from enum import Enum #IntFlag?

'''
一共六种图片信息来源的情况
1.  [Paddle] 直接输出到txt
2.  [Paddle] 直接保存到png的info中
3.  [PaddleLikeWebUI] 仿webui输出到txt
4.  [PaddleLikeWebUI] 仿webui输出到png的info[parameters]中
5.  [WebUI] webui图片
6.  [NAIFU] NAIFU图片
'''
class InfoFormat(Enum):
    Paddle = 1
    WebUI = 4
    NAIFU = 8
    
    PaddleLikeWebUI = 5
    Unknown = 255

# 所有Paddle输出的参数
PRAM_NAME_LIST = (
    'prompt',
    'negative_prompt',
    'height',
    'width',
    'num_inference_steps',
    'guidance_scale',
    'num_images_per_prompt',
    'eta',
    'seed',
    'latents',
    'max_embeddings_multiples',
    'no_boseos_middle',
    'skip_parsing',
    'skip_weighting',
    'epoch_time',
    'sampler',
    'superres_model_name',
    'model_name',
    
    # img2img
    'init_image',
)

#情况3/4 [Paddle]=>[PaddleLikeWebUI]
MAP_PARAM_TO_LABEL = {
    'prompt': '',
    'negative_prompt': 'Negative prompt: ', #用于与webui保持一致
    'num_inference_steps': 'Steps: ', #用于与webui保持一致
    'sampler': 'Sampler: ',
    'guidance_scale': 'CFG scale: ',
    'strength': 'Strength: ',
    'seed': 'Seed: ',
    'width':'width: ',
    'height':'height: ',
}
#情况3/4/5 [PaddleLikeWebUI]/[WebUI]=>[Paddle]
MAP_LAEBL_TO_PARAM = {
    'Prompt': 'prompt',
    'Negative prompt': 'negative_prompt',
    'Steps': 'num_inference_steps',
    'Sampler': 'sampler',
    'CFG Scale': 'guidance_scale',
    'CFG scale': 'guidance_scale',
    'Strength': 'strength',
    'Seed': 'seed',
    'Width':'width',
    'Height':'height',
    #webui
    'Eta': 'eta',
    'Model': 'model_name', #注意model_name可能是webui的模型，不能直接使用
    'Model hash': 'model_hash',
}
# [NAIFU]=>[Paddle]
MAP_NAIFU_TAG_TO_PARAM = {
    'steps': 'num_inference_steps',
    'scale': 'guidance_scale',
    'uc': 'negative_prompt',
}

# 可信的，是文本内容的tag。用于info信息继承
RELIABEL_TAG_LIST = (
    'Title', 
    'Description', 
    'Software', 
    'Source', 
    'Comment', 
    'parameters',
    'original_parameters',
)

# --------------------------------------------------
# 序列化
# --------------------------------------------------

# 输出[PaddleLikeWebUI]的样式
def serialize_to_text(params):
    """
    将参数序列化为文本，以用于保存图片。格式[PaddleLikeWebUI]
    """
    # Todo 剔除无用信息
    labels = MAP_PARAM_TO_LABEL
    info = ''
    for k in labels:
        if k in params: info += labels[k] + str(params[k]) + '\n'
    
    for k in params:
        if k not in labels: info += k + ': ' + str(params[k]) + '\n'
    
    return info
    
    
def serialize_to_pnginfo(params, existing_info = None, mark_paddle = True):
    """
    将参数序列化到图像信息中。格式[PaddleLikeWebUI]
    参数 existing_info 用于继承。
    """
    text = serialize_to_text(params)
    
    dict = {}
    if existing_info is None:
        pass
    elif 'parameters' in existing_info:
        # dict.update(existing_info)
        dict['original_parameters'] = existing_info.pop('parameters')
    elif 'prompt' in existing_info:
        # 如果是[Paddle]，那么将其转换为[PaddleLikeWebUI]，并舍弃掉[Paddle]信息
        dict['original_parameters'] = serialize_to_text(existing_info)
        for k in existing_info:
            if k not in PRAM_NAME_LIST: dict[k] = existing_info[k]
    else:
        #当做[NAIFU]处理(未知的tag不能保证数据类型是iTXt)
        for k in ('Title', 'Description', 'Software', 'Source', 'Comment'):
            if k in existing_info:
                dict[k] = existing_info[k]
    
    if mark_paddle: dict['Software'] = 'PaddleNLP'
    dict['parameters'] = text

    pnginfo = PngImagePlugin.PngInfo()
    for key, val in dict.items():
        pnginfo.add_text(key, str(val))
        
    return pnginfo

def imageinfo_to_pnginfo(info, update_format = True):
    """
    从[Image.info]生成[PngInfo]，用于继承图片信息。
    仅用于超分辨率（highres）的图片保存。
    不认识的信息会被过滤掉。
    """
    dict = {}
    if ('prompt' in info):
        if update_format:
            # 如果是[Paddle]，那么将其转换为[PaddleLikeWebUI]，并舍弃掉[Paddle]信息
            dict['parameters'] = serialize_to_text(info)
            for k in info:
                if k not in PRAM_NAME_LIST: dict[k] = info[k]
        else:
            for k in PRAM_NAME_LIST:
                if k in info: dict[k] = info[k]
                
    # 可信的info
    for k in RELIABEL_TAG_LIST:
        if k in info:
            dict[k] = info[k]

    pnginfo = PngImagePlugin.PngInfo()
    for key, val in dict.items():
        pnginfo.add_text(key, str(val))
        
    return pnginfo
    
# --------------------------------------------------
# 反序列化
# --------------------------------------------------
    
    
def _parse_value(s):
    if s == 'None': return None
    if s == 'False': return False
    if s == 'True': return True
    if re.fullmatch(r'[\d\.]+', s): 
        return int(s) if s.find('.') < 0 else float(s) 
    
    return s
 
# 只支持[Paddle][PaddleLikeWebUI][WebUI]
def _deserialize_from_lines(enumerable, format_presumed = InfoFormat.Unknown):
    dict = {}
    fmt = format_presumed
    
    ln = -1
    name = 'prompt'
    for line in enumerable:
        ln += 1
        line = line.rstrip('\n')
        key, colon, val = line.partition(': ')

        if (ln == 0) and (key == 'prompt'):
            fmt = InfoFormat.Paddle
        
        # 没有冒号分隔
        if colon == '':
            if ln == 0:
                name = 'prompt'
                dict[name] = line
            elif name == 'prompt' or name == 'negative_prompt':
                dict[name] += '\n' + line #追加上一行
            elif line == '':
                pass
            elif re.fullmatch(r'Strength[\d\.]+', name):
                # 兼容之前Strength标签错误
                dict['strength'] = name[8:]
            else:
                #不认识的换行参数
                dict[name] += '\n' + line

        # 有冒号分隔
        elif key in PRAM_NAME_LIST:
            # 1/2原始格式
            name = key
            dict[name] = val
        elif key in MAP_LAEBL_TO_PARAM:
            # 3/4格式
            fmt = InfoFormat.PaddleLikeWebUI
            name = MAP_LAEBL_TO_PARAM[key]
            dict[name] = val
        
        # 发现标签但是不认识
        elif name == 'prompt' or name == 'negative_prompt':
            # prompt下不视为标签
            dict[name] += '\n' + line
        
        # 看着像一个标签
        elif re.fullmatch(r'\w+', name):
            # 当他是个标签
            name = key
            dict[name] = val
        else:
            dict[name] += '\n' + line #追加上一行
        
    
    # 处理webui格式（[WebUI]=>[Paddle]
    if ('num_inference_steps' in dict) and (dict['num_inference_steps'].find(', ') > -1):
        webui_text = dict['num_inference_steps']
        fmt = InfoFormat.WebUI
        webui_text = 'num_inference_steps: '+webui_text
        for pair in webui_text.split(', '):
            key, colon, val = pair.partition(': ')
            key = key if key not in MAP_LAEBL_TO_PARAM else MAP_LAEBL_TO_PARAM[key]
            dict[key] = val
            
        # 处理Size: 768x512
        if ('Size' in dict):
            size = re.split(r'\D',dict.pop('Size'))
            dict['width'] = size[0]
            dict['height'] = size[1]
    
    for k in dict:
        dict[k] = _parse_value(dict[k])
    return (dict,fmt)
    
def deserialize_from_txt(text, format_presumed = InfoFormat.Unknown):
    """ 从一段文本提取参数信息。支持格式[Paddle][PaddleLikeWebUI][WebUI] """
    return _deserialize_from_lines(text.splitlines(), format_presumed)
    
# 直接从图片Info中收集信息[Paddle]
def _collect_from_pnginfo(info):
    dict = {}
    for key in info:
        if key in PRAM_NAME_LIST:
            dict[key] = _parse_value(info[key])
    fmt = InfoFormat.Paddle if 'prompt' in dict else InfoFormat.Unknown
    return (dict,fmt)

# 从Naifu中提取信息 [NAIFU]
def _collect_from_pnginfo_naifu(info):
    if ('Description' not in info) \
    or ('Comment' not in info) \
    or ('Software' not in info):
        return ({}, InfoFormat.Unknown)
    
    try:
        data = json.loads(info['Comment'])
        data['prompt'] = info['Description']
        for key in MAP_NAIFU_TAG_TO_PARAM:
            if key in data:
                data[MAP_NAIFU_TAG_TO_PARAM[key]] = data.pop(key)
        return (data, InfoFormat.NAIFU)
    except JSONDecodeError:
        return ({}, InfoFormat.Unknown)

def deserialize_from_image(image):
    """ 从图片获取参数信息。参数为Image或文件地址。"""
    if isinstance(image, str):
        assert os.path.isfile(image), f'{image}不是可读文件'
        image = Image.open(image)
        
    if 'parameters' in image.info:  #是情况4/5 [PaddleLikeWebUI]
        return deserialize_from_txt(image.info['parameters'], InfoFormat.PaddleLikeWebUI)
    
    # [NAIFU]
    dict, fmt = _collect_from_pnginfo_naifu(image.info)
    if fmt is InfoFormat.NAIFU: return (dict, fmt)
    
    # [Paddle]
    return _collect_from_pnginfo(image.info)
    
def deserialize_from_filename(filename):
    """ 从文本文件或图像文件获取参数信息，优先从其对应的文本文件中提取。参数为文件地址。"""
    txt_path, dot, ext = filename.rpartition('.')
    txt_path += '.txt'
    if os.path.isfile(txt_path):
        with open(txt_path, 'r') as f:
            (dict,fmt) = _deserialize_from_lines(f)
            if (fmt is not InfoFormat.Unknown) or ext.lower() == 'txt':
                return (dict, fmt)
    
    return deserialize_from_image(filename)