diff --git a/fastvideo/v1/dataset/dataloader/schema.py b/fastvideo/v1/dataset/dataloader/schema.py index 22013403c..37dba61df 100644 --- a/fastvideo/v1/dataset/dataloader/schema.py +++ b/fastvideo/v1/dataset/dataloader/schema.py @@ -35,6 +35,9 @@ pyarrow_schema_i2v = pa.schema([ pa.field("clip_feature_bytes", pa.binary()), pa.field("clip_feature_shape", pa.list_(pa.int64())), pa.field("clip_feature_dtype", pa.string()), + pa.field("encoded_first_frame_bytes", pa.binary()), + pa.field("encoded_first_frame_shape", pa.list_(pa.int64())), + pa.field("encoded_first_frame_dtype", pa.string()), # --- Metadata --- pa.field("file_name", pa.string()), pa.field("caption", pa.string()), @@ -126,13 +129,6 @@ pyarrow_schema_t2v = pa.schema([ pyarrow_schema_t2v_validation = 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()), diff --git a/fastvideo/v1/tests/nightly/test_e2e_i2v_overfit_single_sample.py b/fastvideo/v1/tests/nightly/test_e2e_i2v_overfit_single_sample.py new file mode 100644 index 000000000..c53422e57 --- /dev/null +++ b/fastvideo/v1/tests/nightly/test_e2e_i2v_overfit_single_sample.py @@ -0,0 +1,177 @@ +import os +from pathlib import Path +from huggingface_hub import snapshot_download +import shutil +import subprocess +import sys +from fastvideo.v1.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision + +# Import the training pipeline +sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) + +NUM_NODES = "1" +MODEL_PATH = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" + +# preprocessing +DATA_DIR = "data" +LOCAL_RAW_DATA_DIR = Path(os.path.join(DATA_DIR, "cats")) +NUM_GPUS_PER_NODE_PREPROCESSING = "1" +PREPROCESSING_ENTRY_FILE_PATH = "fastvideo/v1/pipelines/preprocess/v1_preprocess.py" + +LOCAL_PREPROCESSED_DATA_DIR = Path(os.path.join(DATA_DIR, "cats_preprocessed_data")) + + +# training +NUM_GPUS_PER_NODE_TRAINING = "8" +TRAINING_ENTRY_FILE_PATH = "fastvideo/v1/training/wan_training_pipeline.py" +LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "combined_parquet_dataset") +LOCAL_VALIDATION_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "validation_parquet_dataset") +LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs")) + +def download_data(): + # create the data dir if it doesn't exist + data_dir = Path(DATA_DIR) + # if data_dir.exists(): + # print(f"Removing existing data directory at {data_dir}") + # shutil.rmtree(data_dir) + + print(f"Creating data directory at {data_dir}") + os.makedirs(data_dir) + + print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...") + try: + # result = snapshot_download( + # repo_id="wlsaidhi/cats-overfit-merged", + # local_dir=str(LOCAL_RAW_DATA_DIR), + # repo_type="dataset", + # resume_download=True, + # token=os.environ.get("HF_TOKEN"), # In case authentication is needed + # ) + print(f"Download completed successfully. Files downloaded to: {result}") + + # Verify the download + if not LOCAL_RAW_DATA_DIR.exists(): + raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist") + + # List downloaded files + print("Downloaded files:") + for file in LOCAL_RAW_DATA_DIR.rglob("*"): + if file.is_file(): + print(f" - {file.relative_to(LOCAL_RAW_DATA_DIR)}") + + except Exception as e: + print(f"Error during download: {str(e)}") + raise + + +def run_preprocessing(): + # Run torchrun command + cmd = [ + "torchrun", + "--nnodes", NUM_NODES, + "--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING, + PREPROCESSING_ENTRY_FILE_PATH, + "--model_path", MODEL_PATH, + "--data_merge_path", os.path.join(LOCAL_RAW_DATA_DIR, "merge_1_sample.txt"), + "--preprocess_video_batch_size", "1", + "--max_height", "480", + "--max_width", "832", + "--num_frames", "77", + "--dataloader_num_workers", "0", + "--output_dir", LOCAL_PREPROCESSED_DATA_DIR, + "--train_fps", "16", + "--validation_dataset_file", os.path.join(LOCAL_RAW_DATA_DIR, "validation_i2v_prompt_1_sample.json"), + "--samples_per_file", "1", + "--flush_frequency", "1", + "--video_length_tolerance_range", "5", + "--dataset", "i2v", + ] + + process = subprocess.run(cmd, check=True) + + +def run_training(): + cmd = [ + "torchrun", + "--nnodes", NUM_NODES, + "--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING, + TRAINING_ENTRY_FILE_PATH, + "--model_path", MODEL_PATH, + "--inference_mode", "False", + "--pretrained_model_name_or_path", MODEL_PATH, + "--data_path", LOCAL_TRAINING_DATA_DIR, + "--validation_preprocessed_path", LOCAL_VALIDATION_DATA_DIR, + "--train_batch_size", "1", + "--num_latent_t", "8", + "--num_gpus", NUM_GPUS_PER_NODE_TRAINING, + "--sp_size", NUM_GPUS_PER_NODE_TRAINING, + "--tp_size", NUM_GPUS_PER_NODE_TRAINING, + "--hsdp_replicate_dim", "1", + "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING, + "--num_gpus", NUM_GPUS_PER_NODE_TRAINING, + "--train_sp_batch_size", "1", + "--dataloader_num_workers", "10", + "--gradient_accumulation_steps", "1", + "--max_train_steps", "901", + "--learning_rate", "1e-5", + "--mixed_precision", "bf16", + "--checkpointing_steps", "6000", + "--validation_steps", "100", + "--validation_sampling_steps", "50", + "--log_validation", + "--checkpoints_total_limit", "3", + "--allow_tf32", + "--ema_start_step", "0", + "--cfg", "0.0", + "--output_dir", LOCAL_OUTPUT_DIR, + "--tracker_project_name", "wan_i2v_finetune_overfit_ci", + "--num_height", "480", + "--num_width", "832", + "--num_frames", "81", + "--validation_guidance_scale", "1.0", + "--num_euler_timesteps", "50", + "--multi_phased_distill_schedule", "4000-1", + "--weight_decay", "0.01", + "--not_apply_cfg_solver", + "--dit_precision", "fp32", + "--max_grad_norm", "1.0", + ] + + print(f"Running training with command: {cmd}") + process = subprocess.run(cmd, check=True) + + +def test_e2e_overfit_single_sample(): + os.environ["WANDB_MODE"] = "online" + + download_data() + run_preprocessing() + run_training() + + reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_v0.mp4") + print(f"reference_video_file: {reference_video_file}") + final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4") + print(f"final_validation_video_file: {final_validation_video_file}") + + + # Ensure both files exist + assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}" + assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}" + + # Compute SSIM + mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision( + reference_video_file, + final_validation_video_file, + use_ms_ssim=True # Using MS-SSIM for better quality assessment + ) + + print("\n===== SSIM Results for Step 900 Validation =====") + print(f"Mean MS-SSIM: {mean_ssim:.4f}") + print(f"Min MS-SSIM: {min_ssim:.4f}") + print(f"Max MS-SSIM: {max_ssim:.4f}") + + assert max_ssim > 0.5, f"Max SSIM is below 0.5: {max_ssim}" + + +if __name__ == "__main__": + test_e2e_overfit_single_sample() \ No newline at end of file