【備忘録メモ】workflow、埋め込みPNGデータ読込みで無限ロードになる場合の対処方法。ComfyUI用PNG形式に含まれるワークフローを抽出したい

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from __future__ import annotations  # Python 3.6ならこの行を消す
import sys, json, zlib
from pathlib import Path
from typing import Dict, List, Union

PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
def _u32(b: bytes) -> int: return int.from_bytes(b, "big", signed=False)
def _w(path: Path, text: str) -> None: path.write_text(text, encoding="utf-8", errors="replace")
def _add(d: Dict[str, Union[str, List[str]]], k: str, v: str) -> None:
    if k in d: d[k] = d[k] + [v] if isinstance(d[k], list) else [d[k], v]
    else: d[k] = v

def extract_png_text_chunks(p: Path) -> Dict[str, Union[str, List[str]]]:
    data = p.read_bytes()
    if len(data) < 8 or data[:8] != PNG_SIGNATURE: raise ValueError("PNGファイルではない可能性があります。")
    pos, out = 8, {}
    while pos + 8 <= len(data):
        ln = _u32(data[pos:pos+4]); ct = data[pos+4:pos+8]; pos += 8
        if pos + ln + 4 > len(data): break
        ck = data[pos:pos+ln]; pos += ln + 4
        if ct == b"IEND": break
        if ct == b"tEXt":
            try:
                nul = ck.index(b"\x00"); key = ck[:nul].decode("latin-1", "replace"); raw = ck[nul+1:]
                try: val = raw.decode("utf-8")
                except UnicodeDecodeError: val = raw.decode("latin-1", "replace")
                _add(out, key, val)
            except ValueError: _add(out, "tEXt_invalid", ck.decode("latin-1", "replace"))
        elif ct == b"zTXt":
            try:
                nul = ck.index(b"\x00"); key = ck[:nul].decode("latin-1", "replace")
                if nul + 2 > len(ck): continue
                m = ck[nul+1]; comp = ck[nul+2:]
                if m != 0: _add(out, key, f"[zTXt] unsupported compression method: {m}"); continue
                try:
                    raw = zlib.decompress(comp)
                    try: val = raw.decode("utf-8")
                    except UnicodeDecodeError: val = raw.decode("latin-1", "replace")
                    _add(out, key, val)
                except Exception as e: _add(out, key, f"[zTXt] decompress error: {e}")
            except ValueError: _add(out, "zTXt_invalid", ck.decode("latin-1", "replace"))
        elif ct == b"iTXt":
            try:
                p0 = 0; nul = ck.index(b"\x00", p0); key = ck[p0:nul].decode("latin-1", "replace"); p0 = nul + 1
                if p0 + 2 > len(ck): continue
                flag, m = ck[p0], ck[p0+1]; p0 += 2
                nul = ck.index(b"\x00", p0); lang = ck[p0:nul].decode("latin-1", "replace"); p0 = nul + 1
                nul = ck.index(b"\x00", p0); tr = ck[p0:nul].decode("utf-8", "replace"); p0 = nul + 1
                tb = ck[p0:]
                if flag == 1:
                    if m != 0: _add(out, key, f"[iTXt] unsupported compression method: {m}"); continue
                    try: tb = zlib.decompress(tb)
                    except Exception as e: _add(out, key, f"[iTXt] decompress error: {e}"); continue
                _add(out, key, tb.decode("utf-8", "replace"))
                if lang: _add(out, f"{key}__language_tag", lang)
                if tr: _add(out, f"{key}__translated_keyword", tr)
            except ValueError: _add(out, "iTXt_invalid", ck.decode("latin-1", "replace"))
    return out

def _maybe_save_special_files(base: Path, meta: Dict[str, Union[str, List[str]]]) -> List[Path]:
    saved: List[Path] = []
    def first(k: str) -> Union[str, None]:
        if k not in meta: return None
        v = meta[k]; return v[0] if isinstance(v, list) and v else (None if isinstance(v, list) else v)
    for k in ("workflow", "prompt", "parameters"):
        v = first(k)
        if not v: continue
        if k == "workflow":
            try:
                obj = json.loads(v); p = base.with_name(base.name + "_workflow.json")
                _w(p, json.dumps(obj, ensure_ascii=False, indent=2)); saved.append(p); continue
            except Exception: pass
        suf = "_prompt.txt" if k == "prompt" else f"_{k}.txt"
        p = base.with_name(base.name + suf); _w(p, v); saved.append(p)
    return saved

def process_one(p: Path) -> bool:
    if not p.exists(): print(f"失敗: ファイルが見つかりません: {p}"); return False
    try: meta = extract_png_text_chunks(p) # PNG検査
    except Exception as e: print(f"失敗: PNGとして読めません: {p} ({e})"); return False
    base = p.with_suffix("")
    lines: List[str] = []
    for k, v in meta.items():
        lines.append(f"===== {k} =====")
        if isinstance(v, list):
            for i, item in enumerate(v, 1): lines.append(f"[{i}]"); lines.append(item); lines.append("")
        else: lines.append(v); lines.append("")
    allp = base.with_name(base.name + "_all.txt"); _w(allp, "\n".join(lines)); print(f"OK: {allp}")
    for sp in _maybe_save_special_files(base, meta): print(f"OK: {sp}")
    if not meta: print(f"注意: メタ情報が見つかりませんでした: {p}")
    return True

def main(argv: List[str]) -> int:
    if len(argv) <= 1: print("使い方: extract_png_json.py <file1.png> [file2.png ...]"); return 1
    ok = sum(1 for a in argv[1:] if process_one(Path(a)))
    ng = (len(argv) - 1) - ok
    print(f"成功: {ok} / 失敗: {ng}")
    return 0 if ng == 0 else 1

if __name__ == "__main__": raise SystemExit(main(sys.argv))

@echo off
setlocal
REM Pythonのパスが通っている前提で実行。通ってないなら python.exe のフルパスに置き換えて。
python "%~dp0extract_png_json.py" %*
pause

@echo off
chcp 65001 >nul
setlocal EnableExtensions DisableDelayedExpansion
if "%~1"=="" exit /b 1
for %%F in (%*) do echo 対象: "%%~fF"
for %%F in (%*) do call :c "%%~fF"||goto end
python "%~dp0extract_png_json.py" %*||goto end
goto end
:c
set "P=%~1"
if not exist "%P%" echo エラー: ファイルが見つからない。ファイル名に%%があると失敗します。%%を消してリネームしてね: "%P%"&exit /b 1
setlocal EnableDelayedExpansion&set "Q=%P%"&if not "!Q:%%=!"=="!Q!" endlocal&goto bad&endlocal
echo(%P%|findstr /R "[\&\|\<\>\^!]" >nul&&goto bad
exit /b 0
:bad
echo エラー: ファイル名に記号が含まれています。記号を除いてリネームしてから再度チャレンジしてね: "%P%"
exit /b 1
:end
pause