diff --git a/examples/inference/basic/basic_self_forcing_causal.py b/examples/inference/basic/basic_self_forcing_causal.py
index e2241a94e..ff5073635 100644
--- a/examples/inference/basic/basic_self_forcing_causal.py
+++ b/examples/inference/basic/basic_self_forcing_causal.py
@@ -8,7 +8,7 @@ def main():
     # model.
     # If a local path is provided, FastVideo will make a best effort
     # attempt to identify the optimal arguments.
-    model_name = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+    model_name = "SFWan2.1-T2V-1.3B-Diffusers/snapshots/bd804886ec4fe04fcc5a3fe17644428b9cee94c0/"
     generator = VideoGenerator.from_pretrained(
         model_name,
         # FastVideo will automatically handle distributed setup
@@ -19,6 +19,7 @@ def main():
     )
 
     sampling_param = SamplingParam.from_pretrained(model_name)
+    sampling_param.num_frames = 81
 
     prompt = (
         "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py
index ffac8dd5b..132e1b5f9 100644
--- a/fastvideo/configs/pipelines/registry.py
+++ b/fastvideo/configs/pipelines/registry.py
@@ -13,7 +13,7 @@ from fastvideo.configs.pipelines.wan import (
     FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config,
     SelfForcingWanT2V480PConfig, Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config,
     Wan2_2_TI2V_5B_Config, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig,
-    WanT2V720PConfig)
+    WanT2V720PConfig, SelfForcingWanT2V480PConfig)
 # isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
@@ -49,6 +49,7 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
     "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(),
     "stepvideo": lambda id: "stepvideo" in id.lower(),
+    "wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
@@ -60,7 +61,8 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
     WanT2V480PConfig,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V480PConfig,
     "wandmdpipeline": FastWan2_1_T2V_480P_Config,
-    "stepvideo": StepVideoT2VConfig
+    "stepvideo": StepVideoT2VConfig,
+    "wancausaldmdpipeline": SelfForcingWanT2V480PConfig,
     # Other fallbacks by architecture
 }
 
diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py
index b7306bd37..7d97b21bc 100644
--- a/fastvideo/entrypoints/video_generator.py
+++ b/fastvideo/entrypoints/video_generator.py
@@ -344,7 +344,9 @@ class VideoGenerator:
                 "size": (target_height, target_width, batch.num_frames),
                 "generation_time": gen_time,
                 "logging_info": logging_info,
-                "trajectory": output_batch.output_trajectory,
+                "trajectory": output_batch.trajectory_latents,
+                "trajectory_timesteps": output_batch.trajectory_timesteps,
+                "trajectory_decoded": output_batch.trajectory_decoded,
             }
 
     def set_lora_adapter(self,
diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py
index dcf4cef9b..16e690d7b 100644
--- a/fastvideo/pipelines/pipeline_batch_info.py
+++ b/fastvideo/pipelines/pipeline_batch_info.py
@@ -151,8 +151,8 @@ class ForwardBatch:
     return_trajectory_latents: bool = False
     return_trajectory_decoded: bool = False
     trajectory_timesteps: list[int] | None = None
-    trajectory_latents: list[tuple[int, torch.Tensor]] | None = None
-    trajectory_decoded: list[tuple[int, torch.Tensor]] | None = None
+    trajectory_latents: torch.Tensor | None = None
+    trajectory_decoded: list[torch.Tensor] | None = None
 
     # Extra parameters that might be needed by specific pipeline implementations
     extra: dict[str, Any] = field(default_factory=dict)
diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py
index c39c0f11c..8d575e7fd 100644
--- a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py
+++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py
@@ -26,7 +26,7 @@ from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
-from fastvideo.utils import shallow_asdict
+from fastvideo.utils import shallow_asdict, save_decoded_latents_as_video
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.preprocess.preprocess_pipeline_base import (
     BasePreprocessPipeline)
@@ -34,7 +34,8 @@ from fastvideo.pipelines.stages import (DenoisingStage, ImageVAEEncodingStage,
                                         InputValidationStage,
                                         LatentPreparationStage,
                                         TextEncodingStage,
-                                        TimestepPreparationStage)
+                                        TimestepPreparationStage,
+                                        DecodingStage)
 
 logger = init_logger(__name__)
 
