Update models/unet.py

models/unet.py

@@ -9,6 +9,25 @@ from einops.layers.torch import Rearrange
 from einops import rearrange
 import matplotlib.pyplot as plt
 import os
+import torch.nn as nn
+
+# Custom LayerNorm class to handle fp16
+class CustomLayerNorm(nn.LayerNorm):
+    def forward(self, x: torch.Tensor):
+        if self.weight.dtype == torch.float32:
+            orig_type = x.dtype
+            ret = super().forward(x.type(torch.float32))
+            return ret.type(orig_type)
+        else:
+            return super().forward(x)
+
+# Function to replace LayerNorm in CLIP model with CustomLayerNorm
+def replace_layer_norm(model):
+    for name, module in model.named_children():
+        if isinstance(module, nn.LayerNorm):
+            setattr(model, name, CustomLayerNorm(module.normalized_shape, elementwise_affine=module.elementwise_affine))
+        else:
+            replace_layer_norm(module)  # Recursively apply to all submodules
 
 
 MONITOR_ATTN = []
@@ -808,7 +827,8 @@ class MotionCLR(nn.Module):
         # text encoder
         self.embed_text = nn.Linear(clip_dim, text_latent_dim)
         self.clip_version = clip_version
-        self.clip_model = self.load_and_freeze_clip(clip_version).cuda()
+        self.clip_model = self.load_and_freeze_clip(clip_version)
+        replace_layer_norm(self.clip_model)
         textTransEncoderLayer = nn.TransformerEncoderLayer(
             d_model=text_latent_dim,
             nhead=text_num_heads,