From c20b2d7d4061c92a739e1a0083ce72aab7f39720 Mon Sep 17 00:00:00 2001
From: nkchocoai <nkai3359@gmail.com>
Date: Thu, 15 Feb 2024 23:03:39 +0900
Subject: [PATCH] Save Image with Metadata

---
 py/capture.py         | 28 ++++++++++++++++++++++++----
 py/defs/captures.py   |  5 +++++
 py/defs/formatters.py | 21 +++++++++++++++++++++
 py/defs/meta.py       |  1 +
 py/nodes/node.py      | 16 ++++++++++------
 5 files changed, 61 insertions(+), 10 deletions(-)
 create mode 100644 py/defs/formatters.py

diff --git a/py/capture.py b/py/capture.py
index 45ddf41..0616862 100644
--- a/py/capture.py
+++ b/py/capture.py
@@ -38,10 +38,10 @@ class Capture:
                         if value is not None:
                             format = field_data.get("format")
                             v = value
-                            if format is not None:
-                                v = format(v)
                             if isinstance(value, list) and len(value) > 0:
                                 v = value[0]
+                            if format is not None:
+                                v = format(v)
                             inputs[meta].append((node_id, v))
         return inputs
 
@@ -52,9 +52,9 @@ class Capture:
         def update_pnginfo_dict(metafield, key):
             x = inputs.get(metafield, [])
             if len(x) > 0:
-                pnginfo_dict[key] = x[-1][1]
+                pnginfo_dict[key] = x[0][1]
 
-        update_pnginfo_dict(MetaField.POSITIVE_PROMPT, "")
+        update_pnginfo_dict(MetaField.POSITIVE_PROMPT, "Positive prompt")
         update_pnginfo_dict(MetaField.NEGATIVE_PROMPT, "Negative prompt")
         update_pnginfo_dict(MetaField.STEPS, "Steps")
 
@@ -76,4 +76,24 @@ class Capture:
         if len(image_widths) > 0 and len(image_heights) > 0:
             pnginfo_dict["Size"] = f"{image_widths[0][1]}x{image_heights[0][1]}"
 
+        update_pnginfo_dict(MetaField.MODEL_NAME, "Model name")
+        update_pnginfo_dict(MetaField.MODEL_HASH, "Model hash")
+
         return pnginfo_dict
+
+    @classmethod
+    def gen_parameters_str(cls, pnginfo_dict):
+        result = pnginfo_dict.get("Positive prompt", "") + "\n"
+        result += "Negative prompt: " + pnginfo_dict.get("Negative prompt", "") + "\n"
+
+        s_list = []
+        pnginfo_dict_without_prompt = {
+            k: v
+            for k, v in pnginfo_dict.items()
+            if k not in {"Positive prompt", "Negative prompt"}
+        }
+        for k, v in pnginfo_dict_without_prompt.items():
+            s = str(v).strip().replace("\n", " ")
+            s_list.append(f"{k}: {s}")
+
+        return result + ", ".join(s_list)
diff --git a/py/defs/captures.py b/py/defs/captures.py
index 4f2f1ed..ef04ef2 100644
--- a/py/defs/captures.py
+++ b/py/defs/captures.py
@@ -1,8 +1,13 @@
 from .meta import MetaField
 from .validators import is_positive_prompt, is_negative_prompt
+from .formatters import calc_model_hash
 
 
 CAPTURE_FIELD_LIST = {
+    "CheckpointLoaderSimple": {
+        MetaField.MODEL_NAME: {"field_name": "ckpt_name"},
+        MetaField.MODEL_HASH: {"field_name": "ckpt_name", "format": calc_model_hash},
+    },
     "EmptyLatentImage": {
         MetaField.IMAGE_WIDTH: {"field_name": "width"},
         MetaField.IMAGE_HEIGHT: {"field_name": "height"},
diff --git a/py/defs/formatters.py b/py/defs/formatters.py
new file mode 100644
index 0000000..1d80d34
--- /dev/null
+++ b/py/defs/formatters.py
@@ -0,0 +1,21 @@
+import hashlib
+
+import folder_paths
+
+cache_model_hash = {}
+
+
+def calc_model_hash(model_name):
+    if model_name in cache_model_hash:
+        return cache_model_hash[model_name]
+
+    filename = folder_paths.get_full_path("checkpoints", model_name)
+    sha256_hash = hashlib.sha256()
+
+    with open(filename, "rb") as f:
+        for byte_block in iter(lambda: f.read(4096), b""):
+            sha256_hash.update(byte_block)
+    model_hash = sha256_hash.hexdigest()[:10]
+
+    cache_model_hash[model_name] = model_hash
+    return model_hash
diff --git a/py/defs/meta.py b/py/defs/meta.py
index d63ca4d..4f3ae56 100644
--- a/py/defs/meta.py
+++ b/py/defs/meta.py
@@ -3,6 +3,7 @@ from enum import Enum
 
 class MetaField(Enum):
     MODEL_NAME = 0
+    MODEL_HASH = 1
     POSITIVE_PROMPT = 10
     NEGATIVE_PROMPT = 11
     SEED = 20
diff --git a/py/nodes/node.py b/py/nodes/node.py
index b8d5d56..4b75073 100644
--- a/py/nodes/node.py
+++ b/py/nodes/node.py
@@ -1,4 +1,5 @@
 import json
+import os
 
 from PIL import Image
 from PIL.PngImagePlugin import PngInfo
@@ -63,7 +64,10 @@ class SaveImageWithMetaData(BaseNode):
             metadata = None
             if not args.disable_metadata:
                 metadata = PngInfo()
-                # metadata.add_text("parameters", comment)
+                if pnginfo_dict:
+                    metadata.add_text(
+                        "parameters", Capture.gen_parameters_str(pnginfo_dict)
+                    )
                 if prompt is not None:
                     metadata.add_text("prompt", json.dumps(prompt))
                 if extra_pnginfo is not None:
@@ -71,11 +75,11 @@ class SaveImageWithMetaData(BaseNode):
                         metadata.add_text(x, json.dumps(extra_pnginfo[x]))
 
             file = f"{filename}_{counter:05}_.png"
-            # img.save(
-            #     os.path.join(full_output_folder, file),
-            #     pnginfo=metadata,
-            #     compress_level=self.compress_level,
-            # )
+            img.save(
+                os.path.join(full_output_folder, file),
+                pnginfo=metadata,
+                compress_level=self.compress_level,
+            )
             results.append(
                 {"filename": file, "subfolder": subfolder, "type": self.type}
             )