diff --git a/fastvideo/data_preprocess/preprocess_text_embeddings.py b/fastvideo/data_preprocess/preprocess_text_embeddings.py new file mode 100644 index 000000000..57c039b22 --- /dev/null +++ b/fastvideo/data_preprocess/preprocess_text_embeddings.py @@ -0,0 +1,156 @@ +import argparse +import torch +from accelerate.logging import get_logger +from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline +from diffusers.utils import export_to_video +import json +import os +import torch.distributed as dist + +logger = get_logger(__name__) +from torch.utils.data import Dataset +from torch.utils.data.distributed import DistributedSampler +from torch.utils.data import DataLoader +from fastvideo.utils.load import load_text_encoder, load_vae +from diffusers.video_processor import VideoProcessor +from tqdm import tqdm +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)) + print("world_size", world_size, "local rank", local_rank) + + 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) + + latents_json_path = os.path.join(args.output_dir, "videos2caption_temp_replace.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=autocast_type): + prompt_embeds, prompt_attention_mask = text_encoder.encode_prompt( + prompt=data["caption"], + ) + 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=fps) + 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/scripts/distill_hunyuan.md b/scripts/distill_hunyuan.md new file mode 100644 index 000000000..1c9c1cd62 --- /dev/null +++ b/scripts/distill_hunyuan.md @@ -0,0 +1,5 @@ +## How to Distill Hunyuan + +python scripts/download_hf.py --repo_id FastVideo/hunyuan --local_dir data/hunyuan --repo_type model +python scripts/download_hf.py --repo_id FastVideo/Hunyuan-Distill-Data --local_dir data/Hunyuan-Distill-Data --repo_type=dataset + diff --git a/scripts/distill_hunyuan.sh b/scripts/distill_hunyuan.sh index 7d83b0727..6c90152e3 100644 --- a/scripts/distill_hunyuan.sh +++ b/scripts/distill_hunyuan.sh @@ -23,7 +23,7 @@ torchrun --nnodes 1 --nproc_per_node 8\ --learning_rate=1e-6\ --mixed_precision="bf16"\ --checkpointing_steps=64\ - --validation_steps 32\ + --validation_steps 64\ --validation_sampling_steps "2,4,8" \ --checkpoints_total_limit 3\ --allow_tf32\ diff --git a/scripts/finetune_data_gen_hunyuan.sh b/scripts/preprocess_hunyuan_data.sh similarity index 58% rename from scripts/finetune_data_gen_hunyuan.sh rename to scripts/preprocess_hunyuan_data.sh index c02e14094..f58df6648 100644 --- a/scripts/finetune_data_gen_hunyuan.sh +++ b/scripts/preprocess_hunyuan_data.sh @@ -4,9 +4,9 @@ MODEL_PATH="data/hunyuan" MODEL_TYPE="hunyuan" DATA_MERGE_PATH="data/Mixkit-All-Clips/merge.txt" OUTPUT_DIR="data/Hunyuan-Mixkit-Data" - +VALIDATION_PATH="assets/prompt.txt" # torchrun --nproc_per_node=$GPU_NUM \ -# ./fastvideo/data_preprocess/finetune_data_VAE.py \ +# fastvideo/data_preprocess/preprocess_vae_latents.py \ # --model_path $MODEL_PATH \ # --data_merge_path $DATA_MERGE_PATH \ # --train_batch_size=1 \ @@ -21,7 +21,16 @@ OUTPUT_DIR="data/Hunyuan-Mixkit-Data" torchrun --nproc_per_node=$GPU_NUM \ - ./fastvideo/data_preprocess/finetune_data_text_encoder.py \ + fastvideo/data_preprocess/preprocess_text_embeddings.py \ --model_type $MODEL_TYPE \ --model_path $MODEL_PATH \ - --output_dir=$OUTPUT_DIR \ No newline at end of file + --output_dir=$OUTPUT_DIR + + + +torchrun --nproc_per_node=1 \ + fastvideo/data_preprocess/preprocess_validation_text_embeddings.py \ + --model_type $MODEL_TYPE \ + --model_path $MODEL_PATH \ + --output_dir=$OUTPUT_DIR \ + --validation_prompt_txt $VALIDATION_PATH \ No newline at end of file