From 3d6db912d17b428c24e1f35568bd41c828bec4b0 Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Fri, 10 May 2024 03:04:47 +0800
Subject: [PATCH] 'update'

---
 README.md                 |  2 +-
 __init__.py               |  1 +
 mz_stylize_photo_core.py  |  9 ++++++++
 mz_stylize_photo_utils.py | 44 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index bc42a94..d8ad8d1 100644
--- a/README.md
+++ b/README.md
@@ -1 +1 @@
-# ComfyUI-StylizePhoto
\ No newline at end of file
+# ComfyUI-StylizePhoto-MZ
\ No newline at end of file
diff --git a/__init__.py b/__init__.py
index e1d8f34..c9e01fb 100644
--- a/__init__.py
+++ b/__init__.py
@@ -45,6 +45,7 @@ class MZ_StylizePhotoKSamplerNode:
                 "denoise": ("FLOAT", {"default": 0.75, "min": 0.0, "max": 1.0, "step": 0.01}),
                 "positive_prompt": ("STRING", {"default": ""}),
                 "negative_prompt": ("STRING", {"default": ""}),
+                "watermark": ("STRING", {"default": "ComfyUI\nGenerated by AI", "multiline": True}),
             },
             "optional":{
                 "model": ("MODEL",),
diff --git a/mz_stylize_photo_core.py b/mz_stylize_photo_core.py
index 9868bc2..3f1a39a 100644
--- a/mz_stylize_photo_core.py
+++ b/mz_stylize_photo_core.py
@@ -102,6 +102,15 @@ def ksampler(kwargs):
     latent = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)[0]
 
     image = VAEDecode().decode(vae, latent)[0]
+
+
+    watermark = kwargs.get("watermark", "")
+    if watermark != "":
+        image = Utils.tensor2pil(image)
+        image = Utils.add_watermark(image, watermark)
+        image = Utils.pil2tensor(image)
+        image = Utils.list_tensor2tensor([image])
+
     return (image,)
 
 
diff --git a/mz_stylize_photo_utils.py b/mz_stylize_photo_utils.py
index b10ee33..8f6e584 100644
--- a/mz_stylize_photo_utils.py
+++ b/mz_stylize_photo_utils.py
@@ -354,6 +354,50 @@ class Utils:
 
         return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS)
 
+    def add_watermark(image, watermark):
+        if watermark == "":
+            return image
+        from PIL import ImageDraw, ImageFont
+
+        font_fullpath = Utils.download_model(
+            {
+                "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=font%2FAlibabaPuHuiTi-2-75-SemiBold.ttf",
+                "output": "font/AlibabaPuHuiTi-2-75-SemiBold.ttf",
+            }
+        )
+
+        watermarks = watermark.split("\n")
+
+        width, height = image.size
+        short_edge = min(width, height)
+        font_size = short_edge // 12
+ 
+            
+        font = ImageFont.truetype(font_fullpath, font_size)
+        draw = ImageDraw.Draw(image) 
+
+        text = watermarks[0]
+        textwidth, textheight = draw.textsize(text, font) 
+        
+        x = (width - textwidth) // 2
+
+        bottom = 10
+
+        y = height - textheight - (textheight * 0.4 + bottom + 8)
+        draw.text((x, y), text, font=font) 
+
+        if len(watermarks) > 1:
+            y1 = y + textheight
+            text = watermarks[1]
+            font_size = int(font_size * 0.4)
+            font = ImageFont.truetype(font_fullpath, font_size)
+            textwidth, textheight = draw.textsize(text, font)
+            x = (width - textwidth) // 2
+            y = y1 - bottom + 4
+            draw.text((x, y), text, font=font)
+
+        return image
+
     def get_device():
         return comfy.model_management.get_torch_device()