add preliminary karras unet 3d for @QuantPrincess at #295

lucidrains · Feb 23, 2024 · 3a3ab1b · 3a3ab1b
1 parent e39893b
commit 3a3ab1b
Show file tree

Hide file tree

Showing 4 changed files with 710 additions and 17 deletions.
diff --git a/denoising_diffusion_pytorch/__init__.py b/denoising_diffusion_pytorch/__init__.py
@@ -13,4 +13,5 @@
     InvSqrtDecayLRSched
 )
 
-from denoising_diffusion_pytorch.karras_unet_1d import KarrasUnet1D
+from denoising_diffusion_pytorch.karras_unet_1d import KarrasUnet1D
+from denoising_diffusion_pytorch.karras_unet_3d import KarrasUnet3D
diff --git a/denoising_diffusion_pytorch/karras_unet_1d.py b/denoising_diffusion_pytorch/karras_unet_1d.py
@@ -693,21 +693,6 @@ def forward(self, x):
 
         return x
 
-# works best with inverse square root decay schedule
-
-def InvSqrtDecayLRSched(
-    optimizer,
-    t_ref = 70000,
-    sigma_ref = 0.01
-):
-    """
-    refer to equation 67 and Table1
-    """
-    def inv_sqrt_decay_fn(t: int):
-        return sigma_ref / sqrt(max(t / t_ref, 1.))
-
-    return LambdaLR(optimizer, lr_lambda = inv_sqrt_decay_fn)
-
 # example
 
 if __name__ == '__main__':