From 23103dffde3a1590e2144146303e714a43da5d8d Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sat, 1 Feb 2025 23:43:22 +0200
Subject: [PATCH] Add ImageNoiseAugmentation

---
 __init__.py    |  1 +
 nodes/nodes.py | 26 +++++++++++++++++++++++++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/__init__.py b/__init__.py
index 791f0fc..8888f3d 100644
--- a/__init__.py
+++ b/__init__.py
@@ -64,6 +64,7 @@ NODE_CONFIG = {
     "ImageGridComposite2x2": {"class": ImageGridComposite2x2, "name": "Image Grid Composite 2x2"},
     "ImageGridComposite3x3": {"class": ImageGridComposite3x3, "name": "Image Grid Composite 3x3"},
     "ImageGridtoBatch": {"class": ImageGridtoBatch, "name": "Image Grid To Batch"},
+    "ImageNoiseAugmentation": {"class": ImageNoiseAugmentation, "name": "Image Noise Augmentation"},
     "ImageNormalize_Neg1_To_1": {"class": ImageNormalize_Neg1_To_1, "name": "Image Normalize -1 to 1"},
     "ImagePass": {"class": ImagePass},
     "ImagePadForOutpaintMasked": {"class": ImagePadForOutpaintMasked, "name": "Image Pad For Outpaint Masked"},
diff --git a/nodes/nodes.py b/nodes/nodes.py
index c66a6d9..cc9809f 100644
--- a/nodes/nodes.py
+++ b/nodes/nodes.py
@@ -2255,4 +2255,28 @@ class LeapfusionHunyuanI2V:
         m = model.clone()
         m.set_model_unet_function_wrapper(outer_wrapper(samples, index))
 
-        return (m,)
\ No newline at end of file
+        return (m,)
+
+class ImageNoiseAugmentation:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "noise_aug_strength": ("FLOAT", {"default": None, "min": 0.0, "max": 100.0, "step": 0.001}),
+                "seed": ("INT", {"default": 123,"min": 0, "max": 0xffffffffffffffff, "step": 1}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "add_noise"
+
+    CATEGORY = "KJNodes/experimental"
+
+    def add_noise(self, image, noise_aug_strength, seed):
+        torch.manual_seed(seed)
+        sigma = torch.ones((image.shape[0],)).to(image.device, image.dtype) * noise_aug_strength
+        image_noise = torch.randn_like(image) * sigma[:, None, None, None]
+        image_noise = torch.where(image==-1, torch.zeros_like(image), image_noise)
+        image_out = image + image_noise
+        return image_out,
\ No newline at end of file