train_camn_audio.py

import os
import shutil
import argparse
import random
import numpy as np
from datetime import datetime
from tqdm import tqdm
import importlib
import copy
import librosa
from pathlib import Path
import json
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.nn.parallel import DistributedDataParallel as DDP
import wandb

from diffusers.optimization import get_scheduler
from omegaconf import OmegaConf

from emage_evaltools.mertic import FGD, BC, L1div
from emage_utils.motion_io import beat_format_load, beat_format_save, MASK_DICT, recover_from_mask, recover_from_mask_ts
import emage_utils.rotation_conversions as rc
from emage_utils import fast_render
from emage_utils.motion_rep_transfer import get_motion_rep_numpy


# ---------------------------------  loss here --------------------------------- #
class GeodesicLoss(nn.Module):
    def __init__(self):
        super(GeodesicLoss, self).__init__()
    def compute_geodesic_distance(self, m1, m2):
        m1 = m1.reshape(-1, 3, 3)
        m2 = m2.reshape(-1, 3, 3)
        m = torch.bmm(m1, m2.transpose(1, 2))
        cos = (m[:, 0, 0] + m[:, 1, 1] + m[:, 2, 2] - 1) / 2
        cos = torch.clamp(cos, min=-1 + 1E-6, max=1-1E-6)
        theta = torch.acos(cos)
        return theta
    def __call__(self, m1, m2, reduction='mean'):
        loss = self.compute_geodesic_distance(m1, m2)
        if reduction == 'mean':
            return loss.mean()
        elif reduction == 'none':
            return loss
        else:
            raise RuntimeError

GeodesicLossFn = GeodesicLoss()


# ---------------------------------  train,val,test fn here --------------------------------- #
def inference_fn(cfg, model, device, test_path, save_path):
    actual_model = model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model
    actual_model.eval()
    torch.set_grad_enabled(False)
    test_list = []
    for data_meta_path in test_path:
        test_list.extend(json.load(open(data_meta_path, "r")))
    test_list = [item for item in test_list if item.get("mode") == "test"]
    seen_ids = set()
    test_list = [item for item in test_list if not (item["video_id"] in seen_ids or seen_ids.add(item["video_id"]))]

    save_list = []
    start_time = time.time()
    total_length = 0
    for test_file in tqdm(test_list, desc="Testing"):
        audio, _ = librosa.load(test_file["audio_path"], sr=cfg.audio_sr)
        audio = torch.from_numpy(audio).to(device).unsqueeze(0)
        speaker_id = torch.zeros(1,1).to(device).long()
        motion_pred = actual_model(audio, speaker_id, seed_frames=4, seed_motion=None)["motion_axis_angle"]
        t = motion_pred.shape[1]
        motion_pred = motion_pred.cpu().numpy().reshape(t, -1)
        beat_format_save(os.path.join(save_path, f"{test_file['video_id']}_output.npz"), motion_pred, upsample=30//cfg.pose_fps)
        save_list.append(
            {
                "audio_path": test_file["audio_path"],
                "motion_path": os.path.join(save_path, f"{test_file['video_id']}_output.npz"),
                "video_id": test_file["video_id"],
            }
        )
        total_length+=t
    time_cost = time.time() - start_time
    print(f"\n cost {time_cost:.2f} seconds to generate {total_length / cfg.pose_fps:.2f} seconds of motion")
    return test_list, save_list


def train_val_fn(cfg, batch, model, device, mode="train", optimizer=None, lr_scheduler=None, fgd_evaluator=None):
    model.train() if mode == "train" else model.eval()
    torch.set_grad_enabled(mode == "train")
    joint_mask = MASK_DICT[cfg.model.joint_mask]
    if mode == "train":
        optimizer.zero_grad()
    motion_gt = batch["motion"].to(device)
    audio = batch["audio"].to(device)
    bs, t, jc = motion_gt.shape
    j = jc // 3
    speaker_id = torch.zeros(bs,1).to(device).long()
    motion_gt = rc.axis_angle_to_rotation_6d(motion_gt.reshape(bs,t,j,3)).reshape(bs, t, j*6)
    motion_pred = model(audio, speaker_id, seed_frames=4, seed_motion=motion_gt, return_axis_angle=False)["motion"]
    motion_pred = rc.rotation_6d_to_matrix(motion_pred.reshape(bs,t,j,6))
    motion_gt = rc.rotation_6d_to_matrix(motion_gt.reshape(bs,t,j,6))
    loss = GeodesicLossFn(motion_pred, motion_gt)
    loss_dict = {"loss": loss}
    all_loss = sum(loss_dict.values())
    loss_dict["all_loss"] = all_loss

    if mode == "train":
        if cfg.solver.max_grad_norm > 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.solver.max_grad_norm)
        all_loss.backward()
        optimizer.step()
        lr_scheduler.step()

    if mode == "val":
        motion_pred = rc.matrix_to_rotation_6d(motion_pred).reshape(bs, t, j*6)
        motion_gt = rc.matrix_to_rotation_6d(motion_gt).reshape(bs, t, j*6)
        padded_pred = recover_from_mask_ts(motion_pred, joint_mask)
        padded_gt = recover_from_mask_ts(motion_gt, joint_mask)
        fgd_evaluator.update(padded_pred, padded_gt)
    return loss_dict


