Add local NovelAI prompt scanner

scan.py

@@ -0,0 +1,228 @@
+import io
+import json
+import zlib
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+ROOT = Path(".")
+OUTPUT = Path("novelai_prompts.json")
+PNG_SIG = b"\x89PNG\r\n\x1a\n"
+ALPHA_SIGS = (b"stealth_pnginfo", b"stealth_pngcomp")
+RGB_SIGS = (b"stealth_rgbinfo", b"stealth_rgbcomp")
+PROMPT_KEYS = ("prompt", "negative_prompt", "v4_prompt", "v4_negative_prompt")
+NAI_KEYS = {
+    "cfg_rescale",
+    "legacy_v3_extend",
+    "noise_schedule",
+    "request_type",
+    "signed_hash",
+    "uc",
+    "v4_negative_prompt",
+    "v4_prompt",
+}
+
+
+def decode_text(data):
+    for enc in ("utf-8", "latin-1"):
+        try:
+            return data.decode(enc).strip()
+        except UnicodeDecodeError:
+            pass
+    return data.decode("utf-8", errors="replace").strip()
+
+
+def parse_json(value):
+    if isinstance(value, dict):
+        return value
+    try:
+        return json.loads(value)
+    except (TypeError, json.JSONDecodeError):
+        return None
+
+
+def parse_png(data):
+    if not data.startswith(PNG_SIG):
+        return None
+    fields = {}
+    index = len(PNG_SIG)
+    while index + 8 <= len(data):
+        size = int.from_bytes(data[index : index + 4], "big")
+        kind = data[index + 4 : index + 8]
+        index += 8
+        content = data[index : min(index + size, len(data))]
+        if kind == b"tEXt" and b"\0" in content:
+            key, text = content.split(b"\0", 1)
+            fields[decode_text(key)] = decode_text(text)
+        elif kind == b"iTXt" and b"\0" in content:
+            parts = content.split(b"\0", 5)
+            if len(parts) == 6:
+                key, flag, method, _, _, text = parts
+                if flag == b"\x01" and method == b"\x00":
+                    try:
+                        text = zlib.decompress(text)
+                    except zlib.error:
+                        text = b""
+                fields[decode_text(key)] = decode_text(text)
+        elif kind == b"zTXt" and b"\0" in content:
+            key, rest = content.split(b"\0", 1)
+            if rest[:1] == b"\x00":
+                try:
+                    fields[decode_text(key)] = decode_text(zlib.decompress(rest[1:]))
+                except zlib.error:
+                    pass
+        index += size + 4
+    return fields or None
+
+
+def byteize(bits):
+    bits = bits[: (bits.shape[0] // 8) * 8]
+    bits = np.bitwise_and(bits, 1).reshape((-1, 8))
+    return np.packbits(bits, axis=1).flatten()
+
+
+def parse_stealth_channel(bits, sigs):
+    raw = bytes(byteize(bits))
+    magic = next((sig for sig in sigs if raw.startswith(sig)), None)
+    if not magic:
+        return None, False
+    try:
+        start = len(magic)
+        bit_len = int.from_bytes(raw[start : start + 4], "big")
+        start += 4
+        if bit_len <= 0 or bit_len % 8 or start + bit_len // 8 > len(raw):
+            return None, True
+        payload = raw[start : start + bit_len // 8]
+        if magic.endswith(b"comp"):
+            payload = zlib.decompress(payload, wbits=31)
+        text = decode_text(payload)
+        fields = parse_json(text)
+        return (fields if isinstance(fields, dict) else {"Stealth": text}), True
+    except (UnicodeDecodeError, ValueError, zlib.error):
+        return None, True
+
+
+def parse_stealth(data):
+    try:
+        image = Image.open(io.BytesIO(data))
+        detected = False
+        if "A" in image.getbands():
+            alpha = np.array(image.getchannel("A")).T.reshape((-1,))
+            fields, found = parse_stealth_channel(alpha, ALPHA_SIGS)
+            if fields:
+                return fields, True, None
+            detected = found
+        rgb = np.array(image.convert("RGB")).transpose((1, 0, 2)).reshape((-1,))
+        fields, found = parse_stealth_channel(rgb, RGB_SIGS)
+        return fields, detected or found, None
+    except Exception as e:
+        return None, False, str(e)
+
+
+def get_field(fields, name, default=None):
+    wanted = name.casefold()
+    return next((value for key, value in fields.items() if key.casefold() == wanted), default)
+
+
+def has_novelai_comment(value):
+    data = parse_json(value)
+    if not isinstance(data, dict):
+        return False
+    keys = set(data)
+    if keys & NAI_KEYS:
+        return True
+    return (
+        isinstance(data.get("prompt"), str)
+        and isinstance(data.get("steps"), int)
+        and ("scale" in data or "sampler" in data)
+    )
+
+
+def is_novelai(fields):
+    if not fields:
+        return False
+    software = str(get_field(fields, "Software", "")).casefold()
+    source = str(get_field(fields, "Source", "")).casefold()
+    title = str(get_field(fields, "Title", ""))
+    comment = get_field(fields, "Comment", "")
+    return (
+        "novelai" in software
+        or "novelai" in source
+        or title == "AI generated image"
+        or has_novelai_comment(comment)
+        or bool({str(key).casefold() for key in fields} & NAI_KEYS)
+    )
+
+
+def prompt_record(source, fields):
+    comment = parse_json(get_field(fields, "Comment"))
+    if not isinstance(comment, dict):
+        comment = fields if has_novelai_comment(fields) else {}
+    return {
+        "source": source,
+        "prompt": get_field(fields, "Description") or comment.get("prompt"),
+        "negative_prompt": comment.get("uc"),
+        "v4_prompt": comment.get("v4_prompt"),
+        "v4_negative_prompt": comment.get("v4_negative_prompt"),
+    }
+
+
+def scan(path):
+    record = {
+        "file": path.relative_to(ROOT).as_posix(),
+        "novelai": False,
+        "prompt_found": False,
+        "stealth_detected": False,
+        "prompts": [],
+        "missing": None,
+        "error": None,
+    }
+    try:
+        data = path.read_bytes()
+        normal = parse_png(data)
+        stealth, detected, error = parse_stealth(data)
+        record["stealth_detected"] = detected
+        record["error"] = error
+        if is_novelai(normal):
+            record["prompts"].append(prompt_record("normal", normal))
+        if is_novelai(stealth):
+            record["prompts"].append(prompt_record("stealth", stealth))
+        record["novelai"] = bool(record["prompts"])
+        record["prompt_found"] = any(
+            any(item.get(key) not in (None, "", {}, []) for key in PROMPT_KEYS)
+            for item in record["prompts"]
+        )
+        if record["novelai"] and not record["prompt_found"]:
+            record["missing"] = "NovelAI metadata found without a prompt"
+        elif detected and not stealth:
+            record["missing"] = "Stealth metadata could not be decoded"
+        elif detected and not record["novelai"]:
+            record["missing"] = "Stealth metadata is not NovelAI"
+        elif not record["novelai"]:
+            record["missing"] = "No NovelAI metadata"
+    except Exception as e:
+        record["missing"] = "Could not scan PNG"
+        record["error"] = str(e)
+    return record
+
+
+def main():
+    paths = sorted(
+        (path for path in ROOT.rglob("*") if path.is_file() and path.suffix.casefold() == ".png"),
+        key=lambda path: path.as_posix().casefold(),
+    )
+    files = [scan(path) for path in paths]
+    found = sum(item["prompt_found"] for item in files)
+    result = {
+        "root": ROOT.resolve().as_posix(),
+        "total": len(files),
+        "found": found,
+        "missing": len(files) - found,
+        "files": files,
+    }
+    OUTPUT.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+    print(f"Scanned {len(files)} PNGs; found {found}; wrote {OUTPUT}")
+
+
+main()