From d2378b6edca0800a83d59b740356901d6bfb7b53 Mon Sep 17 00:00:00 2001
From: A Latent Place <mboms@live.de>
Date: Tue, 4 Jun 2024 19:45:07 +0200
Subject: [PATCH] Added Contrast

---
 yanc.py | 46 +++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 43 insertions(+), 3 deletions(-)

diff --git a/yanc.py b/yanc.py
index 7569e0a..5778c2d 100644
--- a/yanc.py
+++ b/yanc.py
@@ -1446,13 +1446,10 @@ class YANCBrightness:
 
         if mask_opt is not None:
             mask = mask_opt.clone()
-            print_brown(str(mask.shape))
             mask = mask.unsqueeze(1).unsqueeze(1).permute(1, 2, 0, 3)
-            print_cyan(str(mask.shape))
         else:
             mask = torch.ones_like(image)
             mask = permute_tt(mask)
-            print_cyan(str(mask.shape))
 
         img = image.clone()
         img = permute_tt(img)
@@ -1463,6 +1460,47 @@ class YANCBrightness:
         return (img,)
 
 
+# ------------------------------------------------------------------------------------------------------------------ #
+
+
+class YANCContrast:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                {
+                    "image": ("IMAGE",),
+                    "contrast": ("FLOAT", {"default": 1.0, "min": 0, "max": 2, "step": 0.01}),
+                },
+                "optional":
+                {
+                    "mask_opt": ("MASK",),
+                }
+                }
+
+    CATEGORY = yanc_root_name + yanc_sub_image
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "do_it"
+
+    def do_it(self, image, contrast, mask_opt=None):
+
+        if mask_opt is not None:
+            mask = mask_opt.clone()
+            mask = mask.unsqueeze(1).unsqueeze(1).permute(1, 2, 0, 3)
+        else:
+            mask = torch.ones_like(image)
+            mask = permute_tt(mask)
+
+        img = image.clone()
+        img = permute_tt(img)
+        img = F.adjust_contrast(img * mask, contrast)
+        img = img + permute_tt(image) * F.invert(mask)
+        img = permute_ft(img)
+
+        return (img,)
+
+
 # ------------------------------------------------------------------------------------------------------------------ #
 NODE_CLASS_MAPPINGS = {
     # Image
@@ -1474,6 +1512,7 @@ NODE_CLASS_MAPPINGS = {
     "> Load Image From Folder": YANCLoadImageFromFolder,
     "> Normal Map Lighting": YANCNormalMapLighting,
     "> Brightness": YANCBrightness,
+    "> Contrast": YANCContrast,
 
     # Text
     "> Text": YANCText,
@@ -1512,6 +1551,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "> Load Image From Folder": "😼> Load Image From Folder",
     "> Normal Map Lighting": "😼> Normal Map Lighting",
     "> Brightness": "😼> Brightness",
+    "> Contrast": "😼> Contrast",
 
     # Text
     "> Text": "😼> Text",