diff --git a/examples/datasets/crush_smol/download_dataset.sh b/examples/datasets/crush_smol/download_dataset.sh new file mode 100644 index 000000000..f2942b7da --- /dev/null +++ b/examples/datasets/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/datasets/crush_smol/validation.json b/examples/datasets/crush_smol/validation.json new file mode 100644 index 000000000..a7e327a04 --- /dev/null +++ b/examples/datasets/crush_smol/validation.json @@ -0,0 +1,76 @@ +{ + "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 + }, + { + "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.", + "image_path": null, + "video_path": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "The video shows a green and orange object being flattened as if it were under a hydraulic press, with the press moving down and compressing the object.", + "image_path": null, + "video_path": "validation_dataset/1gGQy4nxyUo-Scene-056.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "The video shows a cylindrical object with a cityscape image being flattened as if it were under a hydraulic press. The object is placed on a metal platform, and a large, striped cylinder presses down on it, causing it to collapse and release a liquid inside. The background features a green wall with a yellow and red warning sign.", + "image_path": null, + "video_path": "validation_dataset/1gGQy4nxyUo-Scene-059.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "The video shows a close-up of an orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened.", + "image_path": null, + "video_path": "validation_dataset/EJqsC21GSBY-Scene-059.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press.", + "image_path": null, + "video_path": "validation_dataset/GBSfpTcKegk-Scene-003.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + } + ] +} \ No newline at end of file diff --git a/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-016.mp4 b/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-016.mp4 new file mode 100644 index 000000000..cec24bb69 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-016.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-056.mp4 b/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-056.mp4 new file mode 100644 index 000000000..87d4aed09 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-056.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-059.mp4 b/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-059.mp4 new file mode 100644 index 000000000..a0455d1f9 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-059.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/EJqsC21GSBY-Scene-059.mp4 b/examples/datasets/crush_smol/validation_dataset/EJqsC21GSBY-Scene-059.mp4 new file mode 100644 index 000000000..15c45a810 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/EJqsC21GSBY-Scene-059.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/GBSfpTcKegk-Scene-003.mp4 b/examples/datasets/crush_smol/validation_dataset/GBSfpTcKegk-Scene-003.mp4 new file mode 100644 index 000000000..27cab6878 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/GBSfpTcKegk-Scene-003.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4 b/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4 new file mode 100644 index 000000000..8e416c6b0 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-027.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4 b/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4 new file mode 100644 index 000000000..c2c781b51 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4 differ diff --git a/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4 b/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4 new file mode 100644 index 000000000..abb42a2a8 Binary files /dev/null and b/examples/datasets/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4 differ diff --git a/examples/datasets/crush_smol/validation_single.json b/examples/datasets/crush_smol/validation_single.json new file mode 100644 index 000000000..30c3844a5 --- /dev/null +++ b/examples/datasets/crush_smol/validation_single.json @@ -0,0 +1,13 @@ +{ + "data": [ + { + "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.", + "image_path": null, + "video_path": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4", + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + } + ] +} \ No newline at end of file diff --git a/examples/inference/basic/basic.py b/examples/inference/basic/basic.py index f12a4969e..b987b6fb8 100644 --- a/examples/inference/basic/basic.py +++ b/examples/inference/basic/basic.py @@ -1,6 +1,6 @@ from fastvideo import VideoGenerator -# from fastvideo.configs.sample import SamplingParam +from fastvideo.configs.sample import SamplingParam OUTPUT_PATH = "video_samples" def main(): @@ -17,30 +17,29 @@ def main(): vae_cpu_offload=False, text_encoder_cpu_offload=True, pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" + ti2v_task=True, # image_encoder_cpu_offload=False, ) - # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") + sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") # sampling_param.num_frames = 45 - # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" + sampling_param.image_path = "test.jpg" # Generate videos with the same simple API, regardless of GPU count 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." + "A girl is packing a suitcase when stuff suddently starts flying around the room." ) video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True) # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/") # Generate another video with a different prompt, without reloading the # model! - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic.") - video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True) + # prompt2 = ( + # "A majestic lion strides across the golden savanna, its powerful frame " + # "glistening under the warm afternoon sun. The tall grass ripples gently in " + # "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + # "embodying the raw energy of the wild. Low angle, steady tracking shot, " + # "cinematic.") + # video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True) if __name__ == "__main__": diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh index c7e2c251e..379a8a063 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh @@ -6,8 +6,8 @@ export TOKENIZERS_PARALLELISM=false # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/" -VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" +DATA_DIR="data/crush-smol_processed_t2v_old" +VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json" NUM_GPUS=4 # export CUDA_VISIBLE_DEVICES=4,5 @@ -52,7 +52,7 @@ dataset_args=( validation_args=( --log_validation --validation_dataset_file $VALIDATION_DATASET_FILE - --validation_steps 200 + --validation_steps 50 --validation_sampling_steps "50" --validation_guidance_scale "6.0" ) diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v.sh index 2f1203060..d622c9af0 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v.sh +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v.sh @@ -4,7 +4,7 @@ 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/" +OUTPUT_DIR="data/crush-smol_processed_t2v_old/" torchrun --nproc_per_node=$GPU_NUM \ fastvideo/pipelines/preprocess/v1_preprocess.py \ diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh index 262f39ec9..83f587796 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh @@ -1,6 +1,6 @@ #!/bin/bash -GPU_NUM=2 # 2,4,8 +GPU_NUM=1 # 2,4,8 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DATASET_PATH="data/crush-smol/" OUTPUT_DIR="data/crush-smol_processed_t2v/" @@ -13,7 +13,7 @@ torchrun --nproc_per_node=$GPU_NUM \ --preprocess.dataset_type merged \ --preprocess.dataset_path $DATASET_PATH \ --preprocess.dataset_output_dir $OUTPUT_DIR \ - --preprocess.preprocess_video_batch_size 2 \ + --preprocess.preprocess_video_batch_size 8 \ --preprocess.dataloader_num_workers 0 \ --preprocess.max_height 480 \ --preprocess.max_width 832 \ diff --git a/examples/training/finetune/wan_ti2v/finetune_t2v_to_i2v.sh b/examples/training/finetune/wan_ti2v/finetune_t2v_to_i2v.sh new file mode 100644 index 000000000..f84a6532f --- /dev/null +++ b/examples/training/finetune/wan_ti2v/finetune_t2v_to_i2v.sh @@ -0,0 +1,93 @@ +#!/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="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR="data/crush-smol_processed_t2v_i2v_1_3b" +VALIDATION_DATASET_FILE="examples/datasets/crush_smol/validation.json" +NUM_GPUS=8 +# export CUDA_VISIBLE_DEVICES=4,5 + + +# Training arguments +training_args=( + --tracker_project_name "wan_t2v_i2v_finetune" + --output_dir "checkpoints/wan_t2v_i2v_finetune" + --max_train_steps 5000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 2 + --num_latent_t 20 + --num_height 480 + --num_width 832 + --num_frames 77 + --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 4 + --tp_size 1 + --hsdp_replicate_dim 2 + --hsdp_shard_dim 4 +) + +# 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 + --validation_dataset_file $VALIDATION_DATASET_FILE + --validation_steps 50 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 5e-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" + # --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500" +) + +torchrun \ + --nnodes 1 \ + --nproc_per_node $NUM_GPUS \ + fastvideo/training/wan_t2v_i2v_training_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_ti2v/preprocess_wan_data_t2v.sh b/examples/training/finetune/wan_ti2v/preprocess_wan_data_t2v.sh new file mode 100644 index 000000000..b2a2e5133 --- /dev/null +++ b/examples/training/finetune/wan_ti2v/preprocess_wan_data_t2v.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +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_i2v_1_3b_new/" + +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 "ti2v" \ No newline at end of file diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index 4c1816351..b8c8edf77 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/fastvideo/configs/pipelines/base.py @@ -45,6 +45,8 @@ class PipelineConfig: embedded_cfg_scale: float = 6.0 flow_shift: float | None = None disable_autocast: bool = False + ti2v_task: bool = False + t2v_as_i2v_task: bool = False # Model configuration dit_config: DiTConfig = field(default_factory=DiTConfig) @@ -85,9 +87,6 @@ class PipelineConfig: # DMD parameters dmd_denoising_steps: list[int] | None = field(default=None) - # Wan2.2 TI2V parameters - ti2v_task: bool = False - # Compilation # enable_torch_compile: bool = False @@ -214,6 +213,24 @@ class PipelineConfig: "Comma-separated list of denoising steps (e.g., '1000,757,522')", ) + # TI2V task + parser.add_argument( + f"--{prefix_with_dot}ti2v-task", + action=StoreBoolean, + dest=f"{prefix_with_dot.replace('-', '_')}ti2v_task", + default=PipelineConfig.ti2v_task, + help="Enable TI2V", + ) + + # T2V to I2V task + parser.add_argument( + f"--{prefix_with_dot}t2v-as-i2v-task", + action=StoreBoolean, + dest=f"{prefix_with_dot.replace('-', '_')}t2v_as_i2v_task", + default=PipelineConfig.t2v_as_i2v_task, + help="Enable T2V to I2V task", + ) + # Add VAE configuration arguments from fastvideo.configs.models.vaes.base import VAEConfig VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config") diff --git a/fastvideo/dataset/dataloader/schema.py b/fastvideo/dataset/dataloader/schema.py index ed00fd74b..b1aff2d7d 100644 --- a/fastvideo/dataset/dataloader/schema.py +++ b/fastvideo/dataset/dataloader/schema.py @@ -78,3 +78,36 @@ pyarrow_schema_t2v = pa.schema([ pa.field("duration_sec", pa.float64()), pa.field("fps", pa.float64()), ]) + +pyarrow_schema_t2v_ode_trajectory = pa.schema([ + pa.field("id", pa.string()), + # --- Image/Video VAE latents --- + # Tensors are stored as raw bytes with shape and dtype info for loading + pa.field("vae_latent_bytes", pa.binary()), + # e.g., [C, T, H, W] or [C, H, W] + pa.field("vae_latent_shape", pa.list_(pa.int64())), + # e.g., 'float32' + pa.field("vae_latent_dtype", pa.string()), + # --- Text encoder output tensor --- + # Tensors are stored as raw bytes with shape and dtype info for loading + pa.field("text_embedding_bytes", pa.binary()), + # e.g., [SeqLen, Dim] + pa.field("text_embedding_shape", pa.list_(pa.int64())), + # e.g., 'bfloat16' or 'float32' + pa.field("text_embedding_dtype", pa.string()), + # --- ODE Trajectory --- + pa.field("ode_trajectory_bytes", pa.binary()), + pa.field("ode_trajectory_shape", pa.list_(pa.int64())), + pa.field("ode_trajectory_dtype", pa.string()), + # --- Metadata --- + pa.field("file_name", pa.string()), + pa.field("caption", pa.string()), + pa.field("media_type", pa.string()), # 'image' or 'video' + pa.field("width", pa.int64()), + pa.field("height", pa.int64()), + # -- Video-specific (can be null/default for images) --- + # Number of frames processed (e.g., 1 for image, N for video) + pa.field("num_frames", pa.int64()), + pa.field("duration_sec", pa.float64()), + pa.field("fps", pa.float64()), +]) diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 65d441169..15d7c9495 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -121,14 +121,25 @@ class ComposedPipelineBase(ABC): model_path: str, device: str | None = None, torch_dtype: torch.dtype | None = None, - pipeline_config: str | PipelineConfig | None = None, + pipeline_config: PipelineConfig | None = None, args: argparse.Namespace | None = None, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None, **kwargs) -> "ComposedPipelineBase": """ - Load a pipeline from a pretrained model. + Load a pipeline from a pretrained model. + Few different patterns are supported: + - Only provide model_path: + - This will load the pipeline in inference mode. + - The pipeline will be initialized with the default config. + - The pipeline will be initialized with the default modules. + - The pipeline will be initialized with the default stages. + - The pipeline will be initialized with the default stages. + - override the default config using pipeline_config or args or kwargs + - override the default modules using loaded_modules + - override the pipelineconfig + loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, If provided, loaded_modules will be used instead of loading from config/pretrained weights. """ @@ -136,6 +147,8 @@ class ComposedPipelineBase(ABC): kwargs['model_path'] = model_path fastvideo_args = FastVideoArgs.from_kwargs(**kwargs) + if pipeline_config is not None: + fastvideo_args.pipeline_config = pipeline_config else: assert args is not None, "args must be provided for training mode" fastvideo_args = TrainingArgs.from_cli_args(args) diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_i2v.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_i2v.py index 98ece8a3d..6ff513f53 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_i2v.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_i2v.py @@ -24,7 +24,9 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline): """I2V preprocessing pipeline implementation.""" _required_config_modules = [ - "text_encoder", "tokenizer", "vae", "image_encoder", "image_processor" + "text_encoder", + "tokenizer", + "vae", ] def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index b89489075..ed2ebbbf5 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -132,11 +132,12 @@ class DenoisingStage(PipelineStage): latents = latents[:, :, rank_in_sp_group, :, :, :] batch.latents = latents if batch.image_latent is not None: - image_latent = rearrange(batch.image_latent, - "b c (n t) h w -> b c n t h w", - n=sp_world_size).contiguous() - image_latent = image_latent[:, :, rank_in_sp_group, :, :, :] - batch.image_latent = image_latent + if not fastvideo_args.pipeline_config.ti2v_task: + image_latent = rearrange(batch.image_latent, + "b c (n t) h w -> b c n t h w", + n=sp_world_size).contiguous() + image_latent = image_latent[:, :, rank_in_sp_group, :, :, :] + batch.image_latent = image_latent # Get timesteps and calculate warmup steps timesteps = batch.timesteps # TODO(will): remove this once we add input/output validation for stages @@ -263,11 +264,25 @@ class DenoisingStage(PipelineStage): # Expand latents for I2V latent_model_input = latents.to(target_dtype) - if batch.image_latent is not None: + if batch.image_latent is not None and not fastvideo_args.pipeline_config.t2v_as_i2v_task: assert not fastvideo_args.pipeline_config.ti2v_task, "image latents should not be provided for TI2V task" latent_model_input = torch.cat( [latent_model_input, batch.image_latent], dim=1).to(target_dtype) + elif batch.image_latent is not None and fastvideo_args.pipeline_config.t2v_as_i2v_task: + assert batch.image_latent is not None, "image latents should be provided for T2V to I2V task" + if rank_in_sp_group == 0: + logger.info("latent_model_input.shape: %s", + latent_model_input.shape) + latent_model_input = torch.cat( + [ + batch.image_latent, + latent_model_input[:, :, 1:, :, :], + ], + dim=2).to(target_dtype) + logger.info("latent_model_input.shape: %s", + latent_model_input.shape) + if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None: timestep = torch.stack([t]).to(get_local_torch_device()) temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten() @@ -283,6 +298,13 @@ class DenoisingStage(PipelineStage): assert torch.isnan(latent_model_input).sum() == 0 latent_model_input = self.scheduler.scale_model_input( latent_model_input, t) + if fastvideo_args.pipeline_config.t2v_as_i2v_task: + if rank_in_sp_group == 0: + latent_model_input = torch.cat([ + batch.image_latent, + latent_model_input[:, :, 1:, :, :], + ], + dim=2).to(target_dtype) # Prepare inputs for transformer guidance_expand = ( @@ -400,6 +422,13 @@ class DenoisingStage(PipelineStage): if sp_group: latents = sequence_model_parallel_all_gather(latents, dim=2) + if fastvideo_args.pipeline_config.t2v_as_i2v_task: + latents = torch.cat([ + batch.image_latent, + latents[:, :, 1:, :, :], + ], + dim=2) + # Update batch with final latents batch.latents = latents diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py index 4e817e4ae..d227d76e5 100644 --- a/fastvideo/pipelines/stages/image_encoding.py +++ b/fastvideo/pipelines/stages/image_encoding.py @@ -157,12 +157,21 @@ class ImageVAEEncodingStage(PipelineStage): # (B, C, H, W) -> (B, C, 1, H, W) image = image.unsqueeze(2) - video_condition = torch.cat([ - image, - image.new_zeros(image.shape[0], image.shape[1], num_frames - 1, - image.shape[3], image.shape[4]) - ], - dim=2) + if fastvideo_args.pipeline_config.t2v_as_i2v_task: + # repeat the image self.vae.temporal_compression_ratio times + video_condition = image.repeat(1, 1, + self.vae.temporal_compression_ratio, + 1, 1) + # video_condition = image + logger.info("video_condition.shape: %s", video_condition.shape) + else: + assert False, "should not be here" + video_condition = torch.cat([ + image, + image.new_zeros(image.shape[0], image.shape[1], num_frames - 1, + image.shape[3], image.shape[4]) + ], + dim=2) video_condition = video_condition.to(device=get_local_torch_device(), dtype=torch.float32) @@ -209,6 +218,9 @@ class ImageVAEEncodingStage(PipelineStage): if fastvideo_args.mode == ExecutionMode.PREPROCESS: batch.image_latent = latent_condition + elif fastvideo_args.pipeline_config.t2v_as_i2v_task: + logger.info("latent_condition.shape: %s", latent_condition.shape) + batch.image_latent = latent_condition else: mask_lat_size = torch.ones(1, 1, num_frames, latent_height, latent_width) diff --git a/fastvideo/platforms/__init__.py b/fastvideo/platforms/__init__.py index 2decdbbf6..992d61c7e 100644 --- a/fastvideo/platforms/__init__.py +++ b/fastvideo/platforms/__init__.py @@ -82,8 +82,8 @@ def rocm_platform_plugin() -> str | None: logger.info("ROCm platform is available") finally: amdsmi.amdsmi_shut_down() - except Exception as e: - logger.info("ROCm platform is unavailable: %s", e) + except Exception: + pass return "fastvideo.platforms.rocm.RocmPlatform" if is_rocm else None diff --git a/fastvideo/training/wan_t2v_i2v_training_pipeline.py b/fastvideo/training/wan_t2v_i2v_training_pipeline.py new file mode 100644 index 000000000..882f52ced --- /dev/null +++ b/fastvideo/training/wan_t2v_i2v_training_pipeline.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +import sys +from copy import deepcopy +from typing import Any + +import torch + +from fastvideo.configs.pipelines.base import PipelineConfig +from fastvideo.configs.sample import SamplingParam +from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler) +from fastvideo.pipelines.basic.wan.wan_i2v_pipeline import ( + WanImageToVideoPipeline) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch +from fastvideo.training.training_pipeline import TrainingPipeline +from fastvideo.utils import is_vsa_available, shallow_asdict + +vsa_available = is_vsa_available() + +logger = init_logger(__name__) + + +class WanT2VI2VTrainingPipeline(TrainingPipeline): + """ + A training pipeline for Wan. + """ + _required_config_modules = ["scheduler", "transformer", "vae"] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + self.modules["scheduler"] = FlowUniPCMultistepScheduler( + shift=fastvideo_args.pipeline_config.flow_shift) + + def create_training_stages(self, training_args: TrainingArgs): + """ + May be used in future refactors. + """ + pass + + def set_schemas(self): + self.train_dataset_schema = pyarrow_schema_t2v + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + logger.info("Initializing validation pipeline...") + args_copy = deepcopy(training_args) + + args_copy.inference_mode = True + args_copy.dit_cpu_offload = True + # args_copy.pipeline_config.vae_config.load_encoder = False + # validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained( + pipeline_config = PipelineConfig.from_pretrained( + training_args.model_path) + pipeline_config.vae_config.load_encoder = True + pipeline_config.t2v_as_i2v_task = True + self.validation_pipeline = WanImageToVideoPipeline.from_pretrained( + training_args.model_path, + args=None, + inference_mode=True, + pipeline_config=pipeline_config, + loaded_modules={ + "transformer": self.get_module("transformer"), + }, + required_config_modules=[ + "scheduler", "transformer", "vae", "text_encoder", "tokenizer" + ], + tp_size=training_args.tp_size, + sp_size=training_args.sp_size, + num_gpus=training_args.num_gpus, + dit_cpu_offload=True, + ) + + def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: + batch = next(self.train_loader_iter, None) # type: ignore + if batch is None: + self.current_epoch += 1 + logger.info("Starting epoch %s", self.current_epoch) + # Reset iterator for next epoch + self.train_loader_iter = iter(self.train_dataloader) + # Get first batch of new epoch + batch = next(self.train_loader_iter) + + latents = batch['vae_latent'] + latents = latents[:, :, :self.training_args.num_latent_t] + encoder_hidden_states = batch['text_embedding'] + encoder_attention_mask = batch['text_attention_mask'] + # clip_features = batch['clip_feature'] + # image_latents = batch['first_frame_latent'] + # image_latents = image_latents[:, :, :self.training_args.num_latent_t] + # pil_image = batch['pil_image'] + infos = batch['info_list'] + + training_batch.latents = latents.to(get_local_torch_device(), + dtype=torch.bfloat16) + training_batch.encoder_hidden_states = encoder_hidden_states.to( + get_local_torch_device(), dtype=torch.bfloat16) + training_batch.encoder_attention_mask = encoder_attention_mask.to( + get_local_torch_device(), dtype=torch.bfloat16) + # training_batch.preprocessed_image = pil_image.to( + # get_local_torch_device()) + # training_batch.image_embeds = clip_features.to(get_local_torch_device()) + # training_batch.image_latents = image_latents.to( + # get_local_torch_device()) + training_batch.infos = infos + + return training_batch + + def _prepare_dit_inputs(self, + training_batch: TrainingBatch) -> TrainingBatch: + """Override to properly handle I2V concatenation - call parent first, then concatenate image conditioning.""" + + # First, call parent method to prepare noise, timesteps, etc. for video latents + training_batch = super()._prepare_dit_inputs(training_batch) + + latents = training_batch.latents + logger.info("latents.shape: %s", latents.shape) + first_frame_latent = latents[:, :, 0, :, :] + + logger.info("first_frame_latent.shape: %s", first_frame_latent.shape) + logger.info("training_batch.noisy_model_input.shape: %s", + training_batch.noisy_model_input.shape) + + training_batch.noisy_model_input = torch.cat([ + first_frame_latent.unsqueeze(2), + training_batch.noisy_model_input[:, :, 1:, :, :] + ], + dim=2) + + return training_batch + + def _build_input_kwargs(self, + training_batch: TrainingBatch) -> TrainingBatch: + + # Image Embeds for conditioning + # image_embeds = training_batch.image_embeds + # assert torch.isnan(image_embeds).sum() == 0 + # image_embeds = image_embeds.to(get_local_torch_device(), + # dtype=torch.bfloat16) + # encoder_hidden_states_image = image_embeds + + # NOTE: noisy_model_input already contains concatenated image_latents from _prepare_dit_inputs + training_batch.input_kwargs = { + "hidden_states": + training_batch.noisy_model_input, + "encoder_hidden_states": + training_batch.encoder_hidden_states, + "timestep": + training_batch.timesteps.to(get_local_torch_device(), + dtype=torch.bfloat16), + "encoder_attention_mask": + training_batch.encoder_attention_mask, + # "encoder_hidden_states_image": + # encoder_hidden_states_image, + "return_dict": + False, + } + return training_batch + + def _prepare_validation_batch(self, sampling_param: SamplingParam, + training_args: TrainingArgs, + validation_batch: dict[str, Any], + num_inference_steps: int) -> ForwardBatch: + sampling_param.prompt = validation_batch['prompt'] + sampling_param.height = training_args.num_height + sampling_param.width = training_args.num_width + sampling_param.image_path = validation_batch['video_path'] + sampling_param.num_inference_steps = num_inference_steps + sampling_param.data_type = "video" + assert self.seed is not None + sampling_param.seed = self.seed + + latents_size = [(sampling_param.num_frames - 1) // 4 + 1, + sampling_param.height // 8, sampling_param.width // 8] + n_tokens = latents_size[0] * latents_size[1] * latents_size[2] + temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio + num_frames = (training_args.num_latent_t - + 1) * temporal_compression_factor + 1 + sampling_param.num_frames = num_frames + batch = ForwardBatch( + **shallow_asdict(sampling_param), + latents=None, + generator=torch.Generator(device="cpu").manual_seed(self.seed), + n_tokens=n_tokens, + eta=0.0, + VSA_sparsity=training_args.VSA_sparsity, + ) + + return batch + + +def main(args) -> None: + logger.info("Starting training pipeline...") + + pipeline = WanT2VI2VTrainingPipeline.from_pretrained( + args.pretrained_model_name_or_path, args=args) + args = pipeline.training_args + pipeline.train() + logger.info("Training pipeline done") + + +if __name__ == "__main__": + argv = sys.argv + from fastvideo.fastvideo_args import TrainingArgs + from fastvideo.utils import FlexibleArgumentParser + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + parser = FastVideoArgs.add_cli_args(parser) + args = parser.parse_args() + args.dit_cpu_offload = False + main(args)