@@ -78,6 +79,8 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                            scheduler=self.get_module("scheduler"),
                            pipeline=self,
                        ))
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
 
     def preprocess_video_and_text_and_trajectory(self,
                                                  fastvideo_args: FastVideoArgs,
@@ -121,37 +124,34 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                     valid_data, fastvideo_args)
 
                 batch_captions = valid_data["text"]
-                # for i in range(len(batch_captions)):
-                #     caption = batch_captions[i]
-                #     logger.info(f"===== batch_captions: {caption}")
-                batch = ForwardBatch(
-                    data_type="video",
-                    prompt=batch_captions,
-                    prompt_embeds=[],
-                    prompt_attention_mask=[],
+                # Encode text using the standalone TextEncodingStage API
+                prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text(
+                    batch_captions,
+                    fastvideo_args,
+                    encoder_index=[0],
+                    return_attention_mask=True,
                 )
-                assert hasattr(self, "prompt_encoding_stage")
-                result_batch = self.prompt_encoding_stage(batch, fastvideo_args)
-                prompt_embeds, prompt_attention_mask = result_batch.prompt_embeds[
-                    0], result_batch.prompt_attention_mask[0]
-                assert prompt_embeds.shape[0] == prompt_attention_mask.shape[0]
+                prompt_embeds = prompt_embeds_list[0]
+                prompt_attention_masks = prompt_masks_list[0]
+                assert prompt_embeds.shape[0] == prompt_attention_masks.shape[0]
 
-                # Get sequence lengths from attention masks (number of 1s)
-                seq_lens = prompt_attention_mask.sum(dim=1)
+                # # Get sequence lengths from attention masks (number of 1s)
+                # seq_lens = prompt_attention_mask.sum(dim=1)
 
-                non_padded_embeds = []
-                non_padded_masks = []
+                # non_padded_embeds = []
+                # non_padded_masks = []
 
-                # Process each item in the batch
-                for i in range(prompt_embeds.size(0)):
-                    seq_len = seq_lens[i].item()
-                    # Slice the embeddings and masks to keep only non-padding parts
-                    non_padded_embeds.append(prompt_embeds[i, :seq_len])
-                    non_padded_masks.append(prompt_attention_mask[i, :seq_len])
+                # # Process each item in the batch
+                # for i in range(prompt_embeds.size(0)):
+                #     seq_len = seq_lens[i].item()
+                #     # Slice the embeddings and masks to keep only non-padding parts
+                #     non_padded_embeds.append(prompt_embeds[i, :seq_len])
+                #     non_padded_masks.append(prompt_attention_mask[i, :seq_len])
 
                 # Update the tensors with non-padded versions
-                prompt_embeds = non_padded_embeds
-                prompt_attention_masks = non_padded_masks
+                # prompt_embeds = non_padded_embeds
+                # prompt_attention_masks = non_padded_masks
+                # prompt_embeds = prompt_embeds
 
                 # logger.info(f"===== prompt_embeds: {prompt_embeds[0].shape}")
                 # logger.info(f"===== prompt_attention_masks: {prompt_attention_masks[0].shape}")
@@ -159,8 +159,23 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                 sampling_params = SamplingParam.from_pretrained(
                     args.model_path)
 
+                # encode negative prompt for trajectory collection
+                if sampling_params.guidance_scale > 1 and sampling_params.negative_prompt is not None:
+                    negative_prompt_embeds_list, negative_prompt_masks_list = self.prompt_encoding_stage.encode_text(
+                        sampling_params.negative_prompt,
+                        fastvideo_args,
+                        encoder_index=[0],
+                        return_attention_mask=True,
+                    )
+                    negative_prompt_embed = negative_prompt_embeds_list[0][0]
+                    negative_prompt_attention_mask = negative_prompt_masks_list[0][0]
+                else:
+                    negative_prompt_embed = None
+                    negative_prompt_attention_mask = None
+
                 trajectory_latents = []
                 trajectory_timesteps = []
+                trajectory_decoded = []
                 for i, (prompt_embed, prompt_attention_mask) in enumerate(zip(prompt_embeds, prompt_attention_masks)):
                     prompt_embed = prompt_embed.unsqueeze(0)
                     prompt_attention_mask = prompt_attention_mask.unsqueeze(0)
@@ -185,7 +200,10 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                     )
                     batch.prompt_embeds = [prompt_embed]
                     batch.prompt_attention_mask = [prompt_attention_mask]
