diff --git a/fastvideo/dataset/latent_datasets.py b/fastvideo/dataset/latent_datasets.py index 5e7718b8d..98fda62cc 100644 --- a/fastvideo/dataset/latent_datasets.py +++ b/fastvideo/dataset/latent_datasets.py @@ -4,7 +4,7 @@ import json import os class LatentDataset(Dataset): - def __init__(self, data_merge_path): + def __init__(self, data_merge_path, num_latent_t): # data_merge_path: video_dir, latent_dir, prompt_embed_dir, json_path self.data_merge_path = data_merge_path # read txt @@ -18,14 +18,16 @@ class LatentDataset(Dataset): self.json_path = json_path with open(self.json_path, 'r') as f: self.data_anno = json.load(f) + + self.num_latent_t = num_latent_t def __getitem__(self, idx): latent_file = self.data_anno[idx]["latent_path"] prompt_embed_file = self.data_anno[idx]["prompt_embed_path"] prompt_attention_mask_file = self.data_anno[idx]["prompt_attention_mask"] # load - latent = torch.load(os.path.join(self.latent_dir, latent_file)) - prompt_embed = torch.load(os.path.join(self.prompt_embed_dir, prompt_embed_file)) - prompt_attention_mask = torch.load(os.path.join(self.prompt_attention_mask_dir, prompt_attention_mask_file)) + latent = torch.load(os.path.join(self.latent_dir, latent_file), map_location="cpu", weights_only=True)[:, :self.num_latent_t] + prompt_embed = torch.load(os.path.join(self.prompt_embed_dir, prompt_embed_file), map_location="cpu", weights_only=True) + prompt_attention_mask = torch.load(os.path.join(self.prompt_attention_mask_dir, prompt_attention_mask_file), map_location="cpu", weights_only=True) return latent, prompt_embed, prompt_attention_mask def __len__(self): diff --git a/fastvideo/train.py b/fastvideo/train.py index 5f40890fe..0f6c894d7 100644 --- a/fastvideo/train.py +++ b/fastvideo/train.py @@ -54,6 +54,23 @@ check_min_version("0.24.0") logger = get_logger(__name__) +import sys +import pdb + +class ForkedPdb(pdb.Pdb): + """A Pdb subclass that may be used + from a forked multiprocessing child + + """ + def interaction(self, *args, **kwargs): + _stdin = sys.stdin + try: + sys.stdin = open('/dev/stdin') + pdb.Pdb.interaction(self, *args, **kwargs) + finally: + sys.stdin = _stdin + + @torch.inference_mode() def log_validation(args, transformer,vae, accelerator, weight_dtype, global_step, ema=False): logger.info(f"Running validation....\n") @@ -265,7 +282,7 @@ def main(args): logger.info(f"optimizer: {optimizer}") - train_dataset = LatentDataset(args.data_merge_path) + train_dataset = LatentDataset(args.data_merge_path, args.num_latent_t) train_dataloader = DataLoader( train_dataset, shuffle=True, @@ -426,6 +443,7 @@ def main(args): logit_std=args.logit_std, mode_scale=args.mode_scale, ) + indices = (u * noise_scheduler.config.num_train_timesteps).long() timesteps = noise_scheduler.timesteps[indices].to(device=latents.device) @@ -433,12 +451,11 @@ def main(args): # zt = (1 - texp) * x + texp * z1 sigmas = get_sigmas(timesteps, n_dim=latents.ndim, dtype=latents.dtype) noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise - model_pred = transformer( noisy_model_input, encoder_hidden_states, noise_scheduler.config.num_train_timesteps - timesteps, # the timestep for mochi is invereted - encoder_attention_mask.unsqueeze(1), # B, 1, L + encoder_attention_mask, # B, L return_dict= False )[0] @@ -500,8 +517,10 @@ def main(args): return loss def train_one_step( data_item_, prof_=None): - latents, attn_mask, cond, cond_mask = data_item_ - + latents, cond,attn_mask, cond_mask = data_item_ + # TODO + latents = latents.to(dtype=torch.bfloat16) + cond = cond.to(dtype=torch.bfloat16) latents_mean = ( torch.tensor(vae.config.latents_mean).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype) ) @@ -519,15 +538,16 @@ def main(args): st_idx = iter * args.train_sp_batch_size ed_idx = (iter + 1) * args.train_sp_batch_size encoder_hidden_states=cond[st_idx: ed_idx] + # TODO attention_mask=attn_mask[st_idx: ed_idx] encoder_attention_mask=cond_mask[st_idx: ed_idx] - run(latents[st_idx: ed_idx], encoder_hidden_states,attention_mask, encoder_attention_mask, prof_) + run(latents[st_idx: ed_idx], encoder_hidden_states, encoder_attention_mask, prof_) else: with accelerator.accumulate(transformer): encoder_hidden_states=cond attention_mask=attn_mask, encoder_attention_mask=cond_mask - run(latents, encoder_hidden_states,attention_mask, encoder_attention_mask, prof_) + run(latents, encoder_hidden_states, encoder_attention_mask, prof_) if progress_info.global_step >= args.max_train_steps: return True @@ -560,6 +580,7 @@ if __name__ == "__main__": parser.add_argument("--num_frames", type=int, default=65) 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 @@ -597,6 +618,7 @@ if __name__ == "__main__": default=1.29, help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.", ) + parser.add_argument("--precondition_outputs", action="store_true", help="Whether to precondition the outputs of the model.") # validation & logs parser.add_argument("--validaiton_prompt", nargs='+', help="List of prompts to use for validation.") diff --git a/t2v_debug_multi.sh b/t2v_debug_multi.sh index f91cac25c..f0ad6a91b 100644 --- a/t2v_debug_multi.sh +++ b/t2v_debug_multi.sh @@ -8,6 +8,7 @@ accelerate launch \ --data_merge_path "data/Mochi-Synthetic-Data/merge.txt" \ --gradient_checkpointing \ --train_batch_size=1 \ + --num_latent_t 1 \ --dataloader_num_workers 1 \ --gradient_accumulation_steps=1 \ --max_train_steps=1000000 \