Create pnginfo.py

pnginfo.py

@@ -0,0 +1,127 @@
+from PIL import Image
+import piexif
+import json
+import html
+import re
+
+re_param_code = r'\s*([\w ]+):\s*("(?:\\.|[^\\"])+"|[^,]*)(?:,|$)'
+re_param = re.compile(re_param_code)
+re_imagesize = re.compile(r"^(\d+)x(\d+)$")
+IGNORED_INFO_KEYS = {
+    'jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',
+    'loop', 'background', 'timestamp', 'duration', 'progressive', 'progression',
+    'icc_profile', 'chromaticity', 'photoshop',
+}
+
+def plaintext_to_html(text, classname=None):
+    content = "<br>\n".join(html.escape(x) for x in text.split('\n'))
+
+    return f"<p class='{classname}'>{content}</p>" if classname else f"<p>{content}</p>"
+
+def to_digit(v):
+    try:
+        return float(v)
+    except:
+        return v
+    
+def read_info_from_image(image) -> tuple[str | None, dict]:
+    if image is None:
+        return '', {}
+    elif type(image) == str:
+        image = Image.open(image)
+    items = (image.info or {}).copy()
+    info =''
+    if "exif" in items:
+        exif = piexif.load(items["exif"])
+        exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b'')
+        try:
+            exif_comment = piexif.helper.UserComment.load(exif_comment)
+        except ValueError:
+            exif_comment = exif_comment.decode('utf8', errors="ignore")
+
+        if exif_comment:
+            items['exif comment'] = exif_comment
+            geninfo = exif_comment
+
+    for field in IGNORED_INFO_KEYS:
+        items.pop(field, None)
+
+    if items.get("Software", None) == "NovelAI":
+        json_info = json.loads(items["Comment"])
+
+        geninfo = f"""{items["Description"]}
+Negative prompt: {json_info["uc"]}
+Steps: {json_info["steps"]}, Sampler: {json_info['sampler']}, CFG scale: {json_info["scale"]}, Seed: {json_info["seed"]}, Size: {image.width}x{image.height}, Clip skip: 2, ENSD: 31337"""
+
+        items = {**{'parameters': geninfo}, **items}
+
+        for key, text in items.items():
+            info += f"""
+    <div>
+    <p><b>{plaintext_to_html(str(key))}</b></p>
+    <p>{plaintext_to_html(str(text))}</p>
+    </div>
+    """.strip()+"\n"
+
+        if len(info) == 0:
+            message = "Nothing found in the image."
+            info = f"<div><p>{message}<p></div>"
+    elif 'parameters' in items:
+        res = {}
+
+        prompt = ""
+        negative_prompt = ""
+
+        done_with_prompt = False
+
+        *lines, lastline = items['parameters'].strip().split("\n")
+        if len(re_param.findall(lastline)) < 3:
+            lines.append(lastline)
+            lastline = ''
+
+        for line in lines:
+            line = line.strip()
+            if line.startswith("Negative prompt:"):
+                done_with_prompt = True
+                line = line[16:].strip()
+            if done_with_prompt:
+                negative_prompt += ("" if negative_prompt == "" else "\n") + line
+            else:
+                prompt += ("" if prompt == "" else "\n") + line
+
+        res["prompt"] = prompt
+        res["negative prompt"] = negative_prompt
+
+        for k, v in re_param.findall(lastline):
+            try:
+                if v[0] == '"' and v[-1] == '"':
+                    v = to_digit(json.loads(v))
+
+                m = re_imagesize.match(v)
+                if m is not None:
+                    res[f"{k.lower()}-1"] = to_digit(m.group(1))
+                    res[f"{k.lower()}-2"] = to_digit(m.group(2))
+                else:
+                    res[k.lower()] = v
+            except Exception:
+                print(f"Error parsing \"{k}: {v}\"")
+
+        return items['parameters'], res
+    else:
+        return info, {}
+    return info, json.loads(items['Comment'])
+
+
+def send_paras(*args):
+    if len(args[0]) == 0:
+        return args[1:]
+    items, prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale = args
+    paras = [prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale]
+    for i, keys in enumerate([('prompt',), ('',), ('uc', 'negative prompt'), ('seed',), ('scale', 'cfg scale'), ('width', 'size-1'), ('height', 'size-2'), ('steps',), ('sampler',), ('noise_schedule',), ('sm',), ('sm_dyn',), ('dynamic_thresholding',), ('cfg_rescale',)]):
+        for key in keys:
+            if key in items:
+                paras[i] = items[key]
+    prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale = paras
+    if prompt.endswith(', ' + quality_tags):
+        prompt = prompt[:-(len(quality_tags) + 2)]
+    return prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale