import os import math import wandb import random import logging import inspect import argparse import datetime from pathlib import Path from tqdm.auto import tqdm from einops import rearrange from omegaconf import OmegaConf from typing import Dict, Tuple import torch import torchvision import torch.nn.functional as F import diffusers from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel from diffusers.pipelines import StableDiffusionPipeline from diffusers.optimization import get_scheduler from diffusers.utils import check_min_version from diffusers.utils.import_utils import is_xformers_available import transformers from transformers import CLIPTextModel, CLIPTokenizer from animatediff.models.unet import UNet3DConditionModel from animatediff.pipelines.pipeline_animation import AnimationPipeline from animatediff.utils.util import save_videos_grid, load_diffusers_lora, load_weights from animatediff.utils.lora_handler import LoraHandler from animatediff.utils.lora import extract_lora_child_module from animatediff.utils.dataset import VideoJsonDataset, SingleVideoDataset, \ ImageDataset, VideoFolderDataset, CachedDataset, VID_TYPES from animatediff.utils.configs import get_simple_config from lion_pytorch import Lion augment_text_list = [ "a video of", "a high quality video of", "a good video of", "a nice video of", "a great video of", "a video showing", "video of", "video clip of", "great video of", "cool video of", "best video of", "streamed video of", "excellent video of", "new video of", "new video clip of", "high quality video of", "a video showing of", "a clear video showing", "video clip showing", "a clear video showing", "a nice video showing", "a good video showing", "video, high quality," "high quality, video, video clip,", "nice video, clear quality,", "clear quality video of" ] def create_save_paths(output_dir: str): lora_path = f"{output_dir}/lora" directories = [ output_dir, f"{output_dir}/samples", f"{output_dir}/sanity_check", lora_path ] for directory in directories: os.makedirs(directory, exist_ok=True) return lora_path def get_train_dataset(dataset_types, train_data, tokenizer): def process_folder_of_videos(train_datasets: list, video_folder: str): for video_file in os.listdir(video_folder): is_video = any([video_file.split(".")[-1] in ext for ext in VID_TYPES]) if is_video: train_data["single_video_path"] = f"{video_folder}/{video_file}" train_datasets.append(SingleVideoDataset(**train_data, tokenizer=tokenizer)) train_datasets = [] # Loop through all available datasets, get the name, then add to list of data to process. for DataSet in [VideoJsonDataset, SingleVideoDataset, ImageDataset, VideoFolderDataset]: for dataset in dataset_types: if dataset == DataSet.__getname__(): video_folder = train_data.get("path", "") if os.path.exists(video_folder) and dataset == "folder": process_folder_of_videos( train_datasets, video_folder ) continue train_datasets.append(DataSet(**train_data, tokenizer=tokenizer)) if len(train_datasets) > 0: return train_datasets else: raise ValueError("Dataset type not found: 'json', 'single_video', 'folder', 'image'") def tensor_to_vae_latent(t, vae): video_length = t.shape[1] t = rearrange(t, "b f c h w -> (b f) c h w") latents = vae.encode(t).latent_dist.sample() latents = rearrange(latents, "(b f) c h w -> b c f h w", f=video_length) latents = latents * 0.18215 return latents def get_cached_latent_dir(c_dir): from omegaconf import ListConfig if isinstance(c_dir, str): return os.path.abspath(c_dir) if c_dir is not None else None if isinstance(c_dir, ListConfig): c_dir = OmegaConf.to_object(c_dir) return c_dir return None def handle_cache_latents( should_cache, output_dir, train_dataloader, train_batch_size, vae, cached_latent_dir=None, shuffle=False, minimum_required_frames=16, sampler=None, device='cuda' ): # Cache latents by storing them in VRAM. # Speeds up training and saves memory by not encoding during the train loop. if not should_cache: return None vae_dtype = vae.dtype vae.to(device, dtype=torch.float32) if hasattr(vae, 'enable_slicing'): vae.enable_slicing() cached_latent_dir = get_cached_latent_dir(cached_latent_dir) if cached_latent_dir is None: cache_save_dir = f"{output_dir}/cached_latents" os.makedirs(cache_save_dir, exist_ok=True) for i, batch in enumerate(tqdm(train_dataloader, desc="Caching Latents.")): frames = batch['pixel_values'].shape[1] not_min_frames = frames > 2 and frames < minimum_required_frames not_img_train = (frames == 1 and batch['dataset'] != 'image') if any([not_min_frames, not_img_train]) and minimum_required_frames != 0: print(f""" Batch item at index {i} does not meet required minimum frames. Seeing this error means that some of your video lengths are too short, but training will continue. Minimum Frames: {minimum_required_frames} Batch item frames: Batch index = {i}, Batch Frames = {frames} """ ) continue save_name = f"cached_{i}" full_out_path = f"{cache_save_dir}/{save_name}.pt" pixel_values = batch['pixel_values'].to(device, dtype=torch.float32) batch['pixel_values'] = tensor_to_vae_latent(pixel_values, vae) for k, v in batch.items(): batch[k] = v[0] torch.save(batch, full_out_path) del pixel_values del batch # We do this to avoid fragmentation from casting latents between devices. torch.cuda.empty_cache() else: cache_save_dir = cached_latent_dir # Convert string to list of strings for processing if we have more than. cache_save_dir = ( [cache_save_dir] if not isinstance(cache_save_dir, list) else cache_save_dir ) cached_dataset_list = [] for save_dir in cache_save_dir: cached_dataset = CachedDataset(cache_dir=save_dir) cached_dataset_list.append(cached_dataset) if len(cached_dataset_list) > 1: print(f"Found {len(cached_dataset_list)} cached datasets. Merging...") new_cached_dataset = torch.utils.data.ConcatDataset(cached_dataset_list) else: new_cached_dataset = cached_dataset_list[0] vae.to(dtype=vae_dtype) return torch.utils.data.DataLoader( new_cached_dataset, batch_size=train_batch_size, shuffle=shuffle, num_workers=2, persistent_workers=True, pin_memory=False, sampler=sampler ) def do_sanity_check( batch: Dict, cache_latents: bool, validation_pipeline: AnimationPipeline, device: str, image_finetune: bool=False, output_dir: str = "", dataset_id: int = 0 ): pixel_values, texts = batch['pixel_values'].cpu(), batch["text_prompt"] if cache_latents: pixel_values = validation_pipeline.decode_latents(batch["pixel_values"].to(device)) to_torch = torch.from_numpy(pixel_values) pixel_values = rearrange(to_torch, 'b c f h w -> b f c h w') if not image_finetune: 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, ...] text = f"{str(dataset_id)}_{text}" save_name = f"{'-'.join(text.replace('/', '').split()[:10]) if not text == '' else f'-{idx}'}.mp4" save_videos_grid(pixel_value, f"{output_dir}/sanity_check/{save_name}", rescale=not cache_latents) else: for idx, (pixel_value, text) in enumerate(zip(pixel_values, texts)): pixel_value = pixel_value / 2. + 0.5 text = f"{str(dataset_id)}_{text}" save_name = f"{'-'.join(text.replace('/', '').split()[:10]) if not text == '' else f'-{idx}'}.png" torchvision.utils.save_image(pixel_value, f"{output_dir}/sanity_check/{save_name}") def sample_noise(latents, noise_strength, use_offset_noise=False): b, c, f, *_ = latents.shape noise_latents = torch.randn_like(latents, device=latents.device) if use_offset_noise: offset_noise = torch.randn(b, c, f, 1, 1, device=latents.device) noise_latents = noise_latents + noise_strength * offset_noise return noise_latents def param_optim(model, condition, extra_params=None, is_lora=False, negation=None): extra_params = extra_params if len(extra_params.keys()) > 0 else None return { "model": model, "condition": condition, 'extra_params': extra_params, 'is_lora': is_lora, "negation": negation } def create_optim_params(name='param', params=None, lr=5e-6, extra_params=None): params = { "name": name, "params": params, "lr": lr } if extra_params is not None: for k, v in extra_params.items(): params[k] = v return params def create_optimizer_params(model_list, lr): import itertools optimizer_params = [] for optim in model_list: model, condition, extra_params, is_lora, negation = optim.values() # Check if we are doing LoRA training. if is_lora and condition and isinstance(model, list): params = create_optim_params( params=itertools.chain(*model), extra_params=extra_params ) optimizer_params.append(params) continue if is_lora and condition and not isinstance(model, list): for n, p in model.named_parameters(): if 'lora' in n: params = create_optim_params(n, p, lr, extra_params) optimizer_params.append(params) continue # If this is true, we can train it. if condition: for n, p in model.named_parameters(): should_negate = 'lora' in n and not is_lora if should_negate: continue params = create_optim_params(n, p, lr, extra_params) optimizer_params.append(params) return optimizer_params def scale_loras(lora_list: list, scale: float, step=None, spatial_lora_num=None): # Assumed enumerator if step is not None and spatial_lora_num is not None: process_list = range(0, len(lora_list), spatial_lora_num) else: process_list = lora_list for lora_i in process_list: if step is not None: lora_list[lora_i].scale = scale else: lora_i.scale = scale def get_spatial_latents( batch: Dict, random_hflip_img: int, cache_latents: bool, noisy_latents:torch.Tensor, target: torch.Tensor, timesteps: torch.Tensor, noise_scheduler: DDPMScheduler ): ran_idx = torch.randint(0, batch["pixel_values"].shape[2], (1,)).item() use_hflip = random.uniform(0, 1) < random_hflip_img noisy_latents_input = None target_spatial = None if use_hflip: pixel_values_spatial = torchvision.transforms.functional.hflip( batch["pixel_values"][:, ran_idx, :, :, :] if not cache_latents else\ batch["pixel_values"][:, :, ran_idx, :, :] ).unsqueeze(1) latents_spatial = ( tensor_to_vae_latent(pixel_values_spatial, vae) if not cache_latents else pixel_values_spatial ) noise_spatial = sample_noise(latents_spatial, 0, use_offset_noise=use_offset_noise) noisy_latents_input = noise_scheduler.add_noise(latents_spatial, noise_spatial, timesteps) target_spatial = noise_spatial else: noisy_latents_input = noisy_latents[:, :, ran_idx, :, :] target_spatial = target[:, :, ran_idx, :, :] return noisy_latents_input, target_spatial, use_hflip def create_ad_temporal_loss( model_pred: torch.Tensor, loss_temporal: torch.Tensor, target: torch.Tensor ): beta = 1 alpha = (beta ** 2 + 1) ** 0.5 ran_idx = torch.randint(0, model_pred.shape[2], (1,)).item() model_pred_decent = alpha * model_pred - beta * model_pred[:, :, ran_idx, :, :].unsqueeze(2) target_decent = alpha * target - beta * target[:, :, ran_idx, :, :].unsqueeze(2) loss_ad_temporal = F.mse_loss(model_pred_decent.float(), target_decent.float(), reduction="mean") loss_temporal = loss_temporal + loss_ad_temporal return loss_temporal def main( image_finetune: bool, name: str, use_wandb: bool, output_dir: str, pretrained_model_path: str, train_data: Dict, validation_data: Dict, cfg_random_null_text: bool = True, cfg_random_null_text_ratio: float = 0.1, unet_checkpoint_path: str = "", unet_additional_kwargs: Dict = {}, ema_decay: float = 0.9999, noise_scheduler_kwargs = None, max_train_epoch: int = -1, max_train_steps: int = 100, validation_steps: int = 100, validation_steps_tuple: Tuple = (-1,), learning_rate: float = 3e-5, learning_rate_spatial: float = 1e-4, scale_lr: bool = False, lr_warmup_steps: int = 0, lr_scheduler: str = "constant", 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_training: bool = True, enable_xformers_memory_efficient_attention: bool = True, global_seed: int = 42, is_debug: bool = False, dataset_types: Tuple[str] = ('json'), motion_module_path: str = "", domain_adapter_path: str = "", random_hflip_img: float = -1, use_motion_lora_format: bool = True, single_spatial_lora: bool = False, lora_name: str = "motion_director_lora", lora_rank: int = 8, lora_unet_dropout: float = 0.1, train_temporal_lora: bool = True, target_spatial_modules: str = ["Transformer3DModel"], target_temporal_modules: str = ["TemporalTransformerBlock"], cache_latents: bool = False, cached_latent_dir=None, train_sample_validation: bool = True, device: str = 'cuda', use_text_augmenter: bool = False, use_lion_optim: bool = False, use_offset_noise: bool = False, *args, **kwargs ): check_min_version("0.10.0.dev0") if use_text_augmenter: print("Using random text augmentation") # Initialize distributed training num_processes = 1 seed = global_seed torch.manual_seed(seed) # Logging folder if lora_name != "motion_director_lora": name = lora_name + f"_{name}" date_calendar = datetime.datetime.now().strftime("%Y-%m-%d") date_time = datetime.datetime.now().strftime("-%H-%M-%S") folder_name = "debug" if is_debug else name + date_time output_dir = os.path.join(output_dir, date_calendar, folder_name) if is_debug and os.path.exists(output_dir): os.system(f"rm -rf {output_dir}") *_, config = inspect.getargvalues(inspect.currentframe()) # 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, ) if not is_debug and use_wandb: run = wandb.init(project="animatediff", name=folder_name, config=config) # Handle the output folder creation lora_path = create_save_paths(output_dir) OmegaConf.save(config, os.path.join(output_dir, 'config.yaml')) # Load scheduler, tokenizer and models. noise_scheduler_kwargs.update({"steps_offset": 1}) noise_scheduler = DDIMScheduler(**OmegaConf.to_container(noise_scheduler_kwargs)) del noise_scheduler_kwargs["steps_offset"] noise_scheduler_kwargs['beta_schedule'] = 'scaled_linear' train_noise_scheduler_spatial = DDPMScheduler(**OmegaConf.to_container(noise_scheduler_kwargs)) # AnimateDiff uses a linear schedule for its temporal sampling noise_scheduler_kwargs['beta_schedule'] = 'linear' train_noise_scheduler = DDPMScheduler(**OmegaConf.to_container(noise_scheduler_kwargs)) if kwargs.get("force_spatial_linear_scaling", True): train_noise_scheduler_spatial = train_noise_scheduler 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") if not image_finetune: unet = UNet3DConditionModel.from_pretrained_2d( pretrained_model_path, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(unet_additional_kwargs) ) else: unet = UNet2DConditionModel.from_pretrained(pretrained_model_path, subfolder="unet") # Freeze all models for LoRA training unet.requires_grad_(False) vae.requires_grad_(False) text_encoder.requires_grad_(False) if not use_lion_optim: optimizer = torch.optim.AdamW else: optimizer = Lion learning_rate, learning_rate_spatial = map(lambda lr: lr / 10, (learning_rate, learning_rate_spatial)) adam_weight_decay *= 10 # Enable xformers if enable_xformers_memory_efficient_attention: if is_xformers_available(): unet.enable_xformers_memory_efficient_attention() if kwargs.get("force_temporal_xformers"): for module in unet.modules(): if module.__class__.__name__ == "VersatileAttention": setattr(module, '_use_memory_efficient_attention_xformers', True) else: raise ValueError("xformers is not available. Make sure it is installed correctly") # Enable gradient checkpointing if gradient_checkpointing: unet.enable_gradient_checkpointing() # Move models to GPU vae.to(device) text_encoder.to(device) # Get the training dataset train_dataset = get_train_dataset(dataset_types, train_data, tokenizer) if len(train_dataset) > 0: train_dataset = torch.utils.data.ConcatDataset(train_dataset) else: train_dataset = train_dataset[0] train_dataloader = torch.utils.data.DataLoader( train_dataset, batch_size=train_batch_size, shuffle=False, num_workers=1, pin_memory=True, drop_last=True, ) if cache_latents: torch.multiprocessing.set_start_method('spawn') train_dataloader = handle_cache_latents( cache_latents, output_dir, train_dataloader, train_batch_size, vae, cached_latent_dir=cached_latent_dir, sampler=None, device=device ) # 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) if scale_lr: learning_rate = (learning_rate * gradient_accumulation_steps * train_batch_size * num_processes) # Validation pipeline if not image_finetune: validation_pipeline = AnimationPipeline( unet=unet, vae=vae, tokenizer=tokenizer, text_encoder=text_encoder, scheduler=noise_scheduler, ).to(device) else: validation_pipeline = StableDiffusionPipeline.from_pretrained( pretrained_model_path, unet=unet, vae=vae, tokenizer=tokenizer, text_encoder=text_encoder, scheduler=noise_scheduler, safety_checker=None, ) validation_pipeline = load_weights( validation_pipeline, motion_module_path=motion_module_path, adapter_lora_path=domain_adapter_path, dreambooth_model_path=unet_checkpoint_path ) validation_pipeline.enable_vae_slicing() validation_pipeline.to(device) unet.to(device=device) text_encoder.to(device=device) # Temporal LoRA if train_temporal_lora: # one temporal lora lora_manager_temporal = LoraHandler(use_unet_lora=True, unet_replace_modules=target_temporal_modules) unet_lora_params_temporal, unet_negation_temporal = lora_manager_temporal.add_lora_to_model( True, unet, lora_manager_temporal.unet_replace_modules, 0, lora_path + '/temporal/', r=lora_rank) optimizer_temporal = optimizer( create_optimizer_params([param_optim(unet_lora_params_temporal, True, is_lora=True, extra_params={**{"lr": learning_rate}} )], learning_rate), lr=learning_rate, betas=(adam_beta1, adam_beta2), weight_decay=adam_weight_decay ) lr_scheduler_temporal = get_scheduler( lr_scheduler, optimizer=optimizer_temporal, num_warmup_steps=lr_warmup_steps * gradient_accumulation_steps, num_training_steps=max_train_steps * gradient_accumulation_steps, ) else: lora_manager_temporal = None unet_lora_params_temporal, unet_negation_temporal = [], [] optimizer_temporal = None lr_scheduler_temporal = None # Spatial LoRAs if single_spatial_lora: spatial_lora_num = 1 else: # one spatial lora for each video spatial_lora_num = train_dataset.__len__() lora_managers_spatial = [] unet_lora_params_spatial_list = [] optimizer_spatial_list = [] lr_scheduler_spatial_list = [] for i in range(spatial_lora_num): lora_manager_spatial = LoraHandler(use_unet_lora=True, unet_replace_modules=target_spatial_modules) lora_managers_spatial.append(lora_manager_spatial) unet_lora_params_spatial, unet_negation_spatial = lora_manager_spatial.add_lora_to_model( True, unet, lora_manager_spatial.unet_replace_modules, lora_unet_dropout, lora_path + '/spatial/', r=lora_rank) unet_lora_params_spatial_list.append(unet_lora_params_spatial) optimizer_spatial = optimizer( create_optimizer_params([param_optim(unet_lora_params_spatial, True, is_lora=True, extra_params={**{"lr": learning_rate_spatial}} )], learning_rate_spatial), lr=learning_rate_spatial, betas=(adam_beta1, adam_beta2), weight_decay=adam_weight_decay ) optimizer_spatial_list.append(optimizer_spatial) # Scheduler lr_scheduler_spatial = get_scheduler( lr_scheduler, optimizer=optimizer_spatial, num_warmup_steps=lr_warmup_steps * gradient_accumulation_steps, num_training_steps=max_train_steps * gradient_accumulation_steps, ) lr_scheduler_spatial_list.append(lr_scheduler_spatial) unet_negation_all = unet_negation_spatial + unet_negation_temporal # 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 * num_processes * gradient_accumulation_steps logging.info("***** Running training *****") logging.info(f" Num examples = {len(train_dataset)}") logging.info(f" Num Epochs = {num_train_epochs}") logging.info(f" Instantaneous batch size per device = {train_batch_size}") logging.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") logging.info(f" Gradient Accumulation steps = {gradient_accumulation_steps}") logging.info(f" Total optimization steps = {max_train_steps}") global_step = 0 first_epoch = 0 # Only show the progress bar once on each machine. progress_bar = tqdm(range(global_step, max_train_steps)) progress_bar.set_description("Steps") # Support mixed-precision training scaler = torch.cuda.amp.GradScaler() if mixed_precision_training else None ### <<<< Training <<<< ### for epoch in range(first_epoch, num_train_epochs): unet.train() for step, batch in enumerate(train_dataloader): spatial_scheduler_lr = 0.0 temporal_scheduler_lr = 0.0 # Handle Lora Optimizers & Conditions for optimizer_spatial in optimizer_spatial_list: optimizer_spatial.zero_grad(set_to_none=True) if optimizer_temporal is not None: optimizer_temporal.zero_grad(set_to_none=True) if train_temporal_lora: mask_temporal_lora = False else: mask_temporal_lora = True mask_spatial_lora = random.uniform(0, 1) < 0.2 and not mask_temporal_lora if cfg_random_null_text: batch["text_prompt"] = [name if random.random() > cfg_random_null_text_ratio else "" for name in batch["text_prompt"]] if use_text_augmenter: random.seed() txt_idx = random.randint(0, len(augment_text_list) - 1) augment_text = augment_text_list[txt_idx] batch['text_prompt'] = [ f"{augment_text} {prompt}" for prompt in batch['text_prompt'] ] # Data batch sanity check if epoch == first_epoch and step == 0: for _idx, _batch in enumerate(tqdm(train_dataloader, desc="Dataset sanity check...")): do_sanity_check( _batch, cache_latents, validation_pipeline, device, output_dir=output_dir, dataset_id=_idx ) if _idx > 10: break # Convert videos to latent space pixel_values = batch["pixel_values"].to(device) video_length = pixel_values.shape[2] bsz = pixel_values.shape[0] # Sample a random timestep for each video timesteps = torch.randint(0, train_noise_scheduler.config.num_train_timesteps, (bsz,), device=pixel_values.device) timesteps = timesteps.long() # Add noise to the latents according to the noise magnitude at each timestep # (this is the forward diffusion process) latents = tensor_to_vae_latent(pixel_values, vae) if not cache_latents else pixel_values noise = sample_noise(latents, 0, use_offset_noise=use_offset_noise) target = noise # Get the text embedding for conditioning with torch.no_grad(): prompt_ids = tokenizer( batch['text_prompt'], max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt" ).input_ids.to(pixel_values.device) encoder_hidden_states = text_encoder(prompt_ids)[0] with torch.cuda.amp.autocast(enabled=mixed_precision_training): if mask_spatial_lora: loras = extract_lora_child_module(unet, target_replace_module=target_spatial_modules) scale_loras(loras, 0.) loss_spatial = None else: loras = extract_lora_child_module(unet, target_replace_module=target_spatial_modules) if spatial_lora_num == 1: scale_loras(loras, 1.0) else: scale_loras(loras, 0.) scale_loras(loras, 1.0, step=step, spatial_lora_num=spatial_lora_num) loras = extract_lora_child_module(unet, target_replace_module=target_temporal_modules) if len(loras) > 0: scale_loras(loras, 0.) ### >>>> Spatial LoRA Prediction >>>> ### noisy_latents = train_noise_scheduler_spatial.add_noise(latents, noise, timesteps) noisy_latents_input, target_spatial, use_hflip = get_spatial_latents( batch, random_hflip_img, cache_latents, noisy_latents, target, timesteps, train_noise_scheduler_spatial ) if use_hflip: model_pred_spatial = unet(noisy_latents_input, timesteps, encoder_hidden_states=encoder_hidden_states).sample loss_spatial = F.mse_loss(model_pred_spatial[:, :, 0, :, :].float(), target_spatial[:, :, 0, :, :].float(), reduction="mean") else: model_pred_spatial = unet(noisy_latents_input.unsqueeze(2), timesteps, encoder_hidden_states=encoder_hidden_states).sample loss_spatial = F.mse_loss(model_pred_spatial[:, :, 0, :, :].float(), target_spatial.float(), reduction="mean") if mask_temporal_lora: loras = extract_lora_child_module(unet, target_replace_module=target_temporal_modules) scale_loras(loras, 0.) loss_temporal = None else: loras = extract_lora_child_module(unet, target_replace_module=target_temporal_modules) scale_loras(loras, 1.0) ### >>>> Temporal LoRA Prediction >>>> ### noisy_latents = train_noise_scheduler.add_noise(latents, noise, timesteps) model_pred = unet(noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states).sample loss_temporal = F.mse_loss(model_pred.float(), target.float(), reduction="mean") loss_temporal = create_ad_temporal_loss(model_pred, loss_temporal, target) # Backpropagate if not mask_spatial_lora: scaler.scale(loss_spatial).backward(retain_graph=True) if spatial_lora_num == 1: scaler.step(optimizer_spatial_list[0]) else: # https://github.com/nerfstudio-project/nerfstudio/pull/1919 if any( any(p.grad is not None for p in g["params"]) for g in optimizer_spatial_list[step].param_groups ): scaler.step(optimizer_spatial_list[step]) if not mask_temporal_lora and train_temporal_lora: scaler.scale(loss_temporal).backward() scaler.step(optimizer_temporal) if spatial_lora_num == 1: lr_scheduler_spatial_list[0].step() spatial_scheduler_lr = lr_scheduler_spatial_list[0].get_lr()[0] else: lr_scheduler_spatial_list[step].step() spatial_scheduler_lr = lr_scheduler_spatial_list[step].get_lr()[0] if lr_scheduler_temporal is not None: lr_scheduler_temporal.step() temporal_scheduler_lr = lr_scheduler_temporal.get_lr()[0] scaler.update() progress_bar.update(1) global_step += 1 # Wandb logging if not is_debug and use_wandb: loss = ( loss_temporal if loss_spatial is None else \ loss_temporal + loss_spatial ) wandb.log({"train_loss": loss.item()}, step=global_step) # Save checkpoint if global_step % checkpointing_steps == 0: import copy # We do this to prevent VRAM spiking / increase from the new copy validation_pipeline.to('cpu') lora_manager_spatial.save_lora_weights( model=copy.deepcopy(validation_pipeline), save_path=lora_path+'/spatial', step=global_step, use_safetensors=True, lora_rank=lora_rank, lora_name=lora_name + "_spatial" ) if lora_manager_temporal is not None: lora_manager_temporal.save_lora_weights( model=copy.deepcopy(validation_pipeline), save_path=lora_path+'/temporal', step=global_step, use_safetensors=True, lora_rank=lora_rank, lora_name=lora_name + "_temporal", use_motion_lora_format=use_motion_lora_format ) validation_pipeline.to(device) # Periodically validation if (global_step % validation_steps == 0 or global_step in validation_steps_tuple): samples = [] validation_seed = getattr(validation_data, 'seed', -1) generator = torch.Generator(device=latents.device) generator.manual_seed(global_seed if validation_seed == -1 else validation_seed) if not train_sample_validation: if not isinstance(train_data.sample_size, int): height, width = train_data.sample_size[:2] else: height, width = [train_data.sample_size] * 2 else: if all(['resized_h'in batch, 'resized_w' in batch]): height, width = batch["resized_h"], batch['resized_w'] else: height, width = [512] * 2 prompts = ( validation_data.prompts[:2] if global_step < 1000 and (not image_finetune) \ else validation_data.prompts ) with torch.cuda.amp.autocast(enabled=True): if gradient_checkpointing: unet.disable_gradient_checkpointing() loras = extract_lora_child_module( unet, target_replace_module=target_spatial_modules ) scale_loras(loras, validation_data.spatial_scale) with torch.no_grad(): unet.eval() for idx, prompt in enumerate(prompts): if len(prompt) == 0: prompt = batch['text_prompt'] print(prompt) if not image_finetune: sample = validation_pipeline( prompt, generator = generator, video_length = train_data.sample_n_frames, height = height, width = width, **validation_data, ).videos save_videos_grid(sample, f"{output_dir}/samples/sample-{global_step}/{idx}.gif") samples.append(sample) else: sample = validation_pipeline( prompt, generator = generator, height = height, width = width, num_inference_steps = validation_data.get("num_inference_steps", 25), guidance_scale = validation_data.get("guidance_scale", 8.), ).images[0] sample = torchvision.transforms.functional.to_tensor(sample) samples.append(sample) unet.train() if not image_finetune: samples = torch.concat(samples) save_path = f"{output_dir}/samples/sample-{global_step}.gif" save_videos_grid(samples, save_path) else: samples = torch.stack(samples) save_path = f"{output_dir}/samples/sample-{global_step}.png" torchvision.utils.save_image(samples, save_path, nrow=4) logging.info(f"Saved samples to {save_path}") logs = { "Temporal Loss": loss_temporal.detach().item(), "Temporal LR": temporal_scheduler_lr, "Spatial Loss": loss_spatial.detach().item() if loss_spatial is not None else 0, "Spatial LR": spatial_scheduler_lr } progress_bar.set_postfix(**logs) if gradient_checkpointing: unet.enable_gradient_checkpointing() if global_step >= max_train_steps: break if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--config", type=str, required=True) parser.add_argument("--wandb", action="store_true") args = parser.parse_args() name = Path(args.config).stem config = OmegaConf.load(args.config) if getattr(config, "simple_mode", False): config = get_simple_config(config) main(name=name, use_wandb=args.wandb, **config)