# ---------------------------------  main train loop here --------------------------------- #
def main(cfg):
    seed_everything(cfg.seed)
    os.environ["WANDB_API_KEY"] = cfg.wandb_key
    local_rank = int(os.environ["LOCAL_RANK"]) if "LOCAL_RANK" in os.environ else 0
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    torch.distributed.init_process_group(backend="nccl")
    log_dir = os.path.join(cfg.output_dir, cfg.exp_name)
    experiment_ckpt_dir = os.path.join(log_dir, "checkpoints")
    os.makedirs(experiment_ckpt_dir, exist_ok=True)

    if local_rank == 0 and cfg.validation.wandb:  
        wandb.init(
            project=cfg.wandb_project,
            name=cfg.exp_name,
            entity=cfg.wandb_entity,
            dir=log_dir,
            config=OmegaConf.to_container(cfg)
        )

    # init 
    if cfg.test:
        from models.camn_audio import CamnAudioModel
        model = CamnAudioModel.from_pretrained("H-Liu1997/camn_audio").to(device)
    else:
        model = init_hf_class(cfg.model.name_pyfile, cfg.model.class_name, cfg.model).to(device)
    model = DDP(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)

    # optimizer
    optimizer_cls = torch.optim.Adam
    optimizer = optimizer_cls(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=cfg.solver.learning_rate,
        betas=(cfg.solver.adam_beta1, cfg.solver.adam_beta2),
        weight_decay=cfg.solver.adam_weight_decay,
        eps=cfg.solver.adam_epsilon
    )
    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
    )

    # dataset
    train_dataset = init_class(cfg.data.name_pyfile, cfg.data.class_name, cfg, split='train')
    test_dataset = init_class(cfg.data.name_pyfile, cfg.data.class_name, cfg, split='test')
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
    test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset)
    train_loader = DataLoader(train_dataset, batch_size=cfg.data.train_bs, sampler=train_sampler, drop_last=True, num_workers=8)
    test_loader = DataLoader(test_dataset, batch_size=cfg.data.train_bs, sampler=test_sampler, drop_last=False, num_workers=8)

    # resume
    if cfg.resume_from_checkpoint:
        checkpoint = torch.load(cfg.resume_from_checkpoint, map_location="cpu")
        model.load_state_dict(checkpoint["model_state_dict"])
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        lr_scheduler.load_state_dict(checkpoint["lr_scheduler_state_dict"])
        iteration = checkpoint["iteration"]
    else:  
        iteration = 0
    if cfg.test:
        iteration = 0

    max_epochs = (cfg.solver.max_train_steps // len(train_loader)) + (1 if cfg.solver.max_train_steps % len(train_loader) != 0 else 0)
    start_epoch = iteration // len(train_loader)
    start_step_in_epoch = iteration % len(train_loader)
    fgd_evaluator = FGD(download_path="./emage_evaltools/")
    bc_evaluator = BC(download_path="./emage_evaltools/", sigma=0.3, order=7)
    l1div_evaluator= L1div()
    loss_meters = {}
    loss_meters_val = {}
    best_fgd_val = np.inf
    best_fgd_iteration_val= 0
    best_fgd_test = np.inf
    best_fgd_iteration_test = 0
    
    # train loop
    data_start = time.time()
    for epoch in range(start_epoch, max_epochs):
        train_sampler.set_epoch(epoch)
        pbar = tqdm(train_loader, leave=True)
        for i, batch in enumerate(pbar):
            # for correct resume, if the dataset is very large. since we fixed the seed, we can skip the data
            if i < start_step_in_epoch: 
              iteration += 1
              continue
            
            # test 
            if iteration % cfg.validation.test_steps == 0 and local_rank == 0:
                test_save_path = os.path.join(log_dir, f"test_{iteration}")
                os.makedirs(test_save_path, exist_ok=True)
                with torch.no_grad():
                    test_list, save_list = inference_fn(cfg.model, model, device, cfg.data.test_meta_paths, test_save_path)
                if cfg.validation.evaluation:
                    metrics = evaluation_fn([True]*55, test_list, save_list, fgd_evaluator, bc_evaluator, l1div_evaluator, device)
                if cfg.validation.visualization: visualization_fn(save_list, test_save_path, test_list, only_check_one=True)
                if cfg.validation.evaluation: best_fgd_test, best_fgd_iteration_test =  log_test(model, metrics, iteration, best_fgd_test, best_fgd_iteration_test, cfg, local_rank, experiment_ckpt_dir, test_save_path)
                if cfg.test: return 0

            # validation
            if iteration % cfg.validation.validation_steps == 0:
                loss_meters = {}
                loss_meters_val = {}
                fgd_evaluator.reset()
                pbar_val = tqdm(test_loader, leave=True)

                data_start_val = time.time()  
                for j, batch in enumerate(pbar_val):
                    data_time_val = time.time() - data_start_val
                    with torch.no_grad():
                        val_loss_dict = train_val_fn(cfg, batch, model, device, mode="val", fgd_evaluator=fgd_evaluator)
                    net_time_val = time.time() - data_start_val
                    val_loss_dict["fgd"] = fgd_evaluator.compute() if j == len(test_loader) - 1 else 0
                    log_train_val(cfg, val_loss_dict, local_rank, loss_meters_val, pbar_val, epoch, max_epochs, iteration, net_time_val, data_time_val, optimizer, "Val  ")
                    data_start_val = time.time()
                    if cfg.debug and j > 1: break

                if local_rank == 0:
                    best_fgd_val, best_fgd_iteration_val = save_last_and_best_ckpt(
                        model, optimizer, lr_scheduler, iteration, experiment_ckpt_dir, best_fgd_val, best_fgd_iteration_val, val_loss_dict["fgd"], lower_is_better=True, mertic_name="fgd")

            # train
            data_time = time.time() - data_start
            loss_dict = train_val_fn(cfg, batch, model, device, mode="train", optimizer=optimizer, lr_scheduler=lr_scheduler)
            net_time = time.time() - data_start - data_time
            log_train_val(cfg, loss_dict, local_rank, loss_meters, pbar, epoch, max_epochs, iteration, net_time, data_time, optimizer, "Train")
            data_start = time.time()

            iteration += 1
        start_step_in_epoch = 0
        epoch += 1

    if local_rank == 0 and cfg.validation.wandb:
        wandb.finish()
    torch.distributed.destroy_process_group()


# ---------------------------------  utils fn here --------------------------------- #
def evaluation_fn(joint_mask, gt_list, pred_list, fgd_evaluator, bc_evaluator, l1_evaluator, device):
    fgd_evaluator.reset()
    bc_evaluator.reset()
    l1_evaluator.reset()
    # lvd_evaluator.reset()
    # mse_evaluator.reset()

    for test_file in tqdm(gt_list, desc="Evaluation"):
        # only load selective joints
        pred_file = [item for item in pred_list if item["video_id"] == test_file["video_id"]][0]
        if not pred_file:
            print(f"Missing prediction for {test_file['video_id']}")
            continue
        # print(test_file["motion_path"], pred_file["motion_path"])
        gt_dict = beat_format_load(test_file["motion_path"], joint_mask)
        pred_dict = beat_format_load(pred_file["motion_path"], joint_mask)

        motion_gt = gt_dict["poses"]
        motion_pred = pred_dict["poses"]
        # expressions_gt = gt_dict["expressions"]
        # expressions_pred = pred_dict["expressions"]
        betas = gt_dict["betas"]
        # motion_gt = recover_from_mask(motion_gt, joint_mask) # t1*165
        # motion_pred = recover_from_mask(motion_pred, joint_mask) # t2*165
    
        t = min(motion_gt.shape[0], motion_pred.shape[0])
        motion_gt = motion_gt[:t]
        motion_pred = motion_pred[:t]
        # expressions_gt = expressions_gt[:t]
        # expressions_pred = expressions_pred[:t]
       
        # bc and l1 require position representation
        motion_position_pred = get_motion_rep_numpy(motion_pred, device=device, betas=betas)["position"] # t*55*3
        motion_position_pred = motion_position_pred.reshape(t, -1)
        # ignore the start and end 2s, this may for beat dataset only
        audio_beat = bc_evaluator.load_audio(test_file["audio_path"], t_start=2 * 16000, t_end=int((t-60)/30*16000))
        motion_beat = bc_evaluator.load_motion(motion_position_pred, t_start=60, t_end=t-60, pose_fps=30, without_file=True)
        bc_evaluator.compute(audio_beat, motion_beat, length=t-120, pose_fps=30)
        # audio_beat = bc_evaluator.load_audio(test_file["audio_path"], t_start=0 * 16000, t_end=int((t-0)/30*16000))
        # motion_beat = bc_evaluator.load_motion(motion_position_pred, t_start=0, t_end=t-0, pose_fps=30, without_file=True)
        # bc_evaluator.compute(audio_beat, motion_beat, length=t-0, pose_fps=30)

        l1_evaluator.compute(motion_position_pred)
       
        # face_position_pred = get_motion_rep_numpy(motion_pred, device=device, expressions=expressions_pred, expression_only=True, betas=betas)["vertices"] # t -1
        # face_position_gt = get_motion_rep_numpy(motion_gt, device=device, expressions=expressions_gt, expression_only=True, betas=betas)["vertices"]
        # lvd_evaluator.compute(face_position_pred, face_position_gt)
        # mse_evaluator.compute(face_position_pred, face_position_gt)
       
        # fgd requires rotation 6d representaiton
        motion_gt = torch.from_numpy(motion_gt).to(device).unsqueeze(0)
        motion_pred = torch.from_numpy(motion_pred).to(device).unsqueeze(0)
        motion_gt = rc.axis_angle_to_rotation_6d(motion_gt.reshape(1, t, 55, 3)).reshape(1, t, 55*6)
        motion_pred = rc.axis_angle_to_rotation_6d(motion_pred.reshape(1, t, 55, 3)).reshape(1, t, 55*6)
        fgd_evaluator.update(motion_pred.float(), motion_gt.float())
       
    metrics = {}
    metrics["fgd"] = fgd_evaluator.compute()
    metrics["bc"] = bc_evaluator.avg()
    metrics["l1"] = l1_evaluator.avg()
    # metrics["lvd"] = lvd_evaluator.avg()
    # metrics["mse"] = mse_evaluator.avg()
    return metrics

def visualization_fn(pred_list, save_path, gt_list=None, only_check_one=True):
    if gt_list is None: # single visualization
        for i in range(len(pred_list)):
            fast_render.render_one_sequence(
                pred_list[i]["motion_path"],
                save_path,
                pred_list[i]["audio_path"],
                model_folder="./evaluation/smplx_models/",
            )
            if only_check_one: break
    else: # paired visualization, pad the translation
        for i in range(len(pred_list)):
            npz_pred = np.load(pred_list[i]["motion_path"], allow_pickle=True)
            gt_file = [item for item in gt_list if item["video_id"] == pred_list[i]["video_id"]][0]
            if not gt_file:
                print(f"Missing prediction for {pred_list[i]['video_id']}")
                continue
            npz_gt = np.load(gt_file["motion_path"], allow_pickle=True)
            t  = npz_gt["poses"].shape[0]
            np.savez(
                os.path.join(save_path, f"{pred_list[i]['video_id']}_transpad.npz"),
                betas=npz_pred['betas'][:t],
                poses=npz_pred['poses'][:t],
                expressions=npz_pred['expressions'][:t],
                trans=npz_pred["trans"][:t],
                model='smplx2020',
                gender='neutral',
                mocap_frame_rate=30,
            )
            fast_render.render_one_sequence(
                os.path.join(save_path, f"{pred_list[i]['video_id']}_transpad.npz"),
                gt_file["motion_path"],
                save_path,
                pred_list[i]["audio_path"],
                model_folder="./evaluation/smplx_models/",
            )
            if only_check_one: break
     
def log_test(model, metrics, iteration, best_mertics, best_iteration, cfg, local_rank, experiment_ckpt_dir, video_save_path=None):
    if local_rank == 0:
        print(f"\n Test Results at iteration {iteration}:")
        for key, value in metrics.items():
            print(f"  {key}: {value:.10f}")
        if cfg.validation.wandb:
            for key, value in metrics.items():
                wandb.log({f"test/{key}": value}, step=iteration)
        if cfg.validation.wandb and cfg.validation.visualization:
            videos_to_log = []
            for filename in os.listdir(video_save_path):
                if filename.endswith(".mp4"):
                    videos_to_log.append(wandb.Video(os.path.join(video_save_path, filename)))
            if videos_to_log:
                wandb.log({"test/videos": videos_to_log}, step=iteration)
        if metrics["fgd"] < best_mertics:
            best_mertics = metrics["fgd"]
            best_iteration = iteration
            model.module.save_pretrained(os.path.join(experiment_ckpt_dir, "test_best"))
        # print(metrics, best_mertics, best_iteration)
        message = f"Current Test FGD: {metrics['fgd']:.4f} (Best: {best_mertics:.4f} at iteration {best_iteration})"
        log_metric_with_box(message)
    return best_mertics, best_iteration

def log_metric_with_box(message):
    box_width = len(message) + 2
    border = "-" * box_width
    print(f"\n{border}")
    print(f"|{message}|")
    print(f"{border}\n")

def log_train_val(cfg, loss_dict, local_rank, loss_meters, pbar, epoch, max_epochs, iteration, net_time, data_time, optimizer, ptype="Train"):
    new_loss_dict = {}
    for k, v in loss_dict.items():
        if "fgd" in k: continue
        v_cpu = torch.as_tensor(v).float().cpu().item()
        if k not in loss_meters:
            loss_meters[k] = {"sum":0,"count":0}
        loss_meters[k]["sum"] += v_cpu
        loss_meters[k]["count"] += 1
        new_loss_dict[k] = v_cpu
    mem_used = torch.cuda.memory_reserved() / 1E9
    lr = optimizer.param_groups[0]["lr"]
    loss_str = " ".join([f"{k}: {new_loss_dict[k]:.4f}({loss_meters[k]['sum']/loss_meters[k]['count']:.4f})" for k in new_loss_dict])
    desc = f"{ptype}: Epoch[{epoch}/{max_epochs}] Iter[{iteration}] {loss_str} lr: {lr:.2E} data_time: {data_time:.3f} net_time: {net_time:.3f} mem: {mem_used:.2f}GB"
    pbar.set_description(desc)
    pbar.bar_format = "{desc} {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]"
    if cfg.validation.wandb and local_rank == 0:
        for k, v in new_loss_dict.items():
            wandb.log({f"loss/{ptype}/{k}": v}, step=iteration)

def save_last_and_best_ckpt(model, optimizer, lr_scheduler, iteration, save_dir, previous_best, best_iteration, current, lower_is_better=True, mertic_name="fgd"):
    checkpoint = {
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "lr_scheduler_state_dict": lr_scheduler.state_dict(),
            "iteration": iteration,
        }
    torch.save(checkpoint, os.path.join(save_dir, "last.bin"))
    model.module.save_pretrained(os.path.join(save_dir, "last"))
    if (lower_is_better and current < previous_best) or (not lower_is_better and current > previous_best):
        previous_best = current
        best_iteration = iteration
        shutil.copy(os.path.join(save_dir, "last.bin"), os.path.join(save_dir, "best.bin"))
        model.module.save_pretrained(os.path.join(save_dir, "best"))
    message = f"Current interation {iteration} {mertic_name}: {current:.4f} (Best: {previous_best:.4f} at iteration {best_iteration})"
    log_metric_with_box(message)
    return previous_best, best_iteration

def init_hf_class(module_name, class_name, config, **kwargs):
    module = importlib.import_module(module_name)
    model_class = getattr(module, class_name)
    config_class = model_class.config_class
    config = config_class(config_obj=config)
    instance = model_class(config, **kwargs)
    return instance

def init_class(module_name, class_name, config, **kwargs):
    module = importlib.import_module(module_name)
    model_class = getattr(module, class_name)
    instance = model_class(config, **kwargs)
    return instance

def seed_everything(seed):
    os.environ['PYTHONHASHSEED'] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.enabled = True

def init_env():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, default="./configs/train/stage2.yaml")
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--wandb", action="store_true")
    parser.add_argument("--visualization", action="store_true")
    parser.add_argument("--evaluation", action="store_true")
    parser.add_argument("--test", action="store_true")
    parser.add_argument('overrides', nargs=argparse.REMAINDER)
    args = parser.parse_args()
    config = OmegaConf.load(args.config)
    config.exp_name = os.path.splitext(os.path.basename(args.config))[0]

    if args.overrides: config = OmegaConf.merge(config, OmegaConf.from_dotlist(args.overrides))
    if args.debug:
        config.wandb_project = "debug"
        config.exp_name = "debug"
        config.solver.max_train_steps = 4
    else:
        run_time = datetime.now().strftime("%Y%m%d-%H%M")
        config.exp_name = config.exp_name + "_" + run_time
    if args.wandb:
        config.validation.wandb = True
    if args.visualization:
        config.validation.visualization = True
    if args.evaluation:
        config.validation.evaluation = True
    if args.test:
        config.test = True
    save_dir = os.path.join(config.output_dir, config.exp_name)
    os.makedirs(save_dir, exist_ok=True)
    sanity_check_dir = os.path.join(save_dir, 'sanity_check')
    os.makedirs(sanity_check_dir, exist_ok=True)
    with open(os.path.join(sanity_check_dir, f'{config.exp_name}.yaml'), 'w') as f:
        OmegaConf.save(config, f)
    current_dir = Path.cwd()
    for py_file in current_dir.rglob('*.py'):
        dest_path = Path(sanity_check_dir) / py_file.relative_to(current_dir)
        dest_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy(py_file, dest_path)
    return config

if __name__ == "__main__":
    config = init_env()
    main(config)