X

Client/Scripts/novelai.py

@@ -4,7 +4,7 @@ import os
 import sys
 
 import numpy as np
-from PIL import Image
+from PIL import Image, PngImagePlugin
 
 
 def byteize(alpha):
@@ -54,6 +54,7 @@ def extract_stealth_png(path):
     try:
         json_data = json.loads(json_bytes.decode("utf-8"))
         if "Comment" in json_data and isinstance(json_data["Comment"], str):
+            json_data["_CommentRaw"] = json_data["Comment"]
             json_data["Comment"] = json.loads(json_data["Comment"])
     except json.JSONDecodeError:
         return {"WebUI": json_bytes.decode("utf-8", errors="replace")}
@@ -68,9 +69,29 @@ def main() -> int:
 
     data = extract_stealth_png(path)
     print(json.dumps(data, indent=2, ensure_ascii=True, sort_keys=True))
+    if isinstance(data, dict) and "Comment" in data:
+        comment_json = data.get("_CommentRaw")
+        if not isinstance(comment_json, str):
+            if isinstance(data["Comment"], str):
+                comment_json = data["Comment"]
+            else:
+                comment_json = json.dumps(
+                    data["Comment"],
+                    ensure_ascii=True,
+                    separators=(",", ":"),
+                )
+        reply = input("Create 1x1 metadata PNG? [y/N] ").strip().lower()
+        if reply in ("y", "yes"):
+            base, _ = os.path.splitext(path)
+            out_path = f"{base}_meta.png"
+            img = Image.new("RGBA", (1, 1), (0, 0, 0, 0))
+            meta = PngImagePlugin.PngInfo()
+            meta.add_text("Comment", comment_json)
+            img.save(out_path, pnginfo=meta)
+            print(f"Wrote {out_path}")
     input()
     return 0
 
 
 if __name__ == "__main__":
-    raise SystemExit(main())
+    raise SystemExit(main())