Feature: Add full bfloat16 support.

xora/examples/image_to_video.py

@@ -142,6 +142,12 @@ def main():
         help="Mixed precision in float32 and bfloat16",
     )
 
+    parser.add_argument(
+        "--bfloat16",
+        action="store_true",
+        help="Denoise in bfloat16",
+    )
+
     # Prompts
     parser.add_argument(
         "--prompt",
@@ -176,6 +182,9 @@ def main():
         "PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="tokenizer"
     )
 
+    if args.bfloat16 and unet.dtype != torch.bfloat16:
+        unet = unet.to(torch.bfloat16)
+
     # Use submodels for the pipeline
     submodel_dict = {
         "transformer": unet,

xora/examples/text_to_video.py

@@ -49,6 +49,16 @@ def main():
         required=True,
         help="Path to the directory containing unet, vae, and scheduler subdirectories",
     )
+    parser.add_argument(
+        "--mixed_precision",
+        action="store_true",
+        help="Mixed precision in float32 and bfloat16",
+    )
+    parser.add_argument(
+        "--bfloat16",
+        action="store_true",
+        help="Denoise in bfloat16",
+    )
     args = parser.parse_args()
 
     # Paths for the separate mode directories
@@ -72,6 +82,9 @@ def main():
         "PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="tokenizer"
     )
 
+    if args.bfloat16 and unet.dtype != torch.bfloat16:
+        unet = unet.to(torch.bfloat16)
+
     # Use submodels for the pipeline
     submodel_dict = {
         "transformer": unet,  # using unet for transformer
@@ -115,6 +128,7 @@ def main():
         **sample,
         is_video=True,
         vae_per_channel_normalize=True,
+        mixed_precision=args.mixed_precision,
     ).images
 
     print("Generated images (video frames).")

xora/models/transformers/transformer3d.py

@@ -253,7 +253,7 @@ class Transformer3DModel(ModelMixin, ConfigMixin):
         return fractional_positions
 
     def precompute_freqs_cis(self, indices_grid, spacing="exp"):
-        dtype = self.dtype
+        dtype = torch.float32  # We need full precision in the freqs_cis computation.
         dim = self.inner_dim
         theta = self.positional_embedding_theta
 
@@ -305,7 +305,7 @@ class Transformer3DModel(ModelMixin, ConfigMixin):
             sin_padding = torch.zeros_like(cos_freq[:, :, : dim % 6])
             cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)
             sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)
-        return cos_freq.to(dtype), sin_freq.to(dtype)
+        return cos_freq.to(self.dtype), sin_freq.to(self.dtype)
 
     def forward(
         self,