+                    batch.negative_prompt_embeds = [negative_prompt_embed]
+                    batch.negative_attention_mask = [negative_prompt_attention_mask]
                     batch.return_trajectory_latents = True
+                    batch.return_trajectory_decoded = False
                     batch.height = args.max_height
                     batch.width = args.max_width
                     # batch.num_frames = 81
@@ -204,9 +222,11 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                         result_batch, fastvideo_args)
                     result_batch = self.denoising_stage(result_batch,
                                                         fastvideo_args)
+                    result_batch = self.decoding_stage(result_batch, fastvideo_args)
                     # trajectory_latents = result_batch.trajectory_latents
                     trajectory_latents.append(result_batch.trajectory_latents.cpu())
                     trajectory_timesteps.append(result_batch.trajectory_timesteps.cpu())
+                    trajectory_decoded.append(result_batch.trajectory_decoded)
 
             extra_features["trajectory_latents"] = trajectory_latents
             extra_features["trajectory_timesteps"] = trajectory_timesteps
@@ -214,6 +234,13 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
             logger.info(f"===== trajectory_latents len: {len(trajectory_latents)}")
             logger.info(f"===== trajectory_timesteps: {trajectory_timesteps}")
             logger.info(f"===== trajectory_timesteps len: {len(trajectory_timesteps)}")
+
+            if batch.return_trajectory_decoded:
+                logger.info(f"===== SAVING TRAJECTORY DECODED")
+                for i, decoded_frames in enumerate(trajectory_decoded):
+                    for j, decoded_frame in enumerate(decoded_frames):
+                        logger.info(f"===== SAVING TRAJECTORY DECODED {i} for prompt {batch_captions[i]}")
+                        save_decoded_latents_as_video(decoded_frame, f"decoded_videos/trajectory_decoded_{i}_{j}.mp4", args.train_fps)
             # assert False
             # Prepare batch data for Parquet dataset
             batch_data = []
@@ -345,39 +372,6 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                 frame, height, width, fastvideo_args, generator)
             video_conditions.append(latent)
 
-            # processed_img = frame.to(device="cpu", dtype=torch.float32)
-            # processed_img = processed_img.unsqueeze(0).permute(0, 3, 1,
-            #                                                    2).unsqueeze(2)
-            # # (B, H, W, C) -> (B, C, 1, H, W)
-            # video_condition = processed_img.unsqueeze(2)
-            # video_condition = video_condition.to(
-            #     device=get_local_torch_device(), dtype=torch.float32)
-            # video_conditions.append(video_condition)
-
-        # video_conditions = torch.cat(video_conditions, dim=0)
-
-        # with torch.autocast(device_type="cuda",
-        #                     dtype=torch.float32,
-        #                     enabled=True):
-        #     encoder_outputs = self.get_module("vae").encode(video_conditions)
-
-        # latent_condition = encoder_outputs.mean
-        # if (hasattr(self.get_module("vae"), "shift_factor")
-        #         and self.get_module("vae").shift_factor is not None):
-        #     if isinstance(self.get_module("vae").shift_factor, torch.Tensor):
-        #         latent_condition -= self.get_module("vae").shift_factor.to(
-        #             latent_condition.device, latent_condition.dtype)
-        #     else:
-        #         latent_condition -= self.get_module("vae").shift_factor
-
-        # if isinstance(self.get_module("vae").scaling_factor, torch.Tensor):
-        #     latent_condition = latent_condition * self.get_module(
-        #         "vae").scaling_factor.to(latent_condition.device,
-        #                                  latent_condition.dtype)
-        # else:
-        #     latent_condition = latent_condition * self.get_module(
-        #         "vae").scaling_factor
-
         features["image_condition_latents"] = video_conditions
         features["pil_images"] = pil_images
         return features
