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