fix

examples/inference/basic/basic_self_forcing_causal.py

@@ -21,12 +21,13 @@ 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 "
-        "wide with interest. The playful yet serene atmosphere is complemented by soft "
-        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
-    )
-    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
+    prompts = [
+        "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.",
+        "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The path is narrow as it makes its way between all the plants. the scene is captured from a ground-level angle, following the cat closely, giving a low and intimate perspective. The image is cinematic with warm tones and a grainy texture. The scattered daylight between the leaves and plants above creates a warm contrast, accentuating the cat’s orange fur. The shot is clear and sharp, with a shallow depth of field.",
+    ]
+
+    for prompt in prompts:
+        video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
 
 if __name__ == "__main__":
     main()

examples/training/consistency_finetune/ode_init/finetune_ode_init_16k_vidprom.slurm

@@ -1,5 +1,5 @@
 #!/bin/bash
-#SBATCH --job-name=4e6_ode_vidprom2k
+#SBATCH --job-name=4e6B8_16kFV_no_warp_ode_vidprom
 #SBATCH --partition=main
 #SBATCH --nodes=1
 #SBATCH --ntasks=1
@@ -7,8 +7,8 @@
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
 #SBATCH --mem=1440G
-#SBATCH --output=ode_vidprom2k_outputtest/ode_vidprom2k_4e-6.out
-#SBATCH --error=ode_vidprom2k_outputtest/ode_vidprom2k_4e-6.err
+#SBATCH --output=ode_vidprom16k_no_warp/ode_vidprom8b16k_4e-6.out
+#SBATCH --error=ode_vidprom16k_no_warp/ode_vidprom8b16k_4e-6.err
 #SBATCH --exclusive
 set -e -x
 
@@ -38,7 +38,7 @@ echo "NODE_RANK: $NODE_RANK"
 
 # MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-extended-t2v-1-3b/"
+DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-16k-t2v-1-3b/"
 VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
 NUM_GPUS=2
 # export CUDA_VISIBLE_DEVICES=4,5
