ComfyUI Image prompt extractor

comfy_prompt_extractor.py

import sys
import json
from PIL import Image

# qt
from PyQt5.QtWidgets import QApplication, QMessageBox, QWidget, QTextEdit, QVBoxLayout, QLabel
from PyQt5.QtGui import QIcon, QFont
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QSizePolicy

class PromptType:
    POSITIVE = 0
    NEGATIVE = 1
    UNDEFINED = 2

class Prompt:
    def __init__(self, title, text):
        self.title = title
        self.text = text

    def getPromptType(self):
        if "negative" in self.title.lower():
            return PromptType.NEGATIVE
        elif "positive" in self.title.lower():
            return PromptType.POSITIVE
        else:
            return PromptType.UNDEFINED

    def __str__(self):
        return f"{self.title}: {self.text}"
    
def get_prompt_type_color_string(prompt_type):
    if prompt_type == PromptType.POSITIVE:
        return "#d6feda"
    elif prompt_type == PromptType.NEGATIVE:
        return "#ffd6d6"
    else:
        return "white"

if __name__ == "__main__":
    filepath = sys.argv[1]
    img = Image.open(filepath)
    img.load()

    app = QApplication(sys.argv)

    # Check to see if this is even a ComfyUI image
    if "prompt" not in img.info.keys():
        QMessageBox.critical(None, "Error", "Not a ComfyUI Image", QMessageBox.Ok)
        sys.exit(1)

    prompts = []

    # Extract ComfyUI prompts
    parsed_prompt = json.loads(img.info['prompt'])
    for node_name, node in parsed_prompt.items():
        if "class_type" in node.keys():
            if node["class_type"] == "Text Multiline":
                text = node["inputs"]["text"]
                if isinstance(text, list):
                    continue
                p = Prompt(node["_meta"]["title"], text)
                prompts.append(p)
            elif node["class_type"] == "CLIPTextEncode":
                text = node["inputs"]["text"]
                # If the text is a list, this means it's actually a connection to another node, not a node itself. So we skip it.
                if isinstance(text, list):
                    continue
                p = Prompt(node_name, text)
                prompts.append(p)

    # Quit if no prompts found
    if len(prompts) == 0:
        QMessageBox.critical(None, "Error", "No prompts found", QMessageBox.Ok)
        sys.exit(1)

    # Sort prompts by title
    prompts.sort(key=lambda x: x.title)
    
    # Setup the window
    app.setStyle("Fusion")
    window = QWidget()
    window.setWindowTitle("ComfyUI Prompt Extractor")
    window.setGeometry(100, 100, 640, 480)
    window.setWindowIcon(QIcon("icon.png"))
    window.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
    window.setAttribute(Qt.WA_DeleteOnClose)
    window.setWindowFlags(Qt.WindowStaysOnTopHint | Qt.WindowCloseButtonHint | Qt.WindowMinimizeButtonHint)
    window.setStyleSheet("background-color: white;")
    window.setWindowIcon(QIcon(filepath))

    layout = QVBoxLayout()
    window.setLayout(layout)

    # separate prompts by type
    positive_prompts = [prompt for prompt in prompts if prompt.getPromptType() == PromptType.POSITIVE]
    negative_prompts = [prompt for prompt in prompts if prompt.getPromptType() == PromptType.NEGATIVE]
    other_prompts = [prompt for prompt in prompts if prompt.getPromptType() == PromptType.UNDEFINED]

    # Add prompts to the window
    for prompt in positive_prompts + negative_prompts + other_prompts:
        
        # Create a label for the prompt
        label = QLabel()
        label.setText(f"{prompt.title}")
        label.setStyleSheet("background-color: white; color: black; font-size: 14px;")
        layout.addWidget(label)

        # Create a text edit for the prompt
        textedit = QTextEdit()
        textedit.setText(f"{prompt.text}")
        highlight_color = get_prompt_type_color_string(prompt.getPromptType())
        textedit.setStyleSheet("background-color: {}; color: black; font-size: 14px;".format(highlight_color))
        textedit.setFont(QFont("Roboto Mono", 12))
        textedit.setFontPointSize(12)
        textedit.setReadOnly(True)
        textedit.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        layout.addWidget(textedit)

    window.show()
    app.exec()

Requires pyqt5 and pillow.

Install Roboto Mono if you haven't already