diff --git a/examples/inference/basic/basic_self_forcing_causal.py b/examples/inference/basic/basic_self_forcing_causal.py index 9d2140f23..56ffe527f 100644 --- a/examples/inference/basic/basic_self_forcing_causal.py +++ b/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() diff --git a/examples/training/consistency_finetune/ode_init/finetune_ode_init_70k_vidprom.slurm b/examples/training/consistency_finetune/ode_init/finetune_ode_init_16k_vidprom.slurm similarity index 86% rename from examples/training/consistency_finetune/ode_init/finetune_ode_init_70k_vidprom.slurm rename to examples/training/consistency_finetune/ode_init/finetune_ode_init_16k_vidprom.slurm index 78dc3082e..30bdfe1f6 100644 --- a/examples/training/consistency_finetune/ode_init/finetune_ode_init_70k_vidprom.slurm +++ b/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 ) diff --git a/examples/training/consistency_finetune/ode_init/finetune_ode_init_crush.slurm b/examples/training/consistency_finetune/ode_init/finetune_ode_init_crush.slurm index 38d72de13..0688eefd6 100644 --- a/examples/training/consistency_finetune/ode_init/finetune_ode_init_crush.slurm +++ b/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" ) diff --git a/examples/training/consistency_finetune/ode_init/finetune_ode_init_overfit.sh b/examples/training/consistency_finetune/ode_init/finetune_ode_init_overfit.sh index bbc3b912a..02f7a33af 100644 --- a/examples/training/consistency_finetune/ode_init/finetune_ode_init_overfit.sh +++ b/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 diff --git a/examples/training/consistency_finetune/ode_init/preprocess_wan_data_ode.sh b/examples/training/consistency_finetune/ode_init/preprocess_wan_data_ode.sh index e2777268e..a1b223666 100644 --- a/examples/training/consistency_finetune/ode_init/preprocess_wan_data_ode.sh +++ b/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" diff --git a/fastvideo/configs/pipelines/wan.py b/fastvideo/configs/pipelines/wan.py index d1e27a66c..ce30a8d13 100644 --- a/fastvideo/configs/pipelines/wan.py +++ b/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]) diff --git a/fastvideo/dataset/utils.py b/fastvideo/dataset/utils.py index 943d7acfb..35c651b50 100644 --- a/fastvideo/dataset/utils.py +++ b/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: diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 689f32b78..e50c1801e 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/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, diff --git a/fastvideo/models/utils.py b/fastvideo/models/utils.py index f87551593..7c4347870 100644 --- a/fastvideo/models/utils.py +++ b/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 diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py index 8d575e7fd..7c4a18b9c 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py +++ b/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({ diff --git a/fastvideo/pipelines/preprocess/v1_preprocess.py b/fastvideo/pipelines/preprocess/v1_preprocess.py index 0051b8e68..31dc18a53 100644 --- a/fastvideo/pipelines/preprocess/v1_preprocess.py +++ b/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( diff --git a/fastvideo/training/ode_causal_pipeline.py b/fastvideo/training/ode_causal_pipeline.py index de09bf37a..f9fe16424 100644 --- a/fastvideo/training/ode_causal_pipeline.py +++ b/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) - -