@@ -47,9 +47,9 @@ NUM_GPUS=2
 # Training arguments
 training_args=(
   --tracker_project_name "wan_ode_init"
-  --output_dir "wan_ode_init_vidprom_70k_test"
+  --output_dir "vidprom_8b16k_test_no_warp_4e-6"
   --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "vidprom_70k_wan_ode_init_4e-6_test"
+  --wandb_run_name "vidprom_8b16k_wan_ode_init_4e-6"
   # --resume_from_checkpoint "ode_init_diffusers/"
   --max_train_steps 6001
   --train_batch_size 1
@@ -59,8 +59,8 @@ training_args=(
   --num_height 480
   --num_width 832
   --num_frames 77
-  --warp_denoising_step
-  # --enable_gradient_checkpointing_type "full"
+  --dmd_denoising_steps "1000,750,500,0"
+  --enable_gradient_checkpointing_type "full"
 )
 
 # Parallel arguments
@@ -68,7 +68,7 @@ parallel_args=(
   --num_gpus $NUM_GPUS
   --sp_size 1
   --tp_size 1
-  --hsdp_replicate_dim 2
+  --hsdp_replicate_dim $NUM_GPUS
   --hsdp_shard_dim 1
 )
 
@@ -88,7 +88,7 @@ dataset_args=(
 validation_args=(
   --log_validation
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 100
+  --validation_steps 50
   --validation_sampling_steps "50"
   --validation_guidance_scale "6.0"
 )
@@ -97,7 +97,7 @@ validation_args=(
 optimizer_args=(
   --learning_rate 4e-6
   --mixed_precision "bf16"
-  --checkpointing_steps 2000
+  --checkpointing_steps 500
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

examples/training/consistency_finetune/ode_init/finetune_ode_init_crush.slurm

@@ -40,17 +40,18 @@ echo "NODE_RANK: $NODE_RANK"
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
-NUM_GPUS=8
+NUM_GPUS=2
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
 # Training arguments
 training_args=(
   --tracker_project_name "wan_ode_init"
-  --output_dir "wan_ode_init"
+  --output_dir "wan_ode_init_warp_2"
   --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "fixed_wan_ode_init_5e-6"
+  --wandb_run_name "2warp_fixed_wan_ode_init_5e-6"
   # --resume_from_checkpoint "ode_init_diffusers/"
+  # --warp_denoising_step
   --max_train_steps 6001
   --train_batch_size 1
   --train_sp_batch_size 1
@@ -67,7 +68,7 @@ parallel_args=(
   --num_gpus $NUM_GPUS
   --sp_size 1
   --tp_size 1
-  --hsdp_replicate_dim 8
+  --hsdp_replicate_dim $NUM_GPUS
   --hsdp_shard_dim 1
 )
 
@@ -87,7 +88,7 @@ dataset_args=(
 validation_args=(
   --log_validation
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 100
+  --validation_steps 20
   --validation_sampling_steps "50"
   --validation_guidance_scale "6.0"
 )

examples/training/consistency_finetune/ode_init/finetune_ode_init_overfit.sh

@@ -4,11 +4,12 @@ export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 export TOKENIZERS_PARALLELISM=false
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
 
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/combined_parquet_dataset/"
+DATA_DIR="/mnt/weka/home/hao.zhang/wl/FastVideo2/data/crush-smol_processed_t2v_1_3b_ode_init_single"
 VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
-NUM_GPUS=8
+NUM_GPUS=1
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 
@@ -17,9 +18,10 @@ training_args=(
   --tracker_project_name "wan_ode_init"
   --output_dir "wan_ode_init_crush_smol"
   --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "wan_ode_init_crush_smol"
+  --wandb_run_name "overfitwan_ode_init_crush_smol"
   # --resume_from_checkpoint "ode_init_diffusers/"
   --max_train_steps 2001
+  # --warp_denoising_step
   --train_batch_size 1
   --train_sp_batch_size 1
   --gradient_accumulation_steps 1
@@ -35,7 +37,7 @@ parallel_args=(
   --num_gpus $NUM_GPUS
   --sp_size 1
   --tp_size 1
-  --hsdp_replicate_dim 8
+  --hsdp_replicate_dim 1
   --hsdp_shard_dim 1
 )
 
@@ -55,14 +57,14 @@ dataset_args=(
 validation_args=(
   --log_validation
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 100
+  --validation_steps 20
   --validation_sampling_steps "50"
   --validation_guidance_scale "6.0"
 )
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 8e-6
+  --learning_rate 1e-5
   --mixed_precision "bf16"
   --checkpointing_steps 500
   --weight_decay 1e-4

examples/training/consistency_finetune/ode_init/preprocess_wan_data_ode.sh

@@ -3,8 +3,8 @@
 GPU_NUM=1 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_TYPE="wan"
-DATA_MERGE_PATH="data/crush-smol/merge.txt"
-OUTPUT_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/"
+DATA_MERGE_PATH="data/crush-smol_single/merge.txt"
+OUTPUT_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_single/"
 
 torchrun --nproc_per_node=$GPU_NUM \
     fastvideo/pipelines/preprocess/v1_preprocess.py \
@@ -18,7 +18,7 @@ torchrun --nproc_per_node=$GPU_NUM \
     --dataloader_num_workers 0 \
     --output_dir=$OUTPUT_DIR \
     --train_fps 16 \
-    --samples_per_file 5 \
-    --flush_frequency 5 \
+    --samples_per_file 1 \
+    --flush_frequency 1 \
     --video_length_tolerance_range 5 \
     --preprocess_task "ode_trajectory" 

fastvideo/configs/pipelines/wan.py

@@ -39,7 +39,7 @@ class WanT2V480PConfig(PipelineConfig):
     vae_sp: bool = False
 
     # Denoising stage
-    flow_shift: int = 3
+    flow_shift: float | None = 8.0
 
     # Text encoding stage
     text_encoder_configs: tuple[EncoderConfig, ...] = field(
@@ -146,6 +146,6 @@ class Wan2_2_I2V_A14B_Config(WanT2V480PConfig):
 @dataclass
 class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
     is_causal: bool = True
-    flow_shift: int = 5
+    flow_shift: float | None = 5.0
     dmd_denoising_steps: list[int] | None = field(
         default_factory=lambda: [1000, 750, 500, 250])

fastvideo/dataset/utils.py

@@ -8,7 +8,7 @@ from fastvideo.logger import init_logger
 logger = init_logger(__name__)
 
 
-def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor:
+def pad(t: torch.Tensor, padding_length: int) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Pad or crop an embedding [L, D] to exactly padding_length tokens.
     Return:
@@ -146,10 +146,7 @@ def collate_rows_from_parquet_schema(rows,
             bytes_key = f"{tensor_name}_bytes"
 
             if shape_key in row and bytes_key in row:
-                logger.info("WTF shape_key: %s", shape_key)
-                # logger.info("WTF bytes_key: %s", bytes_key)
                 shape = row[shape_key]
-                logger.info("WTF shape: %s", shape)
                 bytes_data = row[bytes_key]
 
                 if len(bytes_data) == 0:

fastvideo/models/dits/causal_wanvideo.py

@@ -244,19 +244,19 @@ class CausalWanTransformerBlock(nn.Module):
         current_start: int = 0,
         cache_start: int | None = None,
     ) -> torch.Tensor:
-        logger.info("temb.shape: %s", temb.shape)
+        # logger.info("temb.shape: %s", temb.shape)
         num_frames = temb.shape[1]
-        logger.info("first hidden_states.shape: %s", hidden_states.shape)
-        logger.info("num_frames: %s", num_frames)
+        # logger.info("first hidden_states.shape: %s", hidden_states.shape)
+        # logger.info("num_frames: %s", num_frames)
         if hidden_states.dim() == 4:
             hidden_states = hidden_states.squeeze(1)
         frame_seqlen = hidden_states.shape[1] // temb.shape[1]
-        logger.info("frame_seqlen: %s", frame_seqlen)
+        # logger.info("frame_seqlen: %s", frame_seqlen)
         bs, seq_length, _ = hidden_states.shape
         orig_dtype = hidden_states.dtype
         # assert orig_dtype != torch.float32
         e = self.scale_shift_table + temb.float()
-        logger.info("e.shape: %s", e.shape)
+        # logger.info("e.shape: %s", e.shape)
         shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
             6, dim=2)
         assert shift_msa.dtype == torch.float32
@@ -312,7 +312,7 @@ class CausalWanTransformerBlock(nn.Module):
         ff_output = self.ffn(norm_hidden_states)
         hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
         # logger.info("after mlp_residual norm_hidden_states.shape: %s", norm_hidden_states.shape)
-        logger.info("after mlp_residual hidden_states.shape: %s", hidden_states.shape)
+        # logger.info("after mlp_residual hidden_states.shape: %s", hidden_states.shape)
         hidden_states = hidden_states.to(orig_dtype)
 
         return hidden_states
@@ -650,10 +650,10 @@ class CausalWanTransformer3DModel(BaseDiT):
                     block_mask=self.block_mask)
         else:
             for block_index, block in enumerate(self.blocks):
-                logger.info("===== TRAIN block %d", block_index)
-                logger.info("hidden_states.shape: %s", hidden_states.shape)
+                # logger.info("===== TRAIN block %d", block_index)
+                # logger.info("hidden_states.shape: %s", hidden_states.shape)
                 # logger.info("encoder_hidden_states.shape: %s", encoder_hidden_states.shape)
-                logger.info("timestep_proj.shape: %s", timestep_proj.shape)
+                # logger.info("timestep_proj.shape: %s", timestep_proj.shape)
                 # logger.info("freqs_cis.shape: %s", freqs_cis.shape)
                 # logger.info("block_mask.shape: %s", self.block_mask.shape)
                 hidden_states = block(hidden_states, encoder_hidden_states,

fastvideo/models/utils.py

@@ -149,9 +149,6 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor,
     """
     Convert predicted noise to clean latent.
     """
-    logger.info(f"timestep: {timestep.shape}")
-    logger.info(f"noise_input_latent: {noise_input_latent.shape}")
-    logger.info(f"pred_noise: {pred_noise.shape}")
     timestep = timestep.expand(noise_input_latent.shape[0])
     dtype = pred_noise.dtype
     device = pred_noise.device

fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py

@@ -26,20 +26,142 @@ 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, save_decoded_latents_as_video
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.preprocess.preprocess_pipeline_base import (
     BasePreprocessPipeline)
-from fastvideo.pipelines.stages import (DenoisingStage, ImageVAEEncodingStage,
+from fastvideo.pipelines.stages import (DecodingStage, DenoisingStage,
+                                        ImageVAEEncodingStage,
                                         InputValidationStage,
                                         LatentPreparationStage,
                                         TextEncodingStage,
-                                        TimestepPreparationStage,
-                                        DecodingStage)
+                                        TimestepPreparationStage)
+from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict
 
 logger = init_logger(__name__)
 
 
+class FlowMatchScheduler:
+
+    order = 1
+
+    def __init__(self,
+                 num_inference_steps=100,
+                 num_train_timesteps=1000,
+                 shift=3.0,
+                 sigma_max=1.0,
+                 sigma_min=0.003 / 1.002,
+                 inverse_timesteps=False,
+                 extra_one_step=False,
+                 reverse_sigmas=False):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.set_timesteps(num_inference_steps)
+
+    def set_timesteps(self,
+                      num_inference_steps=100,
+                      denoising_strength=1.0,
+                      training=False,
+                      device=None):
+        sigma_start = self.sigma_min + \
+            (self.sigma_max - self.sigma_min) * denoising_strength
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min,
+                                         num_inference_steps + 1)[:-1]
+        else:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min,
+                                         num_inference_steps)
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        self.sigmas = self.shift * self.sigmas / \
+            (1 + (self.shift - 1) * self.sigmas)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        if training:
+            x = self.timesteps
+            y = torch.exp(
+                -2 * ((x - num_inference_steps / 2) / num_inference_steps)**2)
+            y_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * \
+                (num_inference_steps / y_shifted.sum())
+            self.linear_timesteps_weights = bsmntw_weighing
+
+    def step(self,
+             model_output,
+             timestep,
+             sample,
+             to_final=False,
+             return_dict=False,
+             **kwargs):
+        assert return_dict is False
+        assert kwargs == {}
+        self.sigmas = self.sigmas.to(model_output.device)
+        self.timesteps = self.timesteps.to(model_output.device)
+        logger.info('step timestep: %s', timestep)
+        logger.info('step timestep: %s', timestep.shape)
+        # timestep  is [num_frames]
+        # timestep_id = torch.argmin(
+        #     (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        # assert timestep.ndim == 1
+        # assert timestep.shape[0] == 1
+        timestep_id = torch.argmin((self.timesteps - timestep).abs(), dim=0)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
+            sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        return (prev_sample, )
+
+    def scale_model_input(self, sample: torch.Tensor, *args,
+                          **kwargs) -> torch.Tensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.Tensor`):
+                The input sample.
+
+        Returns:
+            `torch.Tensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    def add_noise(self, original_samples, noise, timestep):
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B, C, H, W]
+            - noise: the noise with shape [B, C, H, W]
+            - timestep: the timestep with shape [B]
+        Output: the corrupted latent with shape [B, C, H, W]
+        """
+        self.sigmas = self.sigmas.to(noise.device)
+        self.timesteps = self.timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample.type_as(noise)
+
+    def training_target(self, sample, noise, timestep):
+        target = noise - sample
+        return target
+
+    def training_weight(self, timestep):
+        timestep_id = torch.argmin(
+            (self.timesteps - timestep.to(self.timesteps.device)).abs())
+        weights = self.linear_timesteps_weights[timestep_id]
+        return weights
+
+
 class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
     """ODE Trajectory preprocessing pipeline implementation."""
 
@@ -55,6 +177,22 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
 
     def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
         """Set up pipeline stages with proper dependency injection."""
+        fastvideo_args.pipeline_config.flow_shift = 5
+        logger.info('WTF flow_shift: %s',
+                    fastvideo_args.pipeline_config.flow_shift)
+
+        assert fastvideo_args.pipeline_config.flow_shift == 5
+        # self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
+        # shift=fastvideo_args.pipeline_config.flow_shift)
+        self.modules["scheduler"] = FlowMatchScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift,
+            sigma_min=0.0,
+            extra_one_step=True)
+        self.modules["scheduler"].set_timesteps(num_inference_steps=48,
+                                                denoising_strength=1.0)
+        logger.info('WTF scheduler timesteps: %s',
+                    self.modules["scheduler"].timesteps)
+
         self.add_stage(stage_name="input_validation_stage",
                        stage=InputValidationStage())
         self.add_stage(stage_name="prompt_encoding_stage",
@@ -124,6 +262,7 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                     valid_data, fastvideo_args)
 
                 batch_captions = valid_data["text"]
+                logger.info(f"===== batch_captions: {batch_captions}")
                 # Encode text using the standalone TextEncodingStage API
                 prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text(
                     batch_captions,
@@ -156,8 +295,7 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                 # logger.info(f"===== prompt_embeds: {prompt_embeds[0].shape}")
                 # logger.info(f"===== prompt_attention_masks: {prompt_attention_masks[0].shape}")
 
-                sampling_params = SamplingParam.from_pretrained(
-                    args.model_path)
+                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:
@@ -168,7 +306,8 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                         return_attention_mask=True,
                     )
                     negative_prompt_embed = negative_prompt_embeds_list[0][0]
-                    negative_prompt_attention_mask = negative_prompt_masks_list[0][0]
+                    negative_prompt_attention_mask = negative_prompt_masks_list[
+                        0][0]
                 else:
                     negative_prompt_embed = None
                     negative_prompt_attention_mask = None
@@ -176,12 +315,15 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                 trajectory_latents = []
                 trajectory_timesteps = []
                 trajectory_decoded = []
-                for i, (prompt_embed, prompt_attention_mask) in enumerate(zip(prompt_embeds, prompt_attention_masks)):
+                for i, (prompt_embed, prompt_attention_mask) in enumerate(
+                        zip(prompt_embeds, prompt_attention_masks, strict=False)):
                     prompt_embed = prompt_embed.unsqueeze(0)
                     prompt_attention_mask = prompt_attention_mask.unsqueeze(0)
-                    logger.info(f"what")
+                    logger.info("what")
                     logger.info(f"===== prompt_embed: {prompt_embed.shape}")
-                    logger.info(f"===== prompt_attention_mask: {prompt_attention_mask.shape}")
+                    logger.info(
+                        f"===== prompt_attention_mask: {prompt_attention_mask.shape}"
+                    )
                     # Collect the trajectory data
                     batch = ForwardBatch(
                         **shallow_asdict(sampling_params),
@@ -201,14 +343,17 @@ 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.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_inference_steps = 48
                     # batch.num_frames = 81
                     batch.fps = args.train_fps
-                    batch.guidance_scale = 3.0
+                    batch.guidance_scale = 6.0
                     batch.do_classifier_free_guidance = True
                     # fastvideo_args.pipeline_config.ti2v_task = True
 
@@ -222,25 +367,36 @@ 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)
+                    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_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
-            logger.info(f"===== trajectory_latents: {trajectory_latents[0].shape}")
-            logger.info(f"===== trajectory_latents len: {len(trajectory_latents)}")
+            logger.info(
+                f"===== trajectory_latents: {trajectory_latents[0].shape}")
+            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)}")
+            logger.info(
+                f"===== trajectory_timesteps len: {len(trajectory_timesteps)}")
 
             if batch.return_trajectory_decoded:
-                logger.info(f"===== SAVING TRAJECTORY DECODED")
+                logger.info("===== 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)
+                        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 = []
@@ -271,11 +427,12 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                         else:
                             assert isinstance(value, list)
                             if isinstance(value[idx], torch.Tensor):
-                                logger.info(f"===== value in list: {value[idx].shape}")
-                                sample_extra_features[key] = value[idx].cpu().float().numpy(
-                                )
+                                logger.info(
+                                    f"===== value in list: {value[idx].shape}")
+                                sample_extra_features[key] = value[idx].cpu(
+                                ).float().numpy()
                             else:
-                                logger.info(f"===== value in list: not tensor")
+                                logger.info("===== value in list: not tensor")
                                 sample_extra_features[key] = value[idx]
                             # logger.info(f"===== value: not tensor")
                             # sample_extra_features[key] = value[idx]
@@ -412,9 +569,12 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
         if extra_features and "trajectory_latents" in extra_features:
             trajectory_latents = extra_features["trajectory_latents"]
             record.update({
-                "trajectory_latents_bytes": trajectory_latents.tobytes(),
-                "trajectory_latents_shape": list(trajectory_latents.shape),
-                "trajectory_latents_dtype": str(trajectory_latents.dtype),
+                "trajectory_latents_bytes":
+                trajectory_latents.tobytes(),
+                "trajectory_latents_shape":
+                list(trajectory_latents.shape),
+                "trajectory_latents_dtype":
+                str(trajectory_latents.dtype),
             })
         else:
             record.update({
@@ -426,9 +586,12 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
         if extra_features and "trajectory_timesteps" in extra_features:
             trajectory_timesteps = extra_features["trajectory_timesteps"]
             record.update({
-                "trajectory_timesteps_bytes": trajectory_timesteps.tobytes(),
-                "trajectory_timesteps_shape": list(trajectory_timesteps.shape),
-                "trajectory_timesteps_dtype": str(trajectory_timesteps.dtype),
+                "trajectory_timesteps_bytes":
+                trajectory_timesteps.tobytes(),
+                "trajectory_timesteps_shape":
+                list(trajectory_timesteps.shape),
+                "trajectory_timesteps_dtype":
+                str(trajectory_timesteps.dtype),
             })
         else:
             record.update({

fastvideo/pipelines/preprocess/v1_preprocess.py

@@ -27,6 +27,7 @@ def main(args) -> None:
     kwargs = {
         "vae_precision": "fp32",
         "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True),
+        "flow_shift": 5,
     }
     pipeline_config.update_config_from_dict(kwargs)
     fastvideo_args = FastVideoArgs(

fastvideo/training/ode_causal_pipeline.py

@@ -7,16 +7,17 @@ import torch
 import torch.nn.functional as F
 
 from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory
+from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler)
-from fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline import WanCausalDMDPipeline
+from fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline import (
+    WanCausalDMDPipeline)
 from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (
     clip_grad_norm_while_handling_failing_dtensor_cases)
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.forward_context import set_forward_context
 
 logger = init_logger(__name__)
 
@@ -49,22 +50,39 @@ class ODEInitTrainingPipeline(TrainingPipeline):
         self.vae.requires_grad_(False)
 
         self.timestep_shift = self.training_args.pipeline_config.flow_shift
+        assert self.timestep_shift == 5.0, "flow_shift must be 5.0"
         self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
             shift=self.timestep_shift)
 
-        self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250], dtype=torch.long, device=get_local_torch_device())
+        # logger.info(f"ARG dmd_denoising_steps: {training_args.pipeline_config.dmd_denoising_steps}")
+        logger.info(
+            f"ARG dmd_denoising_steps: {self.training_args.pipeline_config.dmd_denoising_steps}"
+        )
+        self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 0],
+                                                dtype=torch.long,
+                                                device=get_local_torch_device())
+        # self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250], dtype=torch.long, device=get_local_torch_device())
         if training_args.warp_denoising_step:  # Warp the denoising step according to the scheduler time shift
-            timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))).cuda()
+            timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(),
+                                   torch.tensor([0],
+                                                dtype=torch.float32))).cuda()
             logger.info(f"timesteps: {timesteps}")
-            self.dmd_denoising_steps = timesteps[1000 - self.dmd_denoising_steps]
-            logger.info(f"warped self.dmd_denoising_steps: {self.dmd_denoising_steps}")
+            self.dmd_denoising_steps = timesteps[1000 -
+                                                 self.dmd_denoising_steps]
+            logger.info(
+                f"warped self.dmd_denoising_steps: {self.dmd_denoising_steps}")
+            assert False, "warp_denoising_step must be false"
         else:
-            logger.info(f"not warped")
-        self.dmd_denoising_steps = self.dmd_denoising_steps.to(get_local_torch_device())
+            # assert False, "warp_denoising_step must be true"
+            logger.info("not warped")
+        self.dmd_denoising_steps = self.dmd_denoising_steps.to(
+            get_local_torch_device())
 
         logger.info(f"denoising_step_list: {self.dmd_denoising_steps}")
 
-        logger.info("Initialized ODE-init training pipeline with %s denoising steps", len(self.dmd_denoising_steps))
+        logger.info(
+            "Initialized ODE-init training pipeline with %s denoising steps",
+            len(self.dmd_denoising_steps))
         # Cache for nearest trajectory index per DMD step (computed lazily on first batch)
         self._cached_closest_idx_per_dmd = None
         self.num_train_timestep = self.noise_scheduler.num_train_timesteps
@@ -92,7 +110,6 @@ class ODEInitTrainingPipeline(TrainingPipeline):
             pin_cpu_memory=training_args.pin_cpu_memory,
             dit_cpu_offload=True)
 
-
     def _get_next_batch(self, training_batch):  # type: ignore[override]
         batch = next(self.train_loader_iter, None)  # type: ignore
         if batch is None:
@@ -116,7 +133,8 @@ class ODEInitTrainingPipeline(TrainingPipeline):
             pass
         else:
             raise ValueError(
-                f"Unexpected trajectory_latents dim: {trajectory_latents.dim()}")
+                f"Unexpected trajectory_latents dim: {trajectory_latents.dim()}"
+            )
 
         trajectory_timesteps = batch['trajectory_timesteps']
         if trajectory_timesteps.dim() == 3:
@@ -127,7 +145,8 @@ class ODEInitTrainingPipeline(TrainingPipeline):
             pass
         else:
             raise ValueError(
-                f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}")
+                f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}"
+            )
         # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing
         trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5)
 
@@ -139,46 +158,42 @@ class ODEInitTrainingPipeline(TrainingPipeline):
             device, dtype=torch.bfloat16)
         training_batch.infos = infos
 
-        return training_batch, trajectory_latents.to(device, dtype=torch.bfloat16), trajectory_timesteps.to(device)
-    def _get_timestep(self, 
-            min_timestep: int,
-            max_timestep: int,
-            batch_size: int,
-            num_frame: int,
-            num_frame_per_block: int,
-            uniform_timestep: bool = False
-    ) -> torch.Tensor:
+        return training_batch, trajectory_latents.to(
+            device, dtype=torch.bfloat16), trajectory_timesteps.to(device)
+
+    def _get_timestep(self,
+                      min_timestep: int,
+                      max_timestep: int,
+                      batch_size: int,
+                      num_frame: int,
+                      num_frame_per_block: int,
+                      uniform_timestep: bool = False) -> torch.Tensor:
         if uniform_timestep:
-            timestep = torch.randint(
-                min_timestep,
-                max_timestep,
-                [batch_size, 1],
-                device=self.device,
-                dtype=torch.long
-            ).repeat(1, num_frame)
+            timestep = torch.randint(min_timestep,
+                                     max_timestep, [batch_size, 1],
+                                     device=self.device,
+                                     dtype=torch.long).repeat(1, num_frame)
             return timestep
         else:
-            timestep = torch.randint(
-                min_timestep,
-                max_timestep,
-                [batch_size, num_frame],
-                device=self.device,
-                dtype=torch.long
-            )
-            logger.info(f"individual timestep: {timestep}")
+            timestep = torch.randint(min_timestep,
+                                     max_timestep, [batch_size, num_frame],
+                                     device=self.device,
+                                     dtype=torch.long)
+            # logger.info(f"individual timestep: {timestep}")
             # make the noise level the same within every block
-            timestep = timestep.reshape(
-                timestep.shape[0], -1, num_frame_per_block)
+            timestep = timestep.reshape(timestep.shape[0], -1,
+                                        num_frame_per_block)
             timestep[:, :, 1:] = timestep[:, :, 0:1]
             timestep = timestep.reshape(timestep.shape[0], -1)
             return timestep
 
-    def _step_predict_next_latent(self, traj_latents: torch.Tensor,
-                                   traj_timesteps: torch.Tensor,
-                                   encoder_hidden_states: torch.Tensor,
-                                   encoder_attention_mask: torch.Tensor
-                                   ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def _step_predict_next_latent(
+        self, traj_latents: torch.Tensor, traj_timesteps: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_attention_mask: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         device = get_local_torch_device()
+        target_latent = traj_latents[:, -1]
 
         # logger.info(f"traj_latents: {traj_latents.shape}")
         # logger.info(f"traj_timesteps: {traj_timesteps.shape}")
@@ -189,39 +204,42 @@ class ODEInitTrainingPipeline(TrainingPipeline):
         # Lazily cache nearest trajectory index per DMD step based on the (fixed) S timesteps
         if self._cached_closest_idx_per_dmd is None:
             # Use the first sample's trajectory timesteps; assumed identical across batches
-            s_steps = traj_timesteps[0].to(torch.long)  # [S]
-            dmd = cast(torch.Tensor, self.dmd_denoising_steps).to(s_steps.device)  # [K]
+            # s_steps = traj_timesteps[0].to(torch.long)  # [S]
+            # dmd = cast(torch.Tensor, self.dmd_denoising_steps).to(s_steps.device)  # [K]
             # distances_ks: [K, S] = |s_steps - dmd|
-            distances_ks = (s_steps.unsqueeze(0) - dmd.unsqueeze(1)).abs()
-            self._cached_closest_idx_per_dmd = distances_ks.argmin(dim=1).to(torch.long).cpu()  # [K]
+            # distances_ks = (s_steps.unsqueeze(0) - dmd.unsqueeze(1)).abs()
+            # self._cached_closest_idx_per_dmd = distances_ks.argmin(dim=1).to(torch.long).cpu()  # [K]
+            self._cached_closest_idx_per_dmd = torch.tensor(
+                [0, 12, 24, 36], dtype=torch.long).cpu()
             # logger.info(f"self._cached_closest_idx_per_dmd: {self._cached_closest_idx_per_dmd}")
+            # logger.info(f"corresponding timesteps: {self.noise_scheduler.timesteps[self._cached_closest_idx_per_dmd]}")
 
         # logger.info(f"traj_latents: {traj_latents.shape}")
         # Select the K indexes from traj_latents using self._cached_closest_idx_per_dmd
         # traj_latents: [B, S, C, T, H, W], self._cached_closest_idx_per_dmd: [K]
         # Output: [B, K, C, T, H, W]
         relevant_traj_latents = torch.index_select(
-            traj_latents, dim=1, index=self._cached_closest_idx_per_dmd.to(traj_latents.device)
-        )
+            traj_latents,
+            dim=1,
+            index=self._cached_closest_idx_per_dmd.to(traj_latents.device))
         # logger.info(f"relevant_traj_latents: {relevant_traj_latents.shape}")
-        target_latent = relevant_traj_latents[:, -1]
         # assert relevant_traj_latents.shape[0] == 1
 
-        indexes = self._get_timestep( # [B, num_frames]
+        indexes = self._get_timestep(  # [B, num_frames]
             0,
             len(self.dmd_denoising_steps),
             B,
             num_frames,
             3,
-            uniform_timestep=False
-        )
+            uniform_timestep=False)
         # noisy_input = relevant_traj_latents[indexes]
         # logger.info(f"indexes: {indexes.shape}")
         noisy_input = torch.gather(
-            relevant_traj_latents, dim=1,
-            index=indexes.reshape(B, 1, num_frames, 1, 1, 1).expand(
-                -1, -1, -1, num_channels, height, width).to(self.device)
-        ).squeeze(1)
+            relevant_traj_latents,
+            dim=1,
+            index=indexes.reshape(B, 1, num_frames, 1, 1,
+                                  1).expand(-1, -1, -1, num_channels, height,
+                                            width).to(self.device)).squeeze(1)
         # noisy_input = noisy_input.unsqueeze(0)
 
         # # Sample a single DMD step for the whole batch and fetch its cached nearest S-index
@@ -255,8 +273,8 @@ class ODEInitTrainingPipeline(TrainingPipeline):
         }
         # Predict noise and step the scheduler to obtain next latent
         with set_forward_context(current_timestep=timestep,
-                            attn_metadata=None,
-                            forward_batch=None):
+                                 attn_metadata=None,
+                                 forward_batch=None):
             noise_pred = self.transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
             # logger.info(f"noise_pred: {noise_pred.shape}")
         if isinstance(noise_pred, (tuple, list)):
@@ -274,7 +292,7 @@ class ODEInitTrainingPipeline(TrainingPipeline):
         # next_latent_pred = self.modules["scheduler"].step(
         #     noise_pred, t, current_latents, return_dict=False)[0]
         return noise_pred, target_latent, timestep
-    
+
     def train_one_step(self, training_batch):  # type: ignore[override]
         self.transformer.train()
         self.optimizer.zero_grad()
@@ -283,10 +301,9 @@ class ODEInitTrainingPipeline(TrainingPipeline):
 
         # Using cached nearest index per DMD step; computation happens in _step_predict_next_latent
 
-
-
         for _ in range(args.gradient_accumulation_steps):
-            training_batch, traj_latents, traj_timesteps = self._get_next_batch(training_batch)
+            training_batch, traj_latents, traj_timesteps = self._get_next_batch(
+                training_batch)
             text_embeds = training_batch.encoder_hidden_states
             text_attention_mask = training_batch.encoder_attention_mask
             assert traj_latents.shape[0] == 1
@@ -298,7 +315,6 @@ class ODEInitTrainingPipeline(TrainingPipeline):
 
             # Sample per-sample current step i in [0, S-2]
 
-
             # idx = torch.randint(low=0, high=S - 1, size=(B, ),
             #                     device=traj_latents.device)
 
@@ -314,20 +330,20 @@ class ODEInitTrainingPipeline(TrainingPipeline):
             #     t = t.long()
 
             # Forward to predict next latent by stepping scheduler with predicted noise
-            noise_pred, target_latent, t = self._step_predict_next_latent(traj_latents,
-                                                                traj_timesteps,
-                                                              text_embeds,
-                                                              text_attention_mask)
+            noise_pred, target_latent, t = self._step_predict_next_latent(
+                traj_latents, traj_timesteps, text_embeds, text_attention_mask)
 
-            mask = t != 0                                                            
+            mask = t != 0
 
             # Compute loss
-            loss = F.mse_loss(noise_pred[mask], target_latent[mask], reduction="mean")
+            loss = F.mse_loss(noise_pred[mask],
+                              target_latent[mask],
+                              reduction="mean")
             loss = loss / args.gradient_accumulation_steps
 
             with set_forward_context(current_timestep=t,
-                            attn_metadata=None,
-                            forward_batch=None):
+                                     attn_metadata=None,
+                                     forward_batch=None):
                 loss.backward()
             avg_loss = loss.detach().clone()
             training_batch.total_loss += avg_loss.item()
@@ -356,6 +372,7 @@ class ODEInitTrainingPipeline(TrainingPipeline):
 
 def main(args) -> None:
     logger.info("Starting ODE-init training pipeline...")
+    logger.info(f"ARG dmd_denoising_steps: {args.dmd_denoising_steps}")
     pipeline = ODEInitTrainingPipeline.from_pretrained(
         args.pretrained_model_name_or_path, args=args)
     args = pipeline.training_args
@@ -373,5 +390,3 @@ if __name__ == "__main__":
     args = parser.parse_args()
     args.dit_cpu_offload = False
     main(args)
-
-