update

ldm/modules/attention.py

@@ -17,7 +17,8 @@ from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
 try:
     from apex.normalization import FusedRMSNorm as RMSNorm
 except:
-    from dit.norm import RMSNorm
+    # from dit.norm import RMSNorm
+    from diffusers.models.normalization import RMSNorm
 
 
 def exists(val):

vit/vision_transformer.py

@@ -37,7 +37,8 @@ from pdb import set_trace as st
 try:
     from apex.normalization import FusedRMSNorm as RMSNorm
 except:
-    from dit.norm import RMSNorm
+    # from dit.norm import RMSNorm
+    from diffusers.models.normalization import RMSNorm
 
 # from apex.normalization import FusedLayerNorm as LayerNorm