Add ImageNoiseAugmentation

__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"},

nodes/nodes.py

@@ -2255,4 +2255,28 @@ class LeapfusionHunyuanI2V:
         m = model.clone()
         m.set_model_unet_function_wrapper(outer_wrapper(samples, index))
 
-        return (m,)
+        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,