import os import math import wandb import random import time import logging import inspect import argparse import datetime import numpy as np from pathlib import Path from tqdm.auto import tqdm from einops import rearrange, repeat from omegaconf import OmegaConf from typing import Dict, Optional, Tuple import torch import torch.nn.functional as F import diffusers from diffusers import AutoencoderKL, DDIMScheduler from diffusers.optimization import get_scheduler from diffusers.utils import check_min_version from diffusers.utils.import_utils import is_xformers_available from diffusers.training_utils import EMAModel import transformers from transformers import CLIPTextModel, CLIPTokenizer from accelerate import Accelerator, DistributedDataParallelKwargs, InitProcessGroupKwargs from accelerate.logging import get_logger from accelerate.utils import set_seed from consisti2v.data.dataset import WebVid10M, Pexels, JointDataset from consisti2v.models.videoldm_unet import VideoLDMUNet3DConditionModel from consisti2v.pipelines.pipeline_conditional_animation import ConditionalAnimationPipeline from consisti2v.utils.util import save_videos_grid logger = get_logger(__name__, log_level="INFO") def main( name: str, use_wandb: bool, is_image: bool, output_dir: str, pretrained_model_path: str, train_data: Dict, validation_data: Dict, cfg_random_null_text_ratio: float = 0.1, cfg_random_null_img_ratio: float = 0.0, resume_from_checkpoint: Optional[str] = None, unet_additional_kwargs: Dict = {}, use_ema: bool = False, ema_decay: float = 0.9999, noise_scheduler_kwargs = None, max_train_epoch: int = -1, max_train_steps: int = 100, validation_steps: int = 100, learning_rate: float = 3e-5, scale_lr: bool = False, lr_warmup_steps: int = 0, lr_scheduler: str = "constant", trainable_modules: Tuple[str] = (None, ), num_workers: int = 32, train_batch_size: int = 1, adam_beta1: float = 0.9, adam_beta2: float = 0.999, adam_weight_decay: float = 1e-2, adam_epsilon: float = 1e-08, max_grad_norm: float = 1.0, gradient_accumulation_steps: int = 1, gradient_checkpointing: bool = False, checkpointing_epochs: int = 5, checkpointing_steps: int = -1, mixed_precision: Optional[str] = "fp16", enable_xformers_memory_efficient_attention: bool = True, seed: Optional[int] = 42, is_debug: bool = False, ): check_min_version("0.10.0.dev0") *_, config = inspect.getargvalues(inspect.currentframe()) config = {k: v for k, v in config.items() if k != 'config' and k != '_'} ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True if not is_image else False) init_kwargs = InitProcessGroupKwargs(timeout=datetime.timedelta(seconds=3600)) accelerator = Accelerator( gradient_accumulation_steps=gradient_accumulation_steps, mixed_precision=mixed_precision, kwargs_handlers=[ddp_kwargs, init_kwargs], ) if seed is not None: set_seed(seed) # Logging folder folder_name = "debug" if is_debug else name + datetime.datetime.now().strftime("-%Y-%m-%dT%H-%M-%S") output_dir = os.path.join(output_dir, folder_name) if is_debug and os.path.exists(output_dir): os.system(f"rm -rf {output_dir}") # Make one log on every process with the configuration for debugging. logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO, ) logger.info(accelerator.state, main_process_only=False) if accelerator.is_local_main_process: transformers.utils.logging.set_verbosity_warning() diffusers.utils.logging.set_verbosity_info() else: transformers.utils.logging.set_verbosity_error() diffusers.utils.logging.set_verbosity_error() if accelerator.is_main_process and (not is_debug) and use_wandb: project_name = "text_image_to_video" if not is_image else "image_finetune" wandb.init(project=project_name, name=folder_name, config=config) accelerator.wait_for_everyone() # Handle the output folder creation if accelerator.is_main_process: os.makedirs(output_dir, exist_ok=True) os.makedirs(f"{output_dir}/samples", exist_ok=True) os.makedirs(f"{output_dir}/sanity_check", exist_ok=True) os.makedirs(f"{output_dir}/checkpoints", exist_ok=True) OmegaConf.save(config, os.path.join(output_dir, 'config.yaml')) # TODO: change all datasets to fps+duration in the future if train_data.dataset == "pexels": train_data.sample_n_frames = train_data.sample_duration * train_data.sample_fps elif train_data.dataset == "joint": if train_data.sample_duration is not None: train_data.sample_n_frames = train_data.sample_duration * train_data.sample_fps # Load scheduler, tokenizer and models. noise_scheduler = DDIMScheduler(**OmegaConf.to_container(noise_scheduler_kwargs)) vae = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae") tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer") text_encoder = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder") unet = VideoLDMUNet3DConditionModel.from_pretrained( pretrained_model_path, subfolder="unet", variant=unet_additional_kwargs['variant'], use_temporal=True if not is_image else False, temp_pos_embedding=unet_additional_kwargs['temp_pos_embedding'], augment_temporal_attention=unet_additional_kwargs['augment_temporal_attention'], n_frames=train_data.sample_n_frames if not is_image else 2, n_temp_heads=unet_additional_kwargs['n_temp_heads'], first_frame_condition_mode=unet_additional_kwargs['first_frame_condition_mode'], use_frame_stride_condition=unet_additional_kwargs['use_frame_stride_condition'], use_safetensors=True ) # Freeze vae and text_encoder vae.requires_grad_(False) text_encoder.requires_grad_(False) unet.train() if use_ema: ema_unet = VideoLDMUNet3DConditionModel.from_pretrained( pretrained_model_path, subfolder="unet", variant=unet_additional_kwargs['variant'], use_temporal=True if not is_image else False, temp_pos_embedding=unet_additional_kwargs['temp_pos_embedding'], augment_temporal_attention=unet_additional_kwargs['augment_temporal_attention'], n_frames=train_data.sample_n_frames if not is_image else 2, n_temp_heads=unet_additional_kwargs['n_temp_heads'], first_frame_condition_mode=unet_additional_kwargs['first_frame_condition_mode'], use_frame_stride_condition=unet_additional_kwargs['use_frame_stride_condition'], use_safetensors=True ) ema_unet = EMAModel(ema_unet.parameters(), decay=ema_decay, model_cls=VideoLDMUNet3DConditionModel, model_config=ema_unet.config) # Set unet trainable parameters train_all_parameters = False for trainable_module_name in trainable_modules: if trainable_module_name == 'all': unet.requires_grad_(True) train_all_parameters = True break if not train_all_parameters: unet.requires_grad_(False) for name, param in unet.named_parameters(): for trainable_module_name in trainable_modules: if trainable_module_name in name: param.requires_grad = True break # Enable xformers if enable_xformers_memory_efficient_attention and int(torch.__version__.split(".")[0]) < 2: if is_xformers_available(): unet.enable_xformers_memory_efficient_attention() else: raise ValueError("xformers is not available. Make sure it is installed correctly") def save_model_hook(models, weights, output_dir): if accelerator.is_main_process: if use_ema: ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema")) for i, model in enumerate(models): model.save_pretrained(os.path.join(output_dir, "unet")) # make sure to pop weight so that corresponding model is not saved again weights.pop() def load_model_hook(models, input_dir): if use_ema: load_model = EMAModel.from_pretrained(os.path.join(input_dir, "unet_ema"), VideoLDMUNet3DConditionModel) ema_unet.load_state_dict(load_model.state_dict()) ema_unet.to(accelerator.device) del load_model for i in range(len(models)): # pop models so that they are not loaded again model = models.pop() # load diffusers style into model load_model = VideoLDMUNet3DConditionModel.from_pretrained(input_dir, subfolder="unet") model.register_to_config(**load_model.config) model.load_state_dict(load_model.state_dict()) del load_model accelerator.register_save_state_pre_hook(save_model_hook) accelerator.register_load_state_pre_hook(load_model_hook) # Enable gradient checkpointing if gradient_checkpointing: unet.enable_gradient_checkpointing() if scale_lr: learning_rate = (learning_rate * gradient_accumulation_steps * train_batch_size * accelerator.num_processes) trainable_params = list(filter(lambda p: p.requires_grad, unet.parameters())) optimizer = torch.optim.AdamW( trainable_params, lr=learning_rate, betas=(adam_beta1, adam_beta2), weight_decay=adam_weight_decay, eps=adam_epsilon, ) logger.info(f"trainable params number: {len(trainable_params)}") logger.info(f"trainable params scale: {sum(p.numel() for p in trainable_params) / 1e6:.3f} M") # Get the training dataset if train_data['dataset'] == "webvid": train_dataset = WebVid10M(**train_data, is_image=is_image) elif train_data['dataset'] == "pexels": train_dataset = Pexels(**train_data, is_image=is_image) elif train_data['dataset'] == "joint": train_dataset = JointDataset(**train_data, is_image=is_image) else: raise ValueError(f"Unknown dataset {train_data['dataset']}") # DataLoaders creation: train_dataloader = torch.utils.data.DataLoader( train_dataset, shuffle=True, batch_size=train_batch_size, num_workers=num_workers, pin_memory=True, ) # Get the training iteration if max_train_steps == -1: assert max_train_epoch != -1 max_train_steps = max_train_epoch * len(train_dataloader) if checkpointing_steps == -1: assert checkpointing_epochs != -1 checkpointing_steps = checkpointing_epochs * len(train_dataloader) # Scheduler lr_scheduler = get_scheduler( lr_scheduler, optimizer=optimizer, num_warmup_steps=lr_warmup_steps * gradient_accumulation_steps, num_training_steps=max_train_steps * gradient_accumulation_steps, ) # Validation pipeline validation_pipeline = ConditionalAnimationPipeline( unet=unet, vae=vae, tokenizer=tokenizer, text_encoder=text_encoder, scheduler=noise_scheduler, ) validation_pipeline.enable_vae_slicing() # Prepare everything with our `accelerator`. unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( unet, optimizer, train_dataloader, lr_scheduler ) # For mixed precision training we cast the text_encoder and vae weights to half-precision # as these models are only used for inference, keeping weights in full precision is not required. weight_dtype = torch.float32 if accelerator.mixed_precision == "fp16": weight_dtype = torch.float16 elif accelerator.mixed_precision == "bf16": weight_dtype = torch.bfloat16 if use_ema: ema_unet.to(accelerator.device) # Move text_encode and vae to gpu and cast to weight_dtype text_encoder.to(accelerator.device, dtype=weight_dtype) vae.to(accelerator.device, dtype=weight_dtype) # We need to recalculate our total training steps as the size of the training dataloader may have changed. num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps) # Afterwards we recalculate our number of training epochs num_train_epochs = math.ceil(max_train_steps / num_update_steps_per_epoch) # Train! total_batch_size = train_batch_size * accelerator.num_processes * gradient_accumulation_steps logger.info("***** Running training *****") logger.info(f" Num examples = {len(train_dataset)}") logger.info(f" Num Epochs = {num_train_epochs}") logger.info(f" Instantaneous batch size per device = {train_batch_size}") logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") logger.info(f" Gradient Accumulation steps = {gradient_accumulation_steps}") logger.info(f" Total optimization steps = {max_train_steps}") global_step = 0 first_epoch = 0 # Load pretrained unet weights if resume_from_checkpoint is not None: logger.info(f"Resuming from checkpoint: {resume_from_checkpoint}") accelerator.load_state(resume_from_checkpoint) global_step = int(resume_from_checkpoint.split("-")[-1]) initial_global_step = global_step first_epoch = global_step // num_update_steps_per_epoch logger.info(f"global_step: {global_step}") logger.info(f"first_epoch: {first_epoch}") else: initial_global_step = 0 # Only show the progress bar once on each machine. progress_bar = tqdm(range(0, max_train_steps), initial=initial_global_step, desc="Steps", disable=not accelerator.is_main_process) for epoch in range(first_epoch, num_train_epochs): train_loss = 0.0 train_grad_norm = 0.0 data_loading_time = 0.0 prepare_everything_time = 0.0 network_forward_time = 0.0 network_backward_time = 0.0 t0 = time.time() for step, batch in enumerate(train_dataloader): t1 = time.time() if cfg_random_null_text_ratio > 0.0: batch['text'] = [name if random.random() > cfg_random_null_text_ratio else "" for name in batch['text']] # Data batch sanity check if accelerator.is_main_process and epoch == first_epoch and step == 0: pixel_values, texts = batch['pixel_values'].cpu(), batch['text'] pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") for idx, (pixel_value, text) in enumerate(zip(pixel_values, texts)): pixel_value = pixel_value[None, ...] save_videos_grid(pixel_value, f"{output_dir}/sanity_check/{'-'.join(text.replace('/', '').split()[:10]) if not text == '' else f'no_text-{idx}'}.gif", rescale=True) ### >>>> Training >>>> ### with accelerator.accumulate(unet): # Convert videos to latent space pixel_values = batch["pixel_values"].to(weight_dtype) video_length = pixel_values.shape[1] pixel_values = rearrange(pixel_values, "b f c h w -> (b f) c h w") latents = vae.encode(pixel_values).latent_dist latents = latents.sample() latents = rearrange(latents, "(b f) c h w -> b c f h w", f=video_length) latents = latents * vae.config.scaling_factor if unet_additional_kwargs["first_frame_condition_mode"] != "none": # Get first frame latents first_frame_latents = latents[:, :, 0:1, :, :] # Sample noise that we'll add to the latents if unet_additional_kwargs['noise_sampling_method'] == 'vanilla': noise = torch.randn_like(latents) elif unet_additional_kwargs['noise_sampling_method'] == 'pyoco_mixed': noise_alpha_squared = float(unet_additional_kwargs['noise_alpha']) ** 2 shared_noise = torch.randn_like(latents[:, :, 0:1, :, :]) * math.sqrt((noise_alpha_squared) / (1 + noise_alpha_squared)) ind_noise = torch.randn_like(latents) * math.sqrt(1 / (1 + noise_alpha_squared)) noise = shared_noise + ind_noise elif unet_additional_kwargs['noise_sampling_method'] == 'pyoco_progressive': noise_alpha_squared = float(unet_additional_kwargs['noise_alpha']) ** 2 noise = torch.randn_like(latents) ind_noise = torch.randn_like(latents) * math.sqrt(1 / (1 + noise_alpha_squared)) for i in range(1, noise.shape[2]): noise[:, :, i, :, :] = noise[:, :, i - 1, :, :] * math.sqrt((noise_alpha_squared) / (1 + noise_alpha_squared)) + ind_noise[:, :, i, :, :] else: raise ValueError(f"Unknown noise sampling method {unet_additional_kwargs['noise_sampling_method']}") bsz = latents.shape[0] # Sample a random timestep for each video timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device) timesteps = timesteps.long() # Add noise to the latents according to the noise magnitude at each timestep # (this is the forward diffusion process) noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps) if cfg_random_null_img_ratio > 0.0: for i in range(first_frame_latents.shape[0]): if random.random() <= cfg_random_null_img_ratio: first_frame_latents[i, :, :, :, :] = noisy_latents[i, :, 0:1, :, :] # Remove the first noisy latent from the latents if we're conditioning on the first frame if unet_additional_kwargs["first_frame_condition_mode"] != "none": noisy_latents = noisy_latents[:, :, 1:, :, :] # Get the text embedding for conditioning prompt_ids = tokenizer( batch['text'], max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt" ).input_ids.to(latents.device) encoder_hidden_states = text_encoder(prompt_ids)[0] # Get the target for loss depending on the prediction type if noise_scheduler.config.prediction_type == "epsilon": target = noise elif noise_scheduler.config.prediction_type == "v_prediction": target = noise_scheduler.get_velocity(latents, noise, timesteps) else: raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}") timesteps = repeat(timesteps, "b -> b f", f=video_length) timesteps = rearrange(timesteps, "b f -> (b f)") frame_stride = None if unet_additional_kwargs["use_frame_stride_condition"]: frame_stride = batch['stride'].to(latents.device) frame_stride = frame_stride.long() frame_stride = repeat(frame_stride, "b -> b f", f=video_length) frame_stride = rearrange(frame_stride, "b f -> (b f)") t2 = time.time() # Predict the noise residual and compute loss if unet_additional_kwargs["first_frame_condition_mode"] != "none": model_pred = unet(noisy_latents, timesteps, encoder_hidden_states, first_frame_latents=first_frame_latents, frame_stride=frame_stride).sample loss = F.mse_loss(model_pred.float(), target.float()[:, :, 1:, :, :], reduction="mean") else: model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean") t3 = time.time() avg_loss = accelerator.gather(loss.repeat(train_batch_size)).mean() train_loss += avg_loss.item() / gradient_accumulation_steps # Backpropagate accelerator.backward(loss) if accelerator.sync_gradients: grad_norm = accelerator.clip_grad_norm_(unet.parameters(), max_grad_norm) avg_grad_norm = accelerator.gather(grad_norm.repeat(train_batch_size)).mean() train_grad_norm += avg_grad_norm.item() / gradient_accumulation_steps optimizer.step() lr_scheduler.step() optimizer.zero_grad() t4 = time.time() data_loading_time += (t1 - t0) / gradient_accumulation_steps prepare_everything_time += (t2 - t1) / gradient_accumulation_steps network_forward_time += (t3 - t2) / gradient_accumulation_steps network_backward_time += (t4 - t3) / gradient_accumulation_steps t0 = time.time() ### <<<< Training <<<< ### # Checks if the accelerator has performed an optimization step behind the scenes if accelerator.sync_gradients: if use_ema: ema_unet.step(unet.parameters()) progress_bar.update(1) global_step += 1 # Wandb logging if accelerator.is_main_process and (not is_debug) and use_wandb: wandb.log({"metrics/train_loss": train_loss}, step=global_step) wandb.log({"metrics/train_grad_norm": train_grad_norm}, step=global_step) wandb.log({"profiling/train_data_loading_time": data_loading_time}, step=global_step) wandb.log({"profiling/train_prepare_everything_time": prepare_everything_time}, step=global_step) wandb.log({"profiling/train_network_forward_time": network_forward_time}, step=global_step) wandb.log({"profiling/train_network_backward_time": network_backward_time}, step=global_step) # accelerator.log({"train_loss": train_loss}, step=global_step) train_loss = 0.0 train_grad_norm = 0.0 data_loading_time = 0.0 prepare_everything_time = 0.0 network_forward_time = 0.0 network_backward_time = 0.0 # Save checkpoint if global_step % checkpointing_steps == 0: if accelerator.is_main_process: save_path = os.path.join(output_dir, f"checkpoints/checkpoint-{global_step}") accelerator.save_state(save_path) logger.info(f"Saved state to {save_path} (global_step: {global_step})") # Periodically validation if accelerator.is_main_process and global_step % validation_steps == 0: if use_ema: # Store the UNet parameters temporarily and load the EMA parameters to perform inference. ema_unet.store(unet.parameters()) ema_unet.copy_to(unet.parameters()) samples = [] wandb_samples = [] generator = torch.Generator(device=latents.device) generator.manual_seed(seed) height = train_data.sample_size[0] if not isinstance(train_data.sample_size, int) else train_data.sample_size width = train_data.sample_size[1] if not isinstance(train_data.sample_size, int) else train_data.sample_size prompts = validation_data.prompts first_frame_paths = [None] * len(prompts) if unet_additional_kwargs["first_frame_condition_mode"] != "none": first_frame_paths = validation_data.path_to_first_frames for idx, (prompt, first_frame_path) in enumerate(zip(prompts, first_frame_paths)): sample = validation_pipeline( prompt, generator = generator, video_length = train_data.sample_n_frames if not is_image else 2, height = height, width = width, first_frame_paths = first_frame_path, noise_sampling_method = unet_additional_kwargs['noise_sampling_method'], noise_alpha = float(unet_additional_kwargs['noise_alpha']), **validation_data, ).videos save_videos_grid(sample, f"{output_dir}/samples/sample-{global_step}/{idx}.gif") samples.append(sample) numpy_sample = (sample.squeeze(0).permute(1, 0, 2, 3) * 255).cpu().numpy().astype(np.uint8) wandb_video = wandb.Video(numpy_sample, fps=8, caption=prompt) wandb_samples.append(wandb_video) if (not is_debug) and use_wandb: val_title = 'val_videos' wandb.log({val_title: wandb_samples}, step=global_step) samples = torch.concat(samples) save_path = f"{output_dir}/samples/sample-{global_step}.gif" save_videos_grid(samples, save_path) logger.info(f"Saved samples to {save_path}") if use_ema: # Switch back to the original UNet parameters. ema_unet.restore(unet.parameters()) logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} progress_bar.set_postfix(**logs) if accelerator.is_main_process and (not is_debug) and use_wandb: wandb.log({"metrics/train_lr": lr_scheduler.get_last_lr()[0]}, step=global_step) if global_step >= max_train_steps: break # Create the pipeline using the trained modules and save it. accelerator.wait_for_everyone() if accelerator.is_main_process: unet = accelerator.unwrap_model(unet) pipeline = ConditionalAnimationPipeline( text_encoder=text_encoder, vae=vae, unet=unet, tokenizer=tokenizer, scheduler=noise_scheduler, ) pipeline.save_pretrained(f"{output_dir}/final_checkpoint") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--config", type=str, required=True) parser.add_argument("--name", "-n", type=str, default="") parser.add_argument("--wandb", action="store_true") parser.add_argument("optional_args", nargs='*', default=[]) args = parser.parse_args() name = args.name + "_" + Path(args.config).stem config = OmegaConf.load(args.config) if args.optional_args: modified_config = OmegaConf.from_dotlist(args.optional_args) config = OmegaConf.merge(config, modified_config) main(name=name, use_wandb=args.wandb, **config)