Saving generation metadata

[sfisher]

I find that I save images from the generated results quite often and then when I go back to see how I created the images, they're just PNG files without information about the prompt or anything else.

So . . . I modified the code in my own installation to save the png generation metadata into the png. It's a bit of a workaround since whatever image format is normally used, I don't think really allows saving metadata without using Pillow (PIL) or some other library. I had trouble doing it all in memory and getting the internal format into the pillow format of bytes, so my hacky workaround is to save to a temp file and then use pillow to open the temp file and re-save to the final name with pillow and metadata instead.

If there is general interest, I might be able to make something less hacky or someone who knows more than me could modify to be better. One issue is that Pillow uses binaries, so it's sort of a PITA to include in any zip plugin which I guess would require the binaries to be included for common architectures (Windows x64, Linux x64/arm, MacOS x64/arm for Python 3.10.x) or maybe there is a better way since I don't know all the details of packaging for krita. Or maybe there is a way to manually get into PNG rewriting to include that metadata using binary manipulations of a file directly (which would not require including all those Pillow libraries).

Anyway, I copied the PIL directory from a Python 3.10 install into the plugin directory locally and made these code changes and it works for anyone who wants similar functionality.

My hacky code changes to model.py:

# extra imports at the top
import sys
import os
import tempfile

# Get the path to the directory containing this script
plugin_dir = os.path.dirname(__file__)

# Add the plugin directory to sys.path
sys.path.insert(0, plugin_dir)

from PIL import Image as PilImage, PngImagePlugin

    # change to last line of save_job_results method to
    _save_image_with_metadata(base_image, job, path)
    # replaces the line "base_image.save(path)"

# added method
def _save_image_with_metadata(base_image: Image, job: Job, save_path: str):
    # Step 1: Save base_image to a temporary PNG file
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
        temp_filename = tmpfile.name
    
    try:
        # Save to real file (because base_image.save expects a path)
        base_image.save(temp_filename)
        
        # Load into Pillow
        pil_image = PilImage.open(temp_filename)
        pil_image.load()

    finally:
        # Clean up the temp file
        os.remove(temp_filename)

    # Step 2: Build the metadata text
    params = job.params
    meta = params.metadata

    # Construct a text string similar to A1111
    metadata_text = f"""{meta.get('prompt', '')}
Negative prompt: {meta.get('negative_prompt', '')}
Sampler: {meta.get('sampler', '')}
Seed: {params.seed}
Size: {params.bounds.width}x{params.bounds.height}
Model: {meta.get('checkpoint', '')}
"""

    # Step 3: Save with metadata
    pnginfo = PngImagePlugin.PngInfo()
    pnginfo.add_text("parameters", metadata_text.strip())

    pil_image.save(save_path, pnginfo=pnginfo)

:-)

[sfisher]

If anyone else finds this useful, I believe I have reliable output without having to install any extra external libraries for Pillow by doing some file manipulation directly to add the parameters to the PNG.

I'm not sure I caught every possible edge case or problem, but it does work with UTF-8 characters and some other things I could think of for testing. I guess I can try it out for a while to see if I encounter any problems. I'm fine if the code gets used by whoever or incorporated. I'd also be happy to do a PR, but not sure what things (tests? or ??) I'd need to include to do that.

It should save the information into the PNG files you get when you right click or click the "..." and save a generated image. Helpful for figuring out what you did later if all you can locate is the PNG file and not the original kra file later.

# at the top of the model.py file after  the requirements already there:
import os
import tempfile
import struct
import zlib

# comment out the last line in the method _save_job_result that says "base_image.save(path)" and add the line below
_save_image_with_metadata(base_image, job, path)

# add these three methods at the bottom of the file
def _save_image_with_metadata(base_image: Image, job: Job, save_path: str):
    # Step 1: Save base_image to a temporary PNG file
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
        temp_filename = tmpfile.name
    
    try:
        # Save to real file (because base_image.save expects a path)
        base_image.save(temp_filename)

        # Step 2: Build the metadata text
        params = job.params

        metadata_text = _format_like_automatic1111(params)

        _add_png_itxt(temp_filename, save_path, "parameters", metadata_text)

    finally:
        # Clean up the temp file
        os.remove(temp_filename)

def _format_like_automatic1111(params):
    meta = params.metadata

    prompt = meta.get("prompt", "")
    neg_prompt = meta.get("negative_prompt", "")
    sampler_info = meta.get("sampler", "")
    model = meta.get("checkpoint", "Unknown")
    seed = params.seed
    width = params.bounds.width
    height = params.bounds.height
    strength = meta.get("strength", None)
    loras = meta.get("loras", [])

    # Try to extract sampler, steps, and cfg scale from "sampler"
    import re
    match = re.match(r".*?-\s*(.+?)\s*\((\d+)\s*/\s*([\d.]+)\)", sampler_info)
    if match:
        sampler, steps, cfg_scale = match.groups()
    else:
        sampler, steps, cfg_scale = sampler_info, "Unknown", "Unknown"

    # Embed LoRAs in the prompt
    lora_tags = ""
    for lora in loras:
        if isinstance(lora, dict):
            name = lora.get("name")
            weight = lora.get("weight", 1.0)
        elif isinstance(lora, (list, tuple)) and len(lora) >= 2:
            name, weight = lora[0], lora[1]
        else:
            continue
        lora_tags += f" <lora:{name}:{weight}>"

    full_prompt = f"{prompt.strip()}{lora_tags}"

    # Construct output
    lines = []
    lines.append(f"Prompt: {full_prompt}")
    lines.append(f"Negative prompt: {neg_prompt}")
    lines.append(f"Steps: {steps}, Sampler: {sampler}, CFG scale: {cfg_scale}, Seed: {seed}, Size: {width}x{height}, Model hash: unknown, Model: {model}")
    
    if strength is not None and strength != 1.0:
        lines[-1] += f", Denoising strength: {strength}"

    return "\n".join(lines)
    

def _add_png_itxt(input_path, output_path, keyword, text):
    with open(input_path, 'rb') as f:
        png_data = f.read()

    if png_data[:8] != b'\x89PNG\r\n\x1a\n':
        raise ValueError("Not a valid PNG file")

    offset = 8
    chunks = []
    while offset < len(png_data):
        length = struct.unpack(">I", png_data[offset:offset+4])[0]
        chunk_type = png_data[offset+4:offset+8]
        chunk_data = png_data[offset+8:offset+8+length]
        offset += 12 + length
        chunks.append((chunk_type, chunk_data))

    new_png = bytearray()
    new_png += b'\x89PNG\r\n\x1a\n'

    for chunk_type, chunk_data in chunks:
        new_png += struct.pack(">I", len(chunk_data))
        new_png += chunk_type
        new_png += chunk_data
        new_png += struct.pack(">I", zlib.crc32(chunk_type + chunk_data) & 0xffffffff)

        if chunk_type == b'IHDR':
            # Insert iTXt chunk after IHDR
            keyword_bytes = keyword.encode("latin1")
            text_bytes = text.encode("utf-8")
            itxt_data = (
                keyword_bytes + b'\x00' +
                b'\x00' +     # compression flag: 0 (not compressed)
                b'\x00' +     # compression method: 0
                b'\x00' +     # language tag: empty
                b'\x00' +     # translated keyword: empty
                text_bytes
            )
            crc = zlib.crc32(b'iTXt' + itxt_data) & 0xffffffff
            new_png += struct.pack(">I", len(itxt_data))
            new_png += b'iTXt'
            new_png += itxt_data
            new_png += struct.pack(">I", crc)

    with open(output_path, 'wb') as f:
        f.write(new_png)