commit b766b2b40eb8708026eaa2b42d67bb830508b92a
Author: web3nomad.eth <me@web3nomad.eth>
Date:   Mon Apr 29 14:44:50 2024 +0800

    Initial commit

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..64777da
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+__pycache__/
+.DS_Store
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..012764a
--- /dev/null
+++ b/README.md
@@ -0,0 +1,2 @@
+# ComfyUI Invisible Watermark
+
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..4d4c4f6
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,4 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
diff --git a/core.py b/core.py
new file mode 100644
index 0000000..672d995
--- /dev/null
+++ b/core.py
@@ -0,0 +1,80 @@
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+import math
+
+# Constants
+DEFAULT_ASPECT_RATIO = 0.6
+FONT_COLOR = (0, 0, 0)
+MARGIN = 1.1
+
+
+def get_text_square(txt, font_h, aspect_r=DEFAULT_ASPECT_RATIO):
+    if not txt:
+        return {
+            'lineArr': [],
+            'numLines': 0,
+            'width': 10,
+            'height': 10
+        }
+
+    max_line_len = math.ceil(math.sqrt(len(txt)) * 2)
+    word_arr = txt.split(' ')
+    line_arr = []
+
+    while word_arr:
+        line = ""
+        while word_arr and len(line_arr) < max_line_len:
+            line += word_arr.pop(0) + ' '
+        line_arr.append(line.strip())
+
+    actual_max_line_len = max(len(s) * 2 for s in line_arr)
+    return {
+        'lineArr': line_arr,
+        'numLines': len(line_arr),
+        'width': math.ceil(font_h * aspect_r * actual_max_line_len),
+        'height': font_h * len(line_arr)
+    }
+
+
+def write_on_txt_canvas(width, height, txt_to_embed, font_h, margin=MARGIN, font=None):
+    img = Image.new('RGBA', (width, height), (255, 255, 255, 0))
+    draw = ImageDraw.Draw(img)
+    _font_h = int(height * font_h)
+    _font = ImageFont.truetype(font, _font_h) if font else None
+    text_square = get_text_square(txt_to_embed, _font_h)
+    x_offset = 5
+    y_offset = 5
+
+    width_factor = text_square['width'] * margin
+    height_factor = text_square['height'] * margin
+
+    for i in range(int(height / height_factor)):
+        for j in range(int(width / width_factor)):
+            for k in range(text_square['numLines']):
+                x_pos = x_offset + j * width_factor
+                y_pos = y_offset + i * height_factor + (k + 1) * _font_h
+                draw.text((x_pos, y_pos), text_square['lineArr'][k], fill=FONT_COLOR, font=_font)
+
+    return np.uint8(img)
+
+
+def add_watermark(processing_data, txt_data):
+    mask = txt_data[:, :, 3] != 0
+
+    # 显式将数据类型更改为 uint8
+    processing_data = processing_data.astype(np.uint8)
+
+    # 进行位操作
+    processing_data[mask, :3] = processing_data[mask, :3] & ~1
+    processing_data[~mask, :3] = processing_data[~mask, :3] | 1
+    return processing_data
+
+
+def generate_watermark(input_image, wm_txt, font=None, font_h=0.025):
+    text_image = write_on_txt_canvas(input_image.shape[1],
+                                     input_image.shape[0], txt_to_embed=wm_txt,
+                                     font_h=font_h,
+                                     font=font)
+    result_img_data = add_watermark(input_image, text_image)
+    result_img = Image.fromarray(result_img_data)
+    return result_img
diff --git a/nodes.py b/nodes.py
new file mode 100644
index 0000000..3399465
--- /dev/null
+++ b/nodes.py
@@ -0,0 +1,48 @@
+import base64
+from PIL import Image
+import torch
+import numpy as np
+import io
+
+from .core import generate_watermark
+
+class InvisibleWatermarkEncode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE", ),
+            },
+        }
+
+    RETURN_TYPES = ()
+
+    FUNCTION = "encode"
+
+    OUTPUT_NODE = True
+
+    CATEGORY = "WATERMARK"
+
+    def encode(self, images: list[torch.Tensor]):
+        results = []
+        for image in images:
+            i = 255. * image.cpu().numpy()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+            result_img = generate_watermark(img, "Hello, World!")
+            results.append(result_img)
+        return {"ui": {"images": results}}
+
+
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    "InvisibleWatermarkEncode": InvisibleWatermarkEncode,
+}
+
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "InvisibleWatermarkEncode": "Invisible Watermark Encode",
+}