From 5fdd0593ae3940e7a6331672c25c51e9c757cbd5 Mon Sep 17 00:00:00 2001
From: EllangoK <karun.ellango7@gmail.com>
Date: Sat, 6 May 2023 19:12:29 -0400
Subject: [PATCH] Sepia modes over different channels

---
 .gitignore               |  1 +
 post_processing/sepia.py | 11 +++++++++--
 post_processing_nodes.py | 11 +++++++++--
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/.gitignore b/.gitignore
index f10328e..c08890e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,3 @@
 wip/
 .vscode/*
+*.pyc
diff --git a/post_processing/sepia.py b/post_processing/sepia.py
index fc28c1a..8debe3d 100644
--- a/post_processing/sepia.py
+++ b/post_processing/sepia.py
@@ -15,6 +15,7 @@ class Sepia:
                     "max": 1.0,
                     "step": 0.1
                 }),
+                "mode": (["sepia", "blue-pia", "green-pia"],),
             },
         }
 
@@ -23,12 +24,18 @@ class Sepia:
 
     CATEGORY = "postprocessing/Color Adjustments"
 
-    def sepia(self, image: torch.Tensor, strength: float):
+    def sepia(self, image: torch.Tensor, strength: float, mode: str = "sepia"):
         if strength == 0:
             return (image,)
 
         sepia_weights = torch.tensor([0.2989, 0.5870, 0.1140]).view(1, 1, 1, 3).to(image.device)
-        sepia_filter = torch.tensor([1.0, 0.8, 0.6]).view(1, 1, 1, 3).to(image.device)
+
+        if mode == "sepia":
+            sepia_filter = torch.tensor([1.0, 0.8, 0.6]).view(1, 1, 1, 3).to(image.device)
+        elif mode == "blue-pia":
+            sepia_filter = torch.tensor([0.6, 0.8, 1.0]).view(1, 1, 1, 3).to(image.device)
+        elif mode == "green-pia":
+            sepia_filter = torch.tensor([0.6, 1.0, 0.6]).view(1, 1, 1, 3).to(image.device)
 
         grayscale = torch.sum(image * sepia_weights, dim=-1, keepdim=True)
         sepia = grayscale * sepia_filter
diff --git a/post_processing_nodes.py b/post_processing_nodes.py
index 3548c1e..2d9fb62 100644
--- a/post_processing_nodes.py
+++ b/post_processing_nodes.py
@@ -999,6 +999,7 @@ class Sepia:
                     "max": 1.0,
                     "step": 0.1
                 }),
+                "mode": (["sepia", "blue-pia", "green-pia"],),
             },
         }
 
@@ -1007,12 +1008,18 @@ class Sepia:
 
     CATEGORY = "postprocessing/Color Adjustments"
 
-    def sepia(self, image: torch.Tensor, strength: float):
+    def sepia(self, image: torch.Tensor, strength: float, mode: str = "sepia"):
         if strength == 0:
             return (image,)
 
         sepia_weights = torch.tensor([0.2989, 0.5870, 0.1140]).view(1, 1, 1, 3).to(image.device)
-        sepia_filter = torch.tensor([1.0, 0.8, 0.6]).view(1, 1, 1, 3).to(image.device)
+
+        if mode == "sepia":
+            sepia_filter = torch.tensor([1.0, 0.8, 0.6]).view(1, 1, 1, 3).to(image.device)
+        elif mode == "blue-pia":
+            sepia_filter = torch.tensor([0.6, 0.8, 1.0]).view(1, 1, 1, 3).to(image.device)
+        elif mode == "green-pia":
+            sepia_filter = torch.tensor([0.6, 1.0, 0.6]).view(1, 1, 1, 3).to(image.device)
 
         grayscale = torch.sum(image * sepia_weights, dim=-1, keepdim=True)
         sepia = grayscale * sepia_filter