update

dit/dit_i23d.py

@@ -16,7 +16,7 @@ try:
     from apex.normalization import FusedRMSNorm as RMSNorm
 except:
     from torch.nn import LayerNorm
-    from dit.norm import RMSNorm
+    from diffusers.models.normalization import RMSNorm
 
 # from vit.vit_triplane import XYZPosEmbed
 

dit/dit_models_xformers.py

@@ -29,8 +29,7 @@ try:
     from apex.normalization import FusedLayerNorm as LayerNorm
 except:
     from torch.nn import LayerNorm
-    # from torch.nn import RMSNorm # requires torch2.4
-    from dit.norm import RMSNorm
+    from diffusers.models.normalization import RMSNorm
 
 # from torch.nn import LayerNorm
 # from xformers import triton

dit/norm.py

@@ -1,18 +0,0 @@
-import torch
-
-def rms_norm(x, weight=None, eps=1e-05):
-    output = x / torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
-    return output * weight if weight is not None else output
-
-class RMSNorm(torch.nn.Module):
-
-    def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True, dtype=None, device=None):
-        super().__init__()
-        self.eps = eps
-        if elementwise_affine:
-            self.weight = torch.nn.Parameter(torch.ones(normalized_shape, dtype=dtype, device=device))
-        else:
-            self.register_parameter('weight', None)
-
-    def forward(self, x):
-        return rms_norm(x.float(), self.weight, self.eps).to(dtype=x.dtype)

requirements.txt

@@ -31,4 +31,5 @@ safetensors
 matplotlib
 git+https://github.com/nupurkmr9/vision-aided-gan
 PyMCubes
-trimesh
+trimesh
+diffusers