Feature: Add mixed precision support and direct bfloat16 support.

xora/examples/image_to_video.py

@@ -136,6 +136,12 @@ def main():
         "--frame_rate", type=int, default=25, help="Frame rate for the output video"
     )
 
+    parser.add_argument(
+        "--mixed_precision",
+        action="store_true",
+        help="Mixed precision in float32 and bfloat16",
+    )
+
     # Prompts
     parser.add_argument(
         "--prompt",
@@ -224,6 +230,7 @@ def main():
         is_video=True,
         vae_per_channel_normalize=True,
         conditioning_method=ConditioningMethod.FIRST_FRAME,
+        mixed_precision=args.mixed_precision,
     ).images
 
     # Save output video

xora/models/transformers/transformer3d.py

@@ -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, sin_freq
+        return cos_freq.to(dtype), sin_freq.to(dtype)
 
     def forward(
         self,

xora/pipelines/pipeline_xora_video.py

@@ -9,6 +9,7 @@ from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
+from contextlib import nullcontext
 from diffusers.image_processor import VaeImageProcessor
 from diffusers.models import AutoencoderKL
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
@@ -758,6 +759,7 @@ class XoraVideoPipeline(DiffusionPipeline):
         callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
         clean_caption: bool = True,
         media_items: Optional[torch.FloatTensor] = None,
+        mixed_precision: bool = False,
         **kwargs,
     ) -> Union[ImagePipelineOutput, Tuple]:
         """
@@ -1006,16 +1008,22 @@ class XoraVideoPipeline(DiffusionPipeline):
 
                 if conditioning_mask is not None:
                     current_timestep = current_timestep * (1 - conditioning_mask)
+                # Choose the appropriate context manager based on `mixed_precision`
+                if mixed_precision:
+                    context_manager = torch.autocast("cuda", dtype=torch.bfloat16)
+                else:
+                    context_manager = nullcontext()  # Dummy context manager
 
                 # predict noise model_output
-                noise_pred = self.transformer(
-                    latent_model_input.to(self.transformer.dtype),
-                    indices_grid,
-                    encoder_hidden_states=prompt_embeds.to(self.transformer.dtype),
-                    encoder_attention_mask=prompt_attention_mask,
-                    timestep=current_timestep,
-                    return_dict=False,
-                )[0]
+                with context_manager:
+                    noise_pred = self.transformer(
+                        latent_model_input.to(self.transformer.dtype),
+                        indices_grid,
+                        encoder_hidden_states=prompt_embeds.to(self.transformer.dtype),
+                        encoder_attention_mask=prompt_attention_mask,
+                        timestep=current_timestep,
+                        return_dict=False,
+                    )[0]
 
                 # perform guidance
                 if do_classifier_free_guidance: