From d36c2091916b14437cdbe321a29f4bf76e4eb56e Mon Sep 17 00:00:00 2001 From: Wei Zhou <69577934+JerryZhou54@users.noreply.github.com> Date: Fri, 23 May 2025 02:06:54 -0400 Subject: [PATCH] Will/training (#425) Co-authored-by: Kevin Lin <42618777+kevin314@users.noreply.github.com> Co-authored-by: William Lin Co-authored-by: Will Lin --- distill_v1.sh | 46 + .../preprocess_text_embeddings_v1.py | 199 ++++ .../preprocess_vae_latents_v1.py | 151 +++ ...reprocess_validation_text_embeddings_v1.py | 115 +++ fastvideo/distill/solver.py | 3 +- .../models/mochi_hf/mochi_latents_utils.py | 13 +- fastvideo/utils/checkpoint.py | 40 +- .../base_device_communicator.py | 21 + fastvideo/v1/entrypoints/data_preprocessor.py | 29 + fastvideo/v1/fastvideo_args.py | 1 + fastvideo/v1/models/loader/fsdp_load.py | 1 + .../v1/pipelines/composed_pipeline_base.py | 150 ++- fastvideo/v1/pipelines/distill_pipeline.py | 303 ++++++ fastvideo/v1/pipelines/training_pipeline.py | 635 ++++++++++++ .../v1/pipelines/wan/wan_latent_pipeline.py | 19 + fastvideo/v1/pipelines/wan/wan_pipeline.py | 2 +- fastvideo/v1_distill.py | 953 ++++++++++++++++++ .../dataset_preparation/prepare_json_file.py | 66 +- scripts/dataset_preparation/resize_videos.py | 14 +- scripts/inference/v1_preprocess_wan.sh | 27 + scripts/preprocess/preprocess_wan_data_v1.sh | 30 + train_v1.sh | 48 + 22 files changed, 2819 insertions(+), 47 deletions(-) create mode 100644 distill_v1.sh create mode 100644 fastvideo/data_preprocess/preprocess_text_embeddings_v1.py create mode 100644 fastvideo/data_preprocess/preprocess_vae_latents_v1.py create mode 100644 fastvideo/data_preprocess/preprocess_validation_text_embeddings_v1.py create mode 100644 fastvideo/v1/entrypoints/data_preprocessor.py create mode 100644 fastvideo/v1/pipelines/distill_pipeline.py create mode 100644 fastvideo/v1/pipelines/training_pipeline.py create mode 100644 fastvideo/v1/pipelines/wan/wan_latent_pipeline.py create mode 100644 fastvideo/v1_distill.py create mode 100644 scripts/inference/v1_preprocess_wan.sh create mode 100644 scripts/preprocess/preprocess_wan_data_v1.sh create mode 100644 train_v1.sh diff --git a/distill_v1.sh b/distill_v1.sh new file mode 100644 index 000000000..f7171eb5c --- /dev/null +++ b/distill_v1.sh @@ -0,0 +1,46 @@ +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online + +DATA_DIR=./data +# 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 + # --gradient_checkpointing\ + # --pretrained_model_name_or_path hunyuanvideo-community/HunyuanVideo \ + # --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ +torchrun --nnodes 1 --nproc_per_node 4\ + fastvideo/v1/pipelines/training_pipeline.py\ + --inference_mode False\ + --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ + --cache_dir "/home/ray/.cache"\ + --data_json_path "$DATA_DIR/HD-Mixkit-Finetune-Hunyuan/videos2caption.json"\ + --validation_prompt_dir "$DATA_DIR/HD-Mixkit-Finetune-Hunyuan/validation"\ + --train_batch_size=1\ + --num_latent_t 1 \ + --sp_size 4 \ + --tp_size 4 \ + --train_sp_batch_size 1\ + --dataloader_num_workers 4\ + --gradient_accumulation_steps=1\ + --max_train_steps=320\ + --learning_rate=1e-6\ + --mixed_precision="bf16"\ + --checkpointing_steps=64\ + --validation_steps 64\ + --validation_sampling_steps "2,4,8" \ + --checkpoints_total_limit 3\ + --allow_tf32\ + --ema_start_step 0\ + --cfg 0.0\ + --log_validation\ + --output_dir="$DATA_DIR/outputs/hy_phase1_shift17_bs_16_HD"\ + --tracker_project_name Hunyuan_Distill \ + --num_height 720 \ + --num_width 1280 \ + --num_frames 125 \ + --shift 17 \ + --validation_guidance_scale "1.0" \ + --num_euler_timesteps 50 \ + --multi_phased_distill_schedule "4000-1" \ + --not_apply_cfg_solver \ + --master_weight_type "bf16" \ No newline at end of file diff --git a/fastvideo/data_preprocess/preprocess_text_embeddings_v1.py b/fastvideo/data_preprocess/preprocess_text_embeddings_v1.py new file mode 100644 index 000000000..db33cdfe0 --- /dev/null +++ b/fastvideo/data_preprocess/preprocess_text_embeddings_v1.py @@ -0,0 +1,199 @@ +import argparse +import json +import os + +import torch +import torch.distributed as dist +from accelerate.logging import get_logger +from diffusers.utils import export_to_video +from diffusers.video_processor import VideoProcessor +from torch.utils.data import DataLoader, Dataset +from torch.utils.data.distributed import DistributedSampler +from tqdm import tqdm + +# from fastvideo.utils.load import load_text_encoder, load_vae + +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.loader.component_loader import VAELoader, TextEncoderLoader, TokenizerLoader +from fastvideo.v1.utils import maybe_download_model +from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group +from fastvideo.v1.fastvideo_args import FastVideoArgs +from fastvideo.v1.configs.models.vaes import WanVAEConfig +from fastvideo.v1.configs.models.encoders.t5 import T5Config + +logger = get_logger(__name__) + +class T5dataset(Dataset): + + def __init__( + self, + json_path, + vae_debug, + ): + self.json_path = json_path + self.vae_debug = vae_debug + with open(self.json_path, "r") as f: + train_dataset = json.load(f) + self.train_dataset = sorted(train_dataset, key=lambda x: x["latent_path"]) + + def __getitem__(self, idx): + caption = self.train_dataset[idx]["caption"] + filename = self.train_dataset[idx]["latent_path"].split(".")[0] + length = self.train_dataset[idx]["length"] + if self.vae_debug: + latents = torch.load( + os.path.join(args.output_dir, "latent", self.train_dataset[idx]["latent_path"]), + map_location="cpu", + ) + else: + latents = [] + + return dict(caption=caption, latents=latents, filename=filename, length=length) + + def __len__(self): + return len(self.train_dataset) + + +def main(args): + local_rank = int(os.getenv("RANK", 0)) + world_size = int(os.getenv("WORLD_SIZE", 1)) + rank = int(os.getenv("RANK", 0)) + init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank) + initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size) + print("world_size", world_size, "local rank", local_rank) + device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(device) + world_group = get_world_group() + + # device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + # torch.cuda.set_device(local_rank) + # if not dist.is_initialized(): + # dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank) + + videoprocessor = VideoProcessor(vae_scale_factor=8) + os.makedirs(args.output_dir, exist_ok=True) + os.makedirs(os.path.join(args.output_dir, "video"), exist_ok=True) + os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True) + os.makedirs(os.path.join(args.output_dir, "prompt_embed"), exist_ok=True) + os.makedirs(os.path.join(args.output_dir, "prompt_attention_mask"), exist_ok=True) + + vae_precision = "fp16" + text_encoder_precision = "fp32" + fastvideo_args = FastVideoArgs(model_path=args.model_path, + use_cpu_offload=False, + vae_precision=vae_precision, + text_encoder_precisions=(text_encoder_precision,)) + fastvideo_args.device = device + fastvideo_args.device_str = f"cuda:{local_rank}" + + # fastvideo_args.dit_config = HunyuanVideoConfig() + fastvideo_args.vae_config = WanVAEConfig() + fastvideo_args.text_encoder_configs = (T5Config(),) + + # vae_loader = VAELoader() + # vae = vae_loader.load_vae() + text_encoder_loader = TextEncoderLoader() + tokenizer_loader = TokenizerLoader() + + model_path = args.model_path + path = maybe_download_model(model_path) + # PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" + ENCODER_PATH = os.path.join(path, "text_encoder") + TOKENIZER_PATH = os.path.join(path, "tokenizer") + print(ENCODER_PATH) + text_encoder = text_encoder_loader.load(ENCODER_PATH, "text_encoder", fastvideo_args) + tokenizer = tokenizer_loader.load(TOKENIZER_PATH, "tokenizer", fastvideo_args) + + latents_json_path = os.path.join(args.output_dir, "videos2caption_temp.json") + train_dataset = T5dataset(latents_json_path, args.vae_debug) + # text_encoder = load_text_encoder(args.model_type, args.model_path, device=device) + # vae, autocast_type, fps = load_vae(args.model_type, args.model_path) + # vae.enable_tiling() + sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True) + train_dataloader = DataLoader( + train_dataset, + sampler=sampler, + batch_size=args.train_batch_size, + num_workers=args.dataloader_num_workers, + ) + + json_data = [] + for _, data in tqdm(enumerate(train_dataloader), disable=local_rank != 0): + with torch.inference_mode(): + # with torch.autocast("cuda", dtype=torch.float32): + print(data["caption"]) + text_inputs = tokenizer(data["caption"], **fastvideo_args.text_encoder_configs[0].tokenizer_kwargs).to( + fastvideo_args.device) + input_ids = text_inputs["input_ids"] + attention_mask = text_inputs["attention_mask"] + outputs = text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + from fastvideo.v1.configs.pipelines.wan import t5_postprocess_text + post_process_func = t5_postprocess_text + prompt_embeds = post_process_func(outputs) + prompt_attention_mask = attention_mask + if args.vae_debug: + latents = data["latents"] + video = vae.decode(latents.to(device), return_dict=False)[0] + video = videoprocessor.postprocess_video(video) + for idx, video_name in enumerate(data["filename"]): + prompt_embed_path = os.path.join(args.output_dir, "prompt_embed", video_name + ".pt") + video_path = os.path.join(args.output_dir, "video", video_name + ".mp4") + prompt_attention_mask_path = os.path.join(args.output_dir, "prompt_attention_mask", + video_name + ".pt") + # save latent + torch.save(prompt_embeds[idx], prompt_embed_path) + torch.save(prompt_attention_mask[idx], prompt_attention_mask_path) + print(f"sample {video_name} saved") + if args.vae_debug: + export_to_video(video[idx], video_path, fps=16) + item = {} + item["length"] = int(data["length"][idx]) + item["latent_path"] = video_name + ".pt" + item["prompt_embed_path"] = video_name + ".pt" + item["prompt_attention_mask"] = video_name + ".pt" + item["caption"] = data["caption"][idx] + json_data.append(item) + dist.barrier() + local_data = json_data + gathered_data = [None] * world_size + dist.all_gather_object(gathered_data, local_data) + if local_rank == 0: + # os.remove(latents_json_path) + all_json_data = [item for sublist in gathered_data for item in sublist] + with open(os.path.join(args.output_dir, "videos2caption.json"), "w") as f: + json.dump(all_json_data, f, indent=4) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # dataset & dataloader + parser.add_argument("--model_path", type=str, default="data/mochi") + # parser.add_argument("--model_type", type=str, default="mochi") + # text encoder & vae & diffusion model + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=1, + help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.", + ) + parser.add_argument( + "--train_batch_size", + type=int, + default=1, + help="Batch size (per device) for the training dataloader.", + ) + parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl") + parser.add_argument("--cache_dir", type=str, default="./cache_dir") + parser.add_argument( + "--output_dir", + type=str, + default=None, + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument("--vae_debug", action="store_true") + args = parser.parse_args() + main(args) diff --git a/fastvideo/data_preprocess/preprocess_vae_latents_v1.py b/fastvideo/data_preprocess/preprocess_vae_latents_v1.py new file mode 100644 index 000000000..bb93f463d --- /dev/null +++ b/fastvideo/data_preprocess/preprocess_vae_latents_v1.py @@ -0,0 +1,151 @@ +import argparse +import json +import os + +import torch +# import torch.distributed as dist +# from accelerate.logging import get_logger +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from tqdm import tqdm + +from fastvideo.dataset import getdataset +# from fastvideo.utils.load import load_vae +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.loader.component_loader import VAELoader +from fastvideo.v1.utils import maybe_download_model +from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group +from fastvideo.v1.fastvideo_args import FastVideoArgs +from fastvideo.v1.configs.models.vaes import WanVAEConfig + +logger = init_logger(__name__) + +model_path = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +path = maybe_download_model(model_path) +# PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +VAE_PATH = os.path.join(path, "vae") +print(VAE_PATH) + + +def main(args): + local_rank = int(os.getenv("RANK", 0)) + world_size = int(os.getenv("WORLD_SIZE", 1)) + rank = int(os.getenv("RANK", 0)) + init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank) + initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size) + print("world_size", world_size, "local rank", local_rank) + device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(device) + world_group = get_world_group() + + vae_precision = "fp16" + fastvideo_args = FastVideoArgs(model_path=VAE_PATH, + use_cpu_offload=False, + vae_precision=vae_precision) + fastvideo_args.device = device + # fastvideo_args.dit_config = HunyuanVideoConfig() + fastvideo_args.vae_config = WanVAEConfig() + + train_dataset = getdataset(args) + sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True) + train_dataloader = DataLoader( + train_dataset, + sampler=sampler, + batch_size=args.train_batch_size, + num_workers=args.dataloader_num_workers, + ) + + + # encoder_device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + # torch.cuda.set_device(local_rank) + # if not dist.is_initialized(): + # dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank) + vae_loader = VAELoader() + vae = vae_loader.load(VAE_PATH, "vae", fastvideo_args) + # vae, autocast_type, fps = load_vae(args.model_type, args.model_path) + # vae.enable_tiling() + os.makedirs(args.output_dir, exist_ok=True) + os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True) + + json_data = [] + for _, data in tqdm(enumerate(train_dataloader), disable=local_rank != 0): + with torch.inference_mode(): + with torch.autocast("cuda", dtype=torch.float16): + latents = vae.encode(data["pixel_values"].to(device)).sample() + for idx, video_path in enumerate(data["path"]): + video_name = os.path.basename(video_path).split(".")[0] + latent_path = os.path.join(args.output_dir, "latent", video_name + ".pt") + torch.save(latents[idx].to(torch.bfloat16), latent_path) + item = {} + item["length"] = latents[idx].shape[1] + item["latent_path"] = video_name + ".pt" + item["caption"] = data["text"][idx] + json_data.append(item) + print(f"{video_name} processed") + world_group.barrier() + local_data = json_data + gathered_data = [None] * world_size + for i in range(world_size): + if local_rank == i: + world_group.broadcast_object(local_data, src=i) + else: + gathered_data[i] = world_group.broadcast_object(None, src=i) + gathered_data[local_rank] = json_data + print(gathered_data) + if local_rank == 0: + all_json_data = [item for sublist in gathered_data for item in sublist] + with open(os.path.join(args.output_dir, "videos2caption_temp.json"), "w") as f: + json.dump(all_json_data, f, indent=4) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # dataset & dataloader + parser.add_argument("--model_path", type=str, default="data/mochi") + # parser.add_argument("--model_type", type=str, default="mochi") + parser.add_argument("--data_merge_path", type=str, required=True) + parser.add_argument("--num_frames", type=int, default=163) + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=1, + help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.", + ) + parser.add_argument( + "--train_batch_size", + type=int, + default=16, + help="Batch size (per device) for the training dataloader.", + ) + parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.") + parser.add_argument("--max_height", type=int, default=480) + parser.add_argument("--max_width", type=int, default=848) + parser.add_argument("--video_length_tolerance_range", type=int, default=2.0) + parser.add_argument("--group_frame", action="store_true") # TODO + parser.add_argument("--group_resolution", action="store_true") # TODO + parser.add_argument("--dataset", default="t2v") + parser.add_argument("--train_fps", type=int, default=30) + parser.add_argument("--use_image_num", type=int, default=0) + parser.add_argument("--text_max_length", type=int, default=256) + parser.add_argument("--speed_factor", type=float, default=1.0) + parser.add_argument("--drop_short_ratio", type=float, default=1.0) + # text encoder & vae & diffusion model + parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl") + parser.add_argument("--cache_dir", type=str, default="./cache_dir") + parser.add_argument("--cfg", type=float, default=0.0) + parser.add_argument( + "--output_dir", + type=str, + default=None, + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."), + ) + + args = parser.parse_args() + main(args) diff --git a/fastvideo/data_preprocess/preprocess_validation_text_embeddings_v1.py b/fastvideo/data_preprocess/preprocess_validation_text_embeddings_v1.py new file mode 100644 index 000000000..20c9ae2fc --- /dev/null +++ b/fastvideo/data_preprocess/preprocess_validation_text_embeddings_v1.py @@ -0,0 +1,115 @@ +import argparse +import os + +import torch +# import torch.distributed as dist +from accelerate.logging import get_logger + +# from fastvideo.utils.load import load_text_encoder + +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.loader.component_loader import VAELoader, TextEncoderLoader, TokenizerLoader +from fastvideo.v1.utils import maybe_download_model +from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group +from fastvideo.v1.fastvideo_args import FastVideoArgs +from fastvideo.v1.configs.models.vaes import WanVAEConfig +from fastvideo.v1.configs.models.encoders.t5 import T5Config + +logger = get_logger(__name__) + + +def main(args): + local_rank = int(os.getenv("RANK", 0)) + world_size = int(os.getenv("WORLD_SIZE", 1)) + rank = int(os.getenv("RANK", 0)) + init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank) + initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size) + print("world_size", world_size, "local rank", local_rank) + device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(device) + world_group = get_world_group() + + vae_precision = "fp16" + text_encoder_precision = "fp32" + fastvideo_args = FastVideoArgs(model_path=args.model_path, + use_cpu_offload=False, + vae_precision=vae_precision, + text_encoder_precisions=(text_encoder_precision,)) + fastvideo_args.device = device + fastvideo_args.device_str = f"cuda:{local_rank}" + + # fastvideo_args.dit_config = HunyuanVideoConfig() + fastvideo_args.vae_config = WanVAEConfig() + fastvideo_args.text_encoder_configs = (T5Config(),) + + # vae_loader = VAELoader() + # vae = vae_loader.load_vae() + text_encoder_loader = TextEncoderLoader() + tokenizer_loader = TokenizerLoader() + + model_path = args.model_path + path = maybe_download_model(model_path) + # PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" + ENCODER_PATH = os.path.join(path, "text_encoder") + TOKENIZER_PATH = os.path.join(path, "tokenizer") + print(ENCODER_PATH) + text_encoder = text_encoder_loader.load(ENCODER_PATH, "text_encoder", fastvideo_args) + tokenizer = tokenizer_loader.load(TOKENIZER_PATH, "tokenizer", fastvideo_args) + + # text_encoder = load_text_encoder(args.model_type, args.model_path, device=device) + # autocast_type = torch.float16 if args.model_type == "hunyuan" else torch.bfloat16 + # output_dir/validation/prompt_attention_mask + # output_dir/validation/prompt_embed + os.makedirs(os.path.join(args.output_dir, "validation"), exist_ok=True) + os.makedirs( + os.path.join(args.output_dir, "validation", "prompt_attention_mask"), + exist_ok=True, + ) + os.makedirs(os.path.join(args.output_dir, "validation", "prompt_embed"), exist_ok=True) + + with open(args.validation_prompt_txt, "r", encoding="utf-8") as file: + lines = file.readlines() + prompts = [line.strip() for line in lines] + for prompt in prompts: + with torch.inference_mode(): + # with torch.autocast("cuda", dtype=autocast_type): + text_inputs = tokenizer(prompt, **fastvideo_args.text_encoder_configs[0].tokenizer_kwargs).to( + fastvideo_args.device) + input_ids = text_inputs["input_ids"] + attention_mask = text_inputs["attention_mask"] + outputs = text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + from fastvideo.v1.configs.pipelines.wan import t5_postprocess_text + post_process_func = t5_postprocess_text + prompt_embeds = post_process_func(outputs) + prompt_attention_mask = attention_mask + + file_name = prompt.split(".")[0] + prompt_embed_path = os.path.join(args.output_dir, "validation", "prompt_embed", f"{file_name}.pt") + prompt_attention_mask_path = os.path.join( + args.output_dir, + "validation", + "prompt_attention_mask", + f"{file_name}.pt", + ) + torch.save(prompt_embeds[0], prompt_embed_path) + torch.save(prompt_attention_mask[0], prompt_attention_mask_path) + print(f"sample {file_name} saved") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # dataset & dataloader + parser.add_argument("--model_path", type=str, default="data/mochi") + parser.add_argument("--validation_prompt_txt", type=str) + parser.add_argument( + "--output_dir", + type=str, + default=None, + help="The output directory where the model predictions and checkpoints will be written.", + ) + args = parser.parse_args() + main(args) diff --git a/fastvideo/distill/solver.py b/fastvideo/distill/solver.py index 7f8fec847..34bf42ad6 100644 --- a/fastvideo/distill/solver.py +++ b/fastvideo/distill/solver.py @@ -7,7 +7,7 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.schedulers.scheduling_utils import SchedulerMixin from diffusers.utils import BaseOutput, logging -from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule +# from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -38,6 +38,7 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin): linear_range=0.5, ): if linear_quadratic: + raise NotImplementedError("Linear quadratic schedule is not implemented") linear_steps = int(num_train_timesteps * linear_range) sigmas = linear_quadratic_schedule(num_train_timesteps, linear_quadratic_threshold, linear_steps) sigmas = torch.tensor(sigmas).to(dtype=torch.float32) diff --git a/fastvideo/models/mochi_hf/mochi_latents_utils.py b/fastvideo/models/mochi_hf/mochi_latents_utils.py index c554d56c7..4007f8dda 100644 --- a/fastvideo/models/mochi_hf/mochi_latents_utils.py +++ b/fastvideo/models/mochi_hf/mochi_latents_utils.py @@ -31,7 +31,7 @@ mochi_latents_std = torch.tensor([ mochi_scaling_factor = 1.0 -def normalize_dit_input(model_type, latents): +def normalize_dit_input(model_type, latents, args=None): if model_type == "mochi": latents_mean = mochi_latents_mean.to(latents.device, latents.dtype) latents_std = mochi_latents_std.to(latents.device, latents.dtype) @@ -41,5 +41,16 @@ def normalize_dit_input(model_type, latents): return latents * 0.476986 elif model_type == "hunyuan": return latents * 0.476986 + elif model_type == "wan": + from fastvideo.v1.configs.models.vaes.wanvae import WanVAEConfig + vae_config = WanVAEConfig() + latents_mean = torch.tensor(vae_config.arch_config.latents_mean) + latents_std = 1.0 / torch.tensor(vae_config.arch_config.latents_std) + + + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(device=latents.device) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(device=latents.device) + latents = ((latents.float() - latents_mean) * latents_std).to(latents) + return latents else: raise NotImplementedError(f"model_type {model_type} not supported") diff --git a/fastvideo/utils/checkpoint.py b/fastvideo/utils/checkpoint.py index b624ffee2..ef661a190 100644 --- a/fastvideo/utils/checkpoint.py +++ b/fastvideo/utils/checkpoint.py @@ -11,6 +11,7 @@ from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_ from torch.distributed.fsdp import FullOptimStateDictConfig, FullStateDictConfig from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import StateDictType +import dataclasses from fastvideo.utils.logging_ import main_print @@ -44,13 +45,50 @@ def save_checkpoint_optimizer(model, optimizer, rank, output_dir, step, discrimi optimizer_path = os.path.join(save_dir, "optimizer.pt") torch.save(optim_state, optimizer_path) else: - weight_path = os.path.join(save_dir, "discriminator_pytorch_model.safetensors") + weight_path = os.path.join(save_dstate_dictir, "discriminator_pytorch_model.safetensors") save_file(cpu_state, weight_path) optimizer_path = os.path.join(save_dir, "discriminator_optimizer.pt") torch.save(optim_state, optimizer_path) main_print(f"--> checkpoint saved at step {step}") +def save_checkpoint_v1(transformer, rank, output_dir, step): + # from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +# from torch.distributed.fsdp import StateDictType, FullStateDictConfig + + # Configure FSDP to save full state dict + FSDP.set_state_dict_type( + transformer, + state_dict_type=StateDictType.FULL_STATE_DICT, + state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + ) + + # Now get the state dict + cpu_state = transformer.state_dict() + + # Save it (only on rank 0 since we used rank0_only=True) + # if torch.distributed.get_rank() == 0: + # torch.save(state_dict, "model_checkpoint.pt") + if rank <= 0: + save_dir = os.path.join(output_dir, f"checkpoint-{step}") + os.makedirs(save_dir, exist_ok=True) + # save using safetensors + # weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors") + weight_path = os.path.join(save_dir, "diffusion_pytorch_model.pt") + print(weight_path) + # save_file(cpu_state, weight_path) + torch.save(cpu_state, weight_path) + config_dict = transformer.hf_config + if "dtype" in config_dict: + del config_dict["dtype"] # TODO + config_path = os.path.join(save_dir, "config.json") + # save dict as json + with open(config_path, "w") as f: + json.dump(config_dict, f, indent=4) + main_print(f"--> checkpoint saved at step {step}") + + + def save_checkpoint(transformer, rank, output_dir, step): main_print(f"--> saving checkpoint at step {step}") with FSDP.state_dict_type( diff --git a/fastvideo/v1/distributed/device_communicators/base_device_communicator.py b/fastvideo/v1/distributed/device_communicators/base_device_communicator.py index 81d68f038..769c4a40e 100644 --- a/fastvideo/v1/distributed/device_communicators/base_device_communicator.py +++ b/fastvideo/v1/distributed/device_communicators/base_device_communicator.py @@ -15,6 +15,15 @@ class DistributedAutograd: This class provides custom autograd functions for distributed operations like all_reduce, all_gather, and all_to_all. Each operation is implemented as a static inner class with proper forward and backward implementations. +<<<<<<< HEAD +======= + + The operations are designed to: + 1. Maintain correct gradient flow in distributed settings + 2. Handle memory efficiently + 3. Provide proper error handling and type safety + 4. Support both training and inference modes +>>>>>>> 191d1be3 (Will/training (#425)) """ class AllReduce(torch.autograd.Function): @@ -129,6 +138,10 @@ class DistributedAutograd: output = torch.empty_like(input_t) dist.all_to_all_single(output, input_t, group=group) +<<<<<<< HEAD +======= + torch.cuda.synchronize() +>>>>>>> 191d1be3 (Will/training (#425)) output = output.reshape(seqlen, bs, shard_hc, hs).transpose(0, 1).contiguous() @@ -148,6 +161,10 @@ class DistributedAutograd: output = torch.empty_like(input_t) dist.all_to_all_single(output, input_t, group=group) +<<<<<<< HEAD +======= + torch.cuda.synchronize() +>>>>>>> 191d1be3 (Will/training (#425)) output = output.reshape(hc, shard_seqlen, bs, hs) output = output.transpose(0, 2).contiguous() @@ -156,9 +173,13 @@ class DistributedAutograd: return output else: raise RuntimeError( +<<<<<<< HEAD f"Invalid scatter_dim={scatter_dim}, gather_dim={gather_dim}. " f"Only (scatter_dim=2, gather_dim=1) and (scatter_dim=1, gather_dim=2) are supported." ) +======= + "scatter_dim must be 1 or 2 and gather_dim must be 1 or 2") +>>>>>>> 191d1be3 (Will/training (#425)) @staticmethod def backward( diff --git a/fastvideo/v1/entrypoints/data_preprocessor.py b/fastvideo/v1/entrypoints/data_preprocessor.py new file mode 100644 index 000000000..a3192b3a2 --- /dev/null +++ b/fastvideo/v1/entrypoints/data_preprocessor.py @@ -0,0 +1,29 @@ +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler + +from fastvideo.v1.pipelines.wan.wan_latent_pipeline import WanLatentPipeline + + +def main(): + print("Starting data preprocessor") + pipeline = WanLatentPipeline.from_pretrained( + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers") + + train_dataset = getdataset(args) + sampler = DistributedSampler(train_dataset, + rank=local_rank, + num_replicas=world_size, + shuffle=True) + train_dataloader = DataLoader( + train_dataset, + sampler=sampler, + batch_size=args.train_batch_size, + num_workers=args.dataloader_num_workers, + ) + + for batch in train_dataloader: + pipeline(batch) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py index 3d6a7c670..461f9b77f 100644 --- a/fastvideo/v1/fastvideo_args.py +++ b/fastvideo/v1/fastvideo_args.py @@ -478,6 +478,7 @@ class TrainingArgs(FastVideoArgs): output_dir: str = "" checkpoints_total_limit: int = 0 checkpointing_steps: int = 0 + resume_from_checkpoint: str = "" logging_dir: str = "" # optimizer & scheduler diff --git a/fastvideo/v1/models/loader/fsdp_load.py b/fastvideo/v1/models/loader/fsdp_load.py index 5d10a8738..91eae0eb6 100644 --- a/fastvideo/v1/models/loader/fsdp_load.py +++ b/fastvideo/v1/models/loader/fsdp_load.py @@ -121,6 +121,7 @@ def load_fsdp_model( f"Unexpected param or buffer {n} on meta device.") for p in model.parameters(): p.requires_grad = False + # set_state_dict(model, StateDictType.LOCAL_STATE_DICT) return model diff --git a/fastvideo/v1/pipelines/composed_pipeline_base.py b/fastvideo/v1/pipelines/composed_pipeline_base.py index 4fc34f628..0242d0975 100644 --- a/fastvideo/v1/pipelines/composed_pipeline_base.py +++ b/fastvideo/v1/pipelines/composed_pipeline_base.py @@ -5,19 +5,25 @@ Base class for composed pipelines. This module defines the base class for pipelines that are composed of multiple stages. """ +import argparse import os from abc import ABC, abstractmethod from copy import deepcopy -from typing import Any, Dict, List, Optional, cast +from typing import Any, Dict, List, Optional, Union, cast import torch -from fastvideo.v1.fastvideo_args import FastVideoArgs +from fastvideo.v1.configs.pipelines import (PipelineConfig, + get_pipeline_config_cls_for_name) +from fastvideo.v1.distributed import (init_distributed_environment, + initialize_model_parallel, + model_parallel_is_initialized) +from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs from fastvideo.v1.logger import init_logger from fastvideo.v1.models.loader.component_loader import PipelineComponentLoader from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.v1.pipelines.stages import PipelineStage -from fastvideo.v1.utils import (maybe_download_model, +from fastvideo.v1.utils import (maybe_download_model, shallow_asdict, verify_model_config_and_directory) logger = init_logger(__name__) @@ -39,15 +45,20 @@ class ComposedPipelineBase(ABC): def __init__(self, model_path: str, fastvideo_args: FastVideoArgs, - config: Optional[Dict[str, Any]] = None): + config: Optional[Dict[str, Any]] = None, + required_config_modules: Optional[List[str]] = None): """ Initialize the pipeline. After __init__, the pipeline should be ready to use. The pipeline should be stateless and not hold any batch state. """ + self.fastvideo_args = fastvideo_args self.model_path = model_path self._stages: List[PipelineStage] = [] self._stage_name_mapping: Dict[str, PipelineStage] = {} + if required_config_modules is not None: + self._required_config_modules = required_config_modules + if self._required_config_modules is None: raise NotImplementedError( "Subclass must set _required_config_modules") @@ -59,16 +70,120 @@ class ComposedPipelineBase(ABC): else: self.config = config + self.maybe_init_distributed_environment(fastvideo_args) + # Load modules directly in initialization logger.info("Loading pipeline modules...") self.modules = self.load_modules(fastvideo_args) + if fastvideo_args.training_mode: + if fastvideo_args.log_validation: + self.initialize_validation_pipeline(fastvideo_args) + self.initialize_training_pipeline(fastvideo_args) + self.initialize_pipeline(fastvideo_args) - logger.info("Creating pipeline stages...") - self.create_pipeline_stages(fastvideo_args) + # logger.info("Creating pipeline stages...") + # self.create_pipeline_stages(fastvideo_args) - def get_module(self, module_name: str) -> Any: + if fastvideo_args.training_mode: + logger.info("Creating training pipeline stages...") + self.create_training_stages(fastvideo_args) + else: + logger.info("Creating pipeline stages...") + self.create_pipeline_stages(fastvideo_args) + + def initialize_training_pipeline(self, fastvideo_args: FastVideoArgs): + raise NotImplementedError( + "if training_mode is True, the pipeline must implement this method") + + def initialize_validation_pipeline(self, fastvideo_args: FastVideoArgs): + raise NotImplementedError( + "if log_validation is True, the pipeline must implement this method" + ) + + @classmethod + def from_pretrained(cls, + model_path: str, + device: Optional[str] = None, + torch_dtype: Optional[torch.dtype] = None, + pipeline_config: Optional[ + Union[str + | PipelineConfig]] = None, + args: Optional[argparse.Namespace] = None, + required_config_modules: Optional[List[str]] = None, + **kwargs) -> "ComposedPipelineBase": + config = None + # 1. If users provide a pipeline config, it will override the default pipeline config + if isinstance(pipeline_config, PipelineConfig): + config = pipeline_config + else: + config_cls = get_pipeline_config_cls_for_name(model_path) + if config_cls is not None: + config = config_cls() + if isinstance(pipeline_config, str): + config.load_from_json(pipeline_config) + + # 2. If users also provide some kwargs, it will override the pipeline config. + # The user kwargs shouldn't contain model config parameters! + if config is None: + logger.warning("No config found for model %s, using default config", + model_path) + config_args = kwargs + else: + config_args = shallow_asdict(config) + config_args.update(kwargs) + + if config_args.get("inference_mode"): + fastvideo_args = FastVideoArgs(model_path=model_path, + device_str=device or "cuda" if + torch.cuda.is_available() else "cpu", + **config_args) + else: + assert args is not None, "args must be provided for training mode" + fastvideo_args = TrainingArgs.from_cli_args(args) + + fastvideo_args.model_path = model_path + fastvideo_args.device_str = device or "cuda" if torch.cuda.is_available( + ) else "cpu" + for key, value in config_args.items(): + setattr(fastvideo_args, key, value) + + # fastvideo_args = FastVideoArgs( + # model_path=model_path, + # device_str=device or "cuda" if torch.cuda.is_available() else "cpu", + # **config_args) + fastvideo_args.check_fastvideo_args() + + return cls(model_path, + fastvideo_args, + required_config_modules=required_config_modules) + + def maybe_init_distributed_environment(self, fastvideo_args: FastVideoArgs): + if model_parallel_is_initialized(): + return + local_rank = int(os.environ.get("LOCAL_RANK", -1)) + world_size = int(os.environ.get("WORLD_SIZE", -1)) + rank = int(os.environ.get("RANK", -1)) + + if local_rank == -1 or world_size == -1 or rank == -1: + raise ValueError( + "Local rank, world size, and rank must be set. Use torchrun to launch the script." + ) + + torch.cuda.set_device(local_rank) + init_distributed_environment(world_size=world_size, + rank=rank, + local_rank=local_rank) + initialize_model_parallel( + tensor_model_parallel_size=fastvideo_args.tp_size, + sequence_model_parallel_size=fastvideo_args.sp_size) + device = torch.device(f"cuda:{local_rank}") + fastvideo_args.device = device + + def get_module(self, module_name: str, default_value: Any = None) -> Any: + if module_name not in self.modules: + return default_value return self.modules[module_name] def add_module(self, module_name: str, module: Any): @@ -114,6 +229,19 @@ class ComposedPipelineBase(ABC): """ raise NotImplementedError + # @abstractmethod + # def create_validation_stages(self, fastvideo_args: FastVideoArgs): + # """ + # Create the validation pipeline stages. + # """ + # raise NotImplementedError + + def create_training_stages(self, fastvideo_args: FastVideoArgs): + """ + Create the training pipeline stages. + """ + raise NotImplementedError + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): """ Initialize the pipeline. @@ -140,13 +268,16 @@ class ComposedPipelineBase(ABC): if module_name not in modules_config: raise ValueError( f"model_index.json must contain a {module_name} module") - logger.info("Diffusers config passed sanity checks") # all the component models used by the pipeline + required_modules = self.required_config_modules + logger.info("Loading required modules: %s", required_modules) + modules = {} for module_name, (transformers_or_diffusers, architecture) in modules_config.items(): - if module_name not in self.required_config_modules: + if module_name not in required_modules: + logger.info("Skipping module %s", module_name) continue component_model_path = os.path.join(self.model_path, module_name) module = PipelineComponentLoader.load_module( @@ -163,7 +294,6 @@ class ComposedPipelineBase(ABC): logger.warning("Overwriting module %s", module_name) modules[module_name] = module - required_modules = self.required_config_modules # Check if all required modules were loaded for module_name in required_modules: if module_name not in modules or modules[module_name] is None: diff --git a/fastvideo/v1/pipelines/distill_pipeline.py b/fastvideo/v1/pipelines/distill_pipeline.py new file mode 100644 index 000000000..9796341f5 --- /dev/null +++ b/fastvideo/v1/pipelines/distill_pipeline.py @@ -0,0 +1,303 @@ +import os +import sys +import time +from collections import deque + +import torch +from tqdm.auto import tqdm + +# import torch.distributed as dist +import wandb +from fastvideo.utils.checkpoint import save_checkpoint +from fastvideo.utils.communications import sp_parallel_dataloader_wrapper +from fastvideo.utils.validation import log_validation +from fastvideo.v1.distributed import cleanup_dist_env_and_memory, get_sp_group +from fastvideo.v1.fastvideo_args import FastVideoArgs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.pipelines import ComposedPipelineBase +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.pipelines.stages import (ConditioningStage, DecodingStage, + DenoisingStage, InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage) + +logger = init_logger(__name__) + + +class WanTrainingPipeline(ComposedPipelineBase): # == distill_one_step + _required_config_modules = ["scheduler", "transformer"] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + self.add_stage(stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + )) + + self.add_stage(stage_name="conditioning_stage", + stage=ConditioningStage()) + + self.add_stage(stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"))) + + self.add_stage(stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ): + device = fastvideo_args.device + local_rank = int(os.environ.get("LOCAL_RANK", -1)) + rank = int(os.environ.get("RANK", -1)) + assert rank != -1 + assert local_rank != -1 + sp_group = get_sp_group() + world_size = sp_group.world_size + rank = sp_group.rank + args = fastvideo_args + transformer = self.get_module("transformer") + teacher_transformer = self.get_module("teacher_transformer") + ema_transformer = None + assert not fastvideo_args.use_ema, "ema is not supported now" + assert teacher_transformer is not None + assert transformer is not None + train_dataset = self.train_dataset + train_dataloader = self.train_dataloader + init_steps = self.init_steps + lr_scheduler = self.lr_scheduler + optimizer = self.optimizer + noise_scheduler = self.noise_scheduler + solver = self.solver + noise_random_generator = None + uncond_prompt_embed = self.uncond_prompt_embed + uncond_prompt_mask = self.uncond_prompt_mask + + # Train! + total_batch_size = (world_size * args.gradient_accumulation_steps / + args.sp_size * args.train_sp_batch_size) + logger.info("***** Running training *****") + logger.info(" Num examples = %s", len(train_dataset)) + logger.info(" Dataloader size = %s", len(train_dataloader)) + logger.info(" Num Epochs = %s", args.num_train_epochs) + logger.info(" Resume training from step %s", init_steps) + logger.info(" Instantaneous batch size per device = %s", + args.train_batch_size) + logger.info( + " Total train batch size (w. data & sequence parallel, accumulation) = %s", + total_batch_size) + logger.info(" Gradient Accumulation steps = %s", + args.gradient_accumulation_steps) + logger.info(" Total optimization steps = %s", args.max_train_steps) + logger.info( + " Total training parameters per FSDP shard = %s B", + sum(p.numel() + for p in transformer.parameters() if p.requires_grad) / 1e9) + # print dtype + logger.info(" Master weight dtype: %s", + transformer.parameters().__next__().dtype) + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + assert NotImplementedError( + "resume_from_checkpoint is not supported now.") + # TODO + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=init_steps, + desc="Steps", + # Only show the progress bar once on each machine. + disable=local_rank > 0, + ) + + loader = sp_parallel_dataloader_wrapper( + train_dataloader, + device, + args.train_batch_size, + args.sp_size, + args.train_sp_batch_size, + ) + # loader = self.get_module("train_dataloader") + + step_times = deque(maxlen=100) + + # todo future + for i in range(init_steps): + next(loader) + + # log_validation(args, transformer, device, + # torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False) + def get_num_phases(multi_phased_distill_schedule, step): + # step-phase,step-phase + multi_phases = multi_phased_distill_schedule.split(",") + phase = multi_phases[-1].split("-")[-1] + for step_phases in multi_phases: + phase_step, phase = step_phases.split("-") + if step <= int(phase_step): + return int(phase) + return phase + + for step in range(init_steps + 1, args.max_train_steps + 1): + start_time = time.time() + assert args.multi_phased_distill_schedule is not None + num_phases = get_num_phases(args.multi_phased_distill_schedule, + step) + + loss, grad_norm, pred_norm = self.distill_one_step( + transformer, + args.model_type, + teacher_transformer, + ema_transformer, + optimizer, + lr_scheduler, + loader, + noise_scheduler, + solver, + noise_random_generator, + args.gradient_accumulation_steps, + args.sp_size, + args.max_grad_norm, + uncond_prompt_embed, + uncond_prompt_mask, + args.num_euler_timesteps, + num_phases, + args.not_apply_cfg_solver, + args.distill_cfg, + args.ema_decay, + args.pred_decay_weight, + args.pred_decay_type, + args.hunyuan_teacher_disable_cfg, + ) + + step_time = time.time() - start_time + step_times.append(step_time) + avg_step_time = sum(step_times) / len(step_times) + + progress_bar.set_postfix({ + "loss": f"{loss:.4f}", + "step_time": f"{step_time:.2f}s", + "grad_norm": grad_norm, + "phases": num_phases, + }) + progress_bar.update(1) + if rank <= 0: + wandb.log( + { + "train_loss": + loss, + "learning_rate": + lr_scheduler.get_last_lr()[0], + "step_time": + step_time, + "avg_step_time": + avg_step_time, + "grad_norm": + grad_norm, + "pred_fro_norm": + pred_norm["fro"], # codespell:ignore + "pred_largest_singular_value": + pred_norm["largest singular value"], + "pred_absolute_mean": + pred_norm["absolute mean"], + "pred_absolute_max": + pred_norm["absolute max"], + }, + step=step, + ) + if step % args.checkpointing_steps == 0: + if args.use_lora: + # Save LoRA weights + raise NotImplementedError("lora is not supported now") + # save_lora_checkpoint(transformer, optimizer, rank, + # args.output_dir, step) + else: + # Your existing checkpoint saving code + if args.use_ema: + raise NotImplementedError("ema is not supported now") + save_checkpoint(ema_transformer, rank, args.output_dir, + step) + else: + save_checkpoint(transformer, rank, args.output_dir, + step) + + sp_group.barrier() + if args.log_validation and step % args.validation_steps == 0: + log_validation( + args, + transformer, + device, + torch.bfloat16, + step, + scheduler_type=args.scheduler_type, + shift=args.shift, + num_euler_timesteps=args.num_euler_timesteps, + linear_quadratic_threshold=args.linear_quadratic_threshold, + linear_range=args.linear_range, + ema=False, + ) + if args.use_ema: + log_validation( + args, + ema_transformer, + device, + torch.bfloat16, + step, + scheduler_type=args.scheduler_type, + shift=args.shift, + num_euler_timesteps=args.num_euler_timesteps, + linear_quadratic_threshold=args. + linear_quadratic_threshold, + linear_range=args.linear_range, + ema=True, + ) + + if args.use_lora: + raise NotImplementedError("lora is not supported now") + save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, + args.max_train_steps) + else: + save_checkpoint(transformer, rank, args.output_dir, + args.max_train_steps) + + if get_sp_group(): + cleanup_dist_env_and_memory() + + +def main(args): + logger.info("Starting training pipeline...") + pipeline = WanTrainingPipeline.from_pretrained( + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", args=args) + args = pipeline.fastvideo_args + pipeline.forward(None, args) + logger.info("Training pipeline done") + + +if __name__ == "__main__": + argv = sys.argv + from fastvideo.v1.fastvideo_args import TrainingArgs + from fastvideo.v1.utils import FlexibleArgumentParser + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + parser = FastVideoArgs.add_cli_args(parser) + args = parser.parse_args() + print(args) + main(args) diff --git a/fastvideo/v1/pipelines/training_pipeline.py b/fastvideo/v1/pipelines/training_pipeline.py new file mode 100644 index 000000000..b0ff5be82 --- /dev/null +++ b/fastvideo/v1/pipelines/training_pipeline.py @@ -0,0 +1,635 @@ +import gc +import math +import os +import sys +import time +from abc import ABC, abstractmethod +from collections import deque +from copy import deepcopy + +import imageio +import numpy as np +import torch +import torchvision +from diffusers.optimization import get_scheduler +from einops import rearrange +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from tqdm.auto import tqdm + +# import torch.distributed as dist +import wandb +from fastvideo.dataset.latent_datasets import (LatentDataset, + latent_collate_function) +from fastvideo.distill.solver import EulerSolver +from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_dit_input +from fastvideo.utils.checkpoint import save_checkpoint_v1 +from fastvideo.utils.communications import sp_parallel_dataloader_wrapper +from fastvideo.utils.dataset_utils import LengthGroupedSampler +from fastvideo.v1.configs.sample import SamplingParam +from fastvideo.v1.distributed import cleanup_dist_env_and_memory, get_sp_group +from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.v1.forward_context import set_forward_context +from fastvideo.v1.logger import init_logger +from fastvideo.v1.pipelines import ComposedPipelineBase +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.pipelines.wan.wan_pipeline import WanPipeline + +logger = init_logger(__name__) + + +def compute_density_for_timestep_sampling( + weighting_scheme: str, + batch_size: int, + generator, + logit_mean: float = None, + logit_std: float = None, + mode_scale: float = None, +): + """ + Compute the density for sampling the timesteps when doing SD3 training. + + Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528. + + SD3 paper reference: https://arxiv.org/abs/2403.03206v1. + """ + if weighting_scheme == "logit_normal": + # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$). + u = torch.normal( + mean=logit_mean, + std=logit_std, + size=(batch_size, ), + device="cpu", + generator=generator, + ) + u = torch.nn.functional.sigmoid(u) + elif weighting_scheme == "mode": + u = torch.rand(size=(batch_size, ), device="cpu", generator=generator) + u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2)**2 - 1 + u) + else: + u = torch.rand(size=(batch_size, ), device="cpu", generator=generator) + return u + + +def get_sigmas(noise_scheduler, + device, + timesteps, + n_dim=4, + dtype=torch.float32): + sigmas = noise_scheduler.sigmas.to(device=device, dtype=dtype) + schedule_timesteps = noise_scheduler.timesteps.to(device) + timesteps = timesteps.to(device) + step_indices = [(schedule_timesteps == t).nonzero().item() + for t in timesteps] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + +class TrainingPipeline(ComposedPipelineBase, ABC): + """ + A pipeline for training a model. All training pipelines should inherit from this class. + All reusable components and code should be implemented in this class. + """ + _required_config_modules = ["scheduler", "transformer"] + + def initialize_training_pipeline(self, fastvideo_args: TrainingArgs): + device = fastvideo_args.device + sp_group = get_sp_group() + world_size = sp_group.world_size + rank = sp_group.rank + + # assert isinstance(fastvideo_args, TrainingArgs) + self.modules["transformer"].requires_grad_(True) + self.modules["transformer"].train() + + # self.modules["text_encoder"].requires_grad_(False) + # self.modules["vae"].requires_grad_(False) + # self.modules["tokenizer"].requires_grad_(False) + + # self.modules["teacher_transformer"] = deepcopy( + # self.modules["transformer"]) + # self.modules["teacher_transformer"].requires_grad_(False) + + transformer = self.modules["transformer"] + # teacher_transformer = self.modules["teacher_transformer"] + args = fastvideo_args + + noise_scheduler = self.modules["scheduler"] + solver = EulerSolver( + noise_scheduler.sigmas.numpy()[::-1], + noise_scheduler.config.num_train_timesteps, + euler_timesteps=args.num_euler_timesteps, + ) + solver.to(device) + params_to_optimize = transformer.parameters() + params_to_optimize = list( + filter(lambda p: p.requires_grad, params_to_optimize)) + + optimizer = torch.optim.AdamW( + params_to_optimize, + lr=args.learning_rate, + betas=(0.9, 0.999), + weight_decay=args.weight_decay, + eps=1e-8, + ) + + init_steps = 0 + logger.info("optimizer: %s", optimizer) + + # todo add lr scheduler + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * world_size, + num_training_steps=args.max_train_steps * world_size, + num_cycles=args.lr_num_cycles, + power=args.lr_power, + last_epoch=init_steps - 1, + ) + + train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, + args.cfg) + uncond_prompt_embed = train_dataset.uncond_prompt_embed + uncond_prompt_mask = train_dataset.uncond_prompt_mask + sampler = (LengthGroupedSampler( + args.train_batch_size, + rank=rank, + world_size=world_size, + lengths=train_dataset.lengths, + group_frame=args.group_frame, + group_resolution=args.group_resolution, + ) if (args.group_frame or args.group_resolution) else + DistributedSampler(train_dataset, + rank=rank, + num_replicas=world_size, + shuffle=False)) + + train_dataloader = DataLoader( + train_dataset, + sampler=sampler, + collate_fn=latent_collate_function, + pin_memory=True, + batch_size=args.train_batch_size, + num_workers=args.dataloader_num_workers, + drop_last=True, + ) + self.lr_scheduler = lr_scheduler + self.train_dataset = train_dataset + self.train_dataloader = train_dataloader + self.init_steps = init_steps + self.optimizer = optimizer + self.noise_scheduler = noise_scheduler + self.solver = solver + self.uncond_prompt_embed = uncond_prompt_embed + self.uncond_prompt_mask = uncond_prompt_mask + # self.noise_random_generator = noise_random_generator + + num_update_steps_per_epoch = math.ceil( + len(train_dataloader) / args.gradient_accumulation_steps * + args.sp_size / args.train_sp_batch_size) + args.num_train_epochs = math.ceil(args.max_train_steps / + num_update_steps_per_epoch) + + if rank <= 0: + project = args.tracker_project_name or "fastvideo" + wandb.init(project=project, config=args) + + def initialize_validation_pipeline(self, fastvideo_args: FastVideoArgs): + + required_config_modules = [ + "text_encoder", "tokenizer", "vae", "scheduler" + ] + args = deepcopy(fastvideo_args) + args.inference_mode = True + validation_pipeline = WanPipeline.from_pretrained( + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + args=args, + required_config_modules=required_config_modules) + + self.validation_pipeline = validation_pipeline + + @abstractmethod + def train_one_step(self, transformer, model_type, optimizer, lr_scheduler, + loader, noise_scheduler, noise_random_generator, + gradient_accumulation_steps, sp_size, + precondition_outputs, max_grad_norm, weighting_scheme, + logit_mean, logit_std, mode_scale): + """ + Train one step of the model. + """ + raise NotImplementedError( + "Training pipeline must implement this method") + + def log_validation(self, transformer, fastvideo_args, global_step): + if not fastvideo_args.log_validation: + return + if self.validation_pipeline is None: + raise ValueError("Validation pipeline is not set") + + # Create sampling parameters if not provided + sampling_param = SamplingParam.from_pretrained( + fastvideo_args.model_path) + + # Prepare validation prompts + embe_dir = os.path.join(fastvideo_args.validation_prompt_dir, + "prompt_embed") + mask_dir = os.path.join(fastvideo_args.validation_prompt_dir, + "prompt_attention_mask") + embeds = sorted([f for f in os.listdir(embe_dir)]) + masks = sorted([f for f in os.listdir(mask_dir)]) + num_embeds = len(embeds) + + # Add the transformer to the validation pipeline + self.validation_pipeline.add_module("transformer", transformer) + self.validation_pipeline.denoising_stage.transformer = transformer + + # Process each validation prompt + videos = [] + for i in range(num_embeds): + prompt_embed_path = os.path.join(embe_dir, f"{embeds[i]}") + prompt_mask_path = os.path.join(mask_dir, f"{masks[i]}") + prompt_embeds = torch.load(prompt_embed_path, + map_location="cpu").to( + fastvideo_args.device).unsqueeze(0) + prompt_attention_mask = torch.load( + prompt_mask_path, + map_location="cpu").to(fastvideo_args.device).unsqueeze(0) + + # Calculate sizes + # target_height = align_to(fastvideo_args.num_height, 16) + # target_width = align_to(fastvideo_args.num_width, 16) + 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] + + # Prepare batch for validation + batch = ForwardBatch( + seed=sampling_param.seed, + data_type="video", + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + height=sampling_param.height, + width=sampling_param.width, + num_frames=sampling_param.num_frames, + # num_inference_steps=fastvideo_args.validation_sampling_steps, + num_inference_steps=10, + # guidance_scale=fastvideo_args.validation_guidance_scale, + guidance_scale=sampling_param.guidance_scale, + n_tokens=n_tokens, + eta=0.0, + extra={}, + ) + + # Run validation inference + with torch.inference_mode(): + output_batch = self.validation_pipeline.forward( + batch, fastvideo_args) + samples = output_batch + + # Process outputs + video = rearrange(samples, "b c t h w -> t b c h w") + frames = [] + for x in video: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + frames.append((x * 255).numpy().astype(np.uint8)) + videos.append(frames) + + # Log validation results + if fastvideo_args.global_rank == 0: + video_filenames = [] + for i, video in enumerate(videos): + filename = os.path.join(fastvideo_args.output_dir, + f"validation_step_{global_step}.mp4") + imageio.mimsave(filename, video, fps=fastvideo_args.fps) + video_filenames.append(filename) + + logs = { + "validation_videos": + [wandb.Video(filename) for filename in video_filenames] + } + wandb.log(logs, step=global_step) + + gc.collect() + torch.cuda.empty_cache() + + +# class ValidationPipeline(WanPipeline): +# # we load everything except transformer, since we will use the transformer +# # from the training pipeline +# _required_config_modules = [ +# "text_encoder", "tokenizer", "vae", "scheduler" +# ] + + +class WanTrainingPipeline(TrainingPipeline): + """ + A training pipeline for Wan. + """ + _required_config_modules = ["scheduler", "transformer"] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + pass + + def create_training_stages(self, fastvideo_args: FastVideoArgs): + pass + + def train_one_step( + self, + transformer, + model_type, + optimizer, + lr_scheduler, + loader, + noise_scheduler, + noise_random_generator, + gradient_accumulation_steps, + sp_size, + precondition_outputs, + max_grad_norm, + weighting_scheme, + logit_mean, + logit_std, + mode_scale, + ): + total_loss = 0.0 + optimizer.zero_grad() + for _ in range(gradient_accumulation_steps): + ( + latents, + encoder_hidden_states, + latents_attention_mask, + encoder_attention_mask, + ) = next(loader) + latents = normalize_dit_input(model_type, latents) + batch_size = latents.shape[0] + noise = torch.randn_like(latents) + u = compute_density_for_timestep_sampling( + weighting_scheme=weighting_scheme, + batch_size=batch_size, + generator=noise_random_generator, + logit_mean=logit_mean, + logit_std=logit_std, + mode_scale=mode_scale, + ) + indices = (u * noise_scheduler.config.num_train_timesteps).long() + timesteps = noise_scheduler.timesteps[indices].to( + device=latents.device) + if sp_size > 1: + # Make sure that the timesteps are the same across all sp processes. + sp_group = get_sp_group() + sp_group.broadcast(timesteps, src=0) + sigmas = get_sigmas( + noise_scheduler, + latents.device, + timesteps, + n_dim=latents.ndim, + dtype=latents.dtype, + ) + noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise + with torch.autocast("cuda", dtype=torch.bfloat16): + input_kwargs = { + "hidden_states": noisy_model_input, + "encoder_hidden_states": encoder_hidden_states, + "timestep": timesteps, + "encoder_attention_mask": encoder_attention_mask, # B, L + "return_dict": False, + } + if 'hunyuan' in model_type: + input_kwargs["guidance"] = torch.tensor( + [1000.0], + device=noisy_model_input.device, + dtype=torch.bfloat16) + with set_forward_context(current_timestep=timesteps, + attn_metadata=None): + model_pred = transformer(**input_kwargs)[0] + + if precondition_outputs: + model_pred = noisy_model_input - model_pred * sigmas + if precondition_outputs: + target = latents + else: + target = noise - latents + + loss = (torch.mean((model_pred.float() - target.float())**2) / + gradient_accumulation_steps) + + loss.backward() + + avg_loss = loss.detach().clone() + sp_group = get_sp_group() + sp_group.all_reduce(avg_loss, op=torch.distributed.ReduceOp.AVG) + # dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG) + total_loss += avg_loss.item() + + # TODO(will): clip grad norm + # grad_norm = transformer.clip_grad_norm_(max_grad_norm) + optimizer.step() + lr_scheduler.step() + return total_loss, 0.0 + # return total_loss, grad_norm.item() + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ): + device = fastvideo_args.device + local_rank = int(os.environ.get("LOCAL_RANK", -1)) + rank = int(os.environ.get("RANK", -1)) + assert rank != -1 + assert local_rank != -1 + sp_group = get_sp_group() + world_size = sp_group.world_size + rank = sp_group.rank + args = fastvideo_args + transformer = self.get_module("transformer") + # teacher_transformer = self.get_module("teacher_transformer") + # ema_transformer = None + assert not fastvideo_args.use_ema, "ema is not supported now" + # assert teacher_transformer is not None + assert transformer is not None + train_dataset = self.train_dataset + train_dataloader = self.train_dataloader + init_steps = self.init_steps + lr_scheduler = self.lr_scheduler + optimizer = self.optimizer + noise_scheduler = self.noise_scheduler + solver = self.solver + noise_random_generator = None + uncond_prompt_embed = self.uncond_prompt_embed + uncond_prompt_mask = self.uncond_prompt_mask + + from diffusers import FlowMatchEulerDiscreteScheduler + noise_scheduler = FlowMatchEulerDiscreteScheduler() + + # params_to_optimize = transformer.parameters() + # params_to_optimize = list(filter(lambda p: p.requires_grad, params_to_optimize)) + + # optimizer = torch.optim.AdamW( + # params_to_optimize, + # lr=args.learning_rate, + # betas=(0.9, 0.999), + # weight_decay=args.weight_decay, + # eps=1e-8, + # ) + + # init_steps = 0 + # if args.resume_from_lora_checkpoint: + # transformer, optimizer, init_steps = resume_lora_optimizer(transformer, args.resume_from_lora_checkpoint, + # optimizer) + # main_print(f"optimizer: {optimizer}") + + # lr_scheduler = get_scheduler( + # args.lr_scheduler, + # optimizer=optimizer, + # num_warmup_steps=args.lr_warmup_steps, + # num_training_steps=args.max_train_steps, + # num_cycles=args.lr_num_cycles, + # power=args.lr_power, + # last_epoch=init_steps - 1, + # ) + + # Train! + total_batch_size = (world_size * args.gradient_accumulation_steps / + args.sp_size * args.train_sp_batch_size) + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Dataloader size = {len(train_dataloader)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Resume training from step {init_steps}") + logger.info( + f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info( + f" Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}" + ) + logger.info( + f" Gradient Accumulation steps = {args.gradient_accumulation_steps}" + ) + logger.info(f" Total optimization steps = {args.max_train_steps}") + logger.info( + f" Total training parameters per FSDP shard = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B" + ) + # print dtype + logger.info( + f" Master weight dtype: {transformer.parameters().__next__().dtype}" + ) + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + assert NotImplementedError( + "resume_from_checkpoint is not supported now.") + # TODO + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=init_steps, + desc="Steps", + # Only show the progress bar once on each machine. + disable=local_rank > 0, + ) + + loader = sp_parallel_dataloader_wrapper( + train_dataloader, + device, + args.train_batch_size, + args.sp_size, + args.train_sp_batch_size, + ) + + step_times = deque(maxlen=100) + + # todo future + for i in range(init_steps): + next(loader) + for step in range(init_steps + 1, args.max_train_steps + 1): + start_time = time.perf_counter() + loss, grad_norm = self.train_one_step( + transformer, + # args.model_type, + "wan", + optimizer, + lr_scheduler, + loader, + noise_scheduler, + noise_random_generator, + args.gradient_accumulation_steps, + args.sp_size, + args.precondition_outputs, + args.max_grad_norm, + args.weighting_scheme, + args.logit_mean, + args.logit_std, + args.mode_scale, + ) + + step_time = time.perf_counter() - start_time + step_times.append(step_time) + avg_step_time = sum(step_times) / len(step_times) + + progress_bar.set_postfix({ + "loss": f"{loss:.4f}", + "step_time": f"{step_time:.2f}s", + "grad_norm": grad_norm, + }) + progress_bar.update(1) + if rank <= 0: + wandb.log( + { + "train_loss": loss, + "learning_rate": lr_scheduler.get_last_lr()[0], + "step_time": step_time, + "avg_step_time": avg_step_time, + "grad_norm": grad_norm, + }, + step=step, + ) + if step % args.checkpointing_steps == 0: + if args.use_lora: + raise NotImplementedError("LoRA is not supported now") + # Save LoRA weights + save_lora_checkpoint(transformer, optimizer, rank, + args.output_dir, step, pipe) + else: + # Your existing checkpoint saving code + save_checkpoint_v1(transformer, rank, args.output_dir, step) + transformer.train() + sp_group.barrier() + if args.log_validation and step % args.validation_steps == 0: + self.log_validation(transformer, args, step) + + if args.use_lora: + raise NotImplementedError("LoRA is not supported now") + # save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps, pipe) + else: + save_checkpoint_v1(transformer, rank, args.output_dir, + args.max_train_steps) + + if get_sp_group(): + cleanup_dist_env_and_memory() + + +def main(args): + logger.info("Starting training pipeline...") + + pipeline = WanTrainingPipeline.from_pretrained( + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", args=args) + args = pipeline.fastvideo_args + pipeline.forward(None, args) + logger.info("Training pipeline done") + + +if __name__ == "__main__": + argv = sys.argv + from fastvideo.v1.fastvideo_args import TrainingArgs + from fastvideo.v1.utils import FlexibleArgumentParser + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + parser = FastVideoArgs.add_cli_args(parser) + args = parser.parse_args() + print(args) + main(args) diff --git a/fastvideo/v1/pipelines/wan/wan_latent_pipeline.py b/fastvideo/v1/pipelines/wan/wan_latent_pipeline.py new file mode 100644 index 000000000..659eb8c65 --- /dev/null +++ b/fastvideo/v1/pipelines/wan/wan_latent_pipeline.py @@ -0,0 +1,19 @@ +from fastvideo.v1.fastvideo_args import FastVideoArgs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch + +logger = init_logger(__name__) + + +class WanLatentPipeline(ComposedPipelineBase): + _required_config_modules = ["text_encoder", "tokenizer", "vae"] + + # def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + pass + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs): + logger.info("WAN Latent Pipeline forward") + pass diff --git a/fastvideo/v1/pipelines/wan/wan_pipeline.py b/fastvideo/v1/pipelines/wan/wan_pipeline.py index 8159d5cce..aea126d20 100644 --- a/fastvideo/v1/pipelines/wan/wan_pipeline.py +++ b/fastvideo/v1/pipelines/wan/wan_pipeline.py @@ -48,7 +48,7 @@ class WanPipeline(ComposedPipelineBase): self.add_stage(stage_name="latent_preparation_stage", stage=LatentPreparationStage( scheduler=self.get_module("scheduler"), - transformer=self.get_module("transformer"))) + transformer=self.get_module("transformer", None))) self.add_stage(stage_name="denoising_stage", stage=DenoisingStage( diff --git a/fastvideo/v1_distill.py b/fastvideo/v1_distill.py new file mode 100644 index 000000000..0279d98e7 --- /dev/null +++ b/fastvideo/v1_distill.py @@ -0,0 +1,953 @@ +# !/bin/python3 +# isort: skip_file +import argparse +import math +import os +import time +from collections import deque +from copy import deepcopy + +import torch +import torch.distributed as dist +import wandb +from accelerate.utils import set_seed +from diffusers import FlowMatchEulerDiscreteScheduler +from diffusers.optimization import get_scheduler +from diffusers.utils import check_min_version +from peft import LoraConfig +# from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +# from torch.distributed.fsdp import ShardingStrategy +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from tqdm.auto import tqdm + +from fastvideo.dataset.latent_datasets import (LatentDataset, + latent_collate_function) +from fastvideo.distill.solver import EulerSolver, extract_into_tensor +from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_dit_input +from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule +from fastvideo.utils.checkpoint import (resume_lora_optimizer, save_checkpoint, + save_lora_checkpoint) +from fastvideo.utils.communications import (broadcast, + sp_parallel_dataloader_wrapper) +from fastvideo.utils.dataset_utils import LengthGroupedSampler +from fastvideo.utils.fsdp_util import (apply_fsdp_checkpointing) +# from fastvideo.utils.load import load_transformer +# from fastvideo.utils.parallel_states import (destroy_sequence_parallel_group, get_sequence_parallel_state, +# initialize_sequence_parallel_state) +from fastvideo.utils.validation import log_validation + +from fastvideo.v1.models.loader.component_loader import TransformerLoader +from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel +from fastvideo.v1.configs.models.dits import WanVideoConfig +from fastvideo.v1.fastvideo_args import FastVideoArgs +from fastvideo.v1.utils import maybe_download_model +from fastvideo.v1.forward_context import set_forward_context + +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.31.0") + + +def main_print(content): + if int(os.environ["LOCAL_RANK"]) <= 0: + print(content) + + +# def reshard_fsdp(model): +# for m in FSDP.fsdp_modules(model): +# if m._has_params and m.sharding_strategy is not ShardingStrategy.NO_SHARD: +# torch.distributed.fsdp._runtime_utils._reshard(m, m._handle, True) + + +def get_norm(model_pred, norms, gradient_accumulation_steps): + fro_norm = ( + torch.linalg.matrix_norm(model_pred, ord="fro") / # codespell:ignore + gradient_accumulation_steps) + largest_singular_value = (torch.linalg.matrix_norm(model_pred, ord=2) / + gradient_accumulation_steps) + absolute_mean = torch.mean( + torch.abs(model_pred)) / gradient_accumulation_steps + absolute_max = torch.max( + torch.abs(model_pred)) / gradient_accumulation_steps + dist.all_reduce(fro_norm, op=dist.ReduceOp.AVG) + dist.all_reduce(largest_singular_value, op=dist.ReduceOp.AVG) + dist.all_reduce(absolute_mean, op=dist.ReduceOp.AVG) + norms["fro"] += torch.mean(fro_norm).item() # codespell:ignore + norms["largest singular value"] += torch.mean(largest_singular_value).item() + norms["absolute mean"] += absolute_mean.item() + norms["absolute max"] += absolute_max.item() + + +def distill_one_step( + transformer, + model_type, + teacher_transformer, + ema_transformer, + optimizer, + lr_scheduler, + loader, + noise_scheduler, + solver, + noise_random_generator, + gradient_accumulation_steps, + sp_size, + max_grad_norm, + uncond_prompt_embed, + uncond_prompt_mask, + num_euler_timesteps, + multiphase, + not_apply_cfg_solver, + distill_cfg, + ema_decay, + pred_decay_weight, + pred_decay_type, + hunyuan_teacher_disable_cfg, +): + total_loss = 0.0 + optimizer.zero_grad() + model_pred_norm = { + "fro": 0.0, # codespell:ignore + "largest singular value": 0.0, + "absolute mean": 0.0, + "absolute max": 0.0, + } + for _ in range(gradient_accumulation_steps): + ( + latents, + encoder_hidden_states, + latents_attention_mask, + encoder_attention_mask, + ) = next(loader) + model_input = normalize_dit_input(model_type, latents) + noise = torch.randn_like(model_input) + bsz = model_input.shape[0] + index = torch.randint(0, + num_euler_timesteps, (bsz, ), + device=model_input.device).long() + if sp_size > 1: + broadcast(index) + # Add noise according to flow matching. + # sigmas = get_sigmas(start_timesteps, n_dim=model_input.ndim, dtype=model_input.dtype) + sigmas = extract_into_tensor(solver.sigmas, index, model_input.shape) + sigmas_prev = extract_into_tensor(solver.sigmas_prev, index, + model_input.shape) + + timesteps = (sigmas * + noise_scheduler.config.num_train_timesteps).view(-1) + # if squeeze to [], unsqueeze to [1] + + timesteps_prev = (sigmas_prev * + noise_scheduler.config.num_train_timesteps).view(-1) + noisy_model_input = sigmas * noise + (1.0 - sigmas) * model_input + # Predict the noise residual + with torch.autocast("cuda", dtype=torch.bfloat16): + print(f"--> noisy_model_input.dtype: {noisy_model_input.dtype}") + print(f"--> noisy_model_input.shape: {noisy_model_input.shape}") + print( + f"--> encoder_hidden_states.dtype: {encoder_hidden_states.dtype}" + ) + print( + f"--> encoder_hidden_states.shape: {encoder_hidden_states.shape}" + ) + print(f"--> timesteps.dtype: {timesteps.dtype}") + print(f"--> timesteps.shape: {timesteps.shape}") + print( + f"--> encoder_attention_mask.dtype: {encoder_attention_mask.dtype}" + ) + print( + f"--> encoder_attention_mask.shape: {encoder_attention_mask.shape}" + ) + noisy_model_input = noisy_model_input.to(dtype=torch.bfloat16) + teacher_kwargs = { + "hidden_states": noisy_model_input, + "encoder_hidden_states": encoder_hidden_states, + "timestep": timesteps, + "encoder_attention_mask": encoder_attention_mask, # B, L + "return_dict": False, + } + if hunyuan_teacher_disable_cfg: + teacher_kwargs["guidance"] = torch.tensor( + [1000.0], + device=noisy_model_input.device, + dtype=torch.bfloat16) + # batch = ForwardBatch( + # enable_teacache=False, + # ) + with set_forward_context(current_timestep=0, + attn_metadata=None, + forward_batch=None, + fastvideo_args=None): + model_pred = transformer(**teacher_kwargs)[0] + + print(f"--> model_pred shape: {model_pred.shape}") + huber_c = 0.001 + target = torch.randn_like(model_pred) + loss = (torch.mean( + torch.sqrt((model_pred.float() - target.float())**2 + huber_c**2) - + huber_c) / gradient_accumulation_steps) + loss.backward() + print(f"--> loss: {loss.item()}") + assert False, "stop here" + + # if accelerator.is_main_process: + model_pred, end_index = solver.euler_style_multiphase_pred( + noisy_model_input, model_pred, index, multiphase) + with torch.no_grad(): + w = distill_cfg + with torch.autocast("cuda", dtype=torch.bfloat16): + cond_teacher_output = teacher_transformer( + noisy_model_input, + encoder_hidden_states, + timesteps, + encoder_attention_mask, # B, L + return_dict=False, + )[0].float() + if not_apply_cfg_solver: + uncond_teacher_output = cond_teacher_output + else: + # Get teacher model prediction on noisy_latents and unconditional embedding + with torch.autocast("cuda", dtype=torch.bfloat16): + uncond_teacher_output = teacher_transformer( + noisy_model_input, + uncond_prompt_embed.unsqueeze(0).expand(bsz, -1, -1), + timesteps, + uncond_prompt_mask.unsqueeze(0).expand(bsz, -1), + return_dict=False, + )[0].float() + teacher_output = uncond_teacher_output + w * (cond_teacher_output - + uncond_teacher_output) + x_prev = solver.euler_step(noisy_model_input, teacher_output, index) + + # 20.4.12. Get target LCM prediction on x_prev, w, c, t_n + with torch.no_grad(): + with torch.autocast("cuda", dtype=torch.bfloat16): + if ema_transformer is not None: + target_pred = ema_transformer( + x_prev.float(), + encoder_hidden_states, + timesteps_prev, + encoder_attention_mask, # B, L + return_dict=False, + )[0] + else: + target_pred = transformer( + x_prev.float(), + encoder_hidden_states, + timesteps_prev, + encoder_attention_mask, # B, L + return_dict=False, + )[0] + + target, end_index = solver.euler_style_multiphase_pred( + x_prev, target_pred, index, multiphase, True) + + huber_c = 0.001 + # loss = loss.mean() + loss = (torch.mean( + torch.sqrt((model_pred.float() - target.float())**2 + huber_c**2) - + huber_c) / gradient_accumulation_steps) + if pred_decay_weight > 0: + if pred_decay_type == "l1": + pred_decay_loss = ( + torch.mean(torch.sqrt(model_pred.float()**2)) * + pred_decay_weight / gradient_accumulation_steps) + loss += pred_decay_loss + elif pred_decay_type == "l2": + # essnetially k2? + pred_decay_loss = (torch.mean(model_pred.float()**2) * + pred_decay_weight / + gradient_accumulation_steps) + loss += pred_decay_loss + else: + assert NotImplementedError("pred_decay_type is not implemented") + + # calculate model_pred norm and mean + get_norm(model_pred.detach().float(), model_pred_norm, + gradient_accumulation_steps) + loss.backward() + + avg_loss = loss.detach().clone() + dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG) + total_loss += avg_loss.item() + + # update ema + if ema_transformer is not None: + reshard_fsdp(ema_transformer) + for p_averaged, p_model in zip(ema_transformer.parameters(), + transformer.parameters()): + with torch.no_grad(): + p_averaged.copy_( + torch.lerp(p_averaged.detach(), p_model.detach(), + 1 - ema_decay)) + + grad_norm = transformer.clip_grad_norm_(max_grad_norm) + optimizer.step() + lr_scheduler.step() + + return total_loss, grad_norm.item(), model_pred_norm + + +def main(args): + torch.backends.cuda.matmul.allow_tf32 = True + + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + rank = int(os.environ.get("RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + init_distributed_environment(world_size=world_size, + rank=rank, + local_rank=local_rank) + initialize_model_parallel(tensor_model_parallel_size=args.sp_size, + sequence_model_parallel_size=args.sp_size) + + # If passed along, set the training seed now. On GPU... + if args.seed is not None: + # TODO: t within the same seq parallel group should be the same. Noise should be different. + set_seed(args.seed + rank) + # We use different seeds for the noise generation in each process to ensure that the noise is different in a batch. + noise_random_generator = None + + # Handle the repository creation + if rank <= 0 and args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + # For mixed precision training we cast all non-trainable weights to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + + # Create model: + + print( + f"--> local_rank: {local_rank}, rank: {rank}, world_size: {world_size}") + main_print(f"--> using model pipeline {args.pretrained_model_name_or_path}") + + model_path = maybe_download_model(args.pretrained_model_name_or_path) + main_print(f"--> loading model from {model_path}") + + transformer_path = os.path.join(model_path, "transformer") + main_print(f"--> loading transformer from {transformer_path}") + + precision = torch.float32 if args.master_weight_type == "fp32" else torch.bfloat16 + + precision_str = "fp32" if precision == torch.float32 else "bf16" + print(f"--> precision: {precision_str}") + print(f"--> precision: {precision}") + + # transformer_path = os.path.join(args.pretrained_model_name_or_path, "transformer") + + fastvideo_args = FastVideoArgs(model_path=transformer_path, + use_cpu_offload=False, + precision=precision_str) + # fastvideo_args.dit_config = HunyuanVideoConfig() + fastvideo_args.dit_config = WanVideoConfig() + + fastvideo_args.device = torch.device(f"cuda:{local_rank}") + fastvideo_args.device_str = f"cuda:{local_rank}" + device = fastvideo_args.device + torch.cuda.set_device(device) + + loader = TransformerLoader() + print(f"--> loading transformer to device {device} on rank {rank}") + transformer = loader.load(transformer_path, "", + fastvideo_args).to(device, dtype=precision) + + # teacher_transformer = deepcopy(transformer) + if args.use_ema: + raise NotImplementedError("EMA is not supported for v1 distillation.") + ema_transformer = deepcopy(transformer) + else: + ema_transformer = None + + if args.use_lora: + raise NotImplementedError("LoRA is not supported for v1 distillation.") + assert args.model_type == "mochi", "LoRA is only supported for Mochi model." + transformer.requires_grad_(False) + transformer_lora_config = LoraConfig( + r=args.lora_rank, + lora_alpha=args.lora_alpha, + init_lora_weights=True, + target_modules=["to_k", "to_q", "to_v", "to_out.0"], + ) + transformer.add_adapter(transformer_lora_config) + + main_print( + f" Total training parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6} M" + ) + main_print( + f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}" + ) + + if args.use_lora: + raise NotImplementedError("LoRA is not supported for v1 distillation.") + transformer.config.lora_rank = args.lora_rank + transformer.config.lora_alpha = args.lora_alpha + transformer.config.lora_target_modules = [ + "to_k", "to_q", "to_v", "to_out.0" + ] + transformer._no_split_modules = no_split_modules + fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"]( + transformer) + + main_print("--> model loaded") + + if args.gradient_checkpointing: + apply_fsdp_checkpointing(transformer, no_split_modules, + args.selective_checkpointing) + apply_fsdp_checkpointing(teacher_transformer, no_split_modules, + args.selective_checkpointing) + if args.use_ema: + apply_fsdp_checkpointing(ema_transformer, no_split_modules, + args.selective_checkpointing) + # Set model as trainable. + transformer.train() + transformer.requires_grad_(True) + # teacher_transformer.requires_grad_(False) + if args.use_ema: + ema_transformer.requires_grad_(False) + noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=args.shift) + if args.scheduler_type == "pcm_linear_quadratic": + linear_steps = int(noise_scheduler.config.num_train_timesteps * + args.linear_range) + sigmas = linear_quadratic_schedule( + noise_scheduler.config.num_train_timesteps, + args.linear_quadratic_threshold, + linear_steps, + ) + sigmas = torch.tensor(sigmas).to(dtype=torch.float32) + else: + sigmas = noise_scheduler.sigmas + solver = EulerSolver( + sigmas.numpy()[::-1], + noise_scheduler.config.num_train_timesteps, + euler_timesteps=args.num_euler_timesteps, + ) + solver.to(device) + params_to_optimize = transformer.parameters() + # l = list(params_to_optimize) + # for p in params_to_optimize: + # main_print(type(p)) + # main_print(f"--> p: {p.shape}") + # main_print(f"--> p: {p.dtype}") + # main_print(f"--> p: {p.device}") + # main_print(f"--> p: {p.requires_grad}") + # main_print('=------------------------') + # break + # print(f"--> params_to_optimize: {list(params_to_optimize)}") + + params_to_optimize = list( + filter(lambda p: p.requires_grad, params_to_optimize)) + # print(f"--> params_to_optimize2: {params_to_optimize}") + + optimizer = torch.optim.AdamW( + params_to_optimize, + lr=args.learning_rate, + betas=(0.9, 0.999), + weight_decay=args.weight_decay, + eps=1e-8, + ) + # optimizer = None + + init_steps = 0 + if args.resume_from_lora_checkpoint: + transformer, optimizer, init_steps = resume_lora_optimizer( + transformer, args.resume_from_lora_checkpoint, optimizer) + main_print(f"optimizer: {optimizer}") + + # todo add lr scheduler + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * world_size, + num_training_steps=args.max_train_steps * world_size, + num_cycles=args.lr_num_cycles, + power=args.lr_power, + last_epoch=init_steps - 1, + ) + + train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, + args.cfg) + uncond_prompt_embed = train_dataset.uncond_prompt_embed + uncond_prompt_mask = train_dataset.uncond_prompt_mask + sampler = (LengthGroupedSampler( + args.train_batch_size, + rank=rank, + world_size=world_size, + lengths=train_dataset.lengths, + group_frame=args.group_frame, + group_resolution=args.group_resolution, + ) if (args.group_frame or args.group_resolution) else DistributedSampler( + train_dataset, rank=rank, num_replicas=world_size, shuffle=False)) + + train_dataloader = DataLoader( + train_dataset, + sampler=sampler, + collate_fn=latent_collate_function, + pin_memory=True, + batch_size=args.train_batch_size, + num_workers=args.dataloader_num_workers, + drop_last=True, + ) + + num_update_steps_per_epoch = math.ceil( + len(train_dataloader) / args.gradient_accumulation_steps * + args.sp_size / args.train_sp_batch_size) + args.num_train_epochs = math.ceil(args.max_train_steps / + num_update_steps_per_epoch) + + # if rank <= 0: + # project = args.tracker_project_name or "fastvideo" + # wandb.init(project=project, config=args) + + # Train! + total_batch_size = (world_size * args.gradient_accumulation_steps / + args.sp_size * args.train_sp_batch_size) + main_print("***** Running training *****") + main_print(f" Num examples = {len(train_dataset)}") + main_print(f" Dataloader size = {len(train_dataloader)}") + main_print(f" Num Epochs = {args.num_train_epochs}") + main_print(f" Resume training from step {init_steps}") + main_print( + f" Instantaneous batch size per device = {args.train_batch_size}") + main_print( + f" Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}" + ) + main_print( + f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + main_print(f" Total optimization steps = {args.max_train_steps}") + main_print( + f" Total training parameters per FSDP shard = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B" + ) + # print dtype + main_print( + f" Master weight dtype: {transformer.parameters().__next__().dtype}") + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + assert NotImplementedError( + "resume_from_checkpoint is not supported now.") + # TODO + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=init_steps, + desc="Steps", + # Only show the progress bar once on each machine. + disable=local_rank > 0, + ) + + loader = sp_parallel_dataloader_wrapper( + train_dataloader, + device, + args.train_batch_size, + args.sp_size, + args.train_sp_batch_size, + ) + + step_times = deque(maxlen=100) + + # todo future + for i in range(init_steps): + next(loader) + + # log_validation(args, transformer, device, + # torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False) + def get_num_phases(multi_phased_distill_schedule, step): + # step-phase,step-phase + multi_phases = multi_phased_distill_schedule.split(",") + phase = multi_phases[-1].split("-")[-1] + for step_phases in multi_phases: + phase_step, phase = step_phases.split("-") + if step <= int(phase_step): + return int(phase) + return phase + + for step in range(init_steps + 1, args.max_train_steps + 1): + start_time = time.time() + assert args.multi_phased_distill_schedule is not None + num_phases = get_num_phases(args.multi_phased_distill_schedule, step) + + loss, grad_norm, pred_norm = distill_one_step( + transformer, + args.model_type, + None, # teacher_transformer + ema_transformer, + optimizer, + lr_scheduler, + loader, + noise_scheduler, + solver, + noise_random_generator, + args.gradient_accumulation_steps, + args.sp_size, + args.max_grad_norm, + uncond_prompt_embed, + uncond_prompt_mask, + args.num_euler_timesteps, + num_phases, + args.not_apply_cfg_solver, + args.distill_cfg, + args.ema_decay, + args.pred_decay_weight, + args.pred_decay_type, + args.hunyuan_teacher_disable_cfg, + ) + + step_time = time.time() - start_time + step_times.append(step_time) + avg_step_time = sum(step_times) / len(step_times) + + progress_bar.set_postfix({ + "loss": f"{loss:.4f}", + "step_time": f"{step_time:.2f}s", + "grad_norm": grad_norm, + "phases": num_phases, + }) + progress_bar.update(1) + if rank <= 0: + wandb.log( + { + "train_loss": + loss, + "learning_rate": + lr_scheduler.get_last_lr()[0], + "step_time": + step_time, + "avg_step_time": + avg_step_time, + "grad_norm": + grad_norm, + "pred_fro_norm": + pred_norm["fro"], # codespell:ignore + "pred_largest_singular_value": + pred_norm["largest singular value"], + "pred_absolute_mean": + pred_norm["absolute mean"], + "pred_absolute_max": + pred_norm["absolute max"], + }, + step=step, + ) + if step % args.checkpointing_steps == 0: + if args.use_lora: + # Save LoRA weights + save_lora_checkpoint(transformer, optimizer, rank, + args.output_dir, step) + else: + # Your existing checkpoint saving code + if args.use_ema: + save_checkpoint(ema_transformer, rank, args.output_dir, + step) + else: + save_checkpoint(transformer, rank, args.output_dir, step) + dist.barrier() + if args.log_validation and step % args.validation_steps == 0: + log_validation( + args, + transformer, + device, + torch.bfloat16, + step, + scheduler_type=args.scheduler_type, + shift=args.shift, + num_euler_timesteps=args.num_euler_timesteps, + linear_quadratic_threshold=args.linear_quadratic_threshold, + linear_range=args.linear_range, + ema=False, + ) + if args.use_ema: + log_validation( + args, + ema_transformer, + device, + torch.bfloat16, + step, + scheduler_type=args.scheduler_type, + shift=args.shift, + num_euler_timesteps=args.num_euler_timesteps, + linear_quadratic_threshold=args.linear_quadratic_threshold, + linear_range=args.linear_range, + ema=True, + ) + + if args.use_lora: + save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, + args.max_train_steps) + else: + save_checkpoint(transformer, rank, args.output_dir, + args.max_train_steps) + + if get_sequence_parallel_state(): + destroy_sequence_parallel_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument("--model_type", + type=str, + default="mochi", + help="The type of model to train.") + + # dataset & dataloader + parser.add_argument("--data_json_path", type=str, required=True) + parser.add_argument("--num_height", type=int, default=480) + parser.add_argument("--num_width", type=int, default=848) + parser.add_argument("--num_frames", type=int, default=163) + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=10, + help= + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.", + ) + parser.add_argument( + "--train_batch_size", + type=int, + default=16, + help="Batch size (per device) for the training dataloader.", + ) + parser.add_argument("--num_latent_t", + type=int, + default=28, + help="Number of latent timesteps.") + parser.add_argument("--group_frame", action="store_true") # TODO + parser.add_argument("--group_resolution", action="store_true") # TODO + + # text encoder & vae & diffusion model + parser.add_argument("--pretrained_model_name_or_path", type=str) + # parser.add_argument("--dit_model_name_or_path", type=str) + parser.add_argument("--cache_dir", type=str, default="./cache_dir") + + # diffusion setting + parser.add_argument("--ema_decay", type=float, default=0.95) + parser.add_argument("--ema_start_step", type=int, default=0) + parser.add_argument("--cfg", type=float, default=0.1) + + # validation & logs + parser.add_argument("--validation_prompt_dir", type=str) + parser.add_argument("--validation_sampling_steps", type=str, default="64") + parser.add_argument("--validation_guidance_scale", type=str, default="4.5") + + parser.add_argument("--validation_steps", type=float, default=64) + parser.add_argument("--log_validation", action="store_true") + parser.add_argument("--tracker_project_name", type=str, default=None) + parser.add_argument("--seed", + type=int, + default=None, + help="A seed for reproducible training.") + parser.add_argument( + "--output_dir", + type=str, + default=None, + help= + "The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument( + "--checkpoints_total_limit", + type=int, + default=None, + help=("Max number of checkpoints to store."), + ) + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help= + ("Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" + " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" + " training using `--resume_from_checkpoint`."), + ) + parser.add_argument("--shift", type=float, default=1.0) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help= + ("Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument( + "--resume_from_lora_checkpoint", + type=str, + default=None, + help= + ("Whether training should be resumed from a previous lora checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help= + ("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."), + ) + + # optimizer & scheduler & Training + parser.add_argument("--num_train_epochs", type=int, default=100) + parser.add_argument( + "--max_train_steps", + type=int, + default=None, + help= + "Total number of training steps to perform. If provided, overrides num_train_epochs.", + ) + parser.add_argument( + "--gradient_accumulation_steps", + type=int, + default=1, + help= + "Number of updates steps to accumulate before performing a backward/update pass.", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + parser.add_argument( + "--scale_lr", + action="store_true", + default=False, + help= + "Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", + ) + parser.add_argument( + "--lr_warmup_steps", + type=int, + default=10, + help="Number of steps for the warmup in the lr scheduler.", + ) + parser.add_argument("--max_grad_norm", + default=1.0, + type=float, + help="Max gradient norm.") + parser.add_argument( + "--gradient_checkpointing", + action="store_true", + help= + "Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", + ) + parser.add_argument("--selective_checkpointing", type=float, default=1.0) + parser.add_argument( + "--allow_tf32", + action="store_true", + help= + ("Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument( + "--mixed_precision", + type=str, + default=None, + choices=["no", "fp16", "bf16"], + help= + ("Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + parser.add_argument( + "--use_cpu_offload", + action="store_true", + help= + "Whether to use CPU offload for param & gradient & optimizer states.", + ) + + parser.add_argument("--sp_size", + type=int, + default=1, + help="For sequence parallel") + parser.add_argument( + "--train_sp_batch_size", + type=int, + default=1, + help="Batch size for sequence parallel training", + ) + + parser.add_argument( + "--use_lora", + action="store_true", + default=False, + help="Whether to use LoRA for finetuning.", + ) + parser.add_argument("--lora_alpha", + type=int, + default=256, + help="Alpha parameter for LoRA.") + parser.add_argument("--lora_rank", + type=int, + default=128, + help="LoRA rank parameter. ") + parser.add_argument("--fsdp_sharding_startegy", default="full") + + # lr_scheduler + parser.add_argument( + "--lr_scheduler", + type=str, + default="constant", + help= + ('The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]'), + ) + parser.add_argument("--num_euler_timesteps", type=int, default=100) + parser.add_argument( + "--lr_num_cycles", + type=int, + default=1, + help="Number of cycles in the learning rate scheduler.", + ) + parser.add_argument( + "--lr_power", + type=float, + default=1.0, + help="Power factor of the polynomial scheduler.", + ) + parser.add_argument( + "--not_apply_cfg_solver", + action="store_true", + help="Whether to apply the cfg_solver.", + ) + parser.add_argument("--distill_cfg", + type=float, + default=3.0, + help="Distillation coefficient.") + # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"] + parser.add_argument("--scheduler_type", + type=str, + default="pcm", + help="The scheduler type to use.") + parser.add_argument( + "--linear_quadratic_threshold", + type=float, + default=0.025, + help="Threshold for linear quadratic scheduler.", + ) + parser.add_argument( + "--linear_range", + type=float, + default=0.5, + help="Range for linear quadratic scheduler.", + ) + parser.add_argument("--weight_decay", + type=float, + default=0.001, + help="Weight decay to apply.") + parser.add_argument("--use_ema", + action="store_true", + help="Whether to use EMA.") + parser.add_argument("--multi_phased_distill_schedule", + type=str, + default=None) + parser.add_argument("--pred_decay_weight", type=float, default=0.0) + parser.add_argument("--pred_decay_type", default="l1") + parser.add_argument("--hunyuan_teacher_disable_cfg", action="store_true") + parser.add_argument( + "--master_weight_type", + type=str, + default="fp32", + help="Weight type to use - fp32 or bf16.", + ) + args = parser.parse_args() + main(args) diff --git a/scripts/dataset_preparation/prepare_json_file.py b/scripts/dataset_preparation/prepare_json_file.py index 91b7e663f..7abd6c073 100644 --- a/scripts/dataset_preparation/prepare_json_file.py +++ b/scripts/dataset_preparation/prepare_json_file.py @@ -1,11 +1,11 @@ import json from pathlib import Path - +import csv import cv2 -def get_video_info(video_path, prompt_text): - """Extract video information using OpenCV and corresponding prompt text""" +def get_video_info(video_path, metadata): + """Extract video information using OpenCV and corresponding metadata""" cap = cv2.VideoCapture(str(video_path)) if not cap.isOpened(): @@ -23,60 +23,66 @@ def get_video_info(video_path, prompt_text): return { "path": video_path.name, + "title": metadata.get("Video Title", ""), + "description": metadata.get("Video Description", ""), + "video_url": metadata.get("Video URL", ""), + "download_url": metadata.get("Download URL", ""), "resolution": { "width": width, "height": height }, "fps": fps, "duration": duration, - "cap": [prompt_text] + "cap": [metadata.get("Video Description", "")] } -def read_prompt_file(prompt_path): - """Read and return the content of a prompt file""" +def read_csv_file(csv_path): + """Read and return the content of a CSV file""" try: - with open(prompt_path, 'r', encoding='utf-8') as f: - return f.read().strip() + with open(csv_path, 'r', encoding='utf-8') as f: + reader = csv.DictReader(f) + return list(reader) except Exception as e: - print(f"Error reading prompt file {prompt_path}: {e}") + print(f"Error reading CSV file {csv_path}: {e}") return None -def process_videos_and_prompts(video_dir_path, prompt_dir_path, verbose=False): - """Process videos and their corresponding prompt files +def process_videos_from_csv(video_dir_path, csv_path, verbose=False): + """Process videos using metadata from CSV file Args: video_dir_path (str): Path to directory containing video files - prompt_dir_path (str): Path to directory containing prompt files + csv_path (str): Path to CSV file containing video metadata verbose (bool): Whether to print verbose processing information """ video_dir = Path(video_dir_path) - prompt_dir = Path(prompt_dir_path) + csv_data = read_csv_file(csv_path) processed_data = [] # Ensure directories exist - if not video_dir.exists() or not prompt_dir.exists(): - print(f"Error: One or both directories do not exist:\nVideos: {video_dir}\nPrompts: {prompt_dir}") + if not video_dir.exists(): + print(f"Error: Video directory does not exist: {video_dir}") + return [] + + if csv_data is None: return [] # Process each video file - for video_file in video_dir.glob('*.mp4'): - video_name = video_file.stem - prompt_file = prompt_dir / f"{video_name}.txt" - - # Check if corresponding prompt file exists - if not prompt_file.exists(): - print(f"Warning: No prompt file found for video {video_name}") + for row in csv_data: + video_filename = row.get("Filename") + if not video_filename: continue - # Read prompt content - prompt_text = read_prompt_file(prompt_file) - if prompt_text is None: + video_file = video_dir / video_filename + + # Check if video file exists + if not video_file.exists(): + print(f"Warning: Video file not found: {video_filename}") continue # Process video and add to results - video_info = get_video_info(video_file, prompt_text) + video_info = get_video_info(video_file, row) if video_info: processed_data.append(video_info) @@ -105,9 +111,9 @@ def parse_args(): """Parse command line arguments""" import argparse - parser = argparse.ArgumentParser(description='Process videos and their corresponding prompt files') + parser = argparse.ArgumentParser(description='Process videos using metadata from CSV file') parser.add_argument('--video_dir', '-v', required=True, help='Directory containing video files') - parser.add_argument('--prompt_dir', '-p', required=True, help='Directory containing prompt text files') + parser.add_argument('--csv_path', '-c', required=True, help='Path to CSV file containing video metadata') parser.add_argument('--output_path', '-o', required=True, @@ -121,8 +127,8 @@ if __name__ == "__main__": # Parse command line arguments args = parse_args() - # Process videos and prompts - processed_videos = process_videos_and_prompts(args.video_dir, args.prompt_dir, args.verbose) + # Process videos from CSV + processed_videos = process_videos_from_csv(args.video_dir, args.csv_path, args.verbose) if processed_videos: # Save results diff --git a/scripts/dataset_preparation/resize_videos.py b/scripts/dataset_preparation/resize_videos.py index 5e37d1aeb..05dd82c30 100644 --- a/scripts/dataset_preparation/resize_videos.py +++ b/scripts/dataset_preparation/resize_videos.py @@ -24,9 +24,9 @@ def is_16_9_ratio(width: int, height: int, tolerance: float = 0.1) -> bool: def resize_video(args_tuple): """ Resize a single video file. - args_tuple: (input_file, output_dir, width, height, fps) + args_tuple: (input_file, output_dir, width, height, fps, num_frames) """ - input_file, output_dir, width, height, fps = args_tuple + input_file, output_dir, width, height, fps, num_frames = args_tuple video = None resized = None output_file = output_dir / f"{input_file.name}" @@ -39,6 +39,13 @@ def resize_video(args_tuple): if not is_16_9_ratio(video.w, video.h): return (input_file.name, "skipped", "Not 16:9") + # Calculate target duration based on num_frames and fps + target_duration = num_frames / fps + + # Trim video if it's longer than target duration + if video.duration > target_duration: + video = video.subclip(0, target_duration) + def process_frame(frame): frame_float = frame.astype(float) / 255.0 resized = resize(frame_float, (height, width, 3), mode='reflect', anti_aliasing=True, preserve_range=True) @@ -75,7 +82,7 @@ def process_folder(args): print(f"Target: {args.width}x{args.height} at {args.fps}fps") # Prepare arguments for parallel processing - process_args = [(video_file, output_path, args.width, args.height, args.fps) for video_file in video_files] + process_args = [(video_file, output_path, args.width, args.height, args.fps, args.num_frames) for video_file in video_files] successful = 0 skipped = 0 @@ -115,6 +122,7 @@ def parse_args(): parser.add_argument('--width', type=int, default=1280, help='Target width in pixels (default: 848)') parser.add_argument('--height', type=int, default=720, help='Target height in pixels (default: 480)') parser.add_argument('--fps', type=int, default=30, help='Target frames per second (default: 30)') + parser.add_argument('--num_frames', type=int, default=163, help='Target number of frames (default: 163)') parser.add_argument('--max_workers', type=int, default=4, diff --git a/scripts/inference/v1_preprocess_wan.sh b/scripts/inference/v1_preprocess_wan.sh new file mode 100644 index 000000000..df677218b --- /dev/null +++ b/scripts/inference/v1_preprocess_wan.sh @@ -0,0 +1,27 @@ +#!/bin/bash + +num_gpus=2 +export FASTVIDEO_ATTENTION_BACKEND= +export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers +# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo +# Note that the tp_size and sp_size should be the same and equal to the number +# of GPUs. They are used for different parallel groups. sp_size is used for +# dit model and tp_size is used for encoder models. +torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ + fastvideo/v1/entrypoints/data_preprocessor.py \ + --sp_size $num_gpus \ + --tp_size $num_gpus \ + --height 480 \ + --width 832 \ + --num_frames 77 \ + --num_inference_steps 50 \ + --fps 16 \ + --guidance_scale 3.0 \ + --prompt_path ./assets/prompt.txt \ + --neg_prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \ + --seed 1024 \ + --output_path outputs_video/ \ + --model_path $MODEL_BASE \ + --vae-sp \ + --text-encoder-precision "fp32" \ + --use-cpu-offload \ No newline at end of file diff --git a/scripts/preprocess/preprocess_wan_data_v1.sh b/scripts/preprocess/preprocess_wan_data_v1.sh new file mode 100644 index 000000000..0445248ed --- /dev/null +++ b/scripts/preprocess/preprocess_wan_data_v1.sh @@ -0,0 +1,30 @@ +# export WANDB_MODE="offline" +GPU_NUM=1 # 2,4,8 +# MODEL_PATH="/home/ray/.cache/huggingface/hub/models--Wan-AI--Wan2.1-T2V-1.3B-Diffusers/snapshots/0fad780a534b6463e45facd96134c9f345acfa5b" +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_MERGE_PATH="data/cats_480/merge.txt" +OUTPUT_DIR="data/cats_480_latents/" +VALIDATION_PATH="assets/prompt.txt" + +# torchrun --nproc_per_node=$GPU_NUM \ +# fastvideo/data_preprocess/preprocess_vae_latents_v1.py \ +# --model_path $MODEL_PATH \ +# --data_merge_path $DATA_MERGE_PATH \ +# --train_batch_size=1 \ +# --max_height=480 \ +# --max_width=832 \ +# --num_frames=81 \ +# --dataloader_num_workers 1 \ +# --output_dir=$OUTPUT_DIR \ +# --train_fps 16 + +# torchrun --nproc_per_node=$GPU_NUM \ +# fastvideo/data_preprocess/preprocess_text_embeddings_v1.py \ +# --model_path $MODEL_PATH \ +# --output_dir=$OUTPUT_DIR + +torchrun --nproc_per_node=1 \ + fastvideo/data_preprocess/preprocess_validation_text_embeddings_v1.py \ + --model_path $MODEL_PATH \ + --output_dir=$OUTPUT_DIR \ + --validation_prompt_txt $VALIDATION_PATH \ No newline at end of file diff --git a/train_v1.sh b/train_v1.sh new file mode 100644 index 000000000..07c1c3e56 --- /dev/null +++ b/train_v1.sh @@ -0,0 +1,48 @@ +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online + +DATA_DIR=./data +NUM_GPUS=1 +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 + # --gradient_checkpointing\ + # --pretrained_model_name_or_path hunyuanvideo-community/HunyuanVideo \ + # --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ +torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\ + fastvideo/v1/pipelines/training_pipeline.py\ + --inference_mode False\ + --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ + --cache_dir "/home/ray/.cache"\ + --data_json_path "$DATA_DIR/cats_480_latents/videos2caption.json"\ + --validation_prompt_dir "$DATA_DIR/cats_480_latents/validation"\ + --train_batch_size=4\ + --num_latent_t 8 \ + --sp_size $NUM_GPUS \ + --tp_size $NUM_GPUS \ + --train_sp_batch_size 1\ + --dataloader_num_workers 4\ + --gradient_accumulation_steps=1\ + --max_train_steps=120 \ + --learning_rate=1e-6\ + --mixed_precision="bf16"\ + --checkpointing_steps=64 \ + --validation_steps 2\ + --validation_sampling_steps "2,4,8" \ + --checkpoints_total_limit 3\ + --allow_tf32\ + --ema_start_step 0\ + --cfg 0.0\ + --output_dir="$DATA_DIR/outputs/wan_finetune"\ + --tracker_project_name wan_finetune \ + --num_height 480 \ + --num_width 832 \ + --num_frames 81 \ + --shift 3 \ + --validation_guidance_scale "1.0" \ + --num_euler_timesteps 50 \ + --multi_phased_distill_schedule "4000-1" \ + --weight_decay 0.01 \ + --not_apply_cfg_solver \ + --master_weight_type "bf16" \ No newline at end of file