Stable Diffusion Web UI で作成された PNG ファイルのメタデータ(パラメータ)を C++ で読む

Stable Diffusion Web UI をセットアップしたい方は (Google Colab なき今は) Paperspace を使うのが得策かと・・・(こちらの記事もお読みください)

また、Web UI の簡単な使い方についてはこちらもどうぞ・・・(Google Colabo を無料で使い続けるのは避けるべきかとは思いますが)

---

Qt が便利なのでとりあえず C++ といいつつ Qt 依存になってしまいました(Qt5 および Qt6 で動作します)

このプログラム全体では、コンパイルされた Exe に PNG のパスを渡すと Stable Diffusion Web UI で作成された際のメタデータを表示します。PNG のチャンクデータを順番に読んで "tEXt" または "iTXt" チャンクを探します。
(日本語等が入っていると "iTXt" チャンクになります)

QString getPngInfo(const QString &fileName) に Stable Diffusion Web UI で作成された PNG ファイルのパスを渡すと、呪文(prompt, negatie prompt)とか各種設定情報が文字列として返ってきます。
(その文字列の解析は、今のところ皆さんにおまかせします)

↑作ってみました：
QString getPngInfoJson(const QString &fileName) を呼ぶと各種設定情報を解析してJSONとして返します。
もし、バグがあったら直したいので、不具合のあった PNG ファイルを添付して [e-mail] まで電子メール送っていただければ対応します(google drive に置くとかでもいいです)。その際は、この記事にコメントで同時にご連絡ください。よろしくお願いします。

他の言語版も作りたいと思います。(次は C# 辺り！)

これを使って、なんか凄いソフト作ってくれや！
それでゃ！

pnginfo.cpp

#include <QtCore>

QString getPngInfo(const QString &fileName)
{
    QString result = "";
    //QFileInfo fi(fileName);
    //qDebug() << fi.size();
    QByteArray signature = "\x89PNG\r\n\x1a\n";
    //qDebug() << signature.size();
    int crcSize = 4;
    QFile file(fileName);
    if (!file.open(QIODevice::ReadOnly)) return "";
    QByteArray realSig = file.read(signature.size());
    if (realSig != signature) return "";
    for (;;)
    {
        QByteArray bytes = file.read(8);
        //qDebug() << "bytes.size()" << bytes.size();
        if (bytes.size() != 8) break;
        QBuffer buf(&bytes);
        if (!buf.open(QBuffer::ReadOnly))
        {
            break;
        }
        QByteArray lengthBytes = buf.read(4);
        quint32 length;
        QDataStream stream(lengthBytes);
        stream.setByteOrder(QDataStream::BigEndian);
        stream >> length;
        QByteArray chunkType = buf.read(4);
        if (chunkType == "tEXt")
        {
            qDebug() << chunkType << length;
            QByteArray s = file.read(length);
            //qDebug() << s;
            if (s.startsWith("parameters"))
            {
                s = s.mid(10);
                if (s.size() > 0 && s.at(0) == 0)
                {
                    while (s.size() > 0 && s.at(0) == 0) s = s.mid(1);
                    result = QString::fromLatin1(s);
                    break;
                }
            }
        }
        else if (chunkType == "iTXt")
        {
            qDebug() << chunkType << length;
            QByteArray s = file.read(length);
            //qDebug() << s;
            if (s.startsWith("parameters"))
            {
                s = s.mid(10);
                if (s.size() > 0 && s.at(0) == 0)
                {
                    while (s.size() > 0 && s.at(0) == 0) s = s.mid(1);
                    result = QString::fromUtf8(s);
                    break;
                }
            }
        }
        else
        {
            qDebug() << chunkType << length;
            file.seek(file.pos()+length);
        }
        buf.close();
        file.seek(file.pos()+crcSize);
        //qDebug() << file.pos();
    }
    file.close();
    return result;
}

QString getPngInfoJson(const QString &fileName)
{
    QJsonObject o;
    QString infoString = getPngInfo(fileName);
    QStringList lines = infoString.split("\n");
    int phase = 0;
    QString prompt;
    QString negPrompt;
    foreach(QString line, lines)
    {
        //qDebug() << line;
        QRegularExpression re("(?<pair>[A-Z][a-zA-Z0-9 ]*: [^,]+)(, )?");
        if (phase == 0)
        {
            if (line.startsWith("Negative prompt: "))
            {
                negPrompt = line.mid(17);
                phase = 1;
                continue;
            }
            QRegularExpressionMatch m = re.match(line);
            if (m.hasMatch())
            {
                phase = 2;
                goto phase2;
            }
            prompt += "\n" + line;
            continue;
        }
        if (phase == 1)
        {
            QRegularExpressionMatch m = re.match(line);
            if (!m.hasMatch())
            {
                negPrompt += "\n" + line;
                continue;
            }
            phase = 2;
        }
phase2:
        QRegularExpressionMatchIterator i = re.globalMatch(line);
        QStringList words;
        while (i.hasNext()) {
            QRegularExpressionMatch match = i.next();
            QString word = match.captured("pair");
            words << word;
        }
        if (words.size() > 0)
        {
            foreach(QString word, words)
            {
                QRegularExpression rx1("(?<name>[A-Z][a-zA-Z0-9 ]*): (?<value>[^,]+)(, )?");
                QRegularExpressionMatch m = rx1.match(word);
                if (m.hasMatch())
                {
                    o[m.captured("name")] = m.captured("value");
                }
            }
        }
    }
    o["Prompt"] = prompt.trimmed();
    o["Negative prompt"] = negPrompt.trimmed();
    QJsonDocument result(o);
    return QString::fromUtf8(result.toJson());
}

int main(int argc, char *argv[])
{
    QCoreApplication app(argc, argv);
    qDebug() << app.arguments();
    if (app.arguments().size() != 2) return 1;
    QString result = getPngInfo(app.arguments()[1]);
    qDebug().noquote() << result;
    QString json = getPngInfoJson(app.arguments()[1]);
    qDebug().noquote() << json;
    return 0;
}