update dep

dit/dit_i23d.py

@@ -9,10 +9,15 @@ from pdb import set_trace as st
 
 from ldm.modules.attention import MemoryEfficientCrossAttention
 from .dit_models_xformers import DiT, get_2d_sincos_pos_embed, ImageCondDiTBlock, FinalLayer, CaptionEmbedder, approx_gelu, ImageCondDiTBlockPixelArt, t2i_modulate, ImageCondDiTBlockPixelArtRMSNorm, T2IFinalLayer, ImageCondDiTBlockPixelArtRMSNormNoClip
-from apex.normalization import FusedLayerNorm as LayerNorm
-from apex.normalization import FusedRMSNorm as RMSNorm
 from timm.models.vision_transformer import Mlp
 
+try:
+    from apex.normalization import FusedLayerNorm as LayerNorm
+    from apex.normalization import FusedRMSNorm as RMSNorm
+except:
+    from torch.nn import LayerNorm
+    from dit.norm import RMSNorm
+
 # from vit.vit_triplane import XYZPosEmbed