Resolve timestep mismatch between dit forward and pred_noise_to_pred_video

fastvideo/models/dits/causal_wanvideo.py

@@ -375,7 +375,7 @@ class CausalWanTransformer3DModel(BaseDiT):
 
         # Causal-specific
         self.block_mask = None
-        self.num_frame_per_block = 1
+        self.num_frame_per_block = 3
         self.independent_first_frame = False
 
         self.__post_init__()

fastvideo/models/utils.py

@@ -171,10 +171,13 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor,
     # timestep shape should be [B]
     dtype = pred_noise.dtype
     device = pred_noise.device
-    pred_noise = pred_noise.float().to(device)
-    noise_input_latent = noise_input_latent.float().to(device)
-    sigmas = scheduler.sigmas.float().to(device)
-    timesteps = scheduler.timesteps.float().to(device)
+
+    # Convert to double following Self-Forcing
+    # https://github.com/guandeh17/Self-Forcing/blob/main/utils/wan_wrapper.py#L184
+    pred_noise = pred_noise.double().to(device)
+    noise_input_latent = noise_input_latent.double().to(device)
+    sigmas = scheduler.sigmas.double().to(device)
+    timesteps = scheduler.timesteps.double().to(device)
     timestep_id = torch.argmin(
         (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
     sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)

fastvideo/training/ode_causal_pipeline.py

@@ -282,10 +282,11 @@ class ODEInitTrainingPipeline(TrainingPipeline):
         latent_vis_dict["noisy_input"] = noisy_input.permute(0, 2, 1, 3, 4).detach().clone().cpu()
         latent_vis_dict["x0"] = target_latent.permute(0, 2, 1, 3, 4).detach().clone().cpu()
 
+        model_dtype = next(self.transformer.parameters()).dtype
         input_kwargs = {
             "hidden_states": noisy_input.permute(0, 2, 1, 3, 4),
             "encoder_hidden_states": encoder_hidden_states,
-            "timestep": timestep.to(device, dtype=torch.bfloat16),
+            "timestep": timestep.to(device, dtype=model_dtype),
             "encoder_attention_mask": encoder_attention_mask,
             "return_dict": False,
         }
@@ -302,7 +303,7 @@ class ODEInitTrainingPipeline(TrainingPipeline):
         pred_video = pred_noise_to_pred_video(
             pred_noise=noise_pred.flatten(0, 1),
             noise_input_latent=noisy_input.flatten(0, 1),
-            timestep=timestep.flatten(0, 1),
+            timestep=timestep.to(dtype=model_dtype).flatten(0, 1),
             scheduler=self.modules["scheduler"]).unflatten(
                 0, noise_pred.shape[:2])
         latent_vis_dict["pred_video"] = pred_video.permute(0, 2, 1, 3, 4).detach().clone().cpu()