diff --git a/fastvideo/pipelines/preprocess/v1_preprocess.py b/fastvideo/pipelines/preprocess/v1_preprocess.py
index 68576a702..0051b8e68 100644
--- a/fastvideo/pipelines/preprocess/v1_preprocess.py
+++ b/fastvideo/pipelines/preprocess/v1_preprocess.py
@@ -26,7 +26,7 @@ def main(args) -> None:
     pipeline_config = PipelineConfig.from_pretrained(args.model_path)
     kwargs = {
         "vae_precision": "fp32",
-        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=False),
+        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True),
     }
     pipeline_config.update_config_from_dict(kwargs)
     fastvideo_args = FastVideoArgs(
diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py
index bc49091a3..9ed07b5c7 100644
--- a/fastvideo/pipelines/stages/decoding.py
+++ b/fastvideo/pipelines/stages/decoding.py
@@ -128,13 +128,17 @@ class DecodingStage(PipelineStage):
         # decode trajectory latents if needed
         if batch.return_trajectory_decoded:
             batch.trajectory_decoded = []
+            logger.info(f"batch.trajectory_latents.shape: {batch.trajectory_latents.shape}")
             assert batch.trajectory_latents is not None, "batch should have trajectory latents"
-            for timestep, latent in batch.trajectory_latents:
+            for idx in range(batch.trajectory_latents.shape[1]):
+                # bathc.trajectory_latents is [batch_size, timesteps, channels, frames, height, width]
+                cur_latent = batch.trajectory_latents[:, idx, :, :, :, :]
+                logger.info(f"cur_latent.shape: {cur_latent.shape}")
+                cur_timestep = batch.trajectory_timesteps[idx]
                 logger.info(
-                    f"decoding trajectory latent for timestep: {timestep}")
-                decoded_latent = self.decode(latent, fastvideo_args)
-                batch.trajectory_decoded.append(
-                    (timestep, decoded_latent.cpu().float()))
+                    f"decoding trajectory latent for timestep: {cur_timestep}")
+                decoded_frames = self.decode(cur_latent, fastvideo_args)
+                batch.trajectory_decoded.append(decoded_frames.cpu().float())
 
         # Convert to CPU float32 for compatibility
         frames = frames.cpu().float()
diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py
index 23c75f7a6..d3bf019c9 100644
--- a/fastvideo/pipelines/stages/denoising.py
+++ b/fastvideo/pipelines/stages/denoising.py
@@ -427,6 +427,7 @@ class DenoisingStage(PipelineStage):
                     progress_bar.update()
 
         # Gather results if using sequence parallelism
+        trajectory_tensor: torch.Tensor | None = None
         if trajectory_latents:
             trajectory_tensor = torch.stack(trajectory_latents, dim=1)
         else:
diff --git a/fastvideo/training/ode_causal_pipeline.py b/fastvideo/training/ode_causal_pipeline.py
index 3e83d3bd6..e74bc105f 100644
--- a/fastvideo/training/ode_causal_pipeline.py
+++ b/fastvideo/training/ode_causal_pipeline.py
@@ -152,7 +152,7 @@ class ODEInitTrainingPipeline(TrainingPipeline):
         t = traj_timesteps[batch_indices, nearest_idx]  # [B]
 
         # Scale model input as in inference for consistency with stored trajectories
