Commit a112168d authored by Justin Maier's avatar Justin Maier Committed by AUTOMATIC1111
Browse files

Save parameters as comments when saving as jpg

parent ca5901b5
Loading
Loading
Loading
Loading
+14 −5
Original line number Diff line number Diff line
@@ -15,6 +15,7 @@ import subprocess as sp
import numpy as np
import torch
from PIL import Image, PngImagePlugin
import piexif

import gradio as gr
import gradio.utils
@@ -113,18 +114,26 @@ def save_files(js_data, images, index):
            writer.writerow(["prompt", "seed", "width", "height", "sampler", "cfgs", "steps", "filename", "negative_prompt"])

        filename_base = str(int(time.time() * 1000))
        extension = opts.samples_format.lower()
        for i, filedata in enumerate(images):
            filename = filename_base + ("" if len(images) == 1 else "-" + str(i + 1)) + f".{opts.samples_format}"
            filename = filename_base + ("" if len(images) == 1 else "-" + str(i + 1)) + f".{extension}"
            filepath = os.path.join(opts.outdir_save, filename)

            if filedata.startswith("data:image/png;base64,"):
                filedata = filedata[len("data:image/png;base64,"):]

            image = Image.open(io.BytesIO(base64.decodebytes(filedata.encode('utf-8'))))
            if opts.enable_pnginfo and extension == 'png':
                pnginfo = PngImagePlugin.PngInfo()
                pnginfo.add_text('parameters', infotexts[i])
                image.save(filepath, pnginfo=pnginfo)
            else:
                image.save(filepath, quality=opts.jpeg_quality)

            image = Image.open(io.BytesIO(base64.decodebytes(filedata.encode('utf-8'))))
            image.save(filepath, quality=opts.jpeg_quality, pnginfo=pnginfo)
            if opts.enable_pnginfo and extension in ("jpg", "jpeg", "webp"):
                piexif.insert(piexif.dump({"Exif": {
                    piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(infotexts[i], encoding="unicode")
                }}), filepath)

            filenames.append(filename)