fix accidental eps change

neggles · Jul 13, 2024 · af528f9 · af528f9
1 parent b9c9068
commit af528f9
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/src/neurosis/modules/attention.py b/src/neurosis/modules/attention.py
@@ -11,7 +11,6 @@
 from torch.utils.checkpoint import checkpoint
 
 from neurosis.modules.diffusion.util import zero_module
-from neurosis.modules.layers import Normalize
 
 logger = logging.getLogger(__name__)
 
@@ -153,7 +152,7 @@ def __init__(self, in_channels: int):
         super().__init__()
         self.in_channels = in_channels
 
-        self.norm = Normalize(in_channels)
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
         self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
         self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
         self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
@@ -610,7 +609,7 @@ def __init__(
             context_dim = [None] * depth
 
         self.in_channels = in_channels
-        self.norm = Normalize(in_channels)
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
 
         inner_dim = n_heads * d_head
         if not use_linear:

diff --git a/src/neurosis/modules/diffusion/openaimodel.py b/src/neurosis/modules/diffusion/openaimodel.py
@@ -16,7 +16,6 @@
     timestep_embedding,
     zero_module,
 )
-from neurosis.modules.layers import Normalize
 
 logger = logging.getLogger(__name__)
 
@@ -246,7 +245,7 @@ def __init__(
             padding = kernel_size // 2
 
         self.in_layers = nn.Sequential(
-            Normalize(channels),
+            nn.GroupNorm(32, channels),
             nn.SiLU(),
             conv_nd(dims, channels, self.out_channels, kernel_size, padding=padding),
         )
@@ -280,7 +279,7 @@ def __init__(
             )
 
         self.out_layers = nn.Sequential(
-            Normalize(self.out_channels),
+            nn.GroupNorm(32, self.out_channels),
             nn.SiLU(),
             nn.Dropout(p=dropout),
             zero_module(
@@ -368,7 +367,7 @@ def __init__(
             ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
             self.num_heads = channels // num_head_channels
         self.use_checkpoint = use_checkpoint
-        self.norm = Normalize(channels)
+        self.norm = nn.GroupNorm(channels)
         self.qkv = conv_nd(1, channels, channels * 3, 1)
         if use_new_attention_order:
             # split qkv before split heads
@@ -796,7 +795,7 @@ def __init__(
                 self._feature_size += ch
 
         self.out = nn.Sequential(
-            Normalize(ch),
+            nn.GroupNorm(32, ch),
             nn.SiLU(),
             zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
         )