-        # scaled_latents = self.modules["scheduler"].scale_model_input(current_latents, t)
+        noisy_input = self.modules["scheduler"].scale_model_input(noisy_input, t)
 
         # Prepare inputs for transformer
         input_kwargs = {
@@ -169,6 +169,16 @@ class ODEInitTrainingPipeline(TrainingPipeline):
             noise_pred = self.transformer(**input_kwargs)
         if isinstance(noise_pred, (tuple, list)):
             noise_pred = noise_pred[0]
+
+        from fastvideo.models.utils import pred_noise_to_pred_video
+        noise_pred = pred_noise_to_pred_video(
+            pred_noise=noise_pred.flatten(0, 1),
+            noise_input_latent=noisy_input.flatten(0, 1),
+            timestep=t,
+            scheduler=self.modules["scheduler"]).unflatten(
+                0, noise_pred.shape[:2])
+
+        # noisy_input = pred_noise_to_pred_video(noise_pred, noisy_input, t, self.modules["scheduler"])
         # next_latent_pred = self.modules["scheduler"].step(
         #     noise_pred, t, current_latents, return_dict=False)[0]
         return noise_pred, target_latent, t
diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py
index bdeb81b4f..bd63340b2 100644
--- a/fastvideo/training/training_utils.py
+++ b/fastvideo/training/training_utils.py
@@ -346,10 +346,14 @@ def load_checkpoint(transformer,
     """
     if not os.path.exists(checkpoint_path):
         logger.warning("Checkpoint path %s does not exist", checkpoint_path)
+        assert False
         return 0
 
     # Extract step number from checkpoint path
-    step = int(os.path.basename(checkpoint_path).split('-')[-1])
+    try:
+        step = int(os.path.basename(checkpoint_path).split('-')[-1])
+    except:
+        step = 1
 
     if rank == 0:
         logger.info("Loading checkpoint from step %s", step)
diff --git a/fastvideo/utils.py b/fastvideo/utils.py
index 83a604454..09ac836ae 100644
--- a/fastvideo/utils.py
+++ b/fastvideo/utils.py
@@ -2,6 +2,10 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py
 
 import argparse
+from einops import rearrange
+import torchvision
+import numpy as np
+import imageio
 import ctypes
 import hashlib
 import importlib
@@ -875,3 +879,16 @@ def best_output_size(w, h, dw, dh, expected_area):
         return ow1, oh1
     else:
         return ow2, oh2
+
+
+def save_decoded_latents_as_video(decoded_latents: list[torch.Tensor], output_path: str, fps: int):
+    # Process outputs
+    videos = rearrange(decoded_latents, "b c t h w -> t b c h w")
+    frames = []
+    for x in videos:
+        x = torchvision.utils.make_grid(x, nrow=6)
+        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+        frames.append((x * 255).numpy().astype(np.uint8))
+
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    imageio.mimsave(output_path, frames, fps=fps, format="mp4")
diff --git a/scripts/checkpoint_conversion/wan_to_diffusers.py b/scripts/checkpoint_conversion/wan_to_diffusers.py
index ea52ed687..5e8908356 100644
--- a/scripts/checkpoint_conversion/wan_to_diffusers.py
+++ b/scripts/checkpoint_conversion/wan_to_diffusers.py
@@ -135,8 +135,8 @@ _self_forcing_to_diffusers_param_names_mapping: dict = {
     r"blocks.\1.norm2.\2",
 }
 
-state_dict = load_state_dict_from_file("checkpoints/self_forcing_dmd.pt")
-state_dict = state_dict["generator_ema"]
+state_dict = load_state_dict_from_file("checkpoints/ode_init.pt")
+state_dict = state_dict["generator"]
 new_state_dict = OrderedDict()
 for k, v in state_dict.items():
     new_key = k
@@ -156,6 +156,6 @@ del state_dict
 
 save_torch_state_dict(
     new_state_dict,
-    "new2/",
+    "ode_init_diffusers/",
     max_shard_size="10GB"
 )