Delete train_stage_2.py

train_stage_2.py

@@ -1,773 +0,0 @@
-import argparse
-import copy
-import logging
-import math
-import os
-import os.path as osp
-import random
-import time
-import warnings
-from collections import OrderedDict
-from datetime import datetime
-from pathlib import Path
-from tempfile import TemporaryDirectory
-
-import diffusers
-import mlflow
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint
-import transformers
-from accelerate import Accelerator
-from accelerate.logging import get_logger
-from accelerate.utils import DistributedDataParallelKwargs
-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 einops import rearrange
-from omegaconf import OmegaConf
-from PIL import Image
-from torchvision import transforms
-from tqdm.auto import tqdm
-from transformers import CLIPVisionModelWithProjection
-
-from src.dataset.dance_video import HumanDanceVideoDataset
-from src.models.mutual_self_attention import ReferenceAttentionControl
-from src.models.pose_guider import PoseGuider
-from src.models.unet_2d_condition import UNet2DConditionModel
-from src.models.unet_3d import UNet3DConditionModel
-from src.pipelines.pipeline_pose2vid import Pose2VideoPipeline
-from src.utils.util import (
-    delete_additional_ckpt,
-    import_filename,
-    read_frames,
-    save_videos_grid,
-    seed_everything,
-)
-
-warnings.filterwarnings("ignore")
-
-# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.10.0.dev0")
-
-logger = get_logger(__name__, log_level="INFO")
-
-
-class Net(nn.Module):
-    def __init__(
-        self,
-        reference_unet: UNet2DConditionModel,
-        denoising_unet: UNet3DConditionModel,
-        pose_guider: PoseGuider,
-        reference_control_writer,
-        reference_control_reader,
-    ):
-        super().__init__()
-        self.reference_unet = reference_unet
-        self.denoising_unet = denoising_unet
-        self.pose_guider = pose_guider
-        self.reference_control_writer = reference_control_writer
-        self.reference_control_reader = reference_control_reader
-
-    def forward(
-        self,
-        noisy_latents,
-        timesteps,
-        ref_image_latents,
-        clip_image_embeds,
-        pose_img,
-        uncond_fwd: bool = False,
-    ):
-        pose_cond_tensor = pose_img.to(device="cuda")
-        pose_fea = self.pose_guider(pose_cond_tensor)
-
-        if not uncond_fwd:
-            ref_timesteps = torch.zeros_like(timesteps)
-            self.reference_unet(
-                ref_image_latents,
-                ref_timesteps,
-                encoder_hidden_states=clip_image_embeds,
-                return_dict=False,
-            )
-            self.reference_control_reader.update(self.reference_control_writer)
-
-        model_pred = self.denoising_unet(
-            noisy_latents,
-            timesteps,
-            pose_cond_fea=pose_fea,
-            encoder_hidden_states=clip_image_embeds,
-        ).sample
-
-        return model_pred
-
-
-def compute_snr(noise_scheduler, timesteps):
-    """
-    Computes SNR as per
-    https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849
-    """
-    alphas_cumprod = noise_scheduler.alphas_cumprod
-    sqrt_alphas_cumprod = alphas_cumprod**0.5
-    sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5
-
-    # Expand the tensors.
-    # Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026
-    sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[
-        timesteps
-    ].float()
-    while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):
-        sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]
-    alpha = sqrt_alphas_cumprod.expand(timesteps.shape)
-
-    sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(
-        device=timesteps.device
-    )[timesteps].float()
-    while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape):
-        sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None]
-    sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape)
-
-    # Compute SNR.
-    snr = (alpha / sigma) ** 2
-    return snr
-
-
-def log_validation(
-    vae,
-    image_enc,
-    net,
-    scheduler,
-    accelerator,
-    width,
-    height,
-    clip_length=24,
-    generator=None,
-):
-    logger.info("Running validation... ")
-
-    ori_net = accelerator.unwrap_model(net)
-    reference_unet = ori_net.reference_unet
-    denoising_unet = ori_net.denoising_unet
-    pose_guider = ori_net.pose_guider
-
-    if generator is None:
-        generator = torch.manual_seed(42)
-    tmp_denoising_unet = copy.deepcopy(denoising_unet)
-    tmp_denoising_unet = tmp_denoising_unet.to(dtype=torch.float16)
-
-    pipe = Pose2VideoPipeline(
-        vae=vae,
-        image_encoder=image_enc,
-        reference_unet=reference_unet,
-        denoising_unet=tmp_denoising_unet,
-        pose_guider=pose_guider,
-        scheduler=scheduler,
-    )
-    pipe = pipe.to(accelerator.device)
-
-    test_cases = [
-        (
-            "./configs/inference/ref_images/anyone-3.png",
-            "./configs/inference/pose_videos/anyone-video-1_kps.mp4",
-        ),
-        (
-            "./configs/inference/ref_images/anyone-2.png",
-            "./configs/inference/pose_videos/anyone-video-2_kps.mp4",
-        ),
-    ]
-
-    results = []
-    for test_case in test_cases:
-        ref_image_path, pose_video_path = test_case
-        ref_name = Path(ref_image_path).stem
-        pose_name = Path(pose_video_path).stem
-        ref_image_pil = Image.open(ref_image_path).convert("RGB")
-
-        pose_list = []
-        pose_tensor_list = []
-        pose_images = read_frames(pose_video_path)
-        pose_transform = transforms.Compose(
-            [transforms.Resize((height, width)), transforms.ToTensor()]
-        )
-        for pose_image_pil in pose_images[:clip_length]:
-            pose_tensor_list.append(pose_transform(pose_image_pil))
-            pose_list.append(pose_image_pil)
-
-        pose_tensor = torch.stack(pose_tensor_list, dim=0)  # (f, c, h, w)
-        pose_tensor = pose_tensor.transpose(0, 1)
-
-        pipeline_output = pipe(
-            ref_image_pil,
-            pose_list,
-            width,
-            height,
-            clip_length,
-            20,
-            3.5,
-            generator=generator,
-        )
-        video = pipeline_output.videos
-
-        # Concat it with pose tensor
-        pose_tensor = pose_tensor.unsqueeze(0)
-        video = torch.cat([video, pose_tensor], dim=0)
-
-        results.append({"name": f"{ref_name}_{pose_name}", "vid": video})
-
-    del tmp_denoising_unet
-    del pipe
-    torch.cuda.empty_cache()
-
-    return results
-
-
-def main(cfg):
-    kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
-    accelerator = Accelerator(
-        gradient_accumulation_steps=cfg.solver.gradient_accumulation_steps,
-        mixed_precision=cfg.solver.mixed_precision,
-        log_with="mlflow",
-        project_dir="./mlruns",
-        kwargs_handlers=[kwargs],
-    )
-
-    # 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 passed along, set the training seed now.
-    if cfg.seed is not None:
-        seed_everything(cfg.seed)
-
-    exp_name = cfg.exp_name
-    save_dir = f"{cfg.output_dir}/{exp_name}"
-    if accelerator.is_main_process:
-        if not os.path.exists(save_dir):
-            os.makedirs(save_dir)
-
-    inference_config_path = "./configs/inference/inference_v2.yaml"
-    infer_config = OmegaConf.load(inference_config_path)
-
-    if cfg.weight_dtype == "fp16":
-        weight_dtype = torch.float16
-    elif cfg.weight_dtype == "fp32":
-        weight_dtype = torch.float32
-    else:
-        raise ValueError(
-            f"Do not support weight dtype: {cfg.weight_dtype} during training"
-        )
-
-    sched_kwargs = OmegaConf.to_container(cfg.noise_scheduler_kwargs)
-    if cfg.enable_zero_snr:
-        sched_kwargs.update(
-            rescale_betas_zero_snr=True,
-            timestep_spacing="trailing",
-            prediction_type="v_prediction",
-        )
-    val_noise_scheduler = DDIMScheduler(**sched_kwargs)
-    sched_kwargs.update({"beta_schedule": "scaled_linear"})
-    train_noise_scheduler = DDIMScheduler(**sched_kwargs)
-
-    image_enc = CLIPVisionModelWithProjection.from_pretrained(
-        cfg.image_encoder_path,
-    ).to(dtype=weight_dtype, device="cuda")
-    vae = AutoencoderKL.from_pretrained(cfg.vae_model_path).to(
-        "cuda", dtype=weight_dtype
-    )
-    reference_unet = UNet2DConditionModel.from_pretrained(
-        cfg.base_model_path,
-        subfolder="unet",
-    ).to(device="cuda", dtype=weight_dtype)
-
-    denoising_unet = UNet3DConditionModel.from_pretrained_2d(
-        cfg.base_model_path,
-        cfg.mm_path,
-        subfolder="unet",
-        unet_additional_kwargs=OmegaConf.to_container(
-            infer_config.unet_additional_kwargs
-        ),
-    ).to(device="cuda")
-
-    pose_guider = PoseGuider(
-        conditioning_embedding_channels=320, block_out_channels=(16, 32, 96, 256)
-    ).to(device="cuda", dtype=weight_dtype)
-
-    stage1_ckpt_dir = cfg.stage1_ckpt_dir
-    stage1_ckpt_step = cfg.stage1_ckpt_step
-    denoising_unet.load_state_dict(
-        torch.load(
-            os.path.join(stage1_ckpt_dir, f"denoising_unet-{stage1_ckpt_step}.pth"),
-            map_location="cpu",
-        ),
-        strict=False,
-    )
-    reference_unet.load_state_dict(
-        torch.load(
-            os.path.join(stage1_ckpt_dir, f"reference_unet-{stage1_ckpt_step}.pth"),
-            map_location="cpu",
-        ),
-        strict=False,
-    )
-    pose_guider.load_state_dict(
-        torch.load(
-            os.path.join(stage1_ckpt_dir, f"pose_guider-{stage1_ckpt_step}.pth"),
-            map_location="cpu",
-        ),
-        strict=False,
-    )
-
-    # Freeze
-    vae.requires_grad_(False)
-    image_enc.requires_grad_(False)
-    reference_unet.requires_grad_(False)
-    denoising_unet.requires_grad_(False)
-    pose_guider.requires_grad_(False)
-
-    # Set motion module learnable
-    for name, module in denoising_unet.named_modules():
-        if "motion_modules" in name:
-            for params in module.parameters():
-                params.requires_grad = True
-
-    reference_control_writer = ReferenceAttentionControl(
-        reference_unet,
-        do_classifier_free_guidance=False,
-        mode="write",
-        fusion_blocks="full",
-    )
-    reference_control_reader = ReferenceAttentionControl(
-        denoising_unet,
-        do_classifier_free_guidance=False,
-        mode="read",
-        fusion_blocks="full",
-    )
-
-    net = Net(
-        reference_unet,
-        denoising_unet,
-        pose_guider,
-        reference_control_writer,
-        reference_control_reader,
-    )
-
-    if cfg.solver.enable_xformers_memory_efficient_attention:
-        if is_xformers_available():
-            reference_unet.enable_xformers_memory_efficient_attention()
-            denoising_unet.enable_xformers_memory_efficient_attention()
-        else:
-            raise ValueError(
-                "xformers is not available. Make sure it is installed correctly"
-            )
-
-    if cfg.solver.gradient_checkpointing:
-        reference_unet.enable_gradient_checkpointing()
-        denoising_unet.enable_gradient_checkpointing()
-
-    if cfg.solver.scale_lr:
-        learning_rate = (
-            cfg.solver.learning_rate
-            * cfg.solver.gradient_accumulation_steps
-            * cfg.data.train_bs
-            * accelerator.num_processes
-        )
-    else:
-        learning_rate = cfg.solver.learning_rate
-
-    # Initialize the optimizer
-    if cfg.solver.use_8bit_adam:
-        try:
-            import bitsandbytes as bnb
-        except ImportError:
-            raise ImportError(
-                "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
-            )
-
-        optimizer_cls = bnb.optim.AdamW8bit
-    else:
-        optimizer_cls = torch.optim.AdamW
-
-    trainable_params = list(filter(lambda p: p.requires_grad, net.parameters()))
-    logger.info(f"Total trainable params {len(trainable_params)}")
-    optimizer = optimizer_cls(
-        trainable_params,
-        lr=learning_rate,
-        betas=(cfg.solver.adam_beta1, cfg.solver.adam_beta2),
-        weight_decay=cfg.solver.adam_weight_decay,
-        eps=cfg.solver.adam_epsilon,
-    )
-
-    # Scheduler
-    lr_scheduler = get_scheduler(
-        cfg.solver.lr_scheduler,
-        optimizer=optimizer,
-        num_warmup_steps=cfg.solver.lr_warmup_steps
-        * cfg.solver.gradient_accumulation_steps,
-        num_training_steps=cfg.solver.max_train_steps
-        * cfg.solver.gradient_accumulation_steps,
-    )
-
-    train_dataset = HumanDanceVideoDataset(
-        width=cfg.data.train_width,
-        height=cfg.data.train_height,
-        n_sample_frames=cfg.data.n_sample_frames,
-        sample_rate=cfg.data.sample_rate,
-        img_scale=(1.0, 1.0),
-        data_meta_paths=cfg.data.meta_paths,
-    )
-    train_dataloader = torch.utils.data.DataLoader(
-        train_dataset, batch_size=cfg.data.train_bs, shuffle=True, num_workers=4
-    )
-
-    # Prepare everything with our `accelerator`.
-    (
-        net,
-        optimizer,
-        train_dataloader,
-        lr_scheduler,
-    ) = accelerator.prepare(
-        net,
-        optimizer,
-        train_dataloader,
-        lr_scheduler,
-    )
-
-    # 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) / cfg.solver.gradient_accumulation_steps
-    )
-    # Afterwards we recalculate our number of training epochs
-    num_train_epochs = math.ceil(
-        cfg.solver.max_train_steps / num_update_steps_per_epoch
-    )
-
-    # We need to initialize the trackers we use, and also store our configuration.
-    # The trackers initializes automatically on the main process.
-    if accelerator.is_main_process:
-        run_time = datetime.now().strftime("%Y%m%d-%H%M")
-        accelerator.init_trackers(
-            exp_name,
-            init_kwargs={"mlflow": {"run_name": run_time}},
-        )
-        # dump config file
-        mlflow.log_dict(OmegaConf.to_container(cfg), "config.yaml")
-
-    # Train!
-    total_batch_size = (
-        cfg.data.train_bs
-        * accelerator.num_processes
-        * cfg.solver.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 = {cfg.data.train_bs}")
-    logger.info(
-        f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}"
-    )
-    logger.info(
-        f"  Gradient Accumulation steps = {cfg.solver.gradient_accumulation_steps}"
-    )
-    logger.info(f"  Total optimization steps = {cfg.solver.max_train_steps}")
-    global_step = 0
-    first_epoch = 0
-
-    # Potentially load in the weights and states from a previous save
-    if cfg.resume_from_checkpoint:
-        if cfg.resume_from_checkpoint != "latest":
-            resume_dir = cfg.resume_from_checkpoint
-        else:
-            resume_dir = save_dir
-        # Get the most recent checkpoint
-        dirs = os.listdir(resume_dir)
-        dirs = [d for d in dirs if d.startswith("checkpoint")]
-        dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
-        path = dirs[-1]
-        accelerator.load_state(os.path.join(resume_dir, path))
-        accelerator.print(f"Resuming from checkpoint {path}")
-        global_step = int(path.split("-")[1])
-
-        first_epoch = global_step // num_update_steps_per_epoch
-        resume_step = global_step % num_update_steps_per_epoch
-
-    # Only show the progress bar once on each machine.
-    progress_bar = tqdm(
-        range(global_step, cfg.solver.max_train_steps),
-        disable=not accelerator.is_local_main_process,
-    )
-    progress_bar.set_description("Steps")
-
-    for epoch in range(first_epoch, num_train_epochs):
-        train_loss = 0.0
-        t_data_start = time.time()
-        for step, batch in enumerate(train_dataloader):
-            t_data = time.time() - t_data_start
-            with accelerator.accumulate(net):
-                # Convert videos to latent space
-                pixel_values_vid = batch["pixel_values_vid"].to(weight_dtype)
-                with torch.no_grad():
-                    video_length = pixel_values_vid.shape[1]
-                    pixel_values_vid = rearrange(
-                        pixel_values_vid, "b f c h w -> (b f) c h w"
-                    )
-                    latents = vae.encode(pixel_values_vid).latent_dist.sample()
-                    latents = rearrange(
-                        latents, "(b f) c h w -> b c f h w", f=video_length
-                    )
-                    latents = latents * 0.18215
-
-                noise = torch.randn_like(latents)
-                if cfg.noise_offset > 0:
-                    noise += cfg.noise_offset * torch.randn(
-                        (latents.shape[0], latents.shape[1], 1, 1, 1),
-                        device=latents.device,
-                    )
-                bsz = latents.shape[0]
-                # Sample a random timestep for each video
-                timesteps = torch.randint(
-                    0,
-                    train_noise_scheduler.num_train_timesteps,
-                    (bsz,),
-                    device=latents.device,
-                )
-                timesteps = timesteps.long()
-
-                pixel_values_pose = batch["pixel_values_pose"]  # (bs, f, c, H, W)
-                pixel_values_pose = pixel_values_pose.transpose(
-                    1, 2
-                )  # (bs, c, f, H, W)
-
-                uncond_fwd = random.random() < cfg.uncond_ratio
-                clip_image_list = []
-                ref_image_list = []
-                for batch_idx, (ref_img, clip_img) in enumerate(
-                    zip(
-                        batch["pixel_values_ref_img"],
-                        batch["clip_ref_img"],
-                    )
-                ):
-                    if uncond_fwd:
-                        clip_image_list.append(torch.zeros_like(clip_img))
-                    else:
-                        clip_image_list.append(clip_img)
-                    ref_image_list.append(ref_img)
-
-                with torch.no_grad():
-                    ref_img = torch.stack(ref_image_list, dim=0).to(
-                        dtype=vae.dtype, device=vae.device
-                    )
-                    ref_image_latents = vae.encode(
-                        ref_img
-                    ).latent_dist.sample()  # (bs, d, 64, 64)
-                    ref_image_latents = ref_image_latents * 0.18215
-
-                    clip_img = torch.stack(clip_image_list, dim=0).to(
-                        dtype=image_enc.dtype, device=image_enc.device
-                    )
-                    clip_img = clip_img.to(device="cuda", dtype=weight_dtype)
-                    clip_image_embeds = image_enc(
-                        clip_img.to("cuda", dtype=weight_dtype)
-                    ).image_embeds
-                    clip_image_embeds = clip_image_embeds.unsqueeze(1)  # (bs, 1, d)
-
-                # add noise
-                noisy_latents = train_noise_scheduler.add_noise(
-                    latents, noise, timesteps
-                )
-
-                # Get the target for loss depending on the prediction type
-                if train_noise_scheduler.prediction_type == "epsilon":
-                    target = noise
-                elif train_noise_scheduler.prediction_type == "v_prediction":
-                    target = train_noise_scheduler.get_velocity(
-                        latents, noise, timesteps
-                    )
-                else:
-                    raise ValueError(
-                        f"Unknown prediction type {train_noise_scheduler.prediction_type}"
-                    )
-
-                # ---- Forward!!! -----
-                model_pred = net(
-                    noisy_latents,
-                    timesteps,
-                    ref_image_latents,
-                    clip_image_embeds,
-                    pixel_values_pose,
-                    uncond_fwd=uncond_fwd,
-                )
-
-                if cfg.snr_gamma == 0:
-                    loss = F.mse_loss(
-                        model_pred.float(), target.float(), reduction="mean"
-                    )
-                else:
-                    snr = compute_snr(train_noise_scheduler, timesteps)
-                    if train_noise_scheduler.config.prediction_type == "v_prediction":
-                        # Velocity objective requires that we add one to SNR values before we divide by them.
-                        snr = snr + 1
-                    mse_loss_weights = (
-                        torch.stack(
-                            [snr, cfg.snr_gamma * torch.ones_like(timesteps)], dim=1
-                        ).min(dim=1)[0]
-                        / snr
-                    )
-                    loss = F.mse_loss(
-                        model_pred.float(), target.float(), reduction="none"
-                    )
-                    loss = (
-                        loss.mean(dim=list(range(1, len(loss.shape))))
-                        * mse_loss_weights
-                    )
-                    loss = loss.mean()
-
-                # Gather the losses across all processes for logging (if we use distributed training).
-                avg_loss = accelerator.gather(loss.repeat(cfg.data.train_bs)).mean()
-                train_loss += avg_loss.item() / cfg.solver.gradient_accumulation_steps
-
-                # Backpropagate
-                accelerator.backward(loss)
-                if accelerator.sync_gradients:
-                    accelerator.clip_grad_norm_(
-                        trainable_params,
-                        cfg.solver.max_grad_norm,
-                    )
-                optimizer.step()
-                lr_scheduler.step()
-                optimizer.zero_grad()
-
-            if accelerator.sync_gradients:
-                reference_control_reader.clear()
-                reference_control_writer.clear()
-                progress_bar.update(1)
-                global_step += 1
-                accelerator.log({"train_loss": train_loss}, step=global_step)
-                train_loss = 0.0
-
-                if global_step % cfg.val.validation_steps == 0:
-                    if accelerator.is_main_process:
-                        generator = torch.Generator(device=accelerator.device)
-                        generator.manual_seed(cfg.seed)
-
-                        sample_dicts = log_validation(
-                            vae=vae,
-                            image_enc=image_enc,
-                            net=net,
-                            scheduler=val_noise_scheduler,
-                            accelerator=accelerator,
-                            width=cfg.data.train_width,
-                            height=cfg.data.train_height,
-                            clip_length=cfg.data.n_sample_frames,
-                            generator=generator,
-                        )
-
-                        for sample_id, sample_dict in enumerate(sample_dicts):
-                            sample_name = sample_dict["name"]
-                            vid = sample_dict["vid"]
-                            with TemporaryDirectory() as temp_dir:
-                                out_file = Path(
-                                    f"{temp_dir}/{global_step:06d}-{sample_name}.gif"
-                                )
-                                save_videos_grid(vid, out_file, n_rows=2)
-                                mlflow.log_artifact(out_file)
-
-            logs = {
-                "step_loss": loss.detach().item(),
-                "lr": lr_scheduler.get_last_lr()[0],
-                "td": f"{t_data:.2f}s",
-            }
-            t_data_start = time.time()
-            progress_bar.set_postfix(**logs)
-
-            if global_step >= cfg.solver.max_train_steps:
-                break
-        # save model after each epoch
-        if accelerator.is_main_process:
-            save_path = os.path.join(save_dir, f"checkpoint-{global_step}")
-            delete_additional_ckpt(save_dir, 1)
-            accelerator.save_state(save_path)
-            # save motion module only
-            unwrap_net = accelerator.unwrap_model(net)
-            save_checkpoint(
-                unwrap_net.denoising_unet,
-                save_dir,
-                "motion_module",
-                global_step,
-                total_limit=3,
-            )
-
-    # Create the pipeline using the trained modules and save it.
-    accelerator.wait_for_everyone()
-    accelerator.end_training()
-
-
-def save_checkpoint(model, save_dir, prefix, ckpt_num, total_limit=None):
-    save_path = osp.join(save_dir, f"{prefix}-{ckpt_num}.pth")
-
-    if total_limit is not None:
-        checkpoints = os.listdir(save_dir)
-        checkpoints = [d for d in checkpoints if d.startswith(prefix)]
-        checkpoints = sorted(
-            checkpoints, key=lambda x: int(x.split("-")[1].split(".")[0])
-        )
-
-        if len(checkpoints) >= total_limit:
-            num_to_remove = len(checkpoints) - total_limit + 1
-            removing_checkpoints = checkpoints[0:num_to_remove]
-            logger.info(
-                f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
-            )
-            logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
-
-            for removing_checkpoint in removing_checkpoints:
-                removing_checkpoint = os.path.join(save_dir, removing_checkpoint)
-                os.remove(removing_checkpoint)
-
-    mm_state_dict = OrderedDict()
-    state_dict = model.state_dict()
-    for key in state_dict:
-        if "motion_module" in key:
-            mm_state_dict[key] = state_dict[key]
-
-    torch.save(mm_state_dict, save_path)
-
-
-def decode_latents(vae, latents):
-    video_length = latents.shape[2]
-    latents = 1 / 0.18215 * latents
-    latents = rearrange(latents, "b c f h w -> (b f) c h w")
-    # video = self.vae.decode(latents).sample
-    video = []
-    for frame_idx in tqdm(range(latents.shape[0])):
-        video.append(vae.decode(latents[frame_idx : frame_idx + 1]).sample)
-    video = torch.cat(video)
-    video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length)
-    video = (video / 2 + 0.5).clamp(0, 1)
-    # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
-    video = video.cpu().float().numpy()
-    return video
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--config", type=str, default="./configs/training/stage2.yaml")
-    args = parser.parse_args()
-
-    if args.config[-5:] == ".yaml":
-        config = OmegaConf.load(args.config)
-    elif args.config[-3:] == ".py":
-        config = import_filename(args.config).cfg
-    else:
-        raise ValueError("Do not support this format config file")
-    main(config)