Stable Diffusion image prompts lookup tool

SD_prompts.py

import argparse
import os
import time
import shutil
from pngparser import PngParser

PARAMS_MARKER = b"parameters\x00"
NEGATIVE_PROMPT_MARKER = "Negative prompt: "

def extract_text_from_png(file_path):
    """Extract text from the specified PNG file."""
    png = PngParser(file_path)

    try:
        text = png.chunks[1].data
    except IndexError:
        raise ValueError("The specified PNG does not contain enough chunks (at least 2 required).")
    except AttributeError:
        raise ValueError("The specified chunk does not have readable data.")

    if not text.startswith(PARAMS_MARKER):
        return ""

    return text.split(PARAMS_MARKER)[-1].decode("ascii")

def extract_positive_prompt(text):
    """Extract the positive prompt from the text."""
    return text.split(NEGATIVE_PROMPT_MARKER)[0].strip()

def extract_positive_prompts_from_directory(directory, days_ago=None):
    """Extract positive prompts from all PNG files in the directory."""
    prompts = []
    current_time = time.time()
    days_in_seconds = days_ago * 86400 if days_ago is not None else 0

    for filename in os.listdir(directory):
        if filename.lower().endswith('.png'):
            file_path = os.path.join(directory, filename)
            # Get the file creation time
            creation_time = os.path.getctime(file_path)
            # Check if the file was created within the specified days
            if days_ago is not None and current_time - creation_time > days_in_seconds:
                continue

            positive_prompt = extract_positive_prompt(extract_text_from_png(file_path))
            if not positive_prompt:
                continue

            # Get the creation date in a readable format
            creation_date = time.ctime(creation_time)
            # Add a tuple of (filename, creation date, prompt) to the list
            prompts.append((filename, creation_date, positive_prompt))

    return prompts

def copy_files_to_directory(prompts, source_directory, target_directory, rename = True):
    # Ensure the target directory exists and is clean
    if os.path.exists(target_directory):
        shutil.rmtree(target_directory)  # Remove existing directory
    os.makedirs(target_directory)  # Create a new directory

    # Copy files with positive prompts
    for i, (filename, date, prompt) in enumerate(prompts):
        if rename:
            target_filename = os.path.join(target_directory, "%03d_%s_%d_%s.png" % (i, filename, len(prompt), date.replace(":", "_")))
        else:
            target_filename = filename
        source_filename = os.path.join(source_directory, filename)
        shutil.copy(source_filename, target_filename)
        print(f"Copied: {source_filename} to {target_filename}")

def main():
    # Set up argument parsing
    parser = argparse.ArgumentParser(description='Extract prompts from PNG files.')
    parser.add_argument('input', type=str, help='Path to the PNG file or directory containing PNG files')
    parser.add_argument('--number', type=int, help='Number of files', default=10)
    parser.add_argument('--unique', type=int, help='Only look for unique prompts', default=1)
    parser.add_argument('--rename', type=int, help='Rename files', default=1)
    parser.add_argument('--created-days-ago', type=int, help='Only include files created within the specified number of days ago')
    args = parser.parse_args()

    input_path = args.input
    number = args.number
    unique = args.unique
    rename = args.rename

    if os.path.isdir(input_path):
        # Process all PNG files in the directory, optionally filtered by creation date
        prompts = extract_positive_prompts_from_directory(input_path, args.created_days_ago)

        # Use a set to find unique prompts
        if unique:
            sorted_prompts = []
            prompts_set = set()
            print(f"amount of prompts total: {len(prompts)}")
            for (filename, date, prompt) in prompts:
                if prompt in prompts_set:
                    continue
                prompts_set.add(prompt)
                sorted_prompts.append((filename, date, prompt))
            print(f"amount of prompts unique: {len(sorted_prompts)}")

        # Prepare a list from the set for sorting, if needed
        sorted_prompts = sorted(sorted_prompts, key=lambda x: len(x[2]))  # Sort by length of prompt

        # Get the top 20 shortest unique positive prompts
        print("Top 20 unique positive prompts with the shortest length:")
        for filename, date, prompt in sorted_prompts[:number]:
            print(f"{filename}: {date}: {prompt}")

        # Copy files to the target directory
        source_directory = os.path.realpath(input_path)
        copy_files_to_directory(sorted_prompts[:number], source_directory, target_directory = "myprompts", rename = rename)

    else:
        # Process a single PNG file
        text_info = extract_positive_prompt(extract_text_from_png(input_path))
        print(text_info)

if __name__ == "__main__":
    main()