diff --git a/fastvideo/dataset/latent_datasets.py b/fastvideo/dataset/latent_datasets.py index 0cd894c62..9678c83d1 100644 --- a/fastvideo/dataset/latent_datasets.py +++ b/fastvideo/dataset/latent_datasets.py @@ -10,7 +10,6 @@ class LatentDataset(Dataset): json_path, num_latent_t, cfg_rate, - uncond_prompt_embed_mask_dir="data/Encoder_Overfit_Data/uncond_prompt_embed_mask", ): # data_merge_path: video_dir, latent_dir, prompt_embed_dir, json_path self.json_path = json_path @@ -25,8 +24,10 @@ class LatentDataset(Dataset): # json.load(f) already keeps the order # self.data_anno = sorted(self.data_anno, key=lambda x: x['latent_path']) 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) + # just zero embeddings [256, 4096] + self.uncond_prompt_embed = torch.zeros(256, 4096).to(torch.float32) + # 256 zeros + self.uncond_prompt_mask = torch.zeros(256).bool() 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"] diff --git a/fastvideo/distill.py b/fastvideo/distill.py index 433b26b2e..9f28c1630 100644 --- a/fastvideo/distill.py +++ b/fastvideo/distill.py @@ -93,7 +93,7 @@ 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 train_one_step_mochi(transformer, teacher_transformer, ema_transformer, optimizer, lr_scheduler,loader, noise_scheduler, solver,noise_random_generator, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, num_euler_timesteps, multiphase, not_apply_cfg_solver, distill_cfg): +def train_one_step_mochi(transformer, teacher_transformer, ema_transformer, optimizer, lr_scheduler,loader, noise_scheduler, solver,noise_random_generator, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, num_euler_timesteps, multiphase, not_apply_cfg_solver, distill_cfg, ema_decay): total_loss = 0.0 optimizer.zero_grad() for _ in range(gradient_accumulation_steps): @@ -214,7 +214,7 @@ def train_one_step_mochi(transformer, teacher_transformer, ema_transformer, opti 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 - 0.95)) + 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() @@ -396,7 +396,8 @@ def main(args): ema_transformer.requires_grad_(False) noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=args.shift) if args.scheduler_type == "pcm_linear_quadratic": - sigmas = linear_quadratic_schedule(noise_scheduler.config.num_train_timesteps, args.linear_quadratic_threshold) + 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 @@ -435,7 +436,7 @@ def main(args): last_epoch=init_steps - 1, ) - train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg, args.uncond_prompt_dir) + 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( @@ -504,9 +505,18 @@ def main(args): 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) for step in range(init_steps + 1, args.max_train_steps+1): start_time = time.time() - loss, grad_norm= train_one_step_mochi(transformer,teacher_transformer, ema_transformer, optimizer, lr_scheduler, loader, noise_scheduler,solver, noise_random_generator, args.gradient_accumulation_steps, args.sp_size, args.precondition_outputs, args.max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, args.num_euler_timesteps, args.validation_sampling_steps, args.not_apply_cfg_solver,args.distill_cfg) + num_phases = get_num_phases(args.multi_phased_distill_schedule, step) if args.multi_phased_distill_schedule is not None else args.validation_sampling_steps + loss, grad_norm= train_one_step_mochi(transformer,teacher_transformer, ema_transformer, optimizer, lr_scheduler, loader, noise_scheduler,solver, noise_random_generator, args.gradient_accumulation_steps, args.sp_size, args.precondition_outputs, 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) step_time = time.time() - start_time step_times.append(step_time) @@ -515,7 +525,8 @@ def main(args): progress_bar.set_postfix({ "loss": f"{loss:.4f}", "step_time": f"{step_time:.2f}s", - "grad_norm": grad_norm + "grad_norm": grad_norm, + "phases": num_phases, }) progress_bar.update(1) if rank <= 0: @@ -532,14 +543,17 @@ def main(args): save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, step) else: # Your existing checkpoint saving code - save_checkpoint(transformer, rank, args.output_dir, step) + 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, ema=False) + 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, ema=True) + 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) @@ -576,9 +590,9 @@ if __name__ == "__main__": # validation & logs parser.add_argument("--validation_prompt_dir", type=str) - parser.add_argument("--uncond_prompt_dir", type=str) parser.add_argument("--validation_sampling_steps", type=int, default=64) - parser.add_argument('--validation_guidance_scale', type=float, default=4.5) + 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) @@ -660,7 +674,10 @@ if __name__ == "__main__": # ["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("--ema_decay", type=float, default=0.95, help="EMA decay.") + parser.add_argument("--multi_phased_distill_schedule", type=str, default=None) args = parser.parse_args() main(args) \ No newline at end of file diff --git a/fastvideo/distill/solver.py b/fastvideo/distill/solver.py index 9f7a0ef82..9e873b040 100644 --- a/fastvideo/distill/solver.py +++ b/fastvideo/distill/solver.py @@ -35,10 +35,12 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin): pcm_timesteps: int = 50, linear_quadratic=False, linear_quadratic_threshold=0.025, + linear_range=0.5, ): if linear_quadratic: - sigmas = linear_quadratic_schedule(num_train_timesteps, linear_quadratic_threshold) + 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) else: timesteps = np.linspace( @@ -139,7 +141,6 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin): self.sigmas_ = torch.cat( [self.sigmas_, torch.zeros(1, device=self.sigmas_.device)] ) - self._step_index = None self._begin_index = None diff --git a/fastvideo/distill_adv.py b/fastvideo/distill_adv.py index c8013ce95..681670e7c 100644 --- a/fastvideo/distill_adv.py +++ b/fastvideo/distill_adv.py @@ -459,7 +459,7 @@ def main(args): last_epoch=init_steps - 1, ) - train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg, args.uncond_prompt_dir) + 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( @@ -597,7 +597,6 @@ if __name__ == "__main__": # validation & logs parser.add_argument("--validation_prompt_dir", type=str) - parser.add_argument("--uncond_prompt_dir", type=str) parser.add_argument("--validation_sampling_steps", type=int, default=64) parser.add_argument('--validation_guidance_scale', type=float, default=4.5) parser.add_argument('--validation_steps', type=float, default=64) diff --git a/fastvideo/sample/sample_t2v_mochi.py b/fastvideo/sample/sample_t2v_mochi.py index c0f958672..5d68c5346 100644 --- a/fastvideo/sample/sample_t2v_mochi.py +++ b/fastvideo/sample/sample_t2v_mochi.py @@ -110,7 +110,7 @@ def main(args): scheduler = FlowMatchEulerDiscreteScheduler() else: linear_quadratic = True if "linear_quadratic" in args.scheduler_type else False - scheduler = PCMFMScheduler(1000, args.shift, args.num_euler_timesteps, linear_quadratic) + scheduler = PCMFMScheduler(1000, args.shift, args.num_euler_timesteps, linear_quadratic,args.linear_threshold, args.linear_range) if args.transformer_path is not None: transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path) else: @@ -201,5 +201,7 @@ if __name__ == "__main__": parser.add_argument('--lora_checkpoint_dir', type=str, default=None, help='Path to the directory containing LoRA checkpoints') parser.add_argument("--shift", type=float, default=8.0) parser.add_argument("--num_euler_timesteps", type=int, default=100) + parser.add_argument("--linear_threshold", type=float, default=0.025) + parser.add_argument("--linear_range", type=float, default=0.5) args = parser.parse_args() main(args) diff --git a/fastvideo/utils/validation.py b/fastvideo/utils/validation.py index 9f8f57e15..c1dd461b4 100644 --- a/fastvideo/utils/validation.py +++ b/fastvideo/utils/validation.py @@ -189,7 +189,7 @@ def sample_validation_video( @torch.no_grad() @torch.autocast("cuda", dtype=torch.bfloat16) -def log_validation(args, transformer, device, weight_dtype, global_step, scheduler_type="euler",shift=1.0, num_euler_timesteps=100, linear_quadratic_threshold=0.025, ema=False): +def log_validation(args, transformer, device, weight_dtype, global_step, scheduler_type="euler",shift=1.0, num_euler_timesteps=100, linear_quadratic_threshold=0.025, linear_range=0.5, ema=False): #TODO print(f"Running validation....\n") vae = AutoencoderKLMochi.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", torch_dtype=weight_dtype).to("cuda") @@ -198,72 +198,75 @@ def log_validation(args, transformer, device, weight_dtype, global_step, schedu scheduler = FlowMatchEulerDiscreteScheduler() else: linear_quadraic = True if scheduler_type == "pcm_linear_quadratic" else False - scheduler = PCMFMScheduler(1000, shift, num_euler_timesteps, linear_quadraic, linear_quadratic_threshold) + scheduler = PCMFMScheduler(1000, shift, num_euler_timesteps, linear_quadraic, linear_quadratic_threshold, linear_range) # args.validation_prompt_dir - videos = [] - # prompt_embed are named embed0 to embedN - # check how many embeds are there - num_embeds = len([f for f in os.listdir(args.validation_prompt_dir) if "embed" in f]) - validation_prompt_ids = list(range(num_embeds)) - num_sp_groups = int(os.getenv("WORLD_SIZE", '1')) // nccl_info.sp_size - # pad to multiple of groups - validation_prompt_ids += [0] * (num_sp_groups - num_embeds % num_sp_groups) - num_embeds_per_group = len(validation_prompt_ids) // num_sp_groups - local_prompt_ids = validation_prompt_ids[nccl_info.group_id * num_embeds_per_group: (nccl_info.group_id + 1) * num_embeds_per_group] - for i in local_prompt_ids: - prompt_embed_path = os.path.join(args.validation_prompt_dir, f"embed{i}.pt") - prompt_mask_path = os.path.join(args.validation_prompt_dir, f"mask{i}.pt") - negative_prompt_embed_path = os.path.join(args.uncond_prompt_dir, "embed.pt") - negative_prompt_mask_path = os.path.join(args.uncond_prompt_dir, "mask.pt") - prompt_embeds = torch.load(prompt_embed_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0) - prompt_attention_mask = torch.load(prompt_mask_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0) - negative_prompt_embeds = torch.load(negative_prompt_embed_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0) - negative_prompt_attention_mask = torch.load(negative_prompt_mask_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0) - generator = torch.Generator(device="cuda").manual_seed(12345) - video = sample_validation_video( - transformer, - vae, - scheduler, - scheduler_type=scheduler_type, - num_frames=args.num_frames, - # Peiyuan TODO: remove hardcode - height=480, - width=848, - num_inference_steps=args.validation_sampling_steps, - guidance_scale=args.validation_guidance_scale, - generator=generator, - prompt_embeds = prompt_embeds, - prompt_attention_mask = prompt_attention_mask, - negative_prompt_embeds = negative_prompt_embeds, - negative_prompt_attention_mask = negative_prompt_attention_mask, - )[0] - if nccl_info.rank_within_group == 0: - videos.append(video[0]) - # collect videos from all process to process zero - - gc.collect() - torch.cuda.empty_cache() - # log if main process - torch.distributed.barrier() - all_videos = [None for i in range(int(os.getenv("WORLD_SIZE", '1')))] # remove padded videos - torch.distributed.all_gather_object(all_videos, videos) - if nccl_info.global_rank == 0: - # remove padding - videos = [video for videos in all_videos for video in videos] - videos = videos[:num_embeds] - # linearize all videos - video_filenames = [] - for i, video in enumerate(videos): - filename = os.path.join(args.output_dir, f"validation_step_{global_step}_video_{i}.mp4") - export_to_video(video, filename, fps=30) - video_filenames.append(filename) + validation_guidance_scale_ls = args.validation_guidance_scale.split(",") + validation_guidance_scale_ls = [float(scale) for scale in validation_guidance_scale_ls] + for validation_guidance_scale in validation_guidance_scale_ls: - logs = { - f"{'ema_' if ema else ''}validation": [ - wandb.Video(filename) - for i, filename in enumerate(video_filenames) - ] - } - wandb.log(logs, step=global_step) + videos = [] + # prompt_embed are named embed0 to embedN + # check how many embeds are there + num_embeds = len([f for f in os.listdir(args.validation_prompt_dir) if "embed" in f]) + validation_prompt_ids = list(range(num_embeds)) + num_sp_groups = int(os.getenv("WORLD_SIZE", '1')) // nccl_info.sp_size + # pad to multiple of groups + validation_prompt_ids += [0] * (num_sp_groups - num_embeds % num_sp_groups) + num_embeds_per_group = len(validation_prompt_ids) // num_sp_groups + local_prompt_ids = validation_prompt_ids[nccl_info.group_id * num_embeds_per_group: (nccl_info.group_id + 1) * num_embeds_per_group] + + for i in local_prompt_ids: + prompt_embed_path = os.path.join(args.validation_prompt_dir, f"embed{i}.pt") + prompt_mask_path = os.path.join(args.validation_prompt_dir, f"mask{i}.pt") + prompt_embeds = torch.load(prompt_embed_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0) + prompt_attention_mask = torch.load(prompt_mask_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0) + negative_prompt_embeds = torch.zeros(256, 4096).to(device).to(weight_dtype).unsqueeze(0) + negative_prompt_attention_mask = torch.zeros(256).bool().to(device).unsqueeze(0) + generator = torch.Generator(device="cuda").manual_seed(12345) + video = sample_validation_video( + transformer, + vae, + scheduler, + scheduler_type=scheduler_type, + num_frames=args.num_frames, + # Peiyuan TODO: remove hardcode + height=480, + width=848, + num_inference_steps=args.validation_sampling_steps, + guidance_scale=validation_guidance_scale, + generator=generator, + prompt_embeds = prompt_embeds, + prompt_attention_mask = prompt_attention_mask, + negative_prompt_embeds = negative_prompt_embeds, + negative_prompt_attention_mask = negative_prompt_attention_mask, + )[0] + if nccl_info.rank_within_group == 0: + videos.append(video[0]) + # collect videos from all process to process zero + + gc.collect() + torch.cuda.empty_cache() + # log if main process + torch.distributed.barrier() + all_videos = [None for i in range(int(os.getenv("WORLD_SIZE", '1')))] # remove padded videos + torch.distributed.all_gather_object(all_videos, videos) + if nccl_info.global_rank == 0: + # remove padding + videos = [video for videos in all_videos for video in videos] + videos = videos[:num_embeds] + # linearize all videos + video_filenames = [] + for i, video in enumerate(videos): + filename = os.path.join(args.output_dir, f"validation_step_{global_step}_guidance_{validation_guidance_scale}_video_{i}.mp4") + export_to_video(video, filename, fps=30) + video_filenames.append(filename) + + logs = { + f"{'ema_' if ema else ''}validation_guidance_{validation_guidance_scale}": [ + wandb.Video(filename) + for i, filename in enumerate(video_filenames) + ] + } + wandb.log(logs, step=global_step)