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

[feat]: Add validation logging with SP (#36)

Browse Source
This commit is contained in:
Zhang Peiyuan
2024-11-05 14:46:20 -08:00
parent f0bc297260
commit 52ba538e8e
5 changed files with 302 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/model/pipeline_mochi.py
+1 -1
View File
@@ -631,7 +631,7 @@ class MochiPipeline(DiffusionPipeline):
generator,
latents,
)
world_size, rank = nccl_info.sp_size, nccl_info.global_rank
world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
if get_sequence_parallel_state():
latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
latents = latents[:, :, rank, :, :, :]
fastvideo/train.py
+22 -60
View File
@@ -4,16 +4,14 @@ import math
import os
import shutil
from pathlib import Path
import gc
import wandb
from einops import rearrange
from tqdm import tqdm
from diffusers.training_utils import cast_training_params, compute_density_for_timestep_sampling, compute_loss_weighting_for_sd3
from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, \
destroy_sequence_parallel_group, get_sequence_parallel_state
destroy_sequence_parallel_group, get_sequence_parallel_state, nccl_info
from fastvideo.utils.communications import prepare_sequence_parallel_data, broadcast
from fastvideo.model.mochi_monkey_patches import hf_mochi_add_sp_monkey_patch
from fastvideo.model.mochi_latents_stat import mochi_stat
from fastvideo.utils.validation import log_validation
import time
from torch.utils.data import DataLoader
from copy import deepcopy
@@ -27,10 +25,8 @@ from tqdm.auto import tqdm
import diffusers
from diffusers import (
FlowMatchEulerDiscreteScheduler,
MochiTransformer3DModel,
MochiPipeline,
)
from diffusers.optimization import get_scheduler
from diffusers.utils import check_min_version, is_wandb_available
@@ -60,41 +56,7 @@ class ForkedPdb(pdb.Pdb):
sys.stdin = _stdin
@torch.inference_mode()
def log_validation(args, transformer,accelerator, weight_dtype, global_step, ema=False):
#TODO
logger.info(f"Running validation....\n")
transformer = accelerator.unwrap_model(transformer)
mochi_pipeline = MochiPipeline.from_pretrained(args.pretrained_model_name_or_path, transformer=transformer, torch_dtype=weight_dtype).to(device=accelerator.device)
videos = []
for prompt in args.validaiton_prompt:
logger.info('Processing the ({}) prompt'.format(prompt))
video = mochi_pipeline(
args.prompt,
num_frames=args.num_frames,
height=args.max_height,
width=args.max_width,
num_inference_steps=args.validation_sampling_steps,
guidance_scale=args.guidance_scale,
).images
videos.append(video[0])
# import ipdb;ipdb.set_trace()
gc.collect()
torch.cuda.empty_cache()
videos = torch.stack(videos).numpy()
videos = rearrange(videos, 'b t h w c -> b t c h w')
for tracker in accelerator.trackers:
logs = {
f"{'ema_' if ema else ''}validation": [
wandb.Video(video, caption=f"{i}: {prompt}", fps=24)
for i, (video, prompt) in enumerate(zip(videos, args.validaiton_prompt))
]
}
tracker.log(logs, step=global_step)
del mochi_pipeline
gc.collect()
torch.cuda.empty_cache()
class ProgressInfo:
@@ -484,22 +446,20 @@ def main(args):
# accelerator.print("Syncing gradients")
sync_gradients_info()
if accelerator.is_main_process:
if progress_info.global_step % args.checkpointing_steps == 0:
if args.enable_tracker:
log_validation(args, transformer, accelerator,
weight_dtype, progress_info.global_step)
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer.store(transformer.parameters())
ema_transformer.copy_to(transformer.parameters())
log_validation(args, transformer, accelerator,
weight_dtype, progress_info.global_step, ema=True)
# Switch back to the original UNet parameters.
ema_transformer.restore(transformer.parameters())
# TODO: add indent
#Group ID indicates which SP group the process belongs to. We only need SP group 0 to log validation.
if progress_info.global_step % args.validation_steps == 0 and args.log_validation:
log_validation(args, transformer, accelerator,
weight_dtype, progress_info.global_step)
if args.use_ema:
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
ema_transformer.store(transformer.parameters())
ema_transformer.copy_to(transformer.parameters())
log_validation(args, transformer, accelerator,
weight_dtype, progress_info.global_step, ema=True)
# Switch back to the original UNet parameters.
ema_transformer.restore(transformer.parameters())
return loss
@@ -549,7 +509,7 @@ if __name__ == "__main__":
# dataset & dataloader
parser.add_argument("--data_json_path", type=str, required=True)
parser.add_argument("--num_frames", type=int, default=65)
parser.add_argument("--num_frames", type=int, default=163)
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.")
@@ -589,10 +549,12 @@ if __name__ == "__main__":
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.")
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('--guidance_scale', type=float, default=4.5)
parser.add_argument("--enable_tracker", action="store_true")
parser.add_argument('--validation_guidance_scale', type=float, default=4.5)
parser.add_argument('--validation_steps', type=float, default=4.5)
parser.add_argument("--log_validation", action="store_true")
parser.add_argument("--tracker_project_name", type=str, default=None)
parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
parser.add_argument("--output_dir", type=str, default=None, help="The output directory where the model predictions and checkpoints will be written.")
fastvideo/utils/validation.py
+244
View File
@@ -0,0 +1,244 @@
from typing import Optional, Union, List
import numpy as np
import torch
from einops import rearrange
from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
from fastvideo.utils.communications import all_gather
from diffusers.utils.torch_utils import randn_tensor
from fastvideo.model.pipeline_mochi import linear_quadratic_schedule, retrieve_timesteps
from tqdm import tqdm
from diffusers.video_processor import VideoProcessor
from diffusers import (
FlowMatchEulerDiscreteScheduler,
AutoencoderKLMochi,
)
from diffusers.utils import export_to_video
import os
import wandb
import gc
def prepare_latents(
batch_size,
num_channels_latents,
height,
width,
num_frames,
dtype,
device,
generator,
vae_spatial_scale_factor,
vae_temporal_scale_factor,
):
height = height // vae_spatial_scale_factor
width = width // vae_spatial_scale_factor
num_frames = (num_frames - 1) // vae_temporal_scale_factor + 1
shape = (batch_size, num_channels_latents, num_frames, height, width)
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
return latents
def sample_validation_video(
transformer,
vae,
scheduler,
height: Optional[int] = None,
width: Optional[int] = None,
num_frames: int = 16,
num_inference_steps: int = 28,
timesteps: List[int] = None,
guidance_scale: float = 4.5,
num_videos_per_prompt: Optional[int] = 1,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
prompt_embeds: Optional[torch.Tensor] = None,
prompt_attention_mask: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
output_type: Optional[str] = "pil",
vae_spatial_scale_factor = 8,
vae_temporal_scale_factor = 6,
):
device = vae.device
batch_size = prompt_embeds.shape[0]
do_classifier_free_guidance = guidance_scale > 1.0
if do_classifier_free_guidance:
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
# 4. Prepare latent variables
# TODO: Remove hardcore
num_channels_latents = 12
latents = prepare_latents(
batch_size * num_videos_per_prompt,
num_channels_latents,
height,
width,
num_frames,
prompt_embeds.dtype,
device,
generator,
vae_spatial_scale_factor,
vae_temporal_scale_factor
)
world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
if get_sequence_parallel_state():
latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
latents = latents[:, :, rank, :, :, :]
# 5. Prepare timestep
# from https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77
threshold_noise = 0.025
sigmas = linear_quadratic_schedule(num_inference_steps, threshold_noise)
sigmas = np.array(sigmas)
timesteps, num_inference_steps = retrieve_timesteps(
scheduler,
num_inference_steps,
device,
timesteps,
sigmas,
)
num_warmup_steps = max(len(timesteps) - num_inference_steps * scheduler.order, 0)
# 6. Denoising loop
# with self.progress_bar(total=num_inference_steps) as progress_bar:
# write with tqdm instead
# only enable if nccl_info.global_rank == 0
with tqdm(total=num_inference_steps, disable= nccl_info.global_rank != 0, desc="Validation sampling...") as progress_bar:
for i, t in enumerate(timesteps):
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype)
noise_pred = transformer(
hidden_states=latent_model_input,
encoder_hidden_states=prompt_embeds,
timestep=timestep,
encoder_attention_mask=prompt_attention_mask,
return_dict=False,
)[0]
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents_dtype = latents.dtype
latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0]
if latents.dtype != latents_dtype:
if torch.backends.mps.is_available():
# some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
latents = latents.to(latents_dtype)
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % scheduler.order == 0):
progress_bar.update()
if get_sequence_parallel_state():
latents = all_gather(latents, dim=2)
if output_type == "latent":
video = latents
else:
# unscale/denormalize the latents
# denormalize with the mean and std if available and not None
has_latents_mean = hasattr(vae.config, "latents_mean") and vae.config.latents_mean is not None
has_latents_std = hasattr(vae.config, "latents_std") and vae.config.latents_std is not None
if has_latents_mean and has_latents_std:
latents_mean = (
torch.tensor(vae.config.latents_mean).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
)
latents_std = (
torch.tensor(vae.config.latents_std).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
)
latents = latents * latents_std / vae.config.scaling_factor + latents_mean
else:
latents = latents / vae.config.scaling_factor
video = vae.decode(latents, return_dict=False)[0]
video_processor = VideoProcessor(vae_scale_factor=vae_spatial_scale_factor)
video = video_processor.postprocess_video(video, output_type=output_type)
return (video,)
@torch.inference_mode()
def log_validation(args, transformer,accelerator, weight_dtype, global_step, ema=False):
#TODO
accelerator.print(f"Running validation....\n")
transformer = accelerator.unwrap_model(transformer)
generator = torch.Generator(device=accelerator.device).manual_seed(12345)
vae = AutoencoderKLMochi.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", torch_dtype=weight_dtype).to(accelerator.device)
vae.enable_tiling()
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler", torch_dtype=weight_dtype)
# args.validation_prompt_dir
prompt_embed_path = os.path.join(args.validation_prompt_dir, "embed.pt")
prompt_mask_path = os.path.join(args.validation_prompt_dir, "mask.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=accelerator.device).to(accelerator.device).to(weight_dtype).unsqueeze(0)
prompt_attention_mask = torch.load(prompt_mask_path, map_location=accelerator.device).to(accelerator.device).to(weight_dtype).unsqueeze(0)
negative_prompt_embeds = torch.load(negative_prompt_embed_path, map_location=accelerator.device).to(accelerator.device).to(weight_dtype).unsqueeze(0)
negative_prompt_attention_mask = torch.load(negative_prompt_mask_path, map_location=accelerator.device).to(accelerator.device).to(weight_dtype).unsqueeze(0)
videos = []
video = sample_validation_video(
transformer,
vae,
scheduler,
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]
videos.append(video[0])
# import ipdb;ipdb.set_trace()
gc.collect()
torch.cuda.empty_cache()
# log if main process
if accelerator.is_main_process :
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)
for tracker in accelerator.trackers:
logs = {
f"{'ema_' if ema else ''}validation": [
wandb.Video(filename)
for i, filename in enumerate(video_filenames)
]
}
tracker.log(logs, step=global_step)
del vae
del prompt_embeds
del prompt_attention_mask
del negative_prompt_embeds
del negative_prompt_attention_mask
gc.collect()
torch.cuda.empty_cache()
inference_mochi.sh
+21 -22
View File
@@ -1,35 +1,34 @@
#!/bin/bash
num_gpus=2
num_gpus=4
torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
fastvideo/sample/sample_t2v_mochi.py \
--model_path data/mochi \
--prompts "A hand enters the frame, pulling a sheet of plastic wrap over three balls of dough placed on a wooden surface. The plastic wrap is stretched to cover the dough more securely. The hand adjusts the wrap, ensuring that it is tight and smooth over the dough. The scene focuses on the hand's movements as it secures the edges of the plastic wrap. No new objects appear, and the camera remains stationary, focusing on the action of covering the dough." \
--num_frames 79 \
--height 480 \
--width 848 \
--num_inference_steps 64 \
--guidance_scale 4.5 \
--transformer_path data/outputs/debug_synthetic_uniform/checkpoint-200/transformer\
--output_path outputs_video/overfit_reproduce/uniform_cfg_train_200_cfg_f79_syn \
--seed 12345
num_gpus=2
torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
fastvideo/sample/sample_t2v_mochi.py \
--model_path data/mochi \
--prompt_embed_path data/Encoder_Overfit_Data/prompt_embed/0.pt \
--encoder_attention_mask_path data/Encoder_Overfit_Data/prompt_attention_mask/0.pt \
--prompts "A cute fluffy panda eating Chinese food in a restaurant" \
--num_frames 163 \
--height 480 \
--width 848 \
--num_inference_steps 64 \
--guidance_scale 4.5 \
--transformer_path data/outputs/debug_uniform_28/checkpoint-100/transformer \
--output_path outputs_video/overfit/uniform_cfg_train_100_cfg_f163_sp4 \
--output_path outputs_video/debug \
--seed 12345
num_gpus=2
torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
fastvideo/sample/sample_t2v_mochi.py \
--model_path data/mochi \
--prompt_embed_path data/Encoder_Overfit_Data/prompt_embed/1.pt \
--encoder_attention_mask_path data/Encoder_Overfit_Data/prompt_attention_mask/1.pt \
--num_frames 79 \
--height 480 \
--width 848 \
--num_inference_steps 64 \
--guidance_scale 4.5 \
--transformer_path data/outputs/debug_uniform_14_sp2_spb2_4videos_2K/checkpoint-1000/transformer \
--output_path outputs_video/overfit/uniform_cfg_train_1000_cfg_f79_sp2_spb2_4videos_second \
--seed 12345
t2v_debug_multi.sh
+14 -91
View File
@@ -1,3 +1,4 @@
accelerate launch \
--config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \
fastvideo/train.py \
@@ -8,48 +9,21 @@ accelerate launch \
--gradient_checkpointing \
--train_batch_size=1 \
--num_latent_t 14 \
--sp_size 1 \
--train_sp_batch_size 1 \
--sp_size 2 \
--train_sp_batch_size 2 \
--dataloader_num_workers 1 \
--gradient_accumulation_steps=2 \
--max_train_steps=500 \
--max_train_steps=2000 \
--learning_rate=1e-5 \
--mixed_precision="bf16" \
--weighting_scheme "uniform" \
--checkpointing_steps=100 \
--checkpointing_steps=500 \
--checkpoints_total_limit 3 \
--allow_tf32 \
--ema_start_step 0 \
--cfg 0.1 \
--ema_decay 0.999 \
--output_dir="data/outputs/debug_cfg_uniform"
accelerate launch \
--config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \
fastvideo/train.py \
--seed 42 \
--pretrained_model_name_or_path data/mochi \
--cache_dir "data/.cache" \
--data_json_path "data/synthetic_debug/videos2caption.json" \
--gradient_checkpointing \
--train_batch_size=1 \
--num_latent_t 14 \
--sp_size 1 \
--train_sp_batch_size 1 \
--dataloader_num_workers 1 \
--gradient_accumulation_steps=2 \
--max_train_steps=200 \
--learning_rate=1e-5 \
--mixed_precision="bf16" \
--weighting_scheme "uniform" \
--checkpointing_steps=100 \
--checkpoints_total_limit 3 \
--allow_tf32 \
--ema_start_step 0 \
--cfg 0.1 \
--ema_decay 0.999 \
--output_dir="data/outputs/debug_synthetic_uniform"
--output_dir="data/outputs/debug_uniform_14_sp2_spb2_4videos_2K"
@@ -63,78 +37,27 @@ accelerate launch \
--data_json_path "data/Encoder_Overfit_Data/videos2caption.json" \
--gradient_checkpointing \
--train_batch_size=1 \
--num_latent_t 28 \
--num_latent_t 20 \
--sp_size 4 \
--train_sp_batch_size 1 \
--dataloader_num_workers 1 \
--gradient_accumulation_steps=8 \
--max_train_steps=500 \
--learning_rate=1e-5 \
--mixed_precision="bf16" \
--weighting_scheme "uniform" \
--checkpointing_steps=100 \
--checkpoints_total_limit 3 \
--allow_tf32 \
--ema_start_step 0 \
--cfg 0.1 \
--ema_decay 0.999 \
--output_dir="data/outputs/debug_uniform_28"
accelerate launch \
--config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \
fastvideo/train.py \
--seed 42 \
--pretrained_model_name_or_path data/mochi \
--cache_dir "data/.cache" \
--data_json_path "data/Encoder_Overfit_Data/videos2caption.json" \
--gradient_checkpointing \
--train_batch_size=1 \
--num_latent_t 14 \
--sp_size 2 \
--train_sp_batch_size 2 \
--dataloader_num_workers 1 \
--gradient_accumulation_steps=2 \
--max_train_steps=500 \
--learning_rate=1e-5 \
--learning_rate=0.0 \
--mixed_precision="bf16" \
--weighting_scheme "uniform" \
--checkpointing_steps=100 \
--validation_steps 2 \
--validation_sampling_steps 64 \
--checkpoints_total_limit 3 \
--allow_tf32 \
--ema_start_step 0 \
--cfg 0.1 \
--ema_decay 0.999 \
--output_dir="data/outputs/debug_uniform_14_sp2_spb_2"
accelerate launch \
--config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \
fastvideo/train.py \
--seed 42 \
--pretrained_model_name_or_path data/mochi \
--cache_dir "data/.cache" \
--data_json_path "data/Encoder_Overfit_Data/videos2caption.json" \
--gradient_checkpointing \
--train_batch_size=1 \
--num_latent_t 14 \
--sp_size 2 \
--train_sp_batch_size 2 \
--dataloader_num_workers 1 \
--gradient_accumulation_steps=2 \
--max_train_steps=500 \
--learning_rate=1e-5 \
--mixed_precision="bf16" \
--weighting_scheme "uniform" \
--checkpointing_steps=100 \
--checkpoints_total_limit 3 \
--allow_tf32 \
--ema_start_step 0 \
--cfg 0.1 \
--ema_decay 0.999 \
--output_dir="data/outputs/debug_uniform_14_sp2_spb_2_4videos"
--log_validation \
--validation_prompt_dir "data/Encoder_Overfit_Data/validation_prompt_embed_mask" \
--uncond_prompt_dir "data/Encoder_Overfit_Data/uncond_prompt_embed_mask" \
--output_dir="data/outputs/debug_logvalidation_sp2"