update

app.py

@@ -175,8 +175,8 @@ def main(args):
     # denoise_model.load_state_dict(
     #     dist_util.load_state_dict(args.ddpm_model_path, map_location="cpu"))
     denoise_model.to(dist_util.dev())
-    denoise_model = denoise_model.to(th.bfloat16)
-    auto_encoder = auto_encoder.to(th.bfloat16)
+    # denoise_model = denoise_model.to(th.bfloat16)
+    # auto_encoder = auto_encoder.to(th.bfloat16)
     # if args.use_fp16:
     #     denoise_model.convert_to_fp16()
     denoise_model.eval()

configs/i23d_args.json

@@ -3,7 +3,7 @@
   "triplane_scaling_divider": 0.96806,
   "diffusion_input_size": 32,
   "trainer_name": "flow_matching",
-  "use_amp": true,
+  "use_amp": false,
   "clip_denoised": false,
   "num_samples": 1,
   "num_instances": 10,

dit/dit_i23d.py

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

dit/dit_models_xformers.py

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

dit/norm.py

@@ -0,0 +1,40 @@
+import numbers
+from typing import Dict, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from pdb import set_trace as st
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps: float, elementwise_affine: bool = True):
+        super().__init__()
+
+        self.eps = eps
+
+        if isinstance(dim, numbers.Integral):
+            dim = (dim,)
+
+        self.dim = torch.Size(dim)
+
+        if elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim))
+        else:
+            self.weight = None
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+
+        if self.weight is not None:
+            # convert into half-precision if necessary
+            if self.weight.dtype in [torch.float16, torch.bfloat16]:
+                hidden_states = hidden_states.to(self.weight.dtype)
+            hidden_states = hidden_states * self.weight
+        else:
+            hidden_states = hidden_states.to(input_dtype)
+
+        return hidden_states.to(input_dtype)

ldm/modules/attention.py

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

vit/vision_transformer.py

@@ -38,7 +38,7 @@ try:
     from apex.normalization import FusedRMSNorm as RMSNorm
 except:
     # from dit.norm import RMSNorm
-    from diffusers.models.normalization import RMSNorm
+    from dit.norm import RMSNorm
 
 # from apex.normalization import FusedLayerNorm as LayerNorm