update dep

dit/dit_models_xformers.py

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

dit/norm.py

@@ -6,10 +6,10 @@ def rms_norm(x, weight=None, eps=1e-05):
 
 class RMSNorm(torch.nn.Module):
 
-    def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
+    def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True, dtype=None, device=None):
         super().__init__()
         self.eps = eps
-        if weight:
+        if elementwise_affine:
             self.weight = torch.nn.Parameter(torch.ones(normalized_shape, dtype=dtype, device=device))
         else:
             self.register_parameter('weight', None)

ldm/modules/attention.py

@@ -17,7 +17,7 @@ from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
 try:
     from apex.normalization import FusedRMSNorm as RMSNorm
 except:
-    from torch.nn import RMSNorm # requires torch2.4
+    from dit.norm import RMSNorm
 
 
 def exists(val):

vit/vision_transformer.py

@@ -34,7 +34,11 @@ from .utils import trunc_normal_
 
 from pdb import set_trace as st
 # import apex
-from apex.normalization import FusedRMSNorm as RMSNorm
+try:
+    from apex.normalization import FusedRMSNorm as RMSNorm
+except:
+    from dit.norm import RMSNorm
+
 # from apex.normalization import FusedLayerNorm as LayerNorm
 
 try: