remove scan.py

scan.py

@@ -1,227 +0,0 @@
-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]
-    files = [item for item in files if item["prompt_found"]]
-    result = {
-        "root": ROOT.resolve().as_posix(),
-        "scanned": len(paths),
-        "found": len(files),
-        "files": files,
-    }
-    OUTPUT.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
-    print(f"Scanned {len(paths)} PNGs; found {len(files)}; wrote {OUTPUT}")
-
-
-main()