fixed CustomNoise to normalize correctly on 4d and 5d

2025-02-22 02:38:32 -05:00
parent 0ff33fe29f
commit 7ba4e444f0
1 changed files with 2 additions and 6 deletions
@@ -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: