From 7ba4e444f0416f5cf7ee55d56b1504dea2404ac7 Mon Sep 17 00:00:00 2001
From: spacepxl <spacetug93@gmail.com>
Date: Sat, 22 Feb 2025 02:38:32 -0500
Subject: [PATCH] fixed CustomNoise to normalize correctly on 4d and 5d

---
 nodes.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/nodes.py b/nodes.py
index 6ff7857..b3a9c6a 100644
--- a/nodes.py
+++ b/nodes.py
@@ -2150,12 +2150,8 @@ class CustomNoise:
 
     def get_noise(self, noise):
         noise_latent = noise["samples"].detach().clone()
-        
-        for batch in range(noise_latent.size(0)):
-            for channel in range(noise_latent.size(1)):
-                mean, std = torch.mean(noise_latent[batch, channel]), torch.std(noise_latent[batch, channel])
-                noise_latent[batch, channel] = (noise_latent[batch, channel] - mean) / std
-        
+        std, mean = torch.std_mean(noise_latent, dim=(-2, -1), keepdim=True)
+        noise_latent = (noise_latent - mean) / std
         return (Noise_CustomNoise(noise_latent),)
 
 class ExtractNFrames: