From fdbbd970e75a20a111f7f04afb589e8ea97d5d87 Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Fri, 16 Aug 2024 11:09:54 -0600
Subject: [PATCH] Adjust noise scaling approach

---
 py/utils.py | 65 ++++++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 59 insertions(+), 6 deletions(-)

diff --git a/py/utils.py b/py/utils.py
index b93f38e..73b5b8f 100644
--- a/py/utils.py
+++ b/py/utils.py
@@ -4,6 +4,42 @@ import torch
 
 from comfy.k_diffusion.sampling import to_d
 
+from . import latent
+
+
+# def scale_noise_(
+#     noise,
+#     factor=1.0,
+#     *,
+#     normalized=True,
+#     normalize_dims=(-3, -2, -1),
+# ):
+#     if not normalized or noise.numel() == 0:
+#         return noise.mul_(factor) if factor != 1 else noise
+#     mean, std = (
+#         noise.mean(dim=normalize_dims, keepdim=True),
+#         noise.std(dim=normalize_dims, keepdim=True),
+#     )
+#     return latent.normalize_to_scale(
+#         noise.sub_(mean).div_(std).clamp(-1, 1), -1.0, 1.0, dim=normalize_dims
+#     ).mul_(factor)
+
+
+# def scale_noise(
+#     noise,
+#     factor=1.0,
+#     *,
+#     normalized=True,
+#     normalize_dims=(-3, -2, -1),
+# ):
+#     if not normalized or noise.numel() == 0:
+#         return noise * factor if factor != 1 else noise
+#     mean, std = (
+#         noise.mean(dim=normalize_dims, keepdim=True),
+#         noise.std(dim=normalize_dims, keepdim=True),
+#     )
+#     return (noise - mean).div_(std).mul_(factor)
+
 
 def scale_noise(
     noise,
@@ -13,12 +49,29 @@ def scale_noise(
     normalize_dims=(-3, -2, -1),
 ):
     if not normalized or noise.numel() == 0:
-        return noise.mul_(factor) if factor != 1 else noise
-    mean, std = (
-        noise.mean(dim=normalize_dims, keepdim=True),
-        noise.std(dim=normalize_dims, keepdim=True),
-    )
-    return noise.sub_(mean).div_(std).mul_(factor)
+        return noise * factor if factor != 1 else noise
+    noise = noise / noise.std(dim=normalize_dims, keepdim=True)
+    return noise.sub_(noise.mean(dim=normalize_dims, keepdim=True)).mul_(factor)
+
+
+# def scale_noise(
+#     noise,
+#     factor=1.0,
+#     *,
+#     normalized=True,
+#     normalize_dims=(-3, -2, -1),
+# ):
+#     if not normalized or noise.numel() == 0:
+#         return noise.mul_(factor) if factor != 1 else noise
+#     n = (
+#         torch.nn.LayerNorm(noise.shape[1:])
+#         if normalize_dims == (-3, -2, -1)
+#         else torch.nn.InstanceNorm2d(noise.shape[1])
+#     ).to(noise)
+#     return n(noise) * factor
+#     return latent.normalize_to_scale(
+#         n(noise).clamp_(-1, 1), -1, 1, dim=normalize_dims
+#     ).mul_(factor)
 
 
 def find_first_unsorted(tensor, desc=True):