diff --git a/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/README.md b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/README.md new file mode 100644 index 000000000..5d0324253 --- /dev/null +++ b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/README.md @@ -0,0 +1,16 @@ +# Wan2.1-I2V-1.3B-InP Crush-Smol Example +These are e2e example scripts for finetuning Wan2.1 T2V 1.3B InP on the crush-smol dataset. + +## Execute the following commands from `FastVideo/` to run training: + +### Download crush-smol dataset: + +`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh` + +### Preprocess the videos and captions into latents: + +`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh` + +### Edit the following file and run finetuning: + +`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh` diff --git a/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/distill_i2v.sh b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/distill_i2v.sh new file mode 100644 index 000000000..efb9d3923 --- /dev/null +++ b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/distill_i2v.sh @@ -0,0 +1,111 @@ +#!/bin/bash + +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export TOKENIZERS_PARALLELISM=false +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers" +DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/" +VALIDATION_DATASET_FILE="examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json" +NUM_GPUS=8 +export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde' +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "wan_i2v_distill" + --output_dir "checkpoints/wan_i2v_distill" + --wandb_run_name "wan_i2v_distill" + --max_train_steps 2000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 2 + --num_latent_t 8 + --num_height 480 + --num_width 832 + --num_frames 77 + # --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 + --hsdp_shard_dim 8 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 1 +) + +# Validation arguments +validation_args=( + --log_validation + # --log_visualization + --log_visualization + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 50 + --validation_sampling_steps "3" + --validation_guidance_scale "1.0" # not used for dmd inference +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 4e-6 + --lr_scheduler "constant" + --fake_score_learning_rate 8e-7 + --fake_score_lr_scheduler "constant" + --mixed_precision "bf16" + --training_state_checkpointing_steps 2000 + --weight_only_checkpointing_steps 2000 + --weight_decay 0.01 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.0 + --dit_precision "fp32" + --ema_start_step 0 + --flow_shift 8 + --seed 1000 + # --enable_gradient_checkpointing_type "full" +) + +dmd_args=( + --dmd_denoising_steps '1000,757,522' + --min_timestep_ratio 0.02 + --max_timestep_ratio 0.98 + --generator_update_interval 5 + --simulate_generator_forward + --real_score_guidance_scale 3.5 + --VSA_sparsity 0.8 +) + +# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t +torchrun \ + --nnodes 1 \ + --nproc_per_node $NUM_GPUS \ + fastvideo/training/wan_i2v_distillation_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ + "${dmd_args[@]}" \ No newline at end of file diff --git a/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/download_dataset.sh b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/download_dataset.sh new file mode 100644 index 000000000..f2942b7da --- /dev/null +++ b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/download_dataset.sh @@ -0,0 +1,3 @@ +#!/bin/bash + +python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset" \ No newline at end of file diff --git a/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm new file mode 100644 index 000000000..cfd5226b5 --- /dev/null +++ b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm @@ -0,0 +1,129 @@ +#!/bin/bash +#SBATCH --job-name=i2v +#SBATCH --partition=main +#SBATCH --nodes=4 +#SBATCH --ntasks=4 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=i2v_output/i2v_%j.out +#SBATCH --error=i2v_output/i2v_%j.err +#SBATCH --exclusive +set -e -x + +# Environment Setup +source ~/conda/miniconda/bin/activate +conda activate will-fv + +# Basic Info +export WANDB_MODE="online" +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +# different cache dir for different processes +export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export NODE_RANK=$SLURM_PROCID +nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) +export MASTER_ADDR=${nodes[0]} +export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID +export TOKENIZERS_PARALLELISM=false +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +echo "MASTER_ADDR: $MASTER_ADDR" +echo "NODE_RANK: $NODE_RANK" + +# Configs +NUM_GPUS=8 + +MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers" +DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/" +VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json" +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t + +# Training arguments +training_args=( + --tracker_project_name wan_i2v_finetune + --output_dir "checkpoints/wan_i2v_finetune" + --max_train_steps 2000 + --train_batch_size 2 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 8 + --num_height 480 + --num_width 832 + --num_frames 77 + --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size $NUM_GPUS + --tp_size $NUM_GPUS + --hsdp_replicate_dim $SLURM_JOB_NUM_NODES + --hsdp_shard_dim $NUM_GPUS +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 10 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 100 + --validation_sampling_steps "40" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 1e-5 + --mixed_precision "bf16" + --checkpointing_steps 1000 + --weight_decay 1e-4 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 + --enable_gradient_checkpointing_type "full" +) + +srun torchrun \ +--nnodes $SLURM_JOB_NUM_NODES \ +--nproc_per_node $NUM_GPUS \ +--node_rank $SLURM_PROCID \ +--rdzv_backend=c10d \ +--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + fastvideo/training/wan_i2v_training_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ No newline at end of file diff --git a/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/preprocess_wan_data_i2v.sh b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/preprocess_wan_data_i2v.sh new file mode 100644 index 000000000..2d508c964 --- /dev/null +++ b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/preprocess_wan_data_i2v.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +GPU_NUM=1 # 2,4,8 +MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers" +MODEL_TYPE="wan" +DATA_MERGE_PATH="data/crush-smol/merge.txt" +OUTPUT_DIR="data/crush-smol_processed_i2v_1_3b_inp/" + +torchrun --nproc_per_node=$GPU_NUM \ + fastvideo/pipelines/preprocess/v1_preprocess.py \ + --model_path $MODEL_PATH \ + --data_merge_path $DATA_MERGE_PATH \ + --preprocess_video_batch_size 8 \ + --seed 42 \ + --max_height 480 \ + --max_width 832 \ + --num_frames 77 \ + --dataloader_num_workers 0 \ + --output_dir=$OUTPUT_DIR \ + --train_fps 16 \ + --samples_per_file 8 \ + --flush_frequency 8 \ + --video_length_tolerance_range 5 \ + --preprocess_task "i2v" \ No newline at end of file diff --git a/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json new file mode 100644 index 000000000..592166e98 --- /dev/null +++ b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json @@ -0,0 +1,31 @@ +{ + "data": [ + { + "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.", + "image_path": null, + "video_path": "validation_dataset/yYcK4nANZz4-Scene-034.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.", + "image_path": null, + "video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.", + "image_path": null, + "video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + } + ] +} \ No newline at end of file diff --git a/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4 b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4 new file mode 100644 index 000000000..8e416c6b0 Binary files /dev/null and b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4 differ diff --git a/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4 b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4 new file mode 100644 index 000000000..c2c781b51 Binary files /dev/null and b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4 differ diff --git a/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4 b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4 new file mode 100644 index 000000000..abb42a2a8 Binary files /dev/null and b/examples/distill/Wan2.1-Fun-1.3B-InP/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4 differ diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh index 11c389a76..dd060032c 100644 --- a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh +++ b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh @@ -9,17 +9,19 @@ MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers" DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/" VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json" NUM_GPUS=8 +export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde' # export CUDA_VISIBLE_DEVICES=4,5 # IP=[MASTER NODE IP] # Training arguments training_args=( --tracker_project_name "wan_i2v_finetune" - --output_dir "$DATA_DIR/outputs/wan_i2v_finetune" + --output_dir "outputs/wan_i2v_finetune" + --wandb_run_name "wan_i2v_finetune" --max_train_steps 2000 - --train_batch_size 4 + --train_batch_size 1 --train_sp_batch_size 1 - --gradient_accumulation_steps 1 + --gradient_accumulation_steps 2 --num_latent_t 8 --num_height 480 --num_width 832 @@ -30,10 +32,10 @@ training_args=( # Parallel arguments parallel_args=( --num_gpus $NUM_GPUS - --sp_size 4 - --tp_size 4 - --hsdp_replicate_dim 2 - --hsdp_shard_dim 4 + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 + --hsdp_shard_dim 8 ) # Model arguments @@ -59,7 +61,7 @@ validation_args=( # Optimizer arguments optimizer_args=( - --learning_rate 2e-5 + --learning_rate 1e-5 --mixed_precision "bf16" --checkpointing_steps 2000 --weight_decay 1e-4 diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index cd8fb9938..de68a7d45 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -34,6 +34,7 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = { "Wan-AI/Wan2.2-TI2V-5B-Diffusers": WanT2V720PConfig, "Wan-AI/Wan2.2-T2V-A14B-Diffusers": WanT2V480PConfig, "Wan-AI/Wan2.2-I2V-A14B-Diffusers": WanI2V480PConfig, + "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig, # Add other specific weight variants } diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py index 740cbbde9..5fd4e0ac9 100644 --- a/fastvideo/configs/sample/registry.py +++ b/fastvideo/configs/sample/registry.py @@ -17,6 +17,7 @@ from fastvideo.configs.sample.wan import ( WanI2V_14B_720P_SamplingParam, WanT2V_1_3B_SamplingParam, WanT2V_14B_SamplingParam, + Wan2_1_Fun_1_3B_InP_SamplingParam, ) # isort: on from fastvideo.logger import init_logger @@ -38,6 +39,8 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = { "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam, "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam, + "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": + Wan2_1_Fun_1_3B_InP_SamplingParam, # Add other specific weight variants } diff --git a/fastvideo/configs/sample/wan.py b/fastvideo/configs/sample/wan.py index 23cb0cdbf..70e4e3d9c 100644 --- a/fastvideo/configs/sample/wan.py +++ b/fastvideo/configs/sample/wan.py @@ -107,13 +107,28 @@ class FastWanT2V480PConfig(WanT2V_1_3B_SamplingParam): fps: int = 16 +# ============================================= +# ============= Wan2.1 Fun Models ============= +# ============================================= +@dataclass +class Wan2_1_Fun_1_3B_InP_SamplingParam(SamplingParam): + """Sampling parameters for Wan2.1 Fun 1.3B InP model.""" + height: int = 480 + width: int = 832 + num_frames: int = 81 + fps: int = 16 + negative_prompt: str = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" + guidance_scale: float = 6.0 + num_inference_steps: int = 50 + + # ============================================= # ============= Wan2.2 TI2V Models ============= # ============================================= @dataclass class Wan2_2_Base_SamplingParam(SamplingParam): """Sampling parameters for Wan2.2 TI2V 5B model.""" - negative_prompt: str | None = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" + negative_prompt: str = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" @dataclass diff --git a/fastvideo/pipelines/basic/wan/wan_i2v_dmd_pipeline.py b/fastvideo/pipelines/basic/wan/wan_i2v_dmd_pipeline.py index 1bcc57c93..71922301f 100644 --- a/fastvideo/pipelines/basic/wan/wan_i2v_dmd_pipeline.py +++ b/fastvideo/pipelines/basic/wan/wan_i2v_dmd_pipeline.py @@ -12,12 +12,10 @@ from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.lora_pipeline import LoRAPipeline # isort: off -from fastvideo.pipelines.stages import (ImageEncodingStage, ConditioningStage, - DecodingStage, DmdDenoisingStage, - EncodingStage, InputValidationStage, - LatentPreparationStage, - TextEncodingStage, - TimestepPreparationStage) +from fastvideo.pipelines.stages import ( + ImageEncodingStage, ConditioningStage, DecodingStage, DmdDenoisingStage, + ImageVAEEncodingStage, InputValidationStage, LatentPreparationStage, + TextEncodingStage, TimestepPreparationStage) # isort: on from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler) @@ -67,7 +65,7 @@ class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase): transformer=self.get_module("transformer"))) self.add_stage(stage_name="image_latent_preparation_stage", - stage=EncodingStage(vae=self.get_module("vae"))) + stage=ImageVAEEncodingStage(vae=self.get_module("vae"))) self.add_stage(stage_name="denoising_stage", stage=DmdDenoisingStage( diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 8a2d9f0b6..24aa4ddf7 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -197,6 +197,7 @@ class DenoisingStage(PipelineStage): # Run denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): + logger.info('timesteps: %s', timesteps) # Skip if interrupted if hasattr(self, 'interrupt') and self.interrupt: continue @@ -682,10 +683,19 @@ class DmdDenoisingStage(DenoisingStage): video_raw_latent_shape = latents.shape prompt_embeds = batch.prompt_embeds assert torch.isnan(prompt_embeds[0]).sum() == 0 - timesteps = torch.tensor( - fastvideo_args.pipeline_config.dmd_denoising_steps, - dtype=torch.long, - device=get_local_torch_device()) + timesteps_50 = [ + 999, 993, 986, 979, 971, 964, 956, 948, 940, 931, 922, 913, 904, + 895, 885, 874, 864, 853, 841, 830, 818, 805, 792, 778, 764, 749, + 734, 718, 702, 684, 666, 647, 627, 607, 585, 562, 538, 513, 486, + 458, 428, 396, 363, 328, 290, 249, 206, 160, 111, 57 + ] + timesteps = torch.tensor(timesteps_50, + dtype=torch.long, + device=get_local_torch_device()) + # timesteps = torch.tensor( + # fastvideo_args.pipeline_config.dmd_denoising_steps, + # dtype=torch.long, + # device=get_local_torch_device()) # Handle sequence parallelism if enabled sp_world_size, rank_in_sp_group = get_sp_world_size( diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 970e3225d..c4de3dec4 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -802,7 +802,7 @@ class DistillationPipeline(TrainingPipeline): dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict fake_score_latents_vis_dict = training_batch.fake_score_latent_vis_dict fake_score_log_keys = ['generator_pred_video'] - dmd_log_keys = ['faker_score_pred_video'] + dmd_log_keys = ['faker_score_pred_video', 'real_score_pred_video'] for latent_key in fake_score_log_keys: latents = fake_score_latents_vis_dict[latent_key]