diff --git a/fastvideo/dataset/latent_datasets.py b/fastvideo/dataset/latent_datasets.py index b04d6be53..e53fccbfb 100644 --- a/fastvideo/dataset/latent_datasets.py +++ b/fastvideo/dataset/latent_datasets.py @@ -27,7 +27,7 @@ class LatentDataset(Dataset): self.num_latent_t = num_latent_t self.uncond_prompt_embed = torch.load(os.path.join(uncond_prompt_embed_mask_dir, "embed.pt"), map_location="cpu", weights_only=True) self.uncond_prompt_mask = torch.load(os.path.join(uncond_prompt_embed_mask_dir, "mask.pt"), map_location="cpu", weights_only=True) - self.lengths = [torch.load(os.path.join(self.latent_dir, data_item["latent_path"]), map_location="cpu").shape[1] for data_item in self.data_anno] + self.lengths = [data_item['length'] if "length" in data_item else 1 for data_item in self.data_anno] def __getitem__(self, idx): latent_file = self.data_anno[idx]["latent_path"] prompt_embed_file = self.data_anno[idx]["prompt_embed_path"] diff --git a/fastvideo/dataset/t2v_datasets.py b/fastvideo/dataset/t2v_datasets.py index 5936b2800..49002e21b 100644 --- a/fastvideo/dataset/t2v_datasets.py +++ b/fastvideo/dataset/t2v_datasets.py @@ -141,7 +141,6 @@ class T2V_dataset(Dataset): video = torchvision_video[frame_indices] video = self.transform(video) video = rearrange(video, 't c h w -> c t h w') - video = video.unsqueeze(0) video = video.to(torch.uint8) assert video.dtype == torch.uint8 @@ -181,7 +180,6 @@ class T2V_dataset(Dataset): image = self.transform_topcrop(image) if 'human_images' in image_data['path'] else self.transform(image) # [1 C H W] -> num_img [1 C H W] image = image.transpose(0, 1) # [1 C H W] -> [C 1 H W] - image = image.unsqueeze(0) image = image.float() / 127.5 - 1.0 diff --git a/fastvideo/utils/data_preprocess/finetune_data_T5.py b/fastvideo/utils/data_preprocess/finetune_data_T5.py index 796e73ae7..9dacbedb7 100644 --- a/fastvideo/utils/data_preprocess/finetune_data_T5.py +++ b/fastvideo/utils/data_preprocess/finetune_data_T5.py @@ -7,7 +7,35 @@ 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 +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)) @@ -27,45 +55,47 @@ def main(args): 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.json") - with open(latents_json_path, "r") as f: - train_dataset = json.load(f) - train_dataset = sorted(train_dataset, key=lambda x: x['latent_path']) + train_dataset = T5dataset(latents_json_path, args.vae_debug) + 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 enumerate(train_dataset): - video_name =data['latent_path'].split(".")[0] - if int(video_name) % world_size != local_rank: - continue - try: - with torch.inference_mode(): - with torch.autocast("cuda", dtype=torch.bfloat16): - latent = torch.load(os.path.join(args.output_dir, 'latent', data['latent_path'])) - prompt_embeds, prompt_attention_mask, _, _ = pipe.encode_prompt( - prompt=data['caption'], - ) - 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[0], prompt_embed_path) - torch.save(prompt_attention_mask[0], prompt_attention_mask_path) - print(f"sample {video_name} saved") - video = pipe.vae.decode(latent.unsqueeze(0).to(device), return_dict=False)[0] - video = pipe.video_processor.postprocess_video(video) - export_to_video(video[0], video_path, fps=30) - item = {} - item["latent_path"] = video_name + ".pt" - item["prompt_embed_path"] = video_name + ".pt" - item["prompt_attention_mask"] = video_name + ".pt" - item["caption"] = data['caption'] - json_data.append(item) - except: - print("video out of memory") - continue - dist.barrier() - local_data = json_data - gathered_data = [None] * world_size - dist.all_gather_object(gathered_data, local_data) + json_data = [] + for _, data in enumerate(train_dataloader): + with torch.inference_mode(): + with torch.autocast("cuda", dtype=torch.bfloat16): + prompt_embeds, prompt_attention_mask, _, _ = pipe.encode_prompt( + prompt=data['caption'], + ) + if args.vae_debug: + latents = data['latents'] + video = pipe.vae.decode(latents.to(device), return_dict=False)[0] + video = pipe.video_processor.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=30) + 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] @@ -77,9 +107,11 @@ if __name__ == "__main__": # dataset & dataloader parser.add_argument("--model_path", type=str, default="data/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/utils/data_preprocess/finetune_data_VAE.py b/fastvideo/utils/data_preprocess/finetune_data_VAE.py index 8cbdb7efd..54c3d2e61 100644 --- a/fastvideo/utils/data_preprocess/finetune_data_VAE.py +++ b/fastvideo/utils/data_preprocess/finetune_data_VAE.py @@ -10,6 +10,7 @@ import json import os from diffusers import AutoencoderKLMochi import torch.distributed as dist +from torch.utils.data.distributed import DistributedSampler logger = get_logger(__name__) @@ -27,9 +28,10 @@ def main(args): project_config=accelerator_project_config, ) train_dataset = getdataset(args) + sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True) train_dataloader = DataLoader( train_dataset, - shuffle=True, + sampler=sampler, batch_size=args.train_batch_size, num_workers=args.dataloader_num_workers, ) @@ -46,19 +48,19 @@ def main(args): json_data = [] for _, data in enumerate(train_dataloader): - video_name = os.path.basename(data['path'][0]).split(".")[0] - if int(video_name) % world_size != local_rank: - continue with torch.inference_mode(): with torch.autocast("cuda", dtype=torch.bfloat16): - latents = vae.encode(data['pixel_values'][0].to(encoder_device))['latent_dist'].sample() - latent_path = os.path.join(args.output_dir, "latent", video_name + ".pt") - torch.save(latents[0].to(torch.bfloat16), latent_path) - item = {} - item["latent_path"] = video_name + ".pt" - item["caption"] = data['text'] - json_data.append(item) - print(f"{video_name} processed") + latents = vae.encode(data['pixel_values'].to(encoder_device))['latent_dist'].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") dist.barrier() local_data = json_data gathered_data = [None] * world_size