nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

linear range; fix uncond; multi guidance validation

Browse Source
This commit is contained in:
rlsu9
2024-12-01 04:31:53 +04:00
parent a3d4a1cef4
commit 6d3a71b681
6 changed files with 108 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/dataset/latent_datasets.py
+4 -3
View File
@@ -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"]
fastvideo/distill.py
+28 -11
View File
@@ -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)
fastvideo/distill/solver.py
+3 -2
View File
@@ -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
fastvideo/distill_adv.py
+1 -2
View File
@@ -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)
fastvideo/sample/sample_t2v_mochi.py
+3 -1
View File
@@ -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)
fastvideo/utils/validation.py
+69 -66
View File
@@ -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)