UnifiedGesture: A Unified Gesture Synthesis Model for Multiple Skeletons (ACM MM 2023)
import os
import pdb

import numpy as np

import pdb
import yaml
from pprint import pprint
from easydict import EasyDict
import torch
from configs.parse_args import parse_args
import glob
from models.vqvae import VQVAE_ulr, VQVAE_ulr_2
import torch.nn as nn

args = parse_args()
mydevice = torch.device('cuda:' + args.gpu)

with open(args.config) as f:
    config = yaml.safe_load(f)

for k, v in vars(args).items():
    config[k] = v
pprint(config)

config = EasyDict(config)

config.no_cuda = config.gpu


def main(upper, lower, root, model, save_path, prefix, model_name='vqvae_ulr_2'):

    with torch.no_grad():
        result_upper = []
        result_lower = []
        result_root = []

        x1 = torch.tensor(upper).unsqueeze(0).to(mydevice)
        x2 = torch.tensor(lower).unsqueeze(0).to(mydevice)
        x3 = torch.tensor(root).unsqueeze(0).to(mydevice)

        if model_name == 'vqvae_ulr':
            # zs1, zs2, zs3 = model.module.get_code(x1.float(), x2.float(), x3.float())
            zs1, _ = model.module.get_code(x1.float())
        elif model_name == 'vqvae_ulr_2':
            x2 = torch.cat([x2, x3], dim=2)
            zs1, zs2 = model.module.get_code(x1.float(), x2.float())

        if model_name == 'vqvae_ulr':
            # y1, y2, y3 = model.module.code_2_pose(zs1, zs2, zs3)
            y1 = model.module.code_2_pose(zs1)
        elif model_name == 'vqvae_ulr_2':
            y1, y2 = model.module.code_2_pose(zs1, zs2)

        result_upper.append(y1.squeeze(0).data.cpu().numpy())
        # result_lower.append(y2.squeeze(0).data.cpu().numpy())
        # if model_name == 'vqvae_ulr':
        #     result_root.append(y3.squeeze(0).data.cpu().numpy())

    out_zs1 = np.vstack(zs1[0].squeeze(0).data.cpu().numpy())
    # out_zs2 = np.vstack(zs2[0].squeeze(0).data.cpu().numpy())
    # if model_name == 'vqvae_ulr':
    #     out_zs3 = np.vstack(zs3[0].squeeze(0).data.cpu().numpy())
    #     out_root = np.vstack(result_root)
    out_upper = np.vstack(result_upper)
    # out_lower = np.vstack(result_lower)

    # if model_name == 'vqvae_ulr':
    #     print(out_upper.shape, out_lower.shape, out_root.shape)
    # elif model_name == 'vqvae_ulr_2':
    #     print(out_upper.shape, out_lower.shape)

    np.save(os.path.join(save_path, 'code_upper_' + prefix + '.npy'), out_zs1)
    # np.save(os.path.join(save_path, 'code_lower_' + prefix + '.npy'), out_zs2)
    np.save(os.path.join(save_path, 'generate_upper_' + prefix + '.npy'), np.expand_dims(out_upper.transpose(), axis=0))
    # np.save(os.path.join(save_path, 'generate_lower_' + prefix + '.npy'), np.expand_dims(out_lower.transpose(), axis=0))
    if model_name == 'vqvae_ulr':
        # np.save(os.path.join(save_path, 'code_root_' + prefix + '.npy'), out_zs3)
        # np.save(os.path.join(save_path, 'generate_root_' + prefix + '.npy'), np.expand_dims(out_root.transpose(), axis=0))
        # return np.expand_dims(out_upper.transpose(), axis=0), np.expand_dims(out_lower.transpose(), axis=0), \
        #        np.expand_dims(out_root.transpose(), axis=0)
        np.expand_dims(out_upper.transpose(), axis=0)

    elif model_name == 'vqvae_ulr_2':
        return np.expand_dims(out_upper.transpose(), axis=0), np.expand_dims(out_lower.transpose(), axis=0)


def main_2(upper, model, save_path, prefix, w_grad, code=None):
    if not w_grad:
        with torch.no_grad():
            result_upper = []
            x1 = torch.tensor(upper).unsqueeze(0).to(mydevice)
            if type(code) is np.ndarray:
                print('use code')
                zs1 = [torch.tensor(code).to(mydevice)]
                out_distance = None
            else:
                zs1, distance = model.module.get_code(x1.float())
                out_distance = np.vstack(distance[0].squeeze(0).data.cpu().numpy())

            y1 = model.module.code_2_pose(zs1)
            result_upper.append(y1.squeeze(0).data.cpu().numpy())
    else:
        result_upper = []

        x1 = upper.unsqueeze(0).to(mydevice)
        zs1, distance = model.module.get_code(x1.float())
        out_distance = torch.cat(distance, dim=0)
        y1 = model.module.code_2_pose(zs1)
        result_upper.append(y1.squeeze(0).data.cpu().numpy())

    out_zs1 = np.vstack(zs1[0].squeeze(0).data.cpu().numpy())
    out_upper = np.vstack(result_upper)
    # print(out_upper.shape)
    # np.save(os.path.join(save_path, 'code_upper_' + prefix + '.npy'), out_zs1)
    # np.save(os.path.join(save_path, 'generate_upper_' + prefix + '.npy'), np.expand_dims(out_upper.transpose(), axis=0))
    return np.expand_dims(out_upper.transpose(), axis=0), out_zs1, out_distance



def visualize_code(model, save_path, prefix, code_upper, code_lower=None, code_root=None, model_name='vqvae_ulr_2'):
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    with torch.no_grad():
        result_upper = []
        result_lower = []
        if model_name == 'vqvae_ulr':
            result_root = []
            code_root = torch.tensor(code_root).unsqueeze(0).to(mydevice)

        code_upper = torch.tensor(code_upper).unsqueeze(0).to(mydevice)
        if code_lower:
            code_lower = torch.tensor(code_lower).unsqueeze(0).to(mydevice)

        if model_name == 'vqvae_ulr':
            y1, y2, y3 = model.module.code_2_pose([code_upper], [code_lower], [code_root])
            result_root.append(y3.squeeze(0).data.cpu().numpy())
            out_root = np.vstack(result_root)

        elif model_name == 'vqvae_ulr_2':
            if code_lower:
                y1, y2 = model.module.code_2_pose([code_upper], [code_lower])
            else:
                y1 = model.module.code_pose([code_upper])

        if code_lower:
            result_lower.append(y2.squeeze(0).data.cpu().numpy())
        result_upper.append(y1.squeeze(0).data.cpu().numpy())


    out_upper = np.vstack(result_upper)
    if code_lower:
        out_lower = np.vstack(result_lower)

    if model_name == 'vqvae_ulr':
        print(out_upper.shape, out_lower.shape, out_root.shape)
        np.save(os.path.join(save_path, 'generate_root_' + prefix), np.expand_dims(out_root.transpose(), axis=0))
    elif model_name == 'vqvae_ulr_2':
        if code_lower:
            print(out_upper.shape, out_lower.shape)
            np.save(os.path.join(save_path, 'generate_lower_' + prefix), np.expand_dims(out_lower.transpose(), axis=0))
        else:
            print(out_upper.shape)
        np.save(os.path.join(save_path, 'generate_upper_' + prefix), np.expand_dims(out_upper.transpose(), axis=0))

    if model_name == 'vqvae_ulr':
        return np.expand_dims(out_upper.transpose(), axis=0), np.expand_dims(out_lower.transpose(), axis=0), \
           np.expand_dims(out_root.transpose(), axis=0)
    elif model_name == 'vqvae_ulr_2':
        if code_lower:
            return np.expand_dims(out_upper.transpose(), axis=0), np.expand_dims(out_lower.transpose(), axis=0)
        else:
            return np.expand_dims(out_upper.transpose(), axis=0)


def fold2code(path, save_path):
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    upper_files = sorted(glob.glob(path + "/*_upper.npy"))
    for item in upper_files:
        path = os.path.split(item)[0]
        name = os.path.split(item)[1]
        source_upper = np.load(os.path.join(path, name))[0].transpose()
        source_lower = np.load(os.path.join(path, name.replace('upper', 'lower')))[0].transpose()
        source_root = np.load(os.path.join(path, name.replace('upper', 'root')))[0].transpose()
        prefix = name.replace('_upper.npy', '')
        main(source_upper, source_lower, source_root, model, save_path, prefix, model_name='vqvae_ulr')



if __name__ == '__main__':
    '''
    cd codebook/
    python VisualizeCodebook.py --config=./configs/codebook.yml --train --gpu 0
    '''

    model = VQVAE_ulr(config.VQVAE, 7 * 16)  # n_joints * n_chanels
    model_path = "./result/my_codebook/codebook_checkpoint_best.bin"
    model = nn.DataParallel(model, device_ids=[eval(i) for i in config.no_cuda])
    model = model.to(mydevice)
    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
    model.load_state_dict(checkpoint['model_dict'])
    model = model.eval()

    # # test orignal latent
    # source_upper = np.load('../retargeting/datasets/bvh2upper_lower_root/Trinity/Recording_006_upper.npy')[0].transpose()
    # # print(source_upper.shape)       # (3164, 64)
    # save_path = "./result/inference/Trinity"
    # prefix = 'train_codebook_bvh2upper_lower_root_again'
    # main_2(source_upper, model, save_path, prefix)


    # generate code from latent
    ZEGGS_path = '../retargeting/datasets/Trinity_ZEGGS/bvh2upper_lower_root/ZEGGS/'
    Trinity_path = '../retargeting/datasets/Trinity_ZEGGS/bvh2upper_lower_root/Trinity/'
    ZEGGS_save_path = '../retargeting/datasets/Trinity_ZEGGS/VQVAE_result/ZEGGS'
    Trinity_save_path = '../retargeting/datasets/Trinity_ZEGGS/VQVAE_result/Trinity/'
    fold2code(Trinity_path, Trinity_save_path)
    fold2code(ZEGGS_path, ZEGGS_save_path)



    # test generate code
    # test_file = '../result/inference/Recording_006.npy'
    # path = os.path.split(test_file)[0]
    # name = os.path.split(test_file)[1]
    # code_upper = np.load(os.path.join(path, 'code_upper_' + name))
    # code_lower = np.load(os.path.join(path, 'code_lower_' + name))
    # code_root = np.load(os.path.join(path, 'code_root_' + name))
    # save_path = "../result/inference/Trinity/"

    # visualize_code(model, save_path, name, code_upper, code_lower=None, code_root=None, model_name='vqvae_ulr_2')

codebook/VisualizeCodebook.py
YoungSeng/UnifiedGesture
import configargparse
import argparse

def str2bool(v):
    """ from https://stackoverflow.com/a/43357954/1361529 """
    if isinstance(v, bool):
       return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise configargparse.ArgumentTypeError('Boolean value expected.')


def parse_args():
    # parser = configargparse.ArgParser()
    # parser.add('-c', '--config', required=True, is_config_file=True, help='Config file path')
    #
    # # dataset
    # parser.add("--train_data_path", action="append", required=True)
    # parser.add("--val_data_path", action="append", required=True)
    # parser.add("--data_mean", action="append", type=float, nargs='*')
    # parser.add("--data_std", action="append", type=float, nargs='*')
    # parser.add("--n_poses", type=int, default=50)
    # parser.add("--subdivision_stride", type=int, default=5)
    # parser.add("--motion_resampling_framerate", type=int, default=24)
    # parser.add("--batch_size", type=int, default=50)
    # parser.add("--loader_workers", type=int, default=0)
    # parser.add("--epochs", type=int, default=10)

    parser = argparse.ArgumentParser(description='Codebook')
    parser.add_argument('--config', default='./configs/codebook.yml')
    parser.add_argument('--lr', default=0.00003, type=float)
    # parser.add_argument('--lr', default=params['lr'], type=float)       # 20230420  0.00003 -> 0.00004

    # codebook
    parser.add_argument('--gpu', type=str, default='2')
    parser.add_argument('--no_cuda', type=list, default=['2'])
    # exclusive arguments
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('--train', action='store_true')


    # parser.add("--test_data_path", action="append", required=False)
    # parser.add("--model_save_path", required=True)
    # parser.add("--random_seed", type=int, default=-1)
    #
    # # word embedding
    # parser.add("--wordembed_path", type=str, default=None)
    # parser.add("--wordembed_dim", type=int, default=200)
    #
    # # model
    # parser.add("--model", type=str, required=True)
    #
    #
    # parser.add("--dropout_prob", type=float, default=0.3)
    # parser.add("--n_layers", type=int, default=2)
    # parser.add("--hidden_size", type=int, default=200)
    #
    # parser.add("--n_pre_poses", type=int, default=5)
    #
    # # training
    # parser.add("--learning_rate", type=float, default=0.001)
    # parser.add("--loss_l1_weight", type=float, default=50)
    # parser.add("--loss_cont_weight", type=float, default=0.1)
    # parser.add("--loss_var_weight", type=float, default=0.01)

    args = parser.parse_args()
    return args

codebook/configs/parse_args.py
YoungSeng/UnifiedGesture
""" create data samples """
import lmdb
import math
import numpy as np
import pyarrow


class DataPreprocessor:
    def __init__(self, clip_lmdb_dir, out_lmdb_dir, n_poses, subdivision_stride, pose_resampling_fps, file='ag',
                 select='specific', n_codes=30):
        self.n_codes = n_codes
        self.n_poses = n_poses
        self.subdivision_stride = subdivision_stride
        self.skeleton_resampling_fps = pose_resampling_fps

        self.src_lmdb_env = lmdb.open(clip_lmdb_dir, readonly=True, lock=False)
        with self.src_lmdb_env.begin() as txn:
            self.n_videos = txn.stat()['entries']

        self.audio_sample_length = int(self.n_poses / self.skeleton_resampling_fps * 16000)

        # create db for samples
        map_size = 1024 * 1024 * 9  # in TB
        map_size <<= 20  # in B
        self.dst_lmdb_env = lmdb.open(out_lmdb_dir, map_size=map_size)
        self.n_out_samples = 0
        self.file = file
        self.select = select

    def run(self):
        src_txn = self.src_lmdb_env.begin(write=False)

        # sampling and normalization
        cursor = src_txn.cursor()
        for key, value in cursor:
            video = pyarrow.deserialize(value)
            vid = video['vid']
            clips = video['clips']
            for clip_idx, clip in enumerate(clips):
                if self.select == 'all_speaker':
                    self._sample_from_clip_allspeakers(vid, clip)
                else:
                    self._sample_from_clip(vid, clip)

        # print stats
        with self.dst_lmdb_env.begin() as txn:
            print('no. of samples: ', txn.stat()['entries'])

        # close db
        self.src_lmdb_env.close()
        self.dst_lmdb_env.sync()
        self.dst_lmdb_env.close()

    def _sample_from_clip_allspeakers(self, vid, clip):
        clip_skeleton = clip['poses']
        # clip_audio_raw = clip['audio_raw']

        # divide
        aux_info = []
        sample_skeletons_list = []
        sample_audio_list = []

        # MINLEN = min(len(clip_skeleton), int(len(clip_audio_raw) * 60 / 16000))

        MINLEN = len(clip_skeleton)
        num_subdivision = math.floor(
            (MINLEN - self.n_poses)
            / self.subdivision_stride) + 1  # floor((K - (N+M)) / S) + 1

        for i in range(num_subdivision):
            start_idx = i * self.subdivision_stride
            fin_idx = start_idx + self.n_poses

            sample_skeletons = clip_skeleton[start_idx:fin_idx]
            subdivision_start_time = start_idx / self.skeleton_resampling_fps
            subdivision_end_time = fin_idx / self.skeleton_resampling_fps

            # raw audio
            # audio_start = math.floor(start_idx / len(clip_skeleton) * len(clip_audio_raw))
            # audio_end = audio_start + self.audio_sample_length
            # sample_audio = clip_audio_raw[audio_start:audio_end]

            motion_info = {'vid': vid,
                           'start_frame_no': start_idx,
                           'end_frame_no': fin_idx,
                           'start_time': subdivision_start_time,
                           'end_time': subdivision_end_time}

            sample_skeletons_list.append(sample_skeletons)
            # sample_audio_list.append(sample_audio)
            aux_info.append(motion_info)

        if len(sample_skeletons_list) > 0:
            with self.dst_lmdb_env.begin(write=True) as txn:
                for poses, audio, aux in zip(sample_skeletons_list,
                                                    sample_audio_list, aux_info):
                    poses = np.asarray(poses)

                    # save
                    k = '{:010}'.format(self.n_out_samples).encode('ascii')
                    if self.file == 'ag':
                        v = [poses, audio, aux]
                    elif self.file == 'g':
                        v = [poses, np.array([0]), np.array([0])]
                    v = pyarrow.serialize(v).to_buffer()
                    txn.put(k, v)
                    self.n_out_samples += 1

    def _sample_from_clip(self, vid, clip):
        clip_upper = clip['upper']
        clip_lower = clip['lower']
        clip_root = clip['root']
        clip_root_vel = clip['root_vel']

        # divide
        aux_info = []
        sample_upper_list = []
        sample_lower_list = []
        sample_root_list = []
        sample_root_vel_list = []

        # MINLEN = min(len(clip_skeleton), int(len(clip_audio_raw) * 60 / 16000), len(clip_codes_raw) * 8)
        MINLEN = len(clip_upper)

        num_subdivision = math.floor(
            (MINLEN - self.n_poses)
            / self.subdivision_stride) + 1  # floor((K - (N+M)) / S) + 1

        for i in range(num_subdivision):
            start_idx = i * self.subdivision_stride
            fin_idx = start_idx + self.n_poses

            sample_upper = clip_upper[start_idx:fin_idx]
            sample_lower = clip_lower[start_idx:fin_idx]
            sample_root = clip_root[start_idx:fin_idx]
            sample_root_vel = clip_root_vel[start_idx:fin_idx]
            subdivision_start_time = start_idx / self.skeleton_resampling_fps
            subdivision_end_time = fin_idx / self.skeleton_resampling_fps

            motion_info = {'vid': vid,
                           'start_frame_no': start_idx,
                           'end_frame_no': fin_idx,
                           'start_time': subdivision_start_time,
                           'end_time': subdivision_end_time}

            sample_upper_list.append(sample_upper)
            sample_lower_list.append(sample_lower)
            sample_root_list.append(sample_root)
            sample_root_vel_list.append(sample_root_vel)

            aux_info.append(motion_info)

        if len(sample_upper_list) > 0:
            with self.dst_lmdb_env.begin(write=True) as txn:
                for upper, lower, root, root_vel, aux in zip(sample_upper_list, sample_lower_list, sample_root_list,
                                                   sample_root_vel_list, aux_info):
                    upper = np.asarray(upper)
                    lower = np.asarray(lower)
                    root = np.asarray(root)
                    root_vel = np.asarray(root_vel)

                    # save
                    k = '{:010}'.format(self.n_out_samples).encode('ascii')
                    v = [upper, lower, root, root_vel, aux]
                    v = pyarrow.serialize(v).to_buffer()
                    txn.put(k, v)
                    self.n_out_samples += 1

    @staticmethod
    def get_words_in_time_range(word_list, start_time, end_time):
        words = []

        for word in word_list:
            _, word_s, word_e = word[0], word[1], word[2]

            if word_s >= end_time:
                break

            if word_e <= start_time:
                continue

            words.append(word)

        return words

codebook/data_loader/data_preprocessor.py
YoungSeng/UnifiedGesture
import logging
import os
import pdb

import numpy as np
import lmdb as lmdb
import torch
from torch.nn.utils.rnn import pad_sequence

from torch.utils.data import Dataset
from torch.utils.data.dataloader import default_collate

import sys
[sys.path.append(i) for i in ['.', '..']]

from data_loader.data_preprocessor import DataPreprocessor
import pyarrow


def word_seq_collate_fn(data):
    """ collate function for loading word sequences in variable lengths """
    # sort a list by sequence length (descending order) to use pack_padded_sequence
    data.sort(key=lambda x: len(x[0]), reverse=True)

    # separate source and target sequences
    word_seq, poses_seq, audio, aux_info = zip(*data)

    # merge sequences
    words_lengths = torch.LongTensor([len(x) for x in word_seq])
    word_seq = pad_sequence(word_seq, batch_first=True).long()

    poses_seq = default_collate(poses_seq)
    audio = default_collate(audio)
    aux_info = {key: default_collate([d[key] for d in aux_info]) for key in aux_info[0]}

    return word_seq, words_lengths, poses_seq, audio, aux_info


class TrinityDataset(Dataset):
    def __init__(self, lmdb_dir, n_poses, subdivision_stride, pose_resampling_fps, data_mean=None, data_std=None, model=None,
                 file='ag', select='specific'):
        self.lmdb_dir = lmdb_dir
        self.n_poses = n_poses
        self.subdivision_stride = subdivision_stride
        self.skeleton_resampling_fps = pose_resampling_fps
        self.lang_model = None
        # self.data_mean = np.array(data_mean).squeeze()
        # self.data_std = np.array(data_std).squeeze()
        self.file = file

        logging.info("Reading data '{}'...".format(lmdb_dir))
        preloaded_dir = lmdb_dir + '_cache'
        if not os.path.exists(preloaded_dir):
            data_sampler = DataPreprocessor(lmdb_dir, preloaded_dir, n_poses,
                                            subdivision_stride, pose_resampling_fps, file=file, select=select)
            data_sampler.run()
        else:
            logging.info('Found pre-loaded samples from {}'.format(preloaded_dir))

        # init lmdb
        # map_size = 1024 * 20  # in MB
        # map_size <<= 20  # in B
        self.lmdb_env = lmdb.open(preloaded_dir, readonly=True, lock=False)      # default 10485760
        with self.lmdb_env.begin() as txn:
            self.n_samples = txn.stat()['entries']

    def __len__(self):
        return self.n_samples

    def __getitem__(self, idx):
        with self.lmdb_env.begin(write=False) as txn:
            key = '{:010}'.format(idx).encode('ascii')
            sample = txn.get(key)

            sample = pyarrow.deserialize(sample)
            upper_seq, lower_seq, root_seq, root_vel_seq, aux_info = sample

        # # normalize
        # std = np.clip(self.data_std, a_min=0.01, a_max=None)
        # pose_seq = (pose_seq - self.data_mean) / std

        # to tensors
        upper_seq = torch.from_numpy(upper_seq).reshape((upper_seq.shape[0], -1)).float()
        lower_seq = torch.from_numpy(lower_seq).reshape((lower_seq.shape[0], -1)).float()
        root_seq = torch.from_numpy(root_seq).reshape((root_seq.shape[0], -1)).float()
        root_vel_seq = torch.from_numpy(root_vel_seq).reshape((root_vel_seq.shape[0], -1)).float()

        return upper_seq, lower_seq, root_seq, root_vel_seq, aux_info


if __name__ == '__main__':
    '''
    cd codebook/
    python data_loader/lmdb_data_loader.py --config=./configs/codebook.yml --train --no_cuda 3 --gpu 3
    '''

    from configs.parse_args import parse_args
    import os
    import yaml
    from pprint import pprint
    from easydict import EasyDict
    from torch.utils.data import DataLoader

    args = parse_args()

    with open(args.config) as f:
        config = yaml.safe_load(f)

    for k, v in vars(args).items():
        config[k] = v
    pprint(config)

    args = EasyDict(config)

    
    train_dataset = TrinityDataset(args.train_data_path,
                                   n_poses=args.n_poses,
                                   subdivision_stride=args.subdivision_stride,
                                   pose_resampling_fps=args.motion_resampling_framerate,
                                   data_mean=args.data_mean, data_std=args.data_std)
    train_loader = DataLoader(dataset=train_dataset, batch_size=128,
                              shuffle=True, drop_last=True, num_workers=args.loader_workers, pin_memory=True)

    print(len(train_loader))
    for batch_i, batch in enumerate(train_loader, 0):
        target_vec, aux, code, audio = batch
        print(batch_i)
        pdb.set_trace()
        # print(target_vec.shape, audio.shape)

codebook/data_loader/lmdb_data_loader.py
YoungSeng/UnifiedGesture
import pdb

import numpy as np
import torch as t
import torch.nn as nn
import torch.nn.functional as F
import models.utils.dist_adapter as dist
import sys
[sys.path.append(i) for i in ['.', '..']]
from configs.parse_args import parse_args

args = parse_args()
mydevice = t.device('cuda:' + args.gpu)

class BottleneckBlock(nn.Module):
    def __init__(self, k_bins, emb_width, mu):
        super().__init__()
        self.k_bins = k_bins
        self.emb_width = emb_width
        self.mu = mu
        self.reset_k()
        self.threshold = 1.0

    def reset_k(self):
        self.init = False
        self.k_sum = None
        self.k_elem = None
        self.register_buffer('k', t.zeros(self.k_bins, self.emb_width).to(mydevice))

    def _tile(self, x):
        d, ew = x.shape     # 960, 512
        if d < self.k_bins:
            n_repeats = (self.k_bins + d - 1) // d
            std = 0.01 / np.sqrt(ew)
            x = x.repeat(n_repeats, 1)
            x = x + t.randn_like(x) * std
        return x

    def init_k(self, x):
        mu, emb_width, k_bins = self.mu, self.emb_width, self.k_bins        # mu=0.99, emb_width=512, k_bins=512
        self.init = True
        # init k_w using random vectors from x
        y = self._tile(x)
        _k_rand = y[t.randperm(y.shape[0])][:k_bins]        # (512, 512), a random permutation of integers from 0 to n - 1
        # dist.broadcast(_k_rand, 0)
        self.k = _k_rand
        assert self.k.shape == (k_bins, emb_width)
        self.k_sum = self.k
        self.k_elem = t.ones(k_bins, device=self.k.device)

    def restore_k(self, num_tokens=None, threshold=1.0):
        mu, emb_width, k_bins = self.mu, self.emb_width, self.k_bins
        self.init = True
        assert self.k.shape == (k_bins, emb_width)
        self.k_sum = self.k.clone()
        self.k_elem = t.ones(k_bins, device=self.k.device)
        if num_tokens is not None:
            expected_usage = num_tokens / k_bins
            self.k_elem.data.mul_(expected_usage)
            self.k_sum.data.mul_(expected_usage)
        self.threshold = threshold

    def update_k(self, x, x_l):     # (960, 512), (960)
        mu, emb_width, k_bins = self.mu, self.emb_width, self.k_bins        # mu=0.99, emb_width=512, k_bins=512
        with t.no_grad():
            # Calculate new centres
            x_l_onehot = t.zeros(k_bins, x.shape[0], device=x.device)  # (512(k_bins), 960(N * L))
            x_l_onehot.scatter_(0, x_l.view(1, x.shape[0]), 1)      # (1, 190) -> (512, 960), find which axis

            _k_sum = t.matmul(x_l_onehot, x)  #(512(k_bins), 512(w))
            _k_elem = x_l_onehot.sum(dim=-1)  # (512(k_bins))
            y = self._tile(x)       # (960, 512)
            _k_rand = y[t.randperm(y.shape[0])][:k_bins]        # (512, 512)

            # dist.broadcast(_k_rand, 0)
            # dist.all_reduce(_k_sum)
            # dist.all_reduce(_k_elem)

            # Update centres
            old_k = self.k
            self.k_sum = mu * self.k_sum + (1. - mu) * _k_sum  # w, k_bins
            self.k_elem = mu * self.k_elem + (1. - mu) * _k_elem  # k_bins
            usage = (self.k_elem.view(k_bins, 1) >= self.threshold).float()
            self.k = usage * (self.k_sum.view(k_bins, emb_width) / self.k_elem.view(k_bins, 1)) \
                     + (1 - usage) * _k_rand
            _k_prob = _k_elem / t.sum(_k_elem)  # x_l_onehot.mean(dim=-1)  # prob of each bin
            entropy = -t.sum(_k_prob * t.log(_k_prob + 1e-8))  # entropy ie how diverse
            used_curr = (_k_elem >= self.threshold).sum()
            usage = t.sum(usage)
            dk = t.norm(self.k - old_k) / np.sqrt(np.prod(old_k.shape))
        return dict(entropy=entropy,
                    used_curr=used_curr,
                    usage=usage,
                    dk=dk)

    def preprocess(self, x):
        # NCT -> NTC -> [NT, C]
        x = x.permute(0, 2, 1).contiguous()
        x = x.view(-1, x.shape[-1])  # x_en = (N * L, w), k_j = (w, k_bins)

        if x.shape[-1] == self.emb_width:
            prenorm = t.norm(x - t.mean(x)) / np.sqrt(np.prod(x.shape))     # np.sqrt - product of array elements over a given axis
        elif x.shape[-1] == 2 * self.emb_width:
            x1, x2 = x[...,:self.emb_width], x[...,self.emb_width:]
            prenorm = (t.norm(x1 - t.mean(x1)) / np.sqrt(np.prod(x1.shape))) + (t.norm(x2 - t.mean(x2)) / np.sqrt(np.prod(x2.shape)))

            # Normalise
            x = x1 + x2
        else:
            assert False, f"Expected {x.shape[-1]} to be (1 or 2) * {self.emb_width}"
        return x, prenorm

    def postprocess(self, x_l, x_d, x_shape):
        # [NT, C] -> NTC -> NCT
        N, T = x_shape
        x_d = x_d.view(N, T, -1).permute(0, 2, 1).contiguous()
        x_l = x_l.view(N, T)
        return x_l, x_d

    def quantise(self, x):
        # Calculate latent code x_l
        k_w = self.k.t()        # (512, 512)
        distance = t.sum(x ** 2, dim=-1, keepdim=True) - 2 * t.matmul(x, k_w) + t.sum(k_w ** 2, dim=0, keepdim=True)  # (960(N * L), 512(b))
        min_distance, x_l = t.min(distance, dim=-1)     # (960), (960)
        fit = t.mean(min_distance)
        return x_l, fit, distance

    def dequantise(self, x_l):
        x = F.embedding(x_l, self.k)        # self.k: (512, 512) weighted array
        return x

    def encode(self, x):
        N, width, T = x.shape

        # Preprocess.
        x, prenorm = self.preprocess(x)

        # Quantise
        x_l, fit, distance = self.quantise(x)

        # Postprocess.
        x_l = x_l.view(N, T)
        return x_l, distance

    def decode(self, x_l):
        N, T = x_l.shape
        width = self.emb_width

        # Dequantise
        x_d = self.dequantise(x_l)

        # Postprocess
        x_d = x_d.view(N, T, width).permute(0, 2, 1).contiguous()
        return x_d

    def forward(self, x, update_k=True):
        N, width, T = x.shape       # 32, 512, 30

        # Preprocess
        x, prenorm = self.preprocess(x)     # (960, 512), 0.2888

        # Init k if not inited
        if update_k and not self.init:
            self.init_k(x)

        # Quantise and dequantise through bottleneck

        x_l, fit, _ = self.quantise(x)     # (960), 34.1081
        x_d = self.dequantise(x_l)      # (960, 512)

        # Update embeddings
        if update_k:
            update_metrics = self.update_k(x, x_l)
        else:
            update_metrics = {}

        # Loss
        commit_loss = t.norm(x_d.detach() - x) ** 2 / np.prod(x.shape)      # L2 loss -> L1 loss

        # Passthrough
        x_d = x + (x_d - x).detach()

        # Postprocess
        x_l, x_d = self.postprocess(x_l, x_d, (N,T))
        return x_l, x_d, commit_loss, dict(fit=fit,
                                           pn=prenorm,
                                           **update_metrics)


class Bottleneck(nn.Module):
    def __init__(self, l_bins, emb_width, mu, levels):
        super().__init__()
        self.levels = levels
        level_block = lambda level: BottleneckBlock(l_bins, emb_width, mu)
        self.level_blocks = nn.ModuleList()
        for level in range(self.levels):
            self.level_blocks.append(level_block(level))

    def encode(self, xs):
        zs = [level_block.encode(x) for (level_block, x) in zip(self.level_blocks, xs)]
        return zs

    def decode(self, zs, start_level=0, end_level=None):
        if end_level is None:
            end_level = self.levels
        xs_quantised = [level_block.decode(z) for (level_block, z) in zip(self.level_blocks[start_level:end_level], zs)]
        return xs_quantised

    def forward(self, xs):
        zs, xs_quantised, commit_losses, metrics = [], [], [], []
        for level in range(self.levels):
            level_block = self.level_blocks[level]
            x = xs[level]       # (32, 512, 30)
            z, x_quantised, commit_loss, metric = level_block(x, update_k=self.training)
            '''
            z: (32, 30)
            x_quantised: (32, 512, 30)
            commit_loss: 0.0666
            metric: same as models/vqvae.py
            '''
            zs.append(z)
            if not self.training:
                # Be extra paranoid and make sure the encoder weights can't
                # change from straight-through estimator
                x_quantised = x_quantised.detach()
            xs_quantised.append(x_quantised)
            commit_losses.append(commit_loss)
            if self.training:
                metrics.append(metric)
        return zs, xs_quantised, commit_losses, metrics

class NoBottleneckBlock(nn.Module):
    def restore_k(self):
        pass

class NoBottleneck(nn.Module):
    def __init__(self, levels):
        super().__init__()
        self.level_blocks = nn.ModuleList()
        self.levels = levels
        for level in range(levels):
            self.level_blocks.append(NoBottleneckBlock())

    def encode(self, xs):
        return xs

    def decode(self, zs, start_level=0, end_level=None):
        if end_level is None:
            end_level = self.levels
        return zs

    def forward(self, xs):
        zero = t.zeros(()).to(mydevice)
        commit_losses = [zero for _ in range(self.levels)]
        metrics = [dict(entropy=zero, usage=zero, used_curr=zero, pn=zero, dk=zero) for _ in range(self.levels)]
        return xs, xs, commit_losses, metrics

if __name__ == '__main__':
    '''
    python -m models.bottleneck --config configs/sep_vqvae.yaml --train --no_cuda 2 --gpu 2
    '''
    # x = [t.rand(32, 512, 30)]
    # bottleneck = Bottleneck(512, 512, 0.99, 1).to(mydevice)
    # zs, xs_quantised, commit_losses, quantiser_metrics = bottleneck(x)

    x = t.rand(32, 512, 30)
    model = BottleneckBlock(k_bins=512, emb_width=512, mu=0.99)
    zs, xs_quantised, commit_losses, quantiser_metrics = model(x)

codebook/models/bottleneck.py
YoungSeng/UnifiedGesture
"""
Patch Encoders / Decoders as used by DeepMind in their sonnet repo example:
https://github.com/deepmind/sonnet/blob/v2/examples/vqvae_example.ipynb
"""
import pdb

import torch
from torch import nn, einsum
import torch.nn.functional as F

# -----------------------------------------------------------------------------

class ResBlock(nn.Module):
    def __init__(self, input_channels, channel):
        super().__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(input_channels, channel, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(channel, input_channels, 1),
        )

    def forward(self, x):
        out = self.conv(x)
        out += x
        out = F.relu(out)
        return out


class DeepMindEncoder(nn.Module):

    def __init__(self, input_channels):         # 7 * 16
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(input_channels, input_channels // 2),
            nn.ReLU(True),
        )

    def forward(self, x):       # [8, 30, 16 * 7] -> [8, 16, 30, 7]
        x = self.net(x)     # -> [8, 30, 128]
        x = x.unsqueeze(-1)     # -> [8, 30, 128, 1]
        x = x.permute(0, 2, 1, 3)    # -> [8, 128, 30, 1]
        return x


class DeepMindDecoder(nn.Module):

    def __init__(self, out_channels):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(64, out_channels),
            nn.Tanh()
        )

    def forward(self, x):       # [batch, 64, 30, 1]
        x = x.permute(0, 2, 3, 1)   # -> [batch, 30, 1, 112]
        x = self.net(x)     #
        x = x.reshape(x.shape[0], x.shape[1], x.shape[2]*x.shape[3])  # [8, 30, 112]
        return x

codebook/models/dvq/dvqmodel/deepmind_enc_dec.py
YoungSeng/UnifiedGesture
"""
The critical quantization layers that we sandwich in the middle of the autoencoder
(between the encoder and decoder) that force the representation through a categorical
variable bottleneck and use various tricks (softening / straight-through estimators)
to backpropagate through the sampling process.
"""
import pdb

import torch
from torch import nn, einsum
import torch.nn.functional as F

from scipy.cluster.vq import kmeans2

# -----------------------------------------------------------------------------

class VQVAEQuantize(nn.Module):
    """
    Neural Discrete Representation Learning, van den Oord et al. 2017
    https://arxiv.org/abs/1711.00937

    Follows the original DeepMind implementation
    https://github.com/deepmind/sonnet/blob/v2/sonnet/src/nets/vqvae.py
    https://github.com/deepmind/sonnet/blob/v2/examples/vqvae_example.ipynb
    """
    def __init__(self, num_hiddens, n_embed, embedding_dim):
        super().__init__()

        self.embedding_dim = embedding_dim
        self.n_embed = n_embed

        self.kld_scale = 10.0

        self.proj = nn.Conv2d(num_hiddens, embedding_dim, 1)
        self.embed = nn.Embedding(n_embed, embedding_dim)

        self.register_buffer('data_initialized', torch.zeros(1))

    def forward(self, z):
        B, C, H, W = z.size()

        # project and flatten out space, so (B, C, H, W) -> (B*H*W, C)
        z_e = self.proj(z)
        z_e = z_e.permute(0, 2, 3, 1) # make (B, H, W, C)
        flatten = z_e.reshape(-1, self.embedding_dim)

        # DeepMind def does not do this but I find I have to... ;\
        if self.training and self.data_initialized.item() == 0:
            print('running kmeans!!') # data driven initialization for the embeddings
            rp = torch.randperm(flatten.size(0))
            kd = kmeans2(flatten[rp[:3000]].data.cpu().numpy(), self.n_embed, minit='points')
            self.embed.weight.data.copy_(torch.from_numpy(kd[0]))
            self.data_initialized.fill_(1)
            # TODO: this won't work in multi-GPU setups

        dist = (
            flatten.pow(2).sum(1, keepdim=True)
            - 2 * flatten @ self.embed.weight.t()
            + self.embed.weight.pow(2).sum(1, keepdim=True).t()
        )
        _, ind = (-dist).max(1)
        ind = ind.view(B, H, W)

        # vector quantization cost that trains the embedding vectors
        z_q = self.embed_code(ind) # (B, H, W, C)
        commitment_cost = 0.25
        diff = commitment_cost * (z_q.detach() - z_e).pow(2).mean() + (z_q - z_e.detach()).pow(2).mean()
        diff *= self.kld_scale

        z_q = z_e + (z_q - z_e).detach() # noop in forward pass, straight-through gradient estimator in backward pass
        z_q = z_q.permute(0, 3, 1, 2) # stack encodings into channels again: (B, C, H, W)
        return z_q, diff, ind

    def embed_code(self, embed_id):
        return F.embedding(embed_id, self.embed.weight)


class GumbelQuantize(nn.Module):
    """
    Gumbel Softmax trick quantizer
    Categorical Reparameterization with Gumbel-Softmax, Jang et al. 2016
    https://arxiv.org/abs/1611.01144
    """
    def __init__(self, num_hiddens, n_embed, embedding_dim, straight_through=False):
        super().__init__()

        self.embedding_dim = embedding_dim
        self.n_embed = n_embed

        self.straight_through = straight_through
        self.temperature = 1.0
        self.kld_scale = 5e-4

        self.proj = nn.Conv2d(num_hiddens, n_embed, 1)
        self.embed = nn.Embedding(n_embed, embedding_dim)

    def forward(self, z):

        # force hard = True when we are in eval mode, as we must quantize
        hard = self.straight_through if self.training else True

        logits = self.proj(z)
        soft_one_hot = F.gumbel_softmax(logits, tau=self.temperature, dim=1, hard=hard)
        z_q = einsum('b n h w, n d -> b d h w', soft_one_hot, self.embed.weight)

        # + kl divergence to the prior loss
        qy = F.softmax(logits, dim=1)
        diff = self.kld_scale * torch.sum(qy * torch.log(qy * self.n_embed + 1e-10), dim=1).mean()

        ind = soft_one_hot.argmax(dim=1)
        return z_q, diff, ind

codebook/models/dvq/dvqmodel/quantize.py
YoungSeng/UnifiedGesture
"""
Defines the full (PyTorch Lightning module) VQVAE, which incorporates an
encoder, decoder and a quantize layer in the middle for the discrete bottleneck.
"""

import os
import math
import pdb
from argparse import ArgumentParser

import torch
from torch import nn, einsum
import torch.nn.functional as F

import pytorch_lightning as pl
# from pytorch_lightning.callbacks import ModelCheckpoint

# from data.cifar10 import CIFAR10Data
from dvqmodel.deepmind_enc_dec import DeepMindEncoder, DeepMindDecoder
from dvqmodel.openai_enc_dec import OpenAIEncoder, OpenAIDecoder
from dvqmodel.openai_enc_dec import Conv2d as PatchedConv2d
from dvqmodel.quantize import VQVAEQuantize, GumbelQuantize
from dvqmodel.loss import Normal, LogitLaplace

# -----------------------------------------------------------------------------


class VQVAE(pl.LightningModule):

    def __init__(self, args, input_channels=3):
        super().__init__()
        self.args = args

        # encoder/decoder module pair
        Encoder, Decoder = {
            'deepmind': (DeepMindEncoder, DeepMindDecoder),
            'openai': (OpenAIEncoder, OpenAIDecoder),
        }['deepmind']
        self.encoder = Encoder(input_channels=input_channels, n_hid=args.n_hid)
        self.decoder = Decoder(n_init=args.embedding_dim, n_hid=args.n_hid, output_channels=input_channels)

        # the quantizer module sandwiched between them, +contributes a KL(posterior || prior) loss to ELBO
        QuantizerModule = {
            'vqvae': VQVAEQuantize,
            'gumbel': GumbelQuantize,
        }[args.vq_flavor]
        self.quantizer = QuantizerModule(self.encoder.output_channels, args.num_embeddings, args.embedding_dim)

        # the data reconstruction loss in the ELBO
        ReconLoss = {
            'l2': Normal,
            'logit_laplace': LogitLaplace,
            # todo: add vqgan
        }[args.loss_flavor]
        self.recon_loss = ReconLoss

    def forward(self, x):
        z = self.encoder(x)
        z_q, latent_loss, ind = self.quantizer(z)
        x_hat = self.decoder(z_q)
        return x_hat, latent_loss + F.mse_loss(x, x_hat), ind

    def training_step(self, batch, batch_idx):
        x, y = batch # hate that i have to do this here in the model
        x = self.recon_loss.inmap(x)
        x_hat, latent_loss, ind = self.forward(x)
        recon_loss = self.recon_loss.nll(x, x_hat)
        loss = recon_loss + latent_loss
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch # hate that i have to do this here in the model
        x = self.recon_loss.inmap(x)
        x_hat, latent_loss, ind = self.forward(x)
        recon_loss = self.recon_loss.nll(x, x_hat)
        self.log('val_recon_loss', recon_loss, prog_bar=True)

        # debugging: cluster perplexity. when perplexity == num_embeddings then all clusters are used exactly equally
        encodings = F.one_hot(ind, self.quantizer.n_embed).float().reshape(-1, self.quantizer.n_embed)
        avg_probs = encodings.mean(0)
        perplexity = (-(avg_probs * torch.log(avg_probs + 1e-10)).sum()).exp()
        cluster_use = torch.sum(avg_probs > 0)
        self.log('val_perplexity', perplexity, prog_bar=True)
        self.log('val_cluster_use', cluster_use, prog_bar=True)

    def configure_optimizers(self):

        # separate out all parameters to those that will and won't experience regularizing weight decay
        decay = set()
        no_decay = set()
        whitelist_weight_modules = (torch.nn.Linear, torch.nn.Conv2d, torch.nn.ConvTranspose2d, PatchedConv2d)
        blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.BatchNorm2d, torch.nn.Embedding)
        for mn, m in self.named_modules():
            for pn, p in m.named_parameters():
                fpn = '%s.%s' % (mn, pn) if mn else pn # full param name

                if pn.endswith('bias'):
                    # all biases will not be decayed
                    no_decay.add(fpn)
                elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
                    # weights of whitelist modules will be weight decayed
                    decay.add(fpn)
                elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
                    # weights of blacklist modules will NOT be weight decayed
                    no_decay.add(fpn)

        # validate that we considered every parameter
        param_dict = {pn: p for pn, p in self.named_parameters()}
        inter_params = decay & no_decay
        union_params = decay | no_decay
        assert len(inter_params) == 0, "parameters %s made it into both decay/no_decay sets!" % (str(inter_params), )
        assert len(param_dict.keys() - union_params) == 0, "parameters %s were not separated into either decay/no_decay set!" \
                                                    % (str(param_dict.keys() - union_params), )

        # create the pytorch optimizer object
        optim_groups = [
            {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": 1e-4},
            {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
        ]
        optimizer = torch.optim.AdamW(optim_groups, lr=3e-4, betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-4)
        self.optimizer = optimizer

        return optimizer

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = ArgumentParser(parents=[parent_parser], add_help=False)
        # model type
        parser.add_argument("--vq_flavor", type=str, default='vqvae', choices=['vqvae', 'gumbel'])
        parser.add_argument("--enc_dec_flavor", type=str, default='deepmind', choices=['deepmind', 'openai'])
        parser.add_argument("--loss_flavor", type=str, default='l2', choices=['l2', 'logit_laplace'])
        # model size
        parser.add_argument("--num_embeddings", type=int, default=512, help="vocabulary size; number of possible discrete states")
        parser.add_argument("--embedding_dim", type=int, default=64, help="size of the vector of the embedding of each discrete token")
        parser.add_argument("--n_hid", type=int, default=64, help="number of channels controlling the size of the model")
        return parser


# -----------------------------------------------------------------------------
def cos_anneal(e0, e1, t0, t1, e):
    """ ramp from (e0, t0) -> (e1, t1) through a cosine schedule based on e \in [e0, e1] """
    alpha = max(0, min(1, (e - e0) / (e1 - e0))) # what fraction of the way through are we
    alpha = 1.0 - math.cos(alpha * math.pi/2) # warp through cosine
    t = alpha * t1 + (1 - alpha) * t0 # interpolate accordingly
    return t

"""
These ramps/decays follow DALL-E Appendix A.2 Training https://arxiv.org/abs/2102.12092
"""
'''
class DecayTemperature(pl.Callback):
    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx):
        # The relaxation temperature τ is annealed from 1 to 1/16 over the first 150,000 updates.
        t = cos_anneal(0, 150000, 1.0, 1.0/16, trainer.global_step)
        pl_module.quantizer.temperature = t

class RampBeta(pl.Callback):
    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx):
        # The KL weight β is increased from 0 to 6.6 over the first 5000 updates
        # "We divide the overall loss by 256 × 256 × 3, so that the weight of the KL term
        # becomes β/192, where β is the KL weight."
        # TODO: OpenAI uses 6.6/192 but kinda tricky to do the conversion here... about 5e-4 works for this repo so far... :\
        t = cos_anneal(0, 5000, 0.0, 5e-4, trainer.global_step)
        pl_module.quantizer.kld_scale = t

class DecayLR(pl.Callback):
    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx):
        # The step size is annealed from 1e10−4 to 1.25e10−6 over 1,200,000 updates. I use 3e-4
        t = cos_anneal(0, 1200000, 3e-4, 1.25e-6, trainer.global_step)
        for g in pl_module.optimizer.param_groups:
            g['lr'] = t

def cli_main():
    pl.seed_everything(1337)

    # -------------------------------------------------------------------------
    # arguments...
    parser = ArgumentParser()
    # training related
    parser = pl.Trainer.add_argparse_args(parser)
    # model related
    parser = VQVAE.add_model_specific_args(parser)
    # dataloader related
    parser.add_argument("--data_dir", type=str, default='/apcv/users/akarpathy/cifar10')
    parser.add_argument("--batch_size", type=int, default=128)
    parser.add_argument("--num_workers", type=int, default=8)
    # done!
    args = parser.parse_args()
    # -------------------------------------------------------------------------

    data = CIFAR10Data(args)
    model = VQVAE(args)

    # annealing schedules for lots of constants
    callbacks = []
    callbacks.append(ModelCheckpoint(monitor='val_recon_loss', mode='min'))
    callbacks.append(DecayLR())
    if args.vq_flavor == 'gumbel':
       callbacks.extend([DecayTemperature(), RampBeta()])
    trainer = pl.Trainer.from_argparse_args(args, callbacks=callbacks, max_steps=3000000)

    trainer.fit(model, data)
'''

class VQVAE(nn.Module):
    def __init__(self, input_channels):
        super().__init__()

        # encoder/decoder module pair
        Encoder, Decoder = {
            'deepmind': (DeepMindEncoder, DeepMindDecoder),
            'openai': (OpenAIEncoder, OpenAIDecoder),
        }['deepmind']
        self.encoder = Encoder(input_channels=input_channels)
        self.decoder = Decoder(out_channels=input_channels)

        # the quantizer module sandwiched between them, +contributes a KL(posterior || prior) loss to ELBO
        QuantizerModule = {
            'vqvae': VQVAEQuantize,
            'gumbel': GumbelQuantize,
        }['vqvae']

        self.quantizer = QuantizerModule(input_channels//2, 512, 64)

        # the data reconstruction loss in the ELBO
        ReconLoss = {
            'l2': Normal,
            'logit_laplace': LogitLaplace,
            # todo: add vqgan
        }['l2']
        self.recon_loss = ReconLoss

    def forward(self, x):       # [8, 30, 112]
        z = self.encoder(x)     # [8, 56, 30, 1]
        z_q, latent_loss, ind = self.quantizer(z)       # [batch, 64, 30, 1]
        x_hat = self.decoder(z_q)
        return x_hat, latent_loss + F.mse_loss(x, x_hat), ind


if __name__ == "__main__":
    '''
    pip install attr, pytorch_lightning
    cd models
    python dvq/vqvae_2.py
    '''
    # cli_main()

    model_ = VQVAE(input_channels=16*7)
    x = torch.randn(128, 30, 16*7)
    a, b, c = model_(x)
    pdb.set_trace()


codebook/models/dvq/vqvae_2.py
YoungSeng/UnifiedGesture
import pdb

import torch as t
import torch.nn as nn
from .resnet import Resnet, Resnet1D
from .utils.torch_utils import assert_shape

class EncoderConvBlock(nn.Module):
    def __init__(self, input_emb_width, output_emb_width, down_t,
                 stride_t, width, depth, m_conv,
                 dilation_growth_rate=1, dilation_cycle=None, zero_out=False,
                 res_scale=False):
        super().__init__()
        blocks = []
        filter_t, pad_t = stride_t * 2, stride_t // 2
        if down_t > 0:
            for i in range(down_t):
                block = nn.Sequential(
                    # nn.Conv1d(input_emb_width if i == 0 else width, width, filter_t, stride_t, pad_t, padding_mode='replicate'),
                    nn.Conv1d(input_emb_width if i == 0 else width, width, filter_t, stride_t, pad_t),
                    Resnet1D(width, depth, m_conv, dilation_growth_rate, dilation_cycle, zero_out, res_scale),
                )
                blocks.append(block)
            block = nn.Conv1d(width, output_emb_width, 3, 1, 1)
            # block = nn.Conv1d(width, output_emb_width, 3, 1, 1, padding_mode='replicate')
            blocks.append(block)
        else:       # 20230413
            block = nn.Conv1d(input_emb_width, output_emb_width, 3, 1, 1)
            blocks.append(block)
        self.model = nn.Sequential(*blocks)

    def forward(self, x):
        return self.model(x)

class DecoderConvBock(nn.Module):
    def __init__(self, input_emb_width, output_emb_width, down_t,
                 stride_t, width, depth, m_conv, dilation_growth_rate=1, dilation_cycle=None, zero_out=False, res_scale=False, reverse_decoder_dilation=False, checkpoint_res=False):
        super().__init__()
        blocks = []
        if down_t > 0:
            filter_t, pad_t = stride_t * 2, stride_t // 2
            block = nn.Conv1d(output_emb_width, width, 3, 1, 1)
            # block = nn.Conv1d(output_emb_width, width, 3, 1, 1, padding_mode='replicate')
            blocks.append(block)
            for i in range(down_t):
                block = nn.Sequential(
                    Resnet1D(width, depth, m_conv, dilation_growth_rate, dilation_cycle, zero_out=zero_out, res_scale=res_scale, reverse_dilation=reverse_decoder_dilation, checkpoint_res=checkpoint_res),
                    nn.ConvTranspose1d(width, input_emb_width if i == (down_t - 1) else width, filter_t, stride_t, pad_t)
                )
                blocks.append(block)
        self.model = nn.Sequential(*blocks)

    def forward(self, x):
        return self.model(x)

class Encoder(nn.Module):
    def __init__(self, input_emb_width, output_emb_width, levels, downs_t,
                 strides_t, **block_kwargs):
        super().__init__()
        self.input_emb_width = input_emb_width
        self.output_emb_width = output_emb_width
        self.levels = levels
        self.downs_t = downs_t
        self.strides_t = strides_t

        block_kwargs_copy = dict(**block_kwargs)
        if 'reverse_decoder_dilation' in block_kwargs_copy:
            del block_kwargs_copy['reverse_decoder_dilation']

        level_block = lambda level, down_t, stride_t: EncoderConvBlock(input_emb_width if level == 0 else output_emb_width,
                                                           output_emb_width,
                                                           down_t, stride_t,
                                                           **block_kwargs_copy)
        self.level_blocks = nn.ModuleList()
        iterator = zip(list(range(self.levels)), downs_t, strides_t)
        for level, down_t, stride_t in iterator:
            self.level_blocks.append(level_block(level, down_t, stride_t))

    def forward(self, x):
        N, T = x.shape[0], x.shape[-1]
        emb = self.input_emb_width
        assert_shape(x, (N, emb, T))
        xs = []

        # 64, 32, ...
        iterator = zip(list(range(self.levels)), self.downs_t, self.strides_t)
        for level, down_t, stride_t in iterator:
            level_block = self.level_blocks[level]
            x = level_block(x)
            emb, T = self.output_emb_width, T // (stride_t ** down_t)
            assert_shape(x, (N, emb, T))
            xs.append(x)

        return xs

class Decoder(nn.Module):
    def __init__(self, input_emb_width, output_emb_width, levels, downs_t,
                 strides_t, **block_kwargs):
        super().__init__()
        self.input_emb_width = input_emb_width
        self.output_emb_width = output_emb_width
        self.levels = levels

        self.downs_t = downs_t

        self.strides_t = strides_t

        level_block = lambda level, down_t, stride_t: DecoderConvBock(output_emb_width,
                                                          output_emb_width,
                                                          down_t, stride_t,
                                                          **block_kwargs)
        self.level_blocks = nn.ModuleList()
        iterator = zip(list(range(self.levels)), downs_t, strides_t)
        for level, down_t, stride_t in iterator:
            self.level_blocks.append(level_block(level, down_t, stride_t))

        self.out = nn.Conv1d(output_emb_width, input_emb_width, 3, 1, 1)
        # self.out = nn.Conv1d(output_emb_width, input_emb_width, 3, 1, 1, padding_mode='replicate')
    def forward(self, xs, all_levels=True):
        if all_levels:
            assert len(xs) == self.levels
        else:
            assert len(xs) == 1
        x = xs[-1]
        N, T = x.shape[0], x.shape[-1]
        emb = self.output_emb_width
        assert_shape(x, (N, emb, T))

        # 32, 64 ...
        iterator = reversed(list(zip(list(range(self.levels)), self.downs_t, self.strides_t)))
        for level, down_t, stride_t in iterator:
            level_block = self.level_blocks[level]
            x = level_block(x)
            emb, T = self.output_emb_width, T * (stride_t ** down_t)
            assert_shape(x, (N, emb, T))
            if level != 0 and all_levels:
                x = x + xs[level - 1]

        x = self.out(x)
        return x

codebook/models/encdec.py
YoungSeng/UnifiedGesture
import pdb

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions.normal import Normal
from torch.distributions import kl_divergence

from models.functions import vq, vq_st

def to_scalar(arr):
    if type(arr) == list:
        return [x.item() for x in arr]
    else:
        return arr.item()


def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        try:
            nn.init.xavier_uniform_(m.weight.data)
            m.bias.data.fill_(0)
        except AttributeError:
            print("Skipping initialization of ", classname)


class VAE(nn.Module):
    def __init__(self, input_dim, dim, z_dim):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(input_dim, dim, 4, 2, 1),
            nn.BatchNorm2d(dim),
            nn.ReLU(True),
            nn.Conv2d(dim, dim, 4, 2, 1),
            nn.BatchNorm2d(dim),
            nn.ReLU(True),
            nn.Conv2d(dim, dim, 5, 1, 0),
            nn.BatchNorm2d(dim),
            nn.ReLU(True),
            nn.Conv2d(dim, z_dim * 2, 3, 1, 0),
            nn.BatchNorm2d(z_dim * 2)
        )

        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(z_dim, dim, 3, 1, 0),
            nn.BatchNorm2d(dim),
            nn.ReLU(True),
            nn.ConvTranspose2d(dim, dim, 5, 1, 0),
            nn.BatchNorm2d(dim),
            nn.ReLU(True),
            nn.ConvTranspose2d(dim, dim, 4, 2, 1),
            nn.BatchNorm2d(dim),
            nn.ReLU(True),
            nn.ConvTranspose2d(dim, input_dim, 4, 2, 1),
            nn.Tanh()
        )

        self.apply(weights_init)

    def forward(self, x):
        mu, logvar = self.encoder(x).chunk(2, dim=1)

        q_z_x = Normal(mu, logvar.mul(.5).exp())
        p_z = Normal(torch.zeros_like(mu), torch.ones_like(logvar))
        kl_div = kl_divergence(q_z_x, p_z).sum(1).mean()

        x_tilde = self.decoder(q_z_x.rsample())
        return x_tilde, kl_div


class VQEmbedding(nn.Module):
    def __init__(self, K, D):
        super().__init__()
        self.embedding = nn.Embedding(K, D)
        self.embedding.weight.data.uniform_(-1./K, 1./K)

    def forward(self, z_e_x):
        z_e_x_ = z_e_x.permute(0, 2, 3, 1).contiguous()
        latents = vq(z_e_x_, self.embedding.weight)
        return latents

    def straight_through(self, z_e_x):
        z_e_x_ = z_e_x.permute(0, 2, 3, 1).contiguous()
        z_q_x_, indices = vq_st(z_e_x_, self.embedding.weight.detach())
        z_q_x = z_q_x_.permute(0, 3, 1, 2).contiguous()

        z_q_x_bar_flatten = torch.index_select(self.embedding.weight,
            dim=0, index=indices)
        z_q_x_bar_ = z_q_x_bar_flatten.view_as(z_e_x_)
        z_q_x_bar = z_q_x_bar_.permute(0, 3, 1, 2).contiguous()

        return z_q_x, z_q_x_bar


class ResBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.block = nn.Sequential(
            nn.ReLU(True),
            nn.Conv2d(dim, dim, 3, 1, 1),
            nn.BatchNorm2d(dim),
            nn.ReLU(True),
            nn.Conv2d(dim, dim, 1),
            nn.BatchNorm2d(dim)
        )

    def forward(self, x):
        return x + self.block(x)


class VectorQuantizedVAE(nn.Module):
    def __init__(self, input_dim, dim, K=512):
        super().__init__()

        self.encoder = nn.Sequential(
            nn.Linear(input_dim, dim),
            nn.ReLU(True),
        )

        self.codebook = VQEmbedding(K, dim)

        self.decoder = nn.Sequential(
            nn.Linear(dim, input_dim),
            nn.Tanh()
        )

        self.apply(weights_init)

    def encode(self, x):        # (2, 4800, 112)
        x = x.reshape(x.shape[0], x.shape[1], x.shape[2]//16, 16)
        z_e_x = self.encoder(x)
        z_e_x = z_e_x.permute(0, 3, 2, 1)
        latents = self.codebook(z_e_x)
        return latents

    def decode(self, latents):
        z_q_x = self.codebook.embedding(latents).permute(0, 3, 1, 2)  # (B, D, H, W)
        x_tilde = self.decoder(z_q_x)
        return x_tilde

    def forward(self, x):       # [2, 4800, 112]
        x = x.reshape(x.shape[0], x.shape[1], x.shape[2]//16, 16)
        z_e_x = self.encoder(x)    # [2, 256, 1200, 1]
        z_e_x = z_e_x.permute(0, 3, 2, 1)
        z_q_x_st, z_q_x = self.codebook.straight_through(z_e_x)
        x_tilde = self.decoder(z_q_x_st.permute(0, 3, 2, 1))
        x_tilde = x_tilde.reshape(x_tilde.shape[0], x_tilde.shape[1], -1)
        x = x.reshape(x.shape[0], x.shape[1], -1)
        return x_tilde, F.mse_loss(x_tilde, x) + F.mse_loss(z_q_x, z_e_x.detach()) + F.mse_loss(z_e_x, z_q_x.detach())


class GatedActivation(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x, y = x.chunk(2, dim=1)
        return F.tanh(x) * F.sigmoid(y)


class GatedMaskedConv2d(nn.Module):
    def __init__(self, mask_type, dim, kernel, residual=True, n_classes=10):
        super().__init__()
        assert kernel % 2 == 1, print("Kernel size must be odd")
        self.mask_type = mask_type
        self.residual = residual

        self.class_cond_embedding = nn.Embedding(
            n_classes, 2 * dim
        )

        kernel_shp = (kernel // 2 + 1, kernel)  # (ceil(n/2), n)
        padding_shp = (kernel // 2, kernel // 2)
        self.vert_stack = nn.Conv2d(
            dim, dim * 2,
            kernel_shp, 1, padding_shp
        )

        self.vert_to_horiz = nn.Conv2d(2 * dim, 2 * dim, 1)

        kernel_shp = (1, kernel // 2 + 1)
        padding_shp = (0, kernel // 2)
        self.horiz_stack = nn.Conv2d(
            dim, dim * 2,
            kernel_shp, 1, padding_shp
        )

        self.horiz_resid = nn.Conv2d(dim, dim, 1)

        self.gate = GatedActivation()

    def make_causal(self):
        self.vert_stack.weight.data[:, :, -1].zero_()  # Mask final row
        self.horiz_stack.weight.data[:, :, :, -1].zero_()  # Mask final column

    def forward(self, x_v, x_h, h):
        if self.mask_type == 'A':
            self.make_causal()

        h = self.class_cond_embedding(h)
        h_vert = self.vert_stack(x_v)
        h_vert = h_vert[:, :, :x_v.size(-1), :]
        out_v = self.gate(h_vert + h[:, :, None, None])

        h_horiz = self.horiz_stack(x_h)
        h_horiz = h_horiz[:, :, :, :x_h.size(-2)]
        v2h = self.vert_to_horiz(h_vert)

        out = self.gate(v2h + h_horiz + h[:, :, None, None])
        if self.residual:
            out_h = self.horiz_resid(out) + x_h
        else:
            out_h = self.horiz_resid(out)

        return out_v, out_h


class GatedPixelCNN(nn.Module):
    def __init__(self, input_dim=256, dim=64, n_layers=15, n_classes=10):
        super().__init__()
        self.dim = dim

        # Create embedding layer to embed input
        self.embedding = nn.Embedding(input_dim, dim)

        # Building the PixelCNN layer by layer
        self.layers = nn.ModuleList()

        # Initial block with Mask-A convolution
        # Rest with Mask-B convolutions
        for i in range(n_layers):
            mask_type = 'A' if i == 0 else 'B'
            kernel = 7 if i == 0 else 3
            residual = False if i == 0 else True

            self.layers.append(
                GatedMaskedConv2d(mask_type, dim, kernel, residual, n_classes)
            )

        # Add the output layer
        self.output_conv = nn.Sequential(
            nn.Conv2d(dim, 512, 1),
            nn.ReLU(True),
            nn.Conv2d(512, input_dim, 1)
        )

        self.apply(weights_init)

    def forward(self, x, label):
        shp = x.size() + (-1, )
        x = self.embedding(x.view(-1)).view(shp)  # (B, H, W, C)
        x = x.permute(0, 3, 1, 2)  # (B, C, W, W)

        x_v, x_h = (x, x)
        for i, layer in enumerate(self.layers):
            x_v, x_h = layer(x_v, x_h, label)

        return self.output_conv(x_h)

    def generate(self, label, shape=(8, 8), batch_size=64):
        param = next(self.parameters())
        x = torch.zeros(
            (batch_size, *shape),
            dtype=torch.int64, device=param.device
        )

        for i in range(shape[0]):
            for j in range(shape[1]):
                logits = self.forward(x, label)
                probs = F.softmax(logits[:, :, i, j], -1)
                x.data[:, i, j].copy_(
                    probs.multinomial(1).squeeze().data
                )
        return x


if __name__ == '__main__':
    '''
    python -m models.modules
    '''
    model = VectorQuantizedVAE(16, 256, 512)
    x = torch.randn(2, 4800, 4 * 16)
    x_tilde, loss = model(x)
    # code = model.encode(x)
    pdb.set_trace()

codebook/models/modules.py
YoungSeng/UnifiedGesture
import math
import torch.nn as nn
import models.utils.dist_adapter as dist
from .utils.checkpoint import checkpoint

class ResConvBlock(nn.Module):
    def __init__(self, n_in, n_state):
        super().__init__()
        self.model = nn.Sequential(
            nn.ReLU(),
            nn.Conv2d(n_in, n_state, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(n_state, n_in, 1, 1, 0),
        )

    def forward(self, x):
        return x + self.model(x)

class Resnet(nn.Module):
    def __init__(self, n_in, n_depth, m_conv=1.0):
        super().__init__()
        self.model = nn.Sequential(*[ResConvBlock(n_in, int(m_conv * n_in)) for _ in range(n_depth)])

    def forward(self, x):
        return self.model(x)

class ResConv1DBlock(nn.Module):
    def __init__(self, n_in, n_state, dilation=1, zero_out=False, res_scale=1.0):
        super().__init__()
        padding = dilation
        self.model = nn.Sequential(
            nn.ReLU(),
            nn.Conv1d(n_in, n_state, 3, 1, padding, dilation),
            # nn.Conv1d(n_in, n_state, 3, 1, padding, dilation, padding_mode='replicate'),
            nn.ReLU(),
            nn.Conv1d(n_state, n_in, 1, 1, 0,),
            # nn.Conv1d(n_state, n_in, 1, 1, 0, padding_mode='replicate'),
        )
        if zero_out:
            out = self.model[-1]
            nn.init.zeros_(out.weight)
            nn.init.zeros_(out.bias)
        self.res_scale = res_scale

    def forward(self, x):
        return x + self.res_scale * self.model(x)

class Resnet1D(nn.Module):
    def __init__(self, n_in, n_depth, m_conv=1.0, dilation_growth_rate=1, dilation_cycle=None, zero_out=False, res_scale=False, reverse_dilation=False, checkpoint_res=False):
        super().__init__()
        def _get_depth(depth):
            if dilation_cycle is None:
                return depth
            else:
                return depth % dilation_cycle
        blocks = [ResConv1DBlock(n_in, int(m_conv * n_in),
                                 dilation=dilation_growth_rate ** _get_depth(depth),
                                 zero_out=zero_out,
                                 res_scale=1.0 if not res_scale else 1.0 / math.sqrt(n_depth))
                  for depth in range(n_depth)]
        if reverse_dilation:
            blocks = blocks[::-1]
        self.checkpoint_res = checkpoint_res
        if self.checkpoint_res == 1:
            if dist.get_rank() == 0:
                print("Checkpointing convs")
            self.blocks = nn.ModuleList(blocks)
        else:
            self.model = nn.Sequential(*blocks)

    def forward(self, x):
        if self.checkpoint_res == 1:
            for block in self.blocks:
                x = checkpoint(block, (x, ), block.parameters(), True)
            return x
        else:
            return self.model(x)

codebook/models/resnet.py
YoungSeng/UnifiedGesture
import torch as t
import models.utils.dist_adapter as dist
from tqdm import tqdm
from datetime import date
import os
import sys

import sys
sys.path.append(os.path.dirname(os.path.dirname(sys.path[0])))
from configs.parse_args import parse_args

args = parse_args()
mydevice = t.device('cuda:' + args.gpu)

def def_tqdm(x):
    return tqdm(x, leave=True, file=sys.stdout, bar_format="{n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]")

def get_range(x):
    if dist.get_rank() == 0:
        return def_tqdm(x)
    else:
        return x

def init_logging(hps, local_rank, rank):
    logdir = f"{hps.local_logdir}/{hps.name}"
    if local_rank == 0:
        if not os.path.exists(logdir):
            os.makedirs(logdir)
        with open(logdir + 'argv.txt', 'w') as f:
            f.write(hps.argv + '\n')
        print("Logging to", logdir)
    logger = Logger(logdir, rank)
    metrics = Metrics()
    logger.add_text('hps', str(hps))
    return logger, metrics

def get_name(hps):
    name = ""
    for key, value in hps.items():
        name += f"{key}_{value}_"
    return name

def average_metrics(_metrics):
    metrics = {}
    for _metric in _metrics:
        for key, val in _metric.items():
            if key not in metrics:
                metrics[key] = []
            metrics[key].append(val)
    return {key: sum(vals)//len(vals) for key, vals in metrics.items()}

class Metrics:
    def __init__(self):
        self.sum = {}
        self.n = {}

    def update(self, tag, val, batch):
        # v is average value over batch
        # store total value and total batch, returns dist average
        sum = t.tensor(val * batch).float().to(mydevice)
        n = t.tensor(batch).float().to(mydevice)
        dist.all_reduce(sum)
        dist.all_reduce(n)
        sum = sum.item()
        n = n.item()
        self.sum[tag] = self.sum.get(tag, 0.0) + sum
        self.n[tag] = self.n.get(tag, 0.0) + n
        return sum / n

    def avg(self, tag):
        if tag in self.sum:
            return self.sum[tag] / self.n[tag]
        else:
            return 0.0

    def reset(self):
        self.sum = {}
        self.n = {}

class Logger:
    def __init__(self, logdir, rank):
        if rank == 0:
            from tensorboardX import SummaryWriter
            self.sw = SummaryWriter(f"{logdir}/logs")
        self.iters = 0
        self.rank = rank
        self.works = []
        self.logdir = logdir

    def step(self):
        self.iters += 1

    def flush(self):
        if self.rank == 0:
            self.sw.flush()

    def add_text(self, tag, text):
        if self.rank == 0:
            self.sw.add_text(tag, text, self.iters)

    def add_audios(self, tag, auds, sample_rate=22050, max_len=None, max_log=8):
        if self.rank == 0:
            for i in range(min(len(auds), max_log)):
                if max_len:
                    self.sw.add_audio(f"{i}/{tag}", auds[i][:max_len * sample_rate], self.iters, sample_rate)
                else:
                    self.sw.add_audio(f"{i}/{tag}", auds[i], self.iters, sample_rate)

    def add_audio(self, tag, aud, sample_rate=22050):
        if self.rank == 0:
            self.sw.add_audio(tag, aud, self.iters, sample_rate)

    def add_images(self, tag, img, dataformats="NHWC"):
        if self.rank == 0:
            self.sw.add_images(tag, img, self.iters, dataformats=dataformats)

    def add_image(self, tag, img):
        if self.rank == 0:
            self.sw.add_image(tag, img, self.iters)

    def add_scalar(self, tag, val):
        if self.rank == 0:
            self.sw.add_scalar(tag, val, self.iters)

    def get_range(self, loader):
        if self.rank == 0:
            self.trange = def_tqdm(loader)
        else:
            self.trange = loader
        return enumerate(self.trange)

    def close_range(self):
        if self.rank == 0:
            self.trange.close()

    def set_postfix(self, *args, **kwargs):
        if self.rank == 0:
            self.trange.set_postfix(*args, **kwargs)

    # For logging summaries of varies graph ops
    def add_reduce_scalar(self, tag, layer, val):
        if self.iters % 100 == 0:
            with t.no_grad():
                val = val.float().norm()/float(val.numel())
            work = dist.reduce(val, 0, async_op=True)
            self.works.append((tag, layer, val, work))

    def finish_reduce(self):
        for tag, layer, val, work in self.works:
            work.wait()
            if self.rank == 0:
                val = val.item()/dist.get_world_size()
                self.lw[layer].add_scalar(tag, val, self.iters)
        self.works = []

codebook/models/utils/logger.py
YoungSeng/UnifiedGesture
import pdb

import numpy as np
import torch as t
import torch.nn as nn

from .encdec import Encoder, Decoder, assert_shape
from .bottleneck import NoBottleneck, Bottleneck
from .utils.logger import average_metrics

import sys
[sys.path.append(i) for i in ['.', '..', './models/dvq']]
from configs.parse_args import parse_args

args = parse_args()
mydevice = t.device('cuda:' + args.gpu)

def dont_update(params):
    for param in params:
        param.requires_grad = False

def update(params):
    for param in params:
        param.requires_grad = True

def calculate_strides(strides, downs):
    return [stride ** down for stride, down in zip(strides, downs)]

# def _loss_fn(loss_fn, x_target, x_pred, hps):
#     if loss_fn == 'l1':
#         return t.mean(t.abs(x_pred - x_target)) / hps.bandwidth['l1']
#     elif loss_fn == 'l2':
#         return t.mean((x_pred - x_target) ** 2) / hps.bandwidth['l2']
#     elif loss_fn == 'linf':
#         residual = ((x_pred - x_target) ** 2).reshape(x_target.shape[0], -1)
#         values, _ = t.topk(residual, hps.linf_k, dim=1)
#         return t.mean(values) / hps.bandwidth['l2']
#     elif loss_fn == 'lmix':
#         loss = 0.0
#         if hps.lmix_l1:
#             loss += hps.lmix_l1 * _loss_fn('l1', x_target, x_pred, hps)
#         if hps.lmix_l2:
#             loss += hps.lmix_l2 * _loss_fn('l2', x_target, x_pred, hps)
#         if hps.lmix_linf:
#             loss += hps.lmix_linf * _loss_fn('linf', x_target, x_pred, hps)
#         return loss
#     else:
#         assert False, f"Unknown loss_fn {loss_fn}"
def _loss_fn(x_target, x_pred):
    return t.mean(t.abs(x_pred - x_target)) 


class VQVAE_ulr(nn.Module):
    def __init__(self, hps, input_dim):
        super().__init__()
        self.VQVAE_1 = VQVAE(hps, 16*4)       # 20230419
        self.VQVAE_2 = VQVAE(hps, 16*2)
        self.VQVAE_3 = VQVAE(hps, 16*1)

    def forward(self, x1, x2, x3):
        y1, loss1, _ = self.VQVAE_1(x1)
        y2, loss2, _ = self.VQVAE_2(x2)
        y3, loss3, _ = self.VQVAE_3(x3)
        return y1, y2, y3, loss1, loss2, loss3

    # def get_code(self, x1, x2, x3):
    #     zs1 = self.VQVAE_1.encode(x1)
    #     zs2 = self.VQVAE_2.encode(x2)
    #     zs3 = self.VQVAE_3.encode(x3)
    #     return zs1, zs2, zs3
    #
    # def code_2_pose(self, zs1, zs2, zs3):
    #     z1 = self.VQVAE_1.decode(zs1)
    #     z2 = self.VQVAE_2.decode(zs2)
    #     z3 = self.VQVAE_3.decode(zs3)
    #     return z1, z2, z3

    def get_code(self, x1):         # 20230426 we only use x1(upper)
        zs1, distance_list = self.VQVAE_1.encode(x1)
        return zs1, distance_list

    def code_2_pose(self, zs1):
        z1 = self.VQVAE_1.decode(zs1)
        return z1

    def freeze_drop(self):
        for m in self.modules():
            if isinstance(m, nn.Dropout):
                m.eval()


class VQVAE_ulr_2(nn.Module):
    def __init__(self, hps_1, hps_2):
        super().__init__()
        self.VQVAE_1 = VQVAE(hps_1, 16*4)       # 20230419
        self.VQVAE_2 = VQVAE(hps_2, 16*3)

    def forward(self, x1, x2):
        y1, loss1, _ = self.VQVAE_1(x1)
        y2, loss2, _ = self.VQVAE_2(x2)
        return y1, y2, loss1, loss2

    def get_code(self, x1, x2):
        zs1 = self.VQVAE_1.encode(x1)
        zs2 = self.VQVAE_2.encode(x2)
        return zs1, zs2

    def code_2_pose(self, zs1, zs2):
        z1 = self.VQVAE_1.decode(zs1)
        z2 = self.VQVAE_2.decode(zs2)
        return z1, z2

    def code_pose(self, zs1):
        z1 = self.VQVAE_1.decode(zs1)
        return z1


# from models.dvq.vqvae_2 import VQVAE as vqvae_2
# class VQVAE_ulr_easy(nn.Module):
#     def __init__(self):
#         super().__init__()
#         self.VQVAE_1 = vqvae_2(64)
#         self.VQVAE_2 = vqvae_2(32)
#         self.VQVAE_3 = vqvae_2(16)
#
#     def forward(self, x1, x2, x3):
#         y1, loss1, _ = self.VQVAE_1(x1)
#         y2, loss2, _ = self.VQVAE_2(x2)
#         y3, loss3, _ = self.VQVAE_3(x3)
#         return y1, y2, y3, loss1, loss2, loss3
#
#     def get_code(self, x1, x2, x3):
#         _, _, zs1 = self.VQVAE_1(x1)
#         _, _, zs2 = self.VQVAE_2(x2)
#         _, _, zs3 = self.VQVAE_3(x3)
#         return zs1, zs2, zs3


# from models.modules import VectorQuantizedVAE
# class VQVAE_ulr_easy(nn.Module):
#     def __init__(self):
#         super(VQVAE_ulr_easy, self).__init__()
#         self.VQVAE_1 = VectorQuantizedVAE(16, 256, 512)
#         self.VQVAE_2 = VectorQuantizedVAE(16, 256, 512)
#         self.VQVAE_3 = VectorQuantizedVAE(16, 256, 512)
#
#     def forward(self, x1, x2, x3):
#         y1, loss1 = self.VQVAE_1(x1)
#         y2, loss2 = self.VQVAE_2(x2)
#         y3, loss3 = self.VQVAE_3(x3)
#         return y1, y2, y3, loss1, loss2, loss3
#
#     def get_code(self, x1, x2, x3):
#         zs1 = self.VQVAE_1.encode(x1)
#         zs2 = self.VQVAE_2.encode(x2)
#         zs3 = self.VQVAE_3.encode(x3)
#         return zs1, zs2, zs3

'''
from vector_quantize_pytorch import VectorQuantize as ResidualVQ

class VQVAE_ulr_easy(nn.Module):
    def __init__(self):
        super(VQVAE_ulr_easy, self).__init__()
        self.VQVAE_1 = ResidualVQ(
                                dim = 4 * 16,
                                codebook_size = 512,
                                # num_quantizers = 1,
                                decay=0.8,
                                commitment_weight=1.0,
                                kmeans_init = True,   # set to True
                                kmeans_iters = 10     # number of kmeans iterations to calculate the centroids for the codebook on init
                                )
        self.VQVAE_2 = ResidualVQ(
                                dim = 2 * 16,
                                codebook_size = 512,
                                # num_quantizers = 1,
                                kmeans_init = True,   # set to True
                                kmeans_iters = 10     # number of kmeans iterations to calculate the centroids for the codebook on init
                                )
        self.VQVAE_3 = ResidualVQ(
                                dim = 1 * 16,
                                codebook_size = 512,
                                # num_quantizers = 1,
                                kmeans_init = True,   # set to True
                                kmeans_iters = 10     # number of kmeans iterations to calculate the centroids for the codebook on init
                                )

    def forward(self, x1, x2, x3):
        y1, _, loss1 = self.VQVAE_1(x1)
        y2, _, loss2 = self.VQVAE_2(x2)
        y3, _, loss3 = self.VQVAE_3(x3)
        return y1, y2, y3, loss1, loss2, loss3
'''


class VQVAE(nn.Module):
    def __init__(self, hps, input_dim=72):
        super().__init__()
        self.hps = hps

        input_shape = (hps.sample_length, input_dim)
        levels = hps.levels
        downs_t = hps.downs_t
        strides_t = hps.strides_t
        emb_width = hps.emb_width
        l_bins = hps.l_bins
        mu = hps.l_mu
        commit = hps.commit
        # spectral = hps.spectral
        # multispectral = hps.multispectral
        multipliers = hps.hvqvae_multipliers 
        use_bottleneck = hps.use_bottleneck
        if use_bottleneck:
            print('We use bottleneck!')
        else:
            print('We do not use bottleneck!')
        if not hasattr(hps, 'dilation_cycle'):
            hps.dilation_cycle = None
        block_kwargs = dict(width=hps.width, depth=hps.depth, m_conv=hps.m_conv, \
                        dilation_growth_rate=hps.dilation_growth_rate, \
                        dilation_cycle=hps.dilation_cycle, \
                        reverse_decoder_dilation=hps.vqvae_reverse_decoder_dilation)

        self.sample_length = input_shape[0]
        x_shape, x_channels = input_shape[:-1], input_shape[-1]
        self.x_shape = x_shape

        self.downsamples = calculate_strides(strides_t, downs_t)
        self.hop_lengths = np.cumprod(self.downsamples)
        self.z_shapes = z_shapes = [(x_shape[0] // self.hop_lengths[level],) for level in range(levels)]

        self.levels = levels

        if multipliers is None:
            self.multipliers = [1] * levels
        else:
            assert len(multipliers) == levels, "Invalid number of multipliers"
            self.multipliers = multipliers
        def _block_kwargs(level):
            this_block_kwargs = dict(block_kwargs)
            this_block_kwargs["width"] *= self.multipliers[level]
            this_block_kwargs["depth"] *= self.multipliers[level]
            return this_block_kwargs

        encoder = lambda level: Encoder(x_channels, emb_width, level + 1,
                                        downs_t[:level+1], strides_t[:level+1], **_block_kwargs(level))     # different from supplemental
        decoder = lambda level: Decoder(x_channels, emb_width, level + 1,
                                        downs_t[:level+1], strides_t[:level+1], **_block_kwargs(level))
        self.encoders = nn.ModuleList()
        self.decoders = nn.ModuleList()
        for level in range(levels):
            self.encoders.append(encoder(level))
            self.decoders.append(decoder(level))

        if use_bottleneck:
            self.bottleneck = Bottleneck(l_bins, emb_width, mu, levels)     # 512, 512, 0.99, 1
        else:
            self.bottleneck = NoBottleneck(levels)

        self.downs_t = downs_t
        self.strides_t = strides_t
        self.l_bins = l_bins
        self.commit = commit
        self.reg = hps.reg if hasattr(hps, 'reg') else 0
        self.acc = hps.acc if hasattr(hps, 'acc') else 0
        self.vel = hps.vel if hasattr(hps, 'vel') else 0
        if self.reg is 0:
            print('No motion regularization!')
        # self.spectral = spectral
        # self.multispectral = multispectral

    def preprocess(self, x):
        # x: NTC [-1,1] -> NCT [-1,1]
        assert len(x.shape) == 3
        x = x.permute(0,2,1).float()
        return x

    def postprocess(self, x):
        # x: NTC [-1,1] <- NCT [-1,1]
        x = x.permute(0,2,1)
        return x

    def _decode(self, zs, start_level=0, end_level=None):
        # Decode
        if end_level is None:
            end_level = self.levels
        assert len(zs) == end_level - start_level
        xs_quantised = self.bottleneck.decode(zs, start_level=start_level, end_level=end_level)
        assert len(xs_quantised) == end_level - start_level

        # Use only lowest level
        decoder, x_quantised = self.decoders[start_level], xs_quantised[0:1]
        x_out = decoder(x_quantised, all_levels=False)
        x_out = self.postprocess(x_out)
        return x_out

    def decode(self, zs, start_level=0, end_level=None, bs_chunks=1):
        z_chunks = [t.chunk(z, bs_chunks, dim=0) for z in zs]
        x_outs = []
        for i in range(bs_chunks):
            zs_i = [z_chunk[i] for z_chunk in z_chunks]
            x_out = self._decode(zs_i, start_level=start_level, end_level=end_level)
            x_outs.append(x_out)
        return t.cat(x_outs, dim=0)

    def _encode(self, x, start_level=0, end_level=None):
        # Encode
        if end_level is None:
            end_level = self.levels
        x_in = self.preprocess(x)
        xs = []
        for level in range(self.levels):
            encoder = self.encoders[level]
            x_out = encoder(x_in)
            xs.append(x_out[-1])
        zs = self.bottleneck.encode(xs)
        return zs[start_level:end_level]

    def encode(self, x, start_level=0, end_level=None, bs_chunks=1):
        x_chunks = t.chunk(x, bs_chunks, dim=0)
        zs_list = []
        distance_list = []
        for x_i in x_chunks:
            zs_i = self._encode(x_i, start_level=start_level, end_level=end_level)
            zs_i, distance = zs_i[0]
            zs_list.append([zs_i])
            distance_list.append(distance)
        zs = [t.cat(zs_level_list, dim=0) for zs_level_list in zip(*zs_list)]
        return zs, distance_list

    def sample(self, n_samples):
        zs = [t.randint(0, self.l_bins, size=(n_samples, *z_shape), device=mydevice) for z_shape in self.z_shapes]
        return self.decode(zs)

    def forward(self, x):       # (32, 240, 45)
        metrics = {}

        N = x.shape[0]

        # Encode/Decode
        x_in = self.preprocess(x)       # (32, 45, 240)
        xs = []
        for level in range(self.levels):
            encoder = self.encoders[level]
            x_out = encoder(x_in)
            xs.append(x_out[-1])
        # xs[0]: (32, 512, 30)
        zs, xs_quantised, commit_losses, quantiser_metrics = self.bottleneck(xs)
        '''
        zs[0]: (32, 30)
        xs_quantised[0]: (32, 512, 30)
        commit_losses[0]: 0.0009
        quantiser_metrics[0]: 
            fit 0.4646
            pn 0.0791
            entropy 5.9596
            used_curr 512
            usage 512
            dk 0.0006
        '''
        x_outs = []
        for level in range(self.levels):
            decoder = self.decoders[level]
            x_out = decoder(xs_quantised[level:level+1], all_levels=False)
            assert_shape(x_out, x_in.shape)
            x_outs.append(x_out)
        # x_outs[0]: (32, 45, 240)

        # Loss
        # def _spectral_loss(x_target, x_out, self.hps):
        #     if hps.use_nonrelative_specloss:
        #         sl = spectral_loss(x_target, x_out, self.hps) / hps.bandwidth['spec']
        #     else:
        #         sl = spectral_convergence(x_target, x_out, self.hps)
        #     sl = t.mean(sl)
        #     return sl

        # def _multispectral_loss(x_target, x_out, self.hps):
        #     sl = multispectral_loss(x_target, x_out, self.hps) / hps.bandwidth['spec']
        #     sl = t.mean(sl)
        #     return sl

        recons_loss = t.zeros(()).to(x.device)
        regularization = t.zeros(()).to(x.device)
        velocity_loss = t.zeros(()).to(x.device)
        acceleration_loss = t.zeros(()).to(x.device)
        # spec_loss = t.zeros(()).to(x.device)
        # multispec_loss = t.zeros(()).to(x.device)
        # x_target = audio_postprocess(x.float(), self.hps)
        x_target = x.float()

        for level in reversed(range(self.levels)):
            x_out = self.postprocess(x_outs[level])     # (32, 240, 45)
            # x_out = audio_postprocess(x_out, self.hps)
            

            # this_recons_loss = _loss_fn(loss_fn, x_target, x_out, hps)
            this_recons_loss = _loss_fn(x_target, x_out)
            # this_spec_loss = _spectral_loss(x_target, x_out, hps)
            # this_multispec_loss = _multispectral_loss(x_target, x_out, hps)
            metrics[f'recons_loss_l{level + 1}'] = this_recons_loss
            # metrics[f'spectral_loss_l{level + 1}'] = this_spec_loss
            # metrics[f'multispectral_loss_l{level + 1}'] = this_multispec_loss
            recons_loss += this_recons_loss
            # spec_loss += this_spec_loss
            # multispec_loss += this_multispec_loss
            regularization += t.mean((x_out[:, 2:] + x_out[:, :-2] - 2 * x_out[:, 1:-1])**2)

            velocity_loss +=  _loss_fn( x_out[:, 1:] - x_out[:, :-1], x_target[:, 1:] - x_target[:, :-1])
            acceleration_loss +=  _loss_fn(x_out[:, 2:] + x_out[:, :-2] - 2 * x_out[:, 1:-1], x_target[:, 2:] + x_target[:, :-2] - 2 * x_target[:, 1:-1])
        # if not hasattr(self.)
        commit_loss = sum(commit_losses)
        # loss = recons_loss + self.spectral * spec_loss + self.multispectral * multispec_loss + self.commit * commit_loss
        # pdb.set_trace()
        loss = recons_loss +  commit_loss * self.commit + self.reg * regularization + self.vel * velocity_loss + self.acc * acceleration_loss
        '''     x:-0.8474 ~ 1.1465
        0.2080
        5.5e-5 * 0.02
        0.0011
        0.0163 * 1
        0.0274 * 1
        '''

        with t.no_grad():
            # sc = t.mean(spectral_convergence(x_target, x_out, hps))
            # l2_loss = _loss_fn("l2", x_target, x_out, hps)
            l1_loss = _loss_fn(x_target, x_out)
            
            # linf_loss = _loss_fn("linf", x_target, x_out, hps)

        quantiser_metrics = average_metrics(quantiser_metrics)

        metrics.update(dict(
            recons_loss=recons_loss,
            # spectral_loss=spec_loss,
            # multispectral_loss=multispec_loss,
            # spectral_convergence=sc,
            # l2_loss=l2_loss,
            l1_loss=l1_loss,
            # linf_loss=linf_loss,
            commit_loss=commit_loss,
            regularization=regularization,
            velocity_loss=velocity_loss,
            acceleration_loss=acceleration_loss,
            **quantiser_metrics))

        for key, val in metrics.items():
            metrics[key] = val.detach()

        return x_out, loss, metrics




if __name__ == '__main__':
    '''
    cd codebook/
    python -m models.vqvae --config=./configs/codebook.yml --train --no_cuda 2 --gpu 2
    '''
    import yaml
    from pprint import pprint
    from easydict import EasyDict

    with open(args.config) as f:
        config = yaml.safe_load(f)

    for k, v in vars(args).items():
        config[k] = v
    pprint(config)

    config = EasyDict(config)

    x = t.rand(1, 2200, 16 * 4).to(mydevice)
    model = VQVAE_ulr(config.VQVAE, 16 * 7)     # n_joints * n_chanels

    model = nn.DataParallel(model, device_ids=[eval(i) for i in config.no_cuda])
    model = model.to(mydevice)
    model = model.train()
    # output, loss, metrics = model(x)
    zs, distance = model.module.get_code(x)     # [32, 8]
    pdb.set_trace()

    # from torch import optim
    # model = VQVAE_ulr_easy().to(mydevice)
    # optimizer = optim.Adam(model.parameters(), lr=0.0003, betas=[0.5, 0.999])
    # x1 = t.rand(32, 30, 4 * 16).to(mydevice)
    # x2 = t.rand(32, 30, 2 * 16).to(mydevice)
    # x3 = t.rand(32, 30, 1 * 16).to(mydevice)
    # y1, y2, y3, loss1, loss2, loss3 = model(x1, x2, x3)
    # pdb.set_trace()

codebook/models/vqvae.py
YoungSeng/UnifiedGesture
'''
pip install easydict
pip install configargparse
'''

import pdb
# import nni

import logging
logging.getLogger().setLevel(logging.INFO)

from torch.utils.data import DataLoader
from data_loader.lmdb_data_loader import TrinityDataset
import time
import torch
import torch.nn as nn
import yaml
from pprint import pprint
from easydict import EasyDict
from configs.parse_args import parse_args
from models.vqvae import VQVAE_ulr
import os
from torch import optim
import itertools
from datetime import datetime
from torch.utils.tensorboard import SummaryWriter
from pathlib import Path

# params = nni.get_next_parameter()
args = parse_args()
mydevice = torch.device('cuda:' + args.gpu)


def evaluate_testset(model, test_data_loader):
    start = time.time()
    model = model.eval()
    # tot_euclidean_error = 0
    # tot_eval_nums = 0
    upper_euclidean_errors = []
    lower_euclidean_errors = []
    root_euclidean_errors = []
    joint_channel = 16
    with torch.no_grad():
        for iter_idx, data in enumerate(test_data_loader, 0):
            upper, lower, root, _, _ = data  # （batch, 40, 135）
            upper, lower, root = upper.to(mydevice), lower.to(mydevice), root.to(mydevice)
            # pose_seq_eval = target_vec.to(mydevice)
            b, t, c = upper.size()
            upper_, lower_, root_, _, _, _ = model(upper, lower, root)
            diff_upper = (upper_ - upper).view(b, t, c // joint_channel, joint_channel)
            diff_lower = (lower_ - lower).view(b, t, c // joint_channel // 2, joint_channel)
            diff_root = (root_ - root).view(b, t, c // joint_channel // 4, joint_channel)
            # tot_euclidean_error += torch.mean(torch.sqrt(torch.sum(diff_upper ** 2, dim=3)))
            # tot_euclidean_error += torch.mean(torch.sqrt(torch.sum(diff_lower ** 2, dim=3)))
            # tot_euclidean_error += torch.mean(torch.sqrt(torch.sum(diff_root ** 2, dim=3)))
            # tot_eval_nums += 1
            upper_euclidean_errors.append(torch.mean(torch.sqrt(torch.sum(diff_upper ** 2, dim=3))))
            lower_euclidean_errors.append(torch.mean(torch.sqrt(torch.sum(diff_lower ** 2, dim=3))))
            root_euclidean_errors.append(torch.mean(torch.sqrt(torch.sum(diff_root ** 2, dim=3))))
        # print(tot_euclidean_error / (tot_eval_nums * 1.0))
        print('generation took {:.2} s'.format(time.time() - start))
    model.train()
    return torch.mean(torch.stack(upper_euclidean_errors)).data.cpu().numpy(), \
              torch.mean(torch.stack(lower_euclidean_errors)).data.cpu().numpy(), \
                torch.mean(torch.stack(root_euclidean_errors)).data.cpu().numpy()


def main(args):
    # dataset
    train_dataset = TrinityDataset(args.train_data_path,
                                   n_poses=args.n_poses,
                                   subdivision_stride=args.subdivision_stride,
                                   pose_resampling_fps=args.motion_resampling_framerate)
    train_loader = DataLoader(dataset=train_dataset, batch_size=args.batch_size,
                              shuffle=True, drop_last=True, num_workers=args.loader_workers, pin_memory=True)

    val_dataset = TrinityDataset(args.val_data_path,
                                       n_poses=args.n_poses,
                                       subdivision_stride=args.subdivision_stride,
                                       pose_resampling_fps=args.motion_resampling_framerate)
    test_loader = DataLoader(dataset=val_dataset, batch_size=args.batch_size,
                             shuffle=False, drop_last=True, num_workers=args.loader_workers, pin_memory=True)

    logging.info('len of train loader:{}, len of test loader:{}'.format(len(train_loader), len(test_loader)))

    model = VQVAE_ulr(args.VQVAE, 7 * 16)  # n_joints * n_chanels
    # model = VQVAE_ulr_easy()
    model = nn.DataParallel(model, device_ids=[eval(i) for i in config.no_cuda])
    model = model.to(mydevice)

    best_val_loss = (1e+2, 0)  # value, epoch

    if not os.path.exists(args.model_save_path):
        os.mkdir(args.model_save_path)

    # optimizer = optim.Adam(itertools.chain(model.module.parameters()), lr=args.lr, betas=args.betas)
    optimizer = optim.Adam(model.module.parameters(), lr=args.lr, betas=args.betas)
    schedular = optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.milestones, gamma=args.gamma)

    updates = 0
    total = len(train_loader)

    tb_path = args.name + '_' + str(datetime.now().strftime('%Y%m%d_%H%M%S'))
    tb_writer = SummaryWriter(log_dir=str(Path(args.model_save_path).parent / 'tensorboard_runs' / tb_path))


    for epoch in range(1, args.epochs + 1):
        logging.info(f'Epoch: {epoch}')
        i = 0
        upper_mean, lower_mean, root_mean = evaluate_testset(model, test_loader)
        logging.info('upper mean on validation: {:.3f}, lower mean on validation: {:.3f}, '
                     'root mean on validation: {:.3f}'.format(upper_mean, lower_mean, root_mean))
        # tb_writer.add_scalar('upper mean/validation', upper_mean, epoch)
        total_loss = (upper_mean + lower_mean + root_mean) / 3.0
        is_best = total_loss < best_val_loss[0]
        if is_best:
            logging.info(' *** BEST VALIDATION LOSS : {:.3f}'.format(total_loss))
            best_val_loss = (total_loss, epoch)
        else:
            logging.info(' best validation loss so far: {:.3f} at EPOCH {}'.format(best_val_loss[0], best_val_loss[1]))

        if is_best or (epoch % args.save_per_epochs == 0):
            if is_best:
                save_name = '{}/{}_checkpoint_best.bin'.format(args.model_save_path, args.name)
            else:
                save_name = '{}/{}_checkpoint_{:03d}.bin'.format(args.model_save_path, args.name, epoch)

            torch.save({
                'args': args, "epoch": epoch, 'model_dict': model.state_dict()
            }, save_name)
            logging.info('Saved the checkpoint')

        # train model
        model = model.train()
        start = datetime.now()
        for batch_i, batch in enumerate(train_loader, 0):
            upper, lower, root, _, _ = batch
            upper, lower, root = upper.to(mydevice), lower.to(mydevice), root.to(mydevice)
            # pose_seq = target_vec.to(mydevice)      # (b, 240, 15*3)
            optimizer.zero_grad()
            _, _, _, loss1, loss2, loss3 = model(upper, lower, root)
            # write to tensorboard
            # tb_writer.add_scalar('loss' + '/train', loss, updates)
            # for key in metrics.keys():
            #     tb_writer.add_scalar(key + '/train', metrics[key], updates)
            loss = loss1 + loss2 + loss3
            loss.backward()
            optimizer.step()
            # log
            stats = {'updates': updates, 'loss1': loss1.item(), 'loss2': loss2.item(), 'loss3': loss3.item()}
            stats_str = ' '.join(f'{key}[{val:.8f}]' for key, val in stats.items())
            i += 1
            remaining = str((datetime.now() - start) / i * (total - i))
            remaining = remaining.split('.')[0]
            logging.info(f'> epoch [{epoch}] updates[{i}] {stats_str} eta[{remaining}]')
            # if i == total:
            #     logging.debug('\n')
            #     logging.debug(f'elapsed time: {str(datetime.now() - start).split(".")[0]}')
            updates += 1
        schedular.step()
    tb_writer.close()
    # print best losses
    logging.info('--------- Final best loss values ---------')
    logging.info('diff mean: {:.3f} at EPOCH {}'.format(best_val_loss[0], best_val_loss[1]))


if __name__ == '__main__':
    '''
    pip install nni
    cd codebook/
    windows: ssh -p 22 -L 8080:127.0.0.1:8080 yangsc21@server15.mjrc.ml (f8550a408967d2dc99a180893fe0b007)
    python train.py --config=./configs/codebook.yml --train --gpu 2
    '''

    with open(args.config) as f:
        config = yaml.safe_load(f)

    for k, v in vars(args).items():
        config[k] = v
    pprint(config)

    config = EasyDict(config)
    config.no_cuda = config.gpu
    main(config)

codebook/train.py
YoungSeng/UnifiedGesture
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""Various utilities."""

from hashlib import sha256
from pathlib import Path
import typing as tp

import torch
import torchaudio


def _linear_overlap_add(frames: tp.List[torch.Tensor], stride: int):
    # Generic overlap add, with linear fade-in/fade-out, supporting complex scenario
    # e.g., more than 2 frames per position.
    # The core idea is to use a weight function that is a triangle,
    # with a maximum value at the middle of the segment.
    # We use this weighting when summing the frames, and divide by the sum of weights
    # for each positions at the end. Thus:
    #   - if a frame is the only one to cover a position, the weighting is a no-op.
    #   - if 2 frames cover a position:
    #          ...  ...
    #         /   \/   \
    #        /    /\    \
    #            S  T       , i.e. S offset of second frame starts, T end of first frame.
    # Then the weight function for each one is: (t - S), (T - t), with `t` a given offset.
    # After the final normalization, the weight of the second frame at position `t` is
    # (t - S) / (t - S + (T - t)) = (t - S) / (T - S), which is exactly what we want.
    #
    #   - if more than 2 frames overlap at a given point, we hope that by induction
    #      something sensible happens.
    assert len(frames)
    device = frames[0].device
    dtype = frames[0].dtype
    shape = frames[0].shape[:-1]
    total_size = stride * (len(frames) - 1) + frames[-1].shape[-1]

    frame_length = frames[0].shape[-1]
    t = torch.linspace(0, 1, frame_length + 2, device=device, dtype=dtype)[1: -1]
    weight = 0.5 - (t - 0.5).abs()

    sum_weight = torch.zeros(total_size, device=device, dtype=dtype)
    out = torch.zeros(*shape, total_size, device=device, dtype=dtype)
    offset: int = 0

    for frame in frames:
        frame_length = frame.shape[-1]
        out[..., offset:offset + frame_length] += weight[:frame_length] * frame
        sum_weight[offset:offset + frame_length] += weight[:frame_length]
        offset += stride
    assert sum_weight.min() > 0
    return out / sum_weight


def _get_checkpoint_url(root_url: str, checkpoint: str):
    if not root_url.endswith('/'):
        root_url += '/'
    return root_url + checkpoint


def _check_checksum(path: Path, checksum: str):
    sha = sha256()
    with open(path, 'rb') as file:
        while True:
            buf = file.read(2**20)
            if not buf:
                break
            sha.update(buf)
    actual_checksum = sha.hexdigest()[:len(checksum)]
    if actual_checksum != checksum:
        raise RuntimeError(f'Invalid checksum for file {path}, '
                           f'expected {checksum} but got {actual_checksum}')


def convert_audio(wav: torch.Tensor, sr: int, target_sr: int, target_channels: int):
    assert wav.dim() >= 2, "Audio tensor must have at least 2 dimensions"
    assert wav.shape[-2] in [1, 2], "Audio must be mono or stereo."
    *shape, channels, length = wav.shape
    if target_channels == 1:
        wav = wav.mean(-2, keepdim=True)
    elif target_channels == 2:
        wav = wav.expand(*shape, target_channels, length)
    elif channels == 1:
        wav = wav.expand(target_channels, -1)
    else:
        raise RuntimeError(f"Impossible to convert from {channels} to {target_channels}")
    wav = torchaudio.transforms.Resample(sr, target_sr)(wav)
    return wav


def save_audio(wav: torch.Tensor, path: tp.Union[Path, str],
               sample_rate: int, rescale: bool = False):
    limit = 0.99
    mx = wav.abs().max()
    if rescale:
        wav = wav * min(limit / mx, 1)
    else:
        wav = wav.clamp(-limit, limit)
    torchaudio.save(str(path), wav, sample_rate=sample_rate, encoding='PCM_S', bits_per_sample=16)

conditional_gpt/encodec/utils.py
YoungSeng/UnifiedGesture
import configargparse
import argparse

def str2bool(v):
    """ from https://stackoverflow.com/a/43357954/1361529 """
    if isinstance(v, bool):
       return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise configargparse.ArgumentTypeError('Boolean value expected.')


def parse_args():
    # parser = configargparse.ArgParser()
    # parser.add('-c', '--config', required=True, is_config_file=True, help='Config file path')
    #
    # # dataset
    # parser.add("--train_data_path", action="append", required=True)
    # parser.add("--val_data_path", action="append", required=True)
    # parser.add("--data_mean", action="append", type=float, nargs='*')
    # parser.add("--data_std", action="append", type=float, nargs='*')
    # parser.add("--n_poses", type=int, default=50)
    # parser.add("--subdivision_stride", type=int, default=5)
    # parser.add("--motion_resampling_framerate", type=int, default=24)
    # parser.add("--batch_size", type=int, default=50)
    # parser.add("--loader_workers", type=int, default=0)
    # parser.add("--epochs", type=int, default=10)

    parser = argparse.ArgumentParser(description='Codebook')
    parser.add_argument('--config', default='./configs/codebook.yml')
    parser.add_argument('--lr', default=0.00003, type=float)
    # parser.add_argument('--lr', default=params['lr'], type=float)       # 20230420  0.00003 -> 0.00004

    # codebook
    parser.add_argument('--gpu', type=str, default='2')
    parser.add_argument('--no_cuda', type=list, default=['2'])
    # exclusive arguments
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('--train', action='store_true')


    # parser.add("--test_data_path", action="append", required=False)
    # parser.add("--model_save_path", required=True)
    # parser.add("--random_seed", type=int, default=-1)
    #
    # # word embedding
    # parser.add("--wordembed_path", type=str, default=None)
    # parser.add("--wordembed_dim", type=int, default=200)
    #
    # # model
    # parser.add("--model", type=str, required=True)
    #
    #
    # parser.add("--dropout_prob", type=float, default=0.3)
    # parser.add("--n_layers", type=int, default=2)
    # parser.add("--hidden_size", type=int, default=200)
    #
    # parser.add("--n_pre_poses", type=int, default=5)
    #
    # # training
    # parser.add("--learning_rate", type=float, default=0.001)
    # parser.add("--loss_l1_weight", type=float, default=50)
    # parser.add("--loss_cont_weight", type=float, default=0.1)
    # parser.add("--loss_var_weight", type=float, default=0.01)

    args = parser.parse_args()
    return args

configs/parse_args.py
YoungSeng/UnifiedGesture
""" create data samples """
import lmdb
import math
import numpy as np
import pyarrow


class DataPreprocessor:
    def __init__(self, clip_lmdb_dir, out_lmdb_dir, n_poses, subdivision_stride, pose_resampling_fps):
        self.n_poses = n_poses
        self.subdivision_stride = subdivision_stride
        self.skeleton_resampling_fps = pose_resampling_fps

        self.src_lmdb_env = lmdb.open(clip_lmdb_dir, readonly=True, lock=False)
        with self.src_lmdb_env.begin() as txn:
            self.n_videos = txn.stat()['entries']

        self.audio_sample_length = int(self.n_poses / self.skeleton_resampling_fps * 75)
        self.lower_sample_length = int(self.n_poses * 2)

        # create db for samples
        map_size = 1024 * 1024 * 9  # in TB
        map_size <<= 20  # in B
        self.dst_lmdb_env = lmdb.open(out_lmdb_dir, map_size=map_size)
        self.n_out_samples = 0

    def run(self):
        src_txn = self.src_lmdb_env.begin(write=False)

        # sampling and normalization
        cursor = src_txn.cursor()
        for key, value in cursor:
            video = pyarrow.deserialize(value)
            vid = video['vid']
            clips = video['clips']
            for clip_idx, clip in enumerate(clips):
                self._sample_from_clip(vid, clip)

        # print stats
        with self.dst_lmdb_env.begin() as txn:
            print('no. of samples: ', txn.stat()['entries'])

        # close db
        self.src_lmdb_env.close()
        self.dst_lmdb_env.sync()
        self.dst_lmdb_env.close()

    def _sample_from_clip(self, vid, clip):
        clip_speech = clip['speech']
        clip_upper = clip['upper']
        clip_lower = clip['lower']

        # divide
        aux_info = []
        sample_speech_list = []
        sample_gesture_list = []
        sample_lower_list = []

        # MINLEN = min(len(clip_skeleton), int(len(clip_audio_raw) * 60 / 16000), len(clip_codes_raw) * 8)
        MINLEN = len(clip_upper)

        num_subdivision = math.floor(
            (MINLEN - self.n_poses)
            / self.subdivision_stride) + 1  # floor((K - (N+M)) / S) + 1

        for i in range(num_subdivision):
            start_idx = i * self.subdivision_stride
            fin_idx = start_idx + self.n_poses
            sample_gesture = clip_upper[start_idx:fin_idx]

            subdivision_start_time = start_idx / self.skeleton_resampling_fps
            subdivision_end_time = fin_idx / self.skeleton_resampling_fps

            # raw audio
            audio_start = math.floor(start_idx / len(clip_upper) * len(clip_speech))
            audio_end = audio_start + self.audio_sample_length
            sample_audio = clip_speech[audio_start:audio_end]
            # raw lower
            lower_start = math.floor(start_idx / len(clip_upper) * len(clip_lower))
            lower_end = lower_start + self.lower_sample_length
            sample_lower = clip_lower[lower_start:lower_end]

            motion_info = {'vid': vid,
                           'start_frame_no': start_idx,
                           'end_frame_no': fin_idx,
                           'start_time': subdivision_start_time,
                           'end_time': subdivision_end_time}

            sample_gesture_list.append(sample_gesture)
            sample_speech_list.append(sample_audio)
            sample_lower_list.append(sample_lower)

            aux_info.append(motion_info)

        if len(sample_gesture_list) > 0:
            with self.dst_lmdb_env.begin(write=True) as txn:
                for gesture, lower, speech, aux in zip(sample_gesture_list, sample_lower_list, sample_speech_list, aux_info):
                    speech = np.asarray(speech)
                    lower = np.asarray(lower)
                    gesture = np.asarray(gesture)

                    # save
                    k = '{:010}'.format(self.n_out_samples).encode('ascii')
                    v = [speech, gesture, lower, aux]
                    v = pyarrow.serialize(v).to_buffer()
                    txn.put(k, v)
                    self.n_out_samples += 1

data_loader/data_preprocessor.py
YoungSeng/UnifiedGesture
import logging
import os
import pdb

import numpy as np
import lmdb as lmdb
import torch
from torch.nn.utils.rnn import pad_sequence

from torch.utils.data import Dataset
from torch.utils.data.dataloader import default_collate

import sys
[sys.path.append(i) for i in ['.', '..']]

from data_loader.data_preprocessor import DataPreprocessor
import pyarrow


def word_seq_collate_fn(data):
    """ collate function for loading word sequences in variable lengths """
    # sort a list by sequence length (descending order) to use pack_padded_sequence
    data.sort(key=lambda x: len(x[0]), reverse=True)

    # separate source and target sequences
    word_seq, poses_seq, audio, aux_info = zip(*data)

    # merge sequences
    words_lengths = torch.LongTensor([len(x) for x in word_seq])
    word_seq = pad_sequence(word_seq, batch_first=True).long()

    poses_seq = default_collate(poses_seq)
    audio = default_collate(audio)
    aux_info = {key: default_collate([d[key] for d in aux_info]) for key in aux_info[0]}

    return word_seq, words_lengths, poses_seq, audio, aux_info


class TrinityDataset(Dataset):
    def __init__(self, lmdb_dir, n_poses, subdivision_stride, pose_resampling_fps):
        self.lmdb_dir = lmdb_dir
        self.n_poses = n_poses
        self.subdivision_stride = subdivision_stride
        self.skeleton_resampling_fps = pose_resampling_fps
        self.lang_model = None

        logging.info("Reading data '{}'...".format(lmdb_dir))
        preloaded_dir = lmdb_dir + '_cache'
        if not os.path.exists(preloaded_dir):
            data_sampler = DataPreprocessor(lmdb_dir, preloaded_dir, n_poses,
                                            subdivision_stride, pose_resampling_fps)
            data_sampler.run()
        else:
            logging.info('Found pre-loaded samples from {}'.format(preloaded_dir))

        # init lmdb
        # map_size = 1024 * 20  # in MB
        # map_size <<= 20  # in B
        self.lmdb_env = lmdb.open(preloaded_dir, readonly=True, lock=False)      # default 10485760
        with self.lmdb_env.begin() as txn:
            self.n_samples = txn.stat()['entries']

    def __len__(self):
        return self.n_samples

    def __getitem__(self, idx):
        with self.lmdb_env.begin(write=False) as txn:
            key = '{:010}'.format(idx).encode('ascii')
            sample = txn.get(key)

            sample = pyarrow.deserialize(sample)
            speech_seq, gesture_seq, lower_seq, aux_info = sample

        # to tensors
        speech_seq = torch.from_numpy(speech_seq).long()
        gesture_seq = torch.from_numpy(gesture_seq).long()
        lower_seq = torch.from_numpy(lower_seq).float()

        return speech_seq, gesture_seq, lower_seq, aux_info


if __name__ == '__main__':
    '''
    python data_loader/lmdb_data_loader.py --config=./configs/pretrain.yml --train --no_cuda 3 --gpu 3
    '''

    from configs.parse_args import parse_args
    import os
    import yaml
    from pprint import pprint
    from easydict import EasyDict
    from torch.utils.data import DataLoader

    args = parse_args()

    with open(args.config) as f:
        config = yaml.safe_load(f)

    for k, v in vars(args).items():
        config[k] = v
    pprint(config)

    args = EasyDict(config)

    train_dataset = TrinityDataset(args.train_data_path,
                                   n_poses=args.n_poses,
                                   subdivision_stride=args.subdivision_stride,
                                   pose_resampling_fps=args.motion_resampling_framerate)
    train_loader = DataLoader(dataset=train_dataset, batch_size=128,
                              shuffle=True, drop_last=True, num_workers=args.loader_workers, pin_memory=True)

    print(len(train_loader))        # 247
    for batch_i, batch in enumerate(train_loader, 0):
        speech, gesture, aux = batch        # [128, 300, 2] 0-1023, [128, 15, 3] 0-511
        print(batch_i)
        pdb.set_trace()

data_loader/lmdb_data_loader.py
YoungSeng/UnifiedGesture
# --------------------------------------------------------
# WavLM: Large-Scale Self-Supervised  Pre-training  for Full Stack Speech Processing (https://arxiv.org/abs/2110.13900.pdf)
# Github source: https://github.com/microsoft/unilm/tree/master/wavlm
# Copyright (c) 2021 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Based on fairseq code bases
# https://github.com/pytorch/fairseq
# --------------------------------------------------------

import math
import logging
from typing import List, Optional, Tuple

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import LayerNorm
from modules_WavLM import (
    Fp32GroupNorm,
    Fp32LayerNorm,
    GradMultiply,
    MultiheadAttention,
    SamePad,
    init_bert_params,
    get_activation_fn,
    TransposeLast,
    GLU_Linear,
)

logger = logging.getLogger(__name__)


def compute_mask_indices(
    shape: Tuple[int, int],
    padding_mask: Optional[torch.Tensor],
    mask_prob: float,
    mask_length: int,
    mask_type: str = "static",
    mask_other: float = 0.0,
    min_masks: int = 0,
    no_overlap: bool = False,
    min_space: int = 0,
) -> np.ndarray:
    """
    Computes random mask spans for a given shape

    Args:
        shape: the the shape for which to compute masks.
            should be of size 2 where first element is batch size and 2nd is timesteps
        padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements
        mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by
            number of timesteps divided by length of mask span to mask approximately this percentage of all elements.
            however due to overlaps, the actual number will be smaller (unless no_overlap is True)
        mask_type: how to compute mask lengths
            static = fixed size
            uniform = sample from uniform distribution [mask_other, mask_length*2]
            normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element
            poisson = sample from possion distribution with lambda = mask length
        min_masks: minimum number of masked spans
        no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping
        min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans
    """

    bsz, all_sz = shape
    mask = np.full((bsz, all_sz), False)

    all_num_mask = int(
        # add a random number for probabilistic rounding
        mask_prob * all_sz / float(mask_length)
        + np.random.rand()
    )

    all_num_mask = max(min_masks, all_num_mask)

    mask_idcs = []
    for i in range(bsz):
        if padding_mask is not None:
            sz = all_sz - padding_mask[i].long().sum().item()
            num_mask = int(
                # add a random number for probabilistic rounding
                mask_prob * sz / float(mask_length)
                + np.random.rand()
            )
            num_mask = max(min_masks, num_mask)
        else:
            sz = all_sz
            num_mask = all_num_mask

        if mask_type == "static":
            lengths = np.full(num_mask, mask_length)
        elif mask_type == "uniform":
            lengths = np.random.randint(mask_other, mask_length * 2 + 1, size=num_mask)
        elif mask_type == "normal":
            lengths = np.random.normal(mask_length, mask_other, size=num_mask)
            lengths = [max(1, int(round(x))) for x in lengths]
        elif mask_type == "poisson":
            lengths = np.random.poisson(mask_length, size=num_mask)
            lengths = [int(round(x)) for x in lengths]
        else:
            raise Exception("unknown mask selection " + mask_type)

        if sum(lengths) == 0:
            lengths[0] = min(mask_length, sz - 1)

        if no_overlap:
            mask_idc = []

            def arrange(s, e, length, keep_length):
                span_start = np.random.randint(s, e - length)
                mask_idc.extend(span_start + i for i in range(length))

                new_parts = []
                if span_start - s - min_space >= keep_length:
                    new_parts.append((s, span_start - min_space + 1))
                if e - span_start - keep_length - min_space > keep_length:
                    new_parts.append((span_start + length + min_space, e))
                return new_parts

            parts = [(0, sz)]
            min_length = min(lengths)
            for length in sorted(lengths, reverse=True):
                lens = np.fromiter(
                    (e - s if e - s >= length + min_space else 0 for s, e in parts),
                    np.int,
                )
                l_sum = np.sum(lens)
                if l_sum == 0:
                    break
                probs = lens / np.sum(lens)
                c = np.random.choice(len(parts), p=probs)
                s, e = parts.pop(c)
                parts.extend(arrange(s, e, length, min_length))
            mask_idc = np.asarray(mask_idc)
        else:
            min_len = min(lengths)
            if sz - min_len <= num_mask:
                min_len = sz - num_mask - 1

            mask_idc = np.random.choice(sz - min_len, num_mask, replace=False)

            mask_idc = np.asarray(
                [
                    mask_idc[j] + offset
                    for j in range(len(mask_idc))
                    for offset in range(lengths[j])
                ]
            )

        mask_idcs.append(np.unique(mask_idc[mask_idc < sz]))

    min_len = min([len(m) for m in mask_idcs])
    for i, mask_idc in enumerate(mask_idcs):
        if len(mask_idc) > min_len:
            mask_idc = np.random.choice(mask_idc, min_len, replace=False)
        mask[i, mask_idc] = True

    return mask


class WavLMConfig:
    def __init__(self, cfg=None):
        self.extractor_mode: str = "default"     # mode for feature extractor. default has a single group norm with d groups in the first conv block, whereas layer_norm has layer norms in every block (meant to use with normalize=True)
        self.encoder_layers: int = 12     # num encoder layers in the transformer

        self.encoder_embed_dim: int = 768     # encoder embedding dimension
        self.encoder_ffn_embed_dim: int = 3072     # encoder embedding dimension for FFN
        self.encoder_attention_heads: int = 12     # num encoder attention heads
        self.activation_fn: str = "gelu"     # activation function to use

        self.layer_norm_first: bool = False     # apply layernorm first in the transformer
        self.conv_feature_layers: str = "[(512,10,5)] + [(512,3,2)] * 4 + [(512,2,2)] * 2"     # string describing convolutional feature extraction layers in form of a python list that contains [(dim, kernel_size, stride), ...]
        self.conv_bias: bool = False     # include bias in conv encoder
        self.feature_grad_mult: float = 1.0     # multiply feature extractor var grads by this

        self.normalize: bool = False  # normalize input to have 0 mean and unit variance during training

        # dropouts
        self.dropout: float = 0.1     # dropout probability for the transformer
        self.attention_dropout: float = 0.1     # dropout probability for attention weights
        self.activation_dropout: float = 0.0     # dropout probability after activation in FFN
        self.encoder_layerdrop: float = 0.0     # probability of dropping a tarnsformer layer
        self.dropout_input: float = 0.0     # dropout to apply to the input (after feat extr)
        self.dropout_features: float = 0.0     # dropout to apply to the features (after feat extr)

        # masking
        self.mask_length: int = 10     # mask length)
        self.mask_prob: float = 0.65     # probability of replacing a token with mask
        self.mask_selection: str = "static"     # how to choose mask length
        self.mask_other: float = 0     # secondary mask argument (used for more complex distributions), see help in compute_mask_indicesh
        self.no_mask_overlap: bool = False     # whether to allow masks to overlap
        self.mask_min_space: int = 1     # min space between spans (if no overlap is enabled)

        # channel masking
        self.mask_channel_length: int = 10     # length of the mask for features (channels)
        self.mask_channel_prob: float = 0.0     # probability of replacing a feature with 0
        self.mask_channel_selection: str = "static"     # how to choose mask length for channel masking
        self.mask_channel_other: float = 0     # secondary mask argument (used for more complex distributions), see help in compute_mask_indices
        self.no_mask_channel_overlap: bool = False     # whether to allow channel masks to overlap
        self.mask_channel_min_space: int = 1     # min space between spans (if no overlap is enabled)

        # positional embeddings
        self.conv_pos: int = 128     # number of filters for convolutional positional embeddings
        self.conv_pos_groups: int = 16     # number of groups for convolutional positional embedding

        # relative position embedding
        self.relative_position_embedding: bool = False     # apply relative position embedding
        self.num_buckets: int = 320     # number of buckets for relative position embedding
        self.max_distance: int = 1280     # maximum distance for relative position embedding
        self.gru_rel_pos: bool = False     # apply gated relative position embedding

        if cfg is not None:
            self.update(cfg)

    def update(self, cfg: dict):
        self.__dict__.update(cfg)


class WavLM(nn.Module):
    def __init__(
        self,
        cfg: WavLMConfig,
    ) -> None:
        super().__init__()
        logger.info(f"WavLM Config: {cfg.__dict__}")

        self.cfg = cfg
        feature_enc_layers = eval(cfg.conv_feature_layers)
        self.embed = feature_enc_layers[-1][0]

        self.feature_extractor = ConvFeatureExtractionModel(
            conv_layers=feature_enc_layers,
            dropout=0.0,
            mode=cfg.extractor_mode,
            conv_bias=cfg.conv_bias,
        )

        self.post_extract_proj = (
            nn.Linear(self.embed, cfg.encoder_embed_dim)
            if self.embed != cfg.encoder_embed_dim
            else None
        )

        self.mask_prob = cfg.mask_prob
        self.mask_selection = cfg.mask_selection
        self.mask_other = cfg.mask_other
        self.mask_length = cfg.mask_length
        self.no_mask_overlap = cfg.no_mask_overlap
        self.mask_min_space = cfg.mask_min_space

        self.mask_channel_prob = cfg.mask_channel_prob
        self.mask_channel_selection = cfg.mask_channel_selection
        self.mask_channel_other = cfg.mask_channel_other
        self.mask_channel_length = cfg.mask_channel_length
        self.no_mask_channel_overlap = cfg.no_mask_channel_overlap
        self.mask_channel_min_space = cfg.mask_channel_min_space

        self.dropout_input = nn.Dropout(cfg.dropout_input)
        self.dropout_features = nn.Dropout(cfg.dropout_features)

        self.feature_grad_mult = cfg.feature_grad_mult

        self.mask_emb = nn.Parameter(
            torch.FloatTensor(cfg.encoder_embed_dim).uniform_()
        )

        self.encoder = TransformerEncoder(cfg)
        self.layer_norm = LayerNorm(self.embed)

    def apply_mask(self, x, padding_mask):
        B, T, C = x.shape
        if self.mask_prob > 0:
            mask_indices = compute_mask_indices(
                (B, T),
                padding_mask,
                self.mask_prob,
                self.mask_length,
                self.mask_selection,
                self.mask_other,
                min_masks=2,
                no_overlap=self.no_mask_overlap,
                min_space=self.mask_min_space,
            )
            mask_indices = torch.from_numpy(mask_indices).to(x.device)
            x[mask_indices] = self.mask_emb
        else:
            mask_indices = None

        if self.mask_channel_prob > 0:
            mask_channel_indices = compute_mask_indices(
                (B, C),
                None,
                self.mask_channel_prob,
                self.mask_channel_length,
                self.mask_channel_selection,
                self.mask_channel_other,
                no_overlap=self.no_mask_channel_overlap,
                min_space=self.mask_channel_min_space,
            )
            mask_channel_indices = (
                torch.from_numpy(mask_channel_indices)
                .to(x.device)
                .unsqueeze(1)
                .expand(-1, T, -1)
            )
            x[mask_channel_indices] = 0

        return x, mask_indices

    def forward_padding_mask(
            self, features: torch.Tensor, padding_mask: torch.Tensor,
    ) -> torch.Tensor:
        extra = padding_mask.size(1) % features.size(1)
        if extra > 0:
            padding_mask = padding_mask[:, :-extra]
        padding_mask = padding_mask.view(
            padding_mask.size(0), features.size(1), -1
        )
        padding_mask = padding_mask.all(-1)
        return padding_mask

    def extract_features(
        self,
        source: torch.Tensor,
        padding_mask: Optional[torch.Tensor] = None,
        mask: bool = False,
        ret_conv: bool = False,
        output_layer: Optional[int] = None,
        ret_layer_results: bool = False,
    ):

        if self.feature_grad_mult > 0:
            features = self.feature_extractor(source)
            if self.feature_grad_mult != 1.0:
                features = GradMultiply.apply(features, self.feature_grad_mult)
        else:
            with torch.no_grad():
                features = self.feature_extractor(source)

        features = features.transpose(1, 2)
        features = self.layer_norm(features)

        if padding_mask is not None:
            padding_mask = self.forward_padding_mask(features, padding_mask)

        if self.post_extract_proj is not None:
            features = self.post_extract_proj(features)

        features = self.dropout_input(features)

        if mask:
            x, mask_indices = self.apply_mask(
                features, padding_mask
            )
        else:
            x = features

        # feature: (B, T, D), float
        # target: (B, T), long
        # x: (B, T, D), float
        # padding_mask: (B, T), bool
        # mask_indices: (B, T), bool
        x, layer_results = self.encoder(
            x,
            padding_mask=padding_mask,
            layer=None if output_layer is None else output_layer - 1
        )

        res = {"x": x, "padding_mask": padding_mask, "features": features, "layer_results": layer_results}

        feature = res["features"] if ret_conv else res["x"]
        if ret_layer_results:
            feature = (feature, res["layer_results"])
        return feature, res["padding_mask"]


class ConvFeatureExtractionModel(nn.Module):
    def __init__(
            self,
            conv_layers: List[Tuple[int, int, int]],
            dropout: float = 0.0,
            mode: str = "default",
            conv_bias: bool = False,
            conv_type: str = "default"
    ):
        super().__init__()

        assert mode in {"default", "layer_norm"}

        def block(
                n_in,
                n_out,
                k,
                stride,
                is_layer_norm=False,
                is_group_norm=False,
                conv_bias=False,
        ):
            def make_conv():
                conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias)
                nn.init.kaiming_normal_(conv.weight)
                return conv

            assert (
                           is_layer_norm and is_group_norm
                   ) == False, "layer norm and group norm are exclusive"

            if is_layer_norm:
                return nn.Sequential(
                    make_conv(),
                    nn.Dropout(p=dropout),
                    nn.Sequential(
                        TransposeLast(),
                        Fp32LayerNorm(dim, elementwise_affine=True),
                        TransposeLast(),
                    ),
                    nn.GELU(),
                )
            elif is_group_norm:
                return nn.Sequential(
                    make_conv(),
                    nn.Dropout(p=dropout),
                    Fp32GroupNorm(dim, dim, affine=True),
                    nn.GELU(),
                )
            else:
                return nn.Sequential(make_conv(), nn.Dropout(p=dropout), nn.GELU())

        self.conv_type = conv_type
        if self.conv_type == "default":
            in_d = 1
            self.conv_layers = nn.ModuleList()
            for i, cl in enumerate(conv_layers):
                assert len(cl) == 3, "invalid conv definition: " + str(cl)
                (dim, k, stride) = cl

                self.conv_layers.append(
                    block(
                        in_d,
                        dim,
                        k,
                        stride,
                        is_layer_norm=mode == "layer_norm",
                        is_group_norm=mode == "default" and i == 0,
                        conv_bias=conv_bias,
                    )
                )
                in_d = dim
        elif self.conv_type == "conv2d":
            in_d = 1
            self.conv_layers = nn.ModuleList()
            for i, cl in enumerate(conv_layers):
                assert len(cl) == 3
                (dim, k, stride) = cl

                self.conv_layers.append(
                    torch.nn.Conv2d(in_d, dim, k, stride)
                )
                self.conv_layers.append(torch.nn.ReLU())
                in_d = dim
        elif self.conv_type == "custom":
            in_d = 1
            idim = 80
            self.conv_layers = nn.ModuleList()
            for i, cl in enumerate(conv_layers):
                assert len(cl) == 3
                (dim, k, stride) = cl
                self.conv_layers.append(
                    torch.nn.Conv2d(in_d, dim, k, stride, padding=1)
                )
                self.conv_layers.append(
                    torch.nn.LayerNorm([dim, idim])
                )
                self.conv_layers.append(torch.nn.ReLU())
                in_d = dim
                if (i + 1) % 2 == 0:
                    self.conv_layers.append(
                        torch.nn.MaxPool2d(2, stride=2, ceil_mode=True)
                    )
                    idim = int(math.ceil(idim / 2))
        else:
            pass

    def forward(self, x, mask=None):

        # BxT -> BxCxT
        x = x.unsqueeze(1)
        if self.conv_type == "custom":
            for conv in self.conv_layers:
                if isinstance(conv, nn.LayerNorm):
                    x = x.transpose(1, 2)
                    x = conv(x).transpose(1, 2)
                else:
                    x = conv(x)
            x = x.transpose(2, 3).contiguous()
            x = x.view(x.size(0), -1, x.size(-1))
        else:
            for conv in self.conv_layers:
                x = conv(x)
            if self.conv_type == "conv2d":
                b, c, t, f = x.size()
                x = x.transpose(2, 3).contiguous().view(b, c * f, t)
        return x


class TransformerEncoder(nn.Module):
    def __init__(self, args):
        super().__init__()

        self.dropout = args.dropout
        self.embedding_dim = args.encoder_embed_dim

        self.pos_conv = nn.Conv1d(
            self.embedding_dim,
            self.embedding_dim,
            kernel_size=args.conv_pos,
            padding=args.conv_pos // 2,
            groups=args.conv_pos_groups,
        )
        dropout = 0
        std = math.sqrt((4 * (1.0 - dropout)) / (args.conv_pos * self.embedding_dim))
        nn.init.normal_(self.pos_conv.weight, mean=0, std=std)
        nn.init.constant_(self.pos_conv.bias, 0)

        self.pos_conv = nn.utils.weight_norm(self.pos_conv, name="weight", dim=2)
        self.pos_conv = nn.Sequential(self.pos_conv, SamePad(args.conv_pos), nn.GELU())

        if hasattr(args, "relative_position_embedding"):
            self.relative_position_embedding = args.relative_position_embedding
            self.num_buckets = args.num_buckets
            self.max_distance = args.max_distance
        else:
            self.relative_position_embedding = False
            self.num_buckets = 0
            self.max_distance = 0

        self.layers = nn.ModuleList(
            [
                TransformerSentenceEncoderLayer(
                    embedding_dim=self.embedding_dim,
                    ffn_embedding_dim=args.encoder_ffn_embed_dim,
                    num_attention_heads=args.encoder_attention_heads,
                    dropout=self.dropout,
                    attention_dropout=args.attention_dropout,
                    activation_dropout=args.activation_dropout,
                    activation_fn=args.activation_fn,
                    layer_norm_first=args.layer_norm_first,
                    has_relative_attention_bias=(self.relative_position_embedding and i == 0),
                    num_buckets=self.num_buckets,
                    max_distance=self.max_distance,
                    gru_rel_pos=args.gru_rel_pos,
                )
                for i in range(args.encoder_layers)
            ]
        )

        self.layer_norm_first = args.layer_norm_first
        self.layer_norm = LayerNorm(self.embedding_dim)
        self.layerdrop = args.encoder_layerdrop

        self.apply(init_bert_params)

    def forward(self, x, padding_mask=None, streaming_mask=None, layer=None):
        x, layer_results = self.extract_features(x, padding_mask, streaming_mask, layer)

        if self.layer_norm_first and layer is None:
            x = self.layer_norm(x)

        return x, layer_results

    def extract_features(self, x, padding_mask=None, streaming_mask=None, tgt_layer=None):

        if padding_mask is not None:
            x[padding_mask] = 0

        x_conv = self.pos_conv(x.transpose(1, 2))
        x_conv = x_conv.transpose(1, 2)
        x += x_conv

        if not self.layer_norm_first:
            x = self.layer_norm(x)

        x = F.dropout(x, p=self.dropout, training=self.training)

        # B x T x C -> T x B x C
        x = x.transpose(0, 1)

        layer_results = []
        z = None
        if tgt_layer is not None:
            layer_results.append((x, z))
        r = None
        pos_bias = None
        for i, layer in enumerate(self.layers):
            dropout_probability = np.random.random()
            if not self.training or (dropout_probability > self.layerdrop):
                x, z, pos_bias = layer(x, self_attn_padding_mask=padding_mask, need_weights=False,
                                       self_attn_mask=streaming_mask, pos_bias=pos_bias)
            if tgt_layer is not None:
                layer_results.append((x, z))
            if i == tgt_layer:
                r = x
                break

        if r is not None:
            x = r

        # T x B x C -> B x T x C
        x = x.transpose(0, 1)

        return x, layer_results


class TransformerSentenceEncoderLayer(nn.Module):
    """
    Implements a Transformer Encoder Layer used in BERT/XLM style pre-trained
    models.
    """

    def __init__(
            self,
            embedding_dim: float = 768,
            ffn_embedding_dim: float = 3072,
            num_attention_heads: float = 8,
            dropout: float = 0.1,
            attention_dropout: float = 0.1,
            activation_dropout: float = 0.1,
            activation_fn: str = "relu",
            layer_norm_first: bool = False,
            has_relative_attention_bias: bool = False,
            num_buckets: int = 0,
            max_distance: int = 0,
            rescale_init: bool = False,
            gru_rel_pos: bool = False,
    ) -> None:

        super().__init__()
        # Initialize parameters
        self.embedding_dim = embedding_dim
        self.dropout = dropout
        self.activation_dropout = activation_dropout

        # Initialize blocks
        self.activation_name = activation_fn
        self.activation_fn = get_activation_fn(activation_fn)
        self.self_attn = MultiheadAttention(
            self.embedding_dim,
            num_attention_heads,
            dropout=attention_dropout,
            self_attention=True,
            has_relative_attention_bias=has_relative_attention_bias,
            num_buckets=num_buckets,
            max_distance=max_distance,
            rescale_init=rescale_init,
            gru_rel_pos=gru_rel_pos,
        )

        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(self.activation_dropout)
        self.dropout3 = nn.Dropout(dropout)

        self.layer_norm_first = layer_norm_first

        # layer norm associated with the self attention layer
        self.self_attn_layer_norm = LayerNorm(self.embedding_dim)

        if self.activation_name == "glu":
            self.fc1 = GLU_Linear(self.embedding_dim, ffn_embedding_dim, "swish")
        else:
            self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
        self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)

        # layer norm associated with the position wise feed-forward NN
        self.final_layer_norm = LayerNorm(self.embedding_dim)

    def forward(
            self,
            x: torch.Tensor,
            self_attn_mask: torch.Tensor = None,
            self_attn_padding_mask: torch.Tensor = None,
            need_weights: bool = False,
            pos_bias=None
    ):
        """
        LayerNorm is applied either before or after the self-attention/ffn
        modules similar to the original Transformer imlementation.
        """
        residual = x

        if self.layer_norm_first:
            x = self.self_attn_layer_norm(x)
            x, attn, pos_bias = self.self_attn(
                query=x,
                key=x,
                value=x,
                key_padding_mask=self_attn_padding_mask,
                need_weights=False,
                attn_mask=self_attn_mask,
                position_bias=pos_bias
            )
            x = self.dropout1(x)
            x = residual + x

            residual = x
            x = self.final_layer_norm(x)
            if self.activation_name == "glu":
                x = self.fc1(x)
            else:
                x = self.activation_fn(self.fc1(x))
            x = self.dropout2(x)
            x = self.fc2(x)
            x = self.dropout3(x)
            x = residual + x
        else:
            x, attn, pos_bias = self.self_attn(
                query=x,
                key=x,
                value=x,
                key_padding_mask=self_attn_padding_mask,
                need_weights=need_weights,
                attn_mask=self_attn_mask,
                position_bias=pos_bias
            )

            x = self.dropout1(x)
            x = residual + x

            x = self.self_attn_layer_norm(x)

            residual = x
            if self.activation_name == "glu":
                x = self.fc1(x)
            else:
                x = self.activation_fn(self.fc1(x))
            x = self.dropout2(x)
            x = self.fc2(x)
            x = self.dropout3(x)
            x = residual + x
            x = self.final_layer_norm(x)

        return x, attn, pos_bias


diffusion_latent/WavLM/WavLM.py
YoungSeng/UnifiedGesture
import configargparse
import argparse

def str2bool(v):
    """ from https://stackoverflow.com/a/43357954/1361529 """
    if isinstance(v, bool):
       return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise configargparse.ArgumentTypeError('Boolean value expected.')


def parse_args():

    parser = argparse.ArgumentParser(description='UnifiedGesture')
    parser.add_argument('--config', default='./configs/codebook.yml')

    # UnifiedGesture
    parser.add_argument('--gpu', type=str, default='2')
    parser.add_argument('--no_cuda', type=list, default=['2'])
    parser.add_argument('--save_dir', type=str, default='./save')
    parser.add_argument('--model_path', type=str, default='./experiments')
    parser.add_argument('--audio_path', type=str, default='./experiments')

    args = parser.parse_args()
    args.no_cuda = args.gpu
    return args

diffusion_latent/configs/parse_args.py
YoungSeng/UnifiedGesture
""" create data samples """
import pdb

import lmdb
import math
import numpy as np
import pyarrow


import torch
import torch.nn.functional as F

def wavlm_init(device='cuda:2'):
    import sys
    [sys.path.append(i) for i in ['./WavLM']]
    from WavLM import WavLM, WavLMConfig
    wavlm_model_path = './wavlm_cache/WavLM-Large.pt'
    # load the pre-trained checkpoints
    checkpoint = torch.load(wavlm_model_path, map_location=torch.device('cpu'))
    cfg = WavLMConfig(checkpoint['cfg'])
    model = WavLM(cfg)
    device = torch.device(device)
    model = model.to(device)
    model.load_state_dict(checkpoint['model'])
    model.eval()
    return model


def wav2wavlm(model, wav_input_16khz, device='cuda:2'):
    with torch.no_grad():
        device = torch.device(device)
        # print(wav_input_16khz.shape)
        wav_input_16khz = torch.from_numpy(wav_input_16khz).to(device)
        # wav_input_16khz = wav_input_16khz.to(device).unsqueeze(0)
        rep = model.extract_features(wav_input_16khz)[0]
        del wav_input_16khz
        rep = F.interpolate(rep.transpose(1, 2), size=36, align_corners=True, mode='linear').transpose(1, 2)
        return rep.squeeze().cpu().detach().data.cpu().numpy()


class DataPreprocessor:
    def __init__(self, clip_lmdb_dir, out_lmdb_dir, n_poses, subdivision_stride, pose_resampling_fps, n_codes=30):
        self.n_codes = n_codes
        self.n_poses = n_poses
        self.subdivision_stride = subdivision_stride
        self.skeleton_resampling_fps = pose_resampling_fps

        self.src_lmdb_env = lmdb.open(clip_lmdb_dir, readonly=True, lock=False)
        with self.src_lmdb_env.begin() as txn:
            self.n_videos = txn.stat()['entries']

        self.audio_sample_length = int(self.n_poses / self.skeleton_resampling_fps * 16000)

        # create db for samples
        map_size = 1024 * 1024 * 9  # in TB
        map_size <<= 20  # in B
        self.dst_lmdb_env = lmdb.open(out_lmdb_dir, map_size=map_size)
        self.n_out_samples = 0

        self.model = wavlm_init()

    def run(self):
        src_txn = self.src_lmdb_env.begin(write=False)

        # sampling and normalization
        cursor = src_txn.cursor()
        for key, value in cursor:
            video = pyarrow.deserialize(value)
            vid = video['vid']
            clips = video['clips']
            for clip_idx, clip in enumerate(clips):
                self._sample_from_clip(vid, clip)

        # print stats
        with self.dst_lmdb_env.begin() as txn:
            print('no. of samples: ', txn.stat()['entries'])

        # close db
        self.src_lmdb_env.close()
        self.dst_lmdb_env.sync()
        self.dst_lmdb_env.close()

    def _sample_from_clip(self, vid, clip):
        clip_skeleton = clip['full_body']
        clip_audio_raw = clip['audio_wav']
        clip_styles_raw = clip['style']
        clip_upper_code_raw = clip['upper_code']
        # clip_skeleton_vel = clip_skeleton[:-1] - clip_skeleton[1:]
        # clip_skeleton_acc = clip_skeleton_vel[:-1] - clip_skeleton_vel[1:]
        # clip_skeleton_vel = np.pad(clip_skeleton_vel, ((1, 0), (0, 0)), 'constant')
        # clip_skeleton_acc = np.pad(clip_skeleton_acc, ((2, 0), (0, 0)), 'constant')

        # divide
        aux_info = []
        sample_skeletons_list = []
        # sample_skeletons_list_vel = []
        # sample_skeletons_list_acc = []
        sample_style_list = []
        sample_codes_list = []
        sample_wavlm_list = []

        MINLEN = len(clip_skeleton)

        num_subdivision = math.floor((MINLEN - self.n_poses) / self.subdivision_stride)  # floor((K - (N+M)) / S) + 1

        for i in range(num_subdivision):
            start_idx = i * self.subdivision_stride
            fin_idx = start_idx + self.n_poses

            sample_skeletons = clip_skeleton[start_idx:fin_idx]
            # sample_skeletons_vel = clip_skeleton_vel[start_idx:fin_idx]
            # sample_skeletons_acc = clip_skeleton_acc[start_idx:fin_idx]
            sample_code = clip_upper_code_raw[start_idx//2:fin_idx//2]

            subdivision_start_time = start_idx / self.skeleton_resampling_fps
            subdivision_end_time = fin_idx / self.skeleton_resampling_fps

            # raw audio
            audio_start = math.floor(start_idx / len(clip_skeleton) * len(clip_audio_raw))
            audio_end = audio_start + self.audio_sample_length
            sample_audio = clip_audio_raw[:, audio_start:audio_end]
            sample_wavlm = wav2wavlm(self.model, sample_audio)      # (1, 68266)

            motion_info = {'vid': vid,
                           'start_frame_no': start_idx,
                           'end_frame_no': fin_idx,
                           'start_time': subdivision_start_time,
                           'end_time': subdivision_end_time}

            sample_skeletons_list.append(sample_skeletons)
            # sample_skeletons_list_vel.append(sample_skeletons_vel)
            # sample_skeletons_list_acc.append(sample_skeletons_acc)
            sample_codes_list.append(sample_code)
            sample_wavlm_list.append(sample_wavlm)
            sample_style_list.append(clip_styles_raw)

            aux_info.append(motion_info)

        if len(sample_skeletons_list) > 0:
            with self.dst_lmdb_env.begin(write=True) as txn:
                # for poses, vel, acc, code, style, wavlm, aux in zip(sample_skeletons_list, sample_skeletons_list_vel, sample_skeletons_list_acc, sample_codes_list, sample_style_list, sample_wavlm_list, aux_info):
                for poses, code, style, wavlm, aux in zip(sample_skeletons_list, sample_codes_list, sample_style_list, sample_wavlm_list, aux_info):

                    poses = np.asarray(poses)
                    # vel = np.asarray(vel)
                    # acc = np.asarray(acc)
                    code = np.asarray(code)
                    style = np.asarray(style)
                    wavlm = np.asarray(wavlm)

                    # save
                    k = '{:010}'.format(self.n_out_samples).encode('ascii')
                    # v = [poses, vel, acc, code, style, wavlm, aux]
                    v = [poses, code, style, wavlm]
                    v = pyarrow.serialize(v).to_buffer()
                    txn.put(k, v)
                    self.n_out_samples += 1


diffusion_latent/data_loader/data_preprocessor.py
YoungSeng/UnifiedGesture
# ! pip install configargparse easydict

import logging
import os
import pdb

import numpy as np
import lmdb as lmdb
import torch
from torch.nn.utils.rnn import pad_sequence

from torch.utils.data import Dataset
from torch.utils.data.dataloader import default_collate

import sys

[sys.path.append(i) for i in ['.', '..']]

from data_loader.data_preprocessor import DataPreprocessor
import pyarrow


class TrinityDataset(Dataset):
    def __init__(self, lmdb_dir, n_poses, subdivision_stride, pose_resampling_fps, model=None):
        self.lmdb_dir = lmdb_dir
        self.n_poses = n_poses
        self.subdivision_stride = subdivision_stride
        self.skeleton_resampling_fps = pose_resampling_fps
        self.lang_model = None

        logging.info("Reading data '{}'...".format(lmdb_dir))
        if model is not None:
            if 'WavLM' in model:
                preloaded_dir = lmdb_dir + '_cache_' + model
        else:
            preloaded_dir = lmdb_dir + '_cache'
        if not os.path.exists(preloaded_dir):
            data_sampler = DataPreprocessor(lmdb_dir, preloaded_dir, n_poses,
                                            subdivision_stride, pose_resampling_fps)
            data_sampler.run()
        else:
            logging.info('Found pre-loaded samples from {}'.format(preloaded_dir))

        # init lmdb
        # map_size = 1024 * 20  # in MB
        # map_size <<= 20  # in B
        self.lmdb_env = lmdb.open(preloaded_dir, readonly=True, lock=False)  # default 10485760
        with self.lmdb_env.begin() as txn:
            self.n_samples = txn.stat()['entries']

    def __len__(self):
        return self.n_samples

    def __getitem__(self, idx):
        with self.lmdb_env.begin(write=False) as txn:
            key = '{:010}'.format(idx).encode('ascii')
            sample = txn.get(key)

            sample = pyarrow.deserialize(sample)
            # poses, code, style, wavlm, aux_info = sample
            poses, code, style, wavlm = sample

        # to tensors
        pose_seq = torch.from_numpy(poses).float()
        # vel_seq = torch.from_numpy(vel).float()
        # acc_seq = torch.from_numpy(acc).float()
        # print('pose_seq.shape', pose_seq.shape)
        # pose_seq = torch.cat([pose_seq, vel_seq, acc_seq], dim=-1)
        code = torch.from_numpy(code).float()
        style = torch.from_numpy(style).float()
        wavlm = torch.from_numpy(wavlm).float()

        return pose_seq, code, style, wavlm


if __name__ == '__main__':
    '''
    cd ./diffusion_latent/
    python data_loader/lmdb_data_loader.py --config=./configs/all_data.yml --train --no_cuda 3 --gpu 3
    '''

    from configs.parse_args import parse_args
    import os
    import yaml
    from pprint import pprint
    from easydict import EasyDict
    from torch.utils.data import DataLoader

    args = parse_args()

    with open(args.config) as f:
        config = yaml.safe_load(f)

    for k, v in vars(args).items():
        config[k] = v
    pprint(config)

    args = EasyDict(config)


    val_dataset = TrinityDataset(args.val_data_path,
                                       n_poses=args.n_poses,
                                       subdivision_stride=args.subdivision_stride,
                                       pose_resampling_fps=args.motion_resampling_framerate, model='WavLM_36_aux')


    valid_loader = DataLoader(dataset=val_dataset, batch_size=128,
                              shuffle=True, drop_last=True, num_workers=args.loader_workers, pin_memory=True)

    print(len(valid_loader))
    for batch_i, batch in enumerate(valid_loader, 0):
        pose_seq, code, style, wavlm = batch     # [128, 88, 1141], -,  [128, 6], [128, 70400], [128, 88, 13]
        # [128, 32, 112], [128, 16], [128, 7], [128, 32, 1024]
        print(batch_i)
        pdb.set_trace()
        break

    train_dataset = TrinityDataset(args.train_data_path,
                                   n_poses=args.n_poses,
                                   subdivision_stride=args.subdivision_stride,
                                   pose_resampling_fps=args.motion_resampling_framerate, model='WavLM_36_aux')

    train_loader = DataLoader(dataset=train_dataset, batch_size=128,
                              shuffle=True, drop_last=True, num_workers=args.loader_workers, pin_memory=True)
diffusion_latent/data_loader/lmdb_data_loader.py
YoungSeng/UnifiedGesture
import torch
from torch.utils import data
import numpy as np
import os
from os.path import join as pjoin
import random
import codecs as cs
from tqdm import tqdm
import spacy

from torch.utils.data._utils.collate import default_collate
from data_loaders.humanml.utils.word_vectorizer import WordVectorizer
from data_loaders.humanml.utils.get_opt import get_opt

# import spacy

def collate_fn(batch):
    batch.sort(key=lambda x: x[3], reverse=True)
    return default_collate(batch)


'''For use of training text-2-motion generative model'''
class Text2MotionDataset(data.Dataset):
    def __init__(self, opt, mean, std, split_file, w_vectorizer):
        self.opt = opt
        self.w_vectorizer = w_vectorizer
        self.max_length = 20
        self.pointer = 0
        min_motion_len = 40 if self.opt.dataset_name =='t2m' else 24

        joints_num = opt.joints_num

        data_dict = {}
        id_list = []
        with cs.open(split_file, 'r') as f:
            for line in f.readlines():
                id_list.append(line.strip())

        new_name_list = []
        length_list = []
        for name in tqdm(id_list):
            try:
                motion = np.load(pjoin(opt.motion_dir, name + '.npy'))
                if (len(motion)) < min_motion_len or (len(motion) >= 200):
                    continue
                text_data = []
                flag = False
                with cs.open(pjoin(opt.text_dir, name + '.txt')) as f:
                    for line in f.readlines():
                        text_dict = {}
                        line_split = line.strip().split('#')
                        caption = line_split[0]
                        tokens = line_split[1].split(' ')
                        f_tag = float(line_split[2])
                        to_tag = float(line_split[3])
                        f_tag = 0.0 if np.isnan(f_tag) else f_tag
                        to_tag = 0.0 if np.isnan(to_tag) else to_tag

                        text_dict['caption'] = caption
                        text_dict['tokens'] = tokens
                        if f_tag == 0.0 and to_tag == 0.0:
                            flag = True
                            text_data.append(text_dict)
                        else:
                            try:
                                n_motion = motion[int(f_tag*20) : int(to_tag*20)]
                                if (len(n_motion)) < min_motion_len or (len(n_motion) >= 200):
                                    continue
                                new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
                                while new_name in data_dict:
                                    new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
                                data_dict[new_name] = {'motion': n_motion,
                                                       'length': len(n_motion),
                                                       'text':[text_dict]}
                                new_name_list.append(new_name)
                                length_list.append(len(n_motion))
                            except:
                                print(line_split)
                                print(line_split[2], line_split[3], f_tag, to_tag, name)
                                # break

                if flag:
                    data_dict[name] = {'motion': motion,
                                       'length': len(motion),
                                       'text':text_data}
                    new_name_list.append(name)
                    length_list.append(len(motion))
            except:
                # Some motion may not exist in KIT dataset
                pass


        name_list, length_list = zip(*sorted(zip(new_name_list, length_list), key=lambda x: x[1]))

        if opt.is_train:
            # root_rot_velocity (B, seq_len, 1)
            std[0:1] = std[0:1] / opt.feat_bias
            # root_linear_velocity (B, seq_len, 2)
            std[1:3] = std[1:3] / opt.feat_bias
            # root_y (B, seq_len, 1)
            std[3:4] = std[3:4] / opt.feat_bias
            # ric_data (B, seq_len, (joint_num - 1)*3)
            std[4: 4 + (joints_num - 1) * 3] = std[4: 4 + (joints_num - 1) * 3] / 1.0
            # rot_data (B, seq_len, (joint_num - 1)*6)
            std[4 + (joints_num - 1) * 3: 4 + (joints_num - 1) * 9] = std[4 + (joints_num - 1) * 3: 4 + (
                        joints_num - 1) * 9] / 1.0
            # local_velocity (B, seq_len, joint_num*3)
            std[4 + (joints_num - 1) * 9: 4 + (joints_num - 1) * 9 + joints_num * 3] = std[
                                                                                       4 + (joints_num - 1) * 9: 4 + (
                                                                                                   joints_num - 1) * 9 + joints_num * 3] / 1.0
            # foot contact (B, seq_len, 4)
            std[4 + (joints_num - 1) * 9 + joints_num * 3:] = std[
                                                              4 + (joints_num - 1) * 9 + joints_num * 3:] / opt.feat_bias

            assert 4 + (joints_num - 1) * 9 + joints_num * 3 + 4 == mean.shape[-1]
            np.save(pjoin(opt.meta_dir, 'mean.npy'), mean)
            np.save(pjoin(opt.meta_dir, 'std.npy'), std)

        self.mean = mean
        self.std = std
        self.length_arr = np.array(length_list)
        self.data_dict = data_dict
        self.name_list = name_list
        self.reset_max_len(self.max_length)

    def reset_max_len(self, length):
        assert length <= self.opt.max_motion_length
        self.pointer = np.searchsorted(self.length_arr, length)
        print("Pointer Pointing at %d"%self.pointer)
        self.max_length = length

    def inv_transform(self, data):
        return data * self.std + self.mean

    def __len__(self):
        return len(self.data_dict) - self.pointer

    def __getitem__(self, item):
        idx = self.pointer + item
        data = self.data_dict[self.name_list[idx]]
        motion, m_length, text_list = data['motion'], data['length'], data['text']
        # Randomly select a caption
        text_data = random.choice(text_list)
        caption, tokens = text_data['caption'], text_data['tokens']

        if len(tokens) < self.opt.max_text_len:
            # pad with "unk"
            tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
            sent_len = len(tokens)
            tokens = tokens + ['unk/OTHER'] * (self.opt.max_text_len + 2 - sent_len)
        else:
            # crop
            tokens = tokens[:self.opt.max_text_len]
            tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
            sent_len = len(tokens)
        pos_one_hots = []
        word_embeddings = []
        for token in tokens:
            word_emb, pos_oh = self.w_vectorizer[token]
            pos_one_hots.append(pos_oh[None, :])
            word_embeddings.append(word_emb[None, :])
        pos_one_hots = np.concatenate(pos_one_hots, axis=0)
        word_embeddings = np.concatenate(word_embeddings, axis=0)

        len_gap = (m_length - self.max_length) // self.opt.unit_length

        if self.opt.is_train:
            if m_length != self.max_length:
            # print("Motion original length:%d_%d"%(m_length, len(motion)))
                if self.opt.unit_length < 10:
                    coin2 = np.random.choice(['single', 'single', 'double'])
                else:
                    coin2 = 'single'
                if len_gap == 0 or (len_gap == 1 and coin2 == 'double'):
                    m_length = self.max_length
                    idx = random.randint(0, m_length - self.max_length)
                    motion = motion[idx:idx+self.max_length]
                else:
                    if coin2 == 'single':
                        n_m_length = self.max_length + self.opt.unit_length * len_gap
                    else:
                        n_m_length = self.max_length + self.opt.unit_length * (len_gap - 1)
                    idx = random.randint(0, m_length - n_m_length)
                    motion = motion[idx:idx + self.max_length]
                    m_length = n_m_length
                # print(len_gap, idx, coin2)
        else:
            if self.opt.unit_length < 10:
                coin2 = np.random.choice(['single', 'single', 'double'])
            else:
                coin2 = 'single'

            if coin2 == 'double':
                m_length = (m_length // self.opt.unit_length - 1) * self.opt.unit_length
            elif coin2 == 'single':
                m_length = (m_length // self.opt.unit_length) * self.opt.unit_length
            idx = random.randint(0, len(motion) - m_length)
            motion = motion[idx:idx+m_length]

        "Z Normalization"
        motion = (motion - self.mean) / self.std

        return word_embeddings, pos_one_hots, caption, sent_len, motion, m_length


'''For use of training text motion matching model, and evaluations'''
class Text2MotionDatasetV2(data.Dataset):
    def __init__(self, opt, mean, std, split_file, w_vectorizer):
        self.opt = opt
        self.w_vectorizer = w_vectorizer
        self.max_length = 20
        self.pointer = 0
        self.max_motion_length = opt.max_motion_length
        min_motion_len = 40 if self.opt.dataset_name =='t2m' else 24

        data_dict = {}
        id_list = []
        with cs.open(split_file, 'r') as f:
            for line in f.readlines():
                id_list.append(line.strip())
        id_list = id_list[:100]     # debug

        new_name_list = []
        length_list = []
        for name in tqdm(id_list):
            try:
                motion = np.load(pjoin(opt.motion_dir, name + '.npy'))
                if (len(motion)) < min_motion_len or (len(motion) >= 200):
                    continue
                text_data = []
                flag = False
                with cs.open(pjoin(opt.text_dir, name + '.txt')) as f:
                    for line in f.readlines():
                        text_dict = {}
                        line_split = line.strip().split('#')
                        caption = line_split[0]
                        tokens = line_split[1].split(' ')
                        f_tag = float(line_split[2])
                        to_tag = float(line_split[3])
                        f_tag = 0.0 if np.isnan(f_tag) else f_tag
                        to_tag = 0.0 if np.isnan(to_tag) else to_tag

                        text_dict['caption'] = caption
                        text_dict['tokens'] = tokens
                        if f_tag == 0.0 and to_tag == 0.0:
                            flag = True
                            text_data.append(text_dict)
                        else:
                            try:
                                n_motion = motion[int(f_tag*20) : int(to_tag*20)]
                                if (len(n_motion)) < min_motion_len or (len(n_motion) >= 200):
                                    continue
                                new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
                                while new_name in data_dict:
                                    new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
                                data_dict[new_name] = {'motion': n_motion,
                                                       'length': len(n_motion),
                                                       'text':[text_dict]}
                                new_name_list.append(new_name)
                                length_list.append(len(n_motion))
                            except:
                                print(line_split)
                                print(line_split[2], line_split[3], f_tag, to_tag, name)
                                # break

                if flag:
                    data_dict[name] = {'motion': motion,
                                       'length': len(motion),
                                       'text': text_data}
                    new_name_list.append(name)
                    length_list.append(len(motion))
            except:
                pass

        name_list, length_list = zip(*sorted(zip(new_name_list, length_list), key=lambda x: x[1]))

        self.mean = mean
        self.std = std
        self.length_arr = np.array(length_list)
        self.data_dict = data_dict
        self.name_list = name_list
        self.reset_max_len(self.max_length)

    def reset_max_len(self, length):
        assert length <= self.max_motion_length
        self.pointer = np.searchsorted(self.length_arr, length)
        print("Pointer Pointing at %d"%self.pointer)
        self.max_length = length

    def inv_transform(self, data):
        return data * self.std + self.mean

    def __len__(self):
        return len(self.data_dict) - self.pointer

    def __getitem__(self, item):
        idx = self.pointer + item
        data = self.data_dict[self.name_list[idx]]
        motion, m_length, text_list = data['motion'], data['length'], data['text']
        # Randomly select a caption
        text_data = random.choice(text_list)
        caption, tokens = text_data['caption'], text_data['tokens']

        if len(tokens) < self.opt.max_text_len:
            # pad with "unk"
            tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
            sent_len = len(tokens)
            tokens = tokens + ['unk/OTHER'] * (self.opt.max_text_len + 2 - sent_len)
        else:
            # crop
            tokens = tokens[:self.opt.max_text_len]
            tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
            sent_len = len(tokens)
        pos_one_hots = []
        word_embeddings = []
        for token in tokens:
            word_emb, pos_oh = self.w_vectorizer[token]
            pos_one_hots.append(pos_oh[None, :])
            word_embeddings.append(word_emb[None, :])
        pos_one_hots = np.concatenate(pos_one_hots, axis=0)
        word_embeddings = np.concatenate(word_embeddings, axis=0)

        # Crop the motions in to times of 4, and introduce small variations
        if self.opt.unit_length < 10:
            coin2 = np.random.choice(['single', 'single', 'double'])
        else:
            coin2 = 'single'

        if coin2 == 'double':
            m_length = (m_length // self.opt.unit_length - 1) * self.opt.unit_length
        elif coin2 == 'single':
            m_length = (m_length // self.opt.unit_length) * self.opt.unit_length
        idx = random.randint(0, len(motion) - m_length)
        motion = motion[idx:idx+m_length]

        "Z Normalization"
        motion = (motion - self.mean) / self.std

        if m_length < self.max_motion_length:
            motion = np.concatenate([motion,
                                     np.zeros((self.max_motion_length - m_length, motion.shape[1]))
                                     ], axis=0)
        # print(word_embeddings.shape, motion.shape)
        # print(tokens)
        return word_embeddings, pos_one_hots, caption, sent_len, motion, m_length, '_'.join(tokens)


'''For use of training baseline'''
class Text2MotionDatasetBaseline(data.Dataset):
    def __init__(self, opt, mean, std, split_file, w_vectorizer):
        self.opt = opt
        self.w_vectorizer = w_vectorizer
        self.max_length = 20
        self.pointer = 0
        self.max_motion_length = opt.max_motion_length
        min_motion_len = 40 if self.opt.dataset_name =='t2m' else 24

        data_dict = {}
        id_list = []
        with cs.open(split_file, 'r') as f:
            for line in f.readlines():
                id_list.append(line.strip())
        # id_list = id_list[:200]

        new_name_list = []
        length_list = []
        for name in tqdm(id_list):
            try:
                motion = np.load(pjoin(opt.motion_dir, name + '.npy'))
                if (len(motion)) < min_motion_len or (len(motion) >= 200):
                    continue
                text_data = []
                flag = False
                with cs.open(pjoin(opt.text_dir, name + '.txt')) as f:
                    for line in f.readlines():
                        text_dict = {}
                        line_split = line.strip().split('#')
                        caption = line_split[0]
                        tokens = line_split[1].split(' ')
                        f_tag = float(line_split[2])
                        to_tag = float(line_split[3])
                        f_tag = 0.0 if np.isnan(f_tag) else f_tag
                        to_tag = 0.0 if np.isnan(to_tag) else to_tag

                        text_dict['caption'] = caption
                        text_dict['tokens'] = tokens
                        if f_tag == 0.0 and to_tag == 0.0:
                            flag = True
                            text_data.append(text_dict)
                        else:
                            try:
                                n_motion = motion[int(f_tag*20) : int(to_tag*20)]
                                if (len(n_motion)) < min_motion_len or (len(n_motion) >= 200):
                                    continue
                                new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
                                while new_name in data_dict:
                                    new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
                                data_dict[new_name] = {'motion': n_motion,
                                                       'length': len(n_motion),
                                                       'text':[text_dict]}
                                new_name_list.append(new_name)
                                length_list.append(len(n_motion))
                            except:
                                print(line_split)
                                print(line_split[2], line_split[3], f_tag, to_tag, name)
                                # break

                if flag:
                    data_dict[name] = {'motion': motion,
                                       'length': len(motion),
                                       'text': text_data}
                    new_name_list.append(name)
                    length_list.append(len(motion))
            except:
                pass

        name_list, length_list = zip(*sorted(zip(new_name_list, length_list), key=lambda x: x[1]))

        self.mean = mean
        self.std = std
        self.length_arr = np.array(length_list)
        self.data_dict = data_dict
        self.name_list = name_list
        self.reset_max_len(self.max_length)

    def reset_max_len(self, length):
        assert length <= self.max_motion_length
        self.pointer = np.searchsorted(self.length_arr, length)
        print("Pointer Pointing at %d"%self.pointer)
        self.max_length = length

    def inv_transform(self, data):
        return data * self.std + self.mean

    def __len__(self):
        return len(self.data_dict) - self.pointer

    def __getitem__(self, item):
        idx = self.pointer + item
        data = self.data_dict[self.name_list[idx]]
        motion, m_length, text_list = data['motion'], data['length'], data['text']
        # Randomly select a caption
        text_data = random.choice(text_list)
        caption, tokens = text_data['caption'], text_data['tokens']

        if len(tokens) < self.opt.max_text_len:
            # pad with "unk"
            tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
            sent_len = len(tokens)
            tokens = tokens + ['unk/OTHER'] * (self.opt.max_text_len + 2 - sent_len)
        else:
            # crop
            tokens = tokens[:self.opt.max_text_len]
            tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
            sent_len = len(tokens)
        pos_one_hots = []
        word_embeddings = []
        for token in tokens:
            word_emb, pos_oh = self.w_vectorizer[token]
            pos_one_hots.append(pos_oh[None, :])
            word_embeddings.append(word_emb[None, :])
        pos_one_hots = np.concatenate(pos_one_hots, axis=0)
        word_embeddings = np.concatenate(word_embeddings, axis=0)

        len_gap = (m_length - self.max_length) // self.opt.unit_length

        if m_length != self.max_length:
            # print("Motion original length:%d_%d"%(m_length, len(motion)))
            if self.opt.unit_length < 10:
                coin2 = np.random.choice(['single', 'single', 'double'])
            else:
                coin2 = 'single'
            if len_gap == 0 or (len_gap == 1 and coin2 == 'double'):
                m_length = self.max_length
                s_idx = random.randint(0, m_length - self.max_length)
            else:
                if coin2 == 'single':
                    n_m_length = self.max_length + self.opt.unit_length * len_gap
                else:
                    n_m_length = self.max_length + self.opt.unit_length * (len_gap - 1)
                s_idx = random.randint(0, m_length - n_m_length)
                m_length = n_m_length
        else:
            s_idx = 0

        src_motion = motion[s_idx: s_idx + m_length]
        tgt_motion = motion[s_idx: s_idx + self.max_length]

        "Z Normalization"
        src_motion = (src_motion - self.mean) / self.std
        tgt_motion = (tgt_motion - self.mean) / self.std

        if m_length < self.max_motion_length:
            src_motion = np.concatenate([src_motion,
                                     np.zeros((self.max_motion_length - m_length, motion.shape[1]))
                                     ], axis=0)
        # print(m_length, src_motion.shape, tgt_motion.shape)
        # print(word_embeddings.shape, motion.shape)
        # print(tokens)
        return word_embeddings, caption, sent_len, src_motion, tgt_motion, m_length


class MotionDatasetV2(data.Dataset):
    def __init__(self, opt, mean, std, split_file):
        self.opt = opt
        joints_num = opt.joints_num

        self.data = []
        self.lengths = []
        id_list = []
        with cs.open(split_file, 'r') as f:
            for line in f.readlines():
                id_list.append(line.strip())

        for name in tqdm(id_list):
            try:
                motion = np.load(pjoin(opt.motion_dir, name + '.npy'))
                if motion.shape[0] < opt.window_size:
                    continue
                self.lengths.append(motion.shape[0] - opt.window_size)
                self.data.append(motion)
            except:
                # Some motion may not exist in KIT dataset
                pass

        self.cumsum = np.cumsum([0] + self.lengths)

        if opt.is_train:
            # root_rot_velocity (B, seq_len, 1)
            std[0:1] = std[0:1] / opt.feat_bias
            # root_linear_velocity (B, seq_len, 2)
            std[1:3] = std[1:3] / opt.feat_bias
            # root_y (B, seq_len, 1)
            std[3:4] = std[3:4] / opt.feat_bias
            # ric_data (B, seq_len, (joint_num - 1)*3)
            std[4: 4 + (joints_num - 1) * 3] = std[4: 4 + (joints_num - 1) * 3] / 1.0
            # rot_data (B, seq_len, (joint_num - 1)*6)
            std[4 + (joints_num - 1) * 3: 4 + (joints_num - 1) * 9] = std[4 + (joints_num - 1) * 3: 4 + (
                        joints_num - 1) * 9] / 1.0
            # local_velocity (B, seq_len, joint_num*3)
            std[4 + (joints_num - 1) * 9: 4 + (joints_num - 1) * 9 + joints_num * 3] = std[
                                                                                       4 + (joints_num - 1) * 9: 4 + (
                                                                                                   joints_num - 1) * 9 + joints_num * 3] / 1.0
            # foot contact (B, seq_len, 4)
            std[4 + (joints_num - 1) * 9 + joints_num * 3:] = std[
                                                              4 + (joints_num - 1) * 9 + joints_num * 3:] / opt.feat_bias

            assert 4 + (joints_num - 1) * 9 + joints_num * 3 + 4 == mean.shape[-1]
            np.save(pjoin(opt.meta_dir, 'mean.npy'), mean)
            np.save(pjoin(opt.meta_dir, 'std.npy'), std)

        self.mean = mean
        self.std = std
        print("Total number of motions {}, snippets {}".format(len(self.data), self.cumsum[-1]))

    def inv_transform(self, data):
        return data * self.std + self.mean

    def __len__(self):
        return self.cumsum[-1]

    def __getitem__(self, item):
        if item != 0:
            motion_id = np.searchsorted(self.cumsum, item) - 1
            idx = item - self.cumsum[motion_id] - 1
        else:
            motion_id = 0
            idx = 0
        motion = self.data[motion_id][idx:idx+self.opt.window_size]
        "Z Normalization"
        motion = (motion - self.mean) / self.std

        return motion


class RawTextDataset(data.Dataset):
    def __init__(self, opt, mean, std, text_file, w_vectorizer):
        self.mean = mean
        self.std = std
        self.opt = opt
        self.data_dict = []
        self.nlp = spacy.load('en_core_web_sm')

        with cs.open(text_file) as f:
            for line in f.readlines():
                word_list, pos_list = self.process_text(line.strip())
                tokens = ['%s/%s'%(word_list[i], pos_list[i]) for i in range(len(word_list))]
                self.data_dict.append({'caption':line.strip(), "tokens":tokens})

        self.w_vectorizer = w_vectorizer
        print("Total number of descriptions {}".format(len(self.data_dict)))


    def process_text(self, sentence):
        sentence = sentence.replace('-', '')
        doc = self.nlp(sentence)
        word_list = []
        pos_list = []
        for token in doc:
            word = token.text
            if not word.isalpha():
                continue
            if (token.pos_ == 'NOUN' or token.pos_ == 'VERB') and (word != 'left'):
                word_list.append(token.lemma_)
            else:
                word_list.append(word)
            pos_list.append(token.pos_)
        return word_list, pos_list

    def inv_transform(self, data):
        return data * self.std + self.mean

    def __len__(self):
        return len(self.data_dict)

    def __getitem__(self, item):
        data = self.data_dict[item]
        caption, tokens = data['caption'], data['tokens']

        if len(tokens) < self.opt.max_text_len:
            # pad with "unk"
            tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
            sent_len = len(tokens)
            tokens = tokens + ['unk/OTHER'] * (self.opt.max_text_len + 2 - sent_len)
        else:
            # crop
            tokens = tokens[:self.opt.max_text_len]
            tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
            sent_len = len(tokens)
        pos_one_hots = []
        word_embeddings = []
        for token in tokens:
            word_emb, pos_oh = self.w_vectorizer[token]
            pos_one_hots.append(pos_oh[None, :])
            word_embeddings.append(word_emb[None, :])
        pos_one_hots = np.concatenate(pos_one_hots, axis=0)
        word_embeddings = np.concatenate(word_embeddings, axis=0)

        return word_embeddings, pos_one_hots, caption, sent_len

class TextOnlyDataset(data.Dataset):
    def __init__(self, opt, mean, std, split_file):
        self.mean = mean
        self.std = std
        self.opt = opt
        self.data_dict = []
        self.max_length = 20
        self.pointer = 0
        self.fixed_length = 120


        data_dict = {}
        id_list = []
        with cs.open(split_file, 'r') as f:
            for line in f.readlines():
                id_list.append(line.strip())
        # id_list = id_list[:200]

        new_name_list = []
        length_list = []
        for name in tqdm(id_list):
            try:
                text_data = []
                flag = False
                with cs.open(pjoin(opt.text_dir, name + '.txt')) as f:
                    for line in f.readlines():
                        text_dict = {}
                        line_split = line.strip().split('#')
                        caption = line_split[0]
                        tokens = line_split[1].split(' ')
                        f_tag = float(line_split[2])
                        to_tag = float(line_split[3])
                        f_tag = 0.0 if np.isnan(f_tag) else f_tag
                        to_tag = 0.0 if np.isnan(to_tag) else to_tag

                        text_dict['caption'] = caption
                        text_dict['tokens'] = tokens
                        if f_tag == 0.0 and to_tag == 0.0:
                            flag = True
                            text_data.append(text_dict)
                        else:
                            try:
                                new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
                                while new_name in data_dict:
                                    new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
                                data_dict[new_name] = {'text':[text_dict]}
                                new_name_list.append(new_name)
                            except:
                                print(line_split)
                                print(line_split[2], line_split[3], f_tag, to_tag, name)
                                # break

                if flag:
                    data_dict[name] = {'text': text_data}
                    new_name_list.append(name)
            except:
                pass

        self.length_arr = np.array(length_list)
        self.data_dict = data_dict
        self.name_list = new_name_list

    def inv_transform(self, data):
        return data * self.std + self.mean

    def __len__(self):
        return len(self.data_dict)

    def __getitem__(self, item):
        idx = self.pointer + item
        data = self.data_dict[self.name_list[idx]]
        text_list = data['text']

        # Randomly select a caption
        text_data = random.choice(text_list)
        caption, tokens = text_data['caption'], text_data['tokens']
        return None, None, caption, None, np.array([0]), self.fixed_length, None
        # fixed_length can be set from outside before sampling

# A wrapper class for t2m original dataset for MDM purposes
class HumanML3D(data.Dataset):
    def __init__(self, mode, datapath='./dataset/humanml_opt.txt', split="train", **kwargs):
        self.mode = mode
        
        self.dataset_name = 't2m'
        self.dataname = 't2m'

        # Configurations of T2M dataset and KIT dataset is almost the same
        abs_base_path = f'../motion-diffusion-model'
        # abs_base_path = f'.'

        dataset_opt_path = pjoin(abs_base_path, datapath)
        device = None  # torch.device('cuda:4') # This param is not in use in this context
        opt = get_opt(dataset_opt_path, device)
        opt.meta_dir = pjoin(abs_base_path, opt.meta_dir)
        opt.motion_dir = pjoin(abs_base_path, opt.motion_dir)
        opt.text_dir = pjoin(abs_base_path, opt.text_dir)
        opt.model_dir = pjoin(abs_base_path, opt.model_dir)
        opt.checkpoints_dir = pjoin(abs_base_path, opt.checkpoints_dir)
        opt.data_root = pjoin(abs_base_path, opt.data_root)
        opt.save_root = pjoin(abs_base_path, opt.save_root)
        opt.meta_dir = './dataset'
        self.opt = opt
        print('Loading dataset %s ...' % opt.dataset_name)

        if mode == 'gt':
            # used by T2M models (including evaluators)
            self.mean = np.load(pjoin(opt.meta_dir, f'{opt.dataset_name}_mean.npy'))
            self.std = np.load(pjoin(opt.meta_dir, f'{opt.dataset_name}_std.npy'))
        elif mode in ['train', 'eval', 'text_only']:
            # used by our models
            self.mean = np.load(pjoin(opt.data_root, 'Mean.npy'))
            self.std = np.load(pjoin(opt.data_root, 'Std.npy'))

        if mode == 'eval':
            # used by T2M models (including evaluators)
            # this is to translate their norms to ours
            self.mean_for_eval = np.load(pjoin(opt.meta_dir, f'{opt.dataset_name}_mean.npy'))
            self.std_for_eval = np.load(pjoin(opt.meta_dir, f'{opt.dataset_name}_std.npy'))

        self.split_file = pjoin(opt.data_root, f'{split}.txt')
        if mode == 'text_only':
            self.t2m_dataset = TextOnlyDataset(self.opt, self.mean, self.std, self.split_file)
        else:
            self.w_vectorizer = WordVectorizer(pjoin(abs_base_path, 'glove'), 'our_vab')
            self.t2m_dataset = Text2MotionDatasetV2(self.opt, self.mean, self.std, self.split_file, self.w_vectorizer)
            self.num_actions = 1 # dummy placeholder

        assert len(self.t2m_dataset) > 1, 'You loaded an empty dataset, ' \
                                          'it is probably because your data dir has only texts and no motions.\n' \
                                          'To train and evaluate MDM you should get the FULL data as described ' \
                                          'in the README file.'

    def __getitem__(self, item):
        return self.t2m_dataset.__getitem__(item)

    def __len__(self):
        return self.t2m_dataset.__len__()

# A wrapper class for t2m original dataset for MDM purposes
class KIT(HumanML3D):
    def __init__(self, mode, datapath='./dataset/kit_opt.txt', split="train", **kwargs):
        super(KIT, self).__init__(mode, datapath, split, **kwargs)
diffusion_latent/data_loaders/humanml/data/dataset.py
YoungSeng/UnifiedGesture
import pdb

import torch

def lengths_to_mask(lengths, max_len):
    # max_len = max(lengths)
    mask = torch.arange(max_len, device=lengths.device).expand(len(lengths), max_len) < lengths.unsqueeze(1)
    return mask
    

def collate_tensors(batch):
    dims = batch[0].dim()
    max_size = [max([b.size(i) for b in batch]) for i in range(dims)]
    size = (len(batch),) + tuple(max_size)
    canvas = batch[0].new_zeros(size=size)
    for i, b in enumerate(batch):
        sub_tensor = canvas[i]
        for d in range(dims):
            sub_tensor = sub_tensor.narrow(d, 0, b.size(d))
        sub_tensor.add_(b)
    return canvas


def collate(batch):
    notnone_batches = [b for b in batch if b is not None]
    databatch = [b['inp'] for b in notnone_batches]
    if 'lengths' in notnone_batches[0]:
        lenbatch = [b['lengths'] for b in notnone_batches]
    else:
        lenbatch = [len(b['inp'][0][0]) for b in notnone_batches]


    databatchTensor = collate_tensors(databatch)
    lenbatchTensor = torch.as_tensor(lenbatch)
    maskbatchTensor = lengths_to_mask(lenbatchTensor, databatchTensor.shape[-1]).unsqueeze(1).unsqueeze(1) # unqueeze for broadcasting

    motion = databatchTensor
    cond = {'y': {'mask': maskbatchTensor, 'lengths': lenbatchTensor}}

    if 'text' in notnone_batches[0]:
        textbatch = [b['text'] for b in notnone_batches]
        cond['y'].update({'text': textbatch})

    if 'tokens' in notnone_batches[0]:
        textbatch = [b['tokens'] for b in notnone_batches]
        cond['y'].update({'tokens': textbatch})

    if 'action' in notnone_batches[0]:
        actionbatch = [b['action'] for b in notnone_batches]
        cond['y'].update({'action': torch.as_tensor(actionbatch).unsqueeze(1)})

    # collate action textual names
    if 'action_text' in notnone_batches[0]:
        action_text = [b['action_text']for b in notnone_batches]
        cond['y'].update({'action_text': action_text})

    return motion, cond

# an adapter to our collate func
def t2m_collate(batch):
    # batch.sort(key=lambda x: x[3], reverse=True)
    adapted_batch = [{
        'inp': torch.tensor(b[4].T).float().unsqueeze(1), # [seqlen, J] -> [J, 1, seqlen]
        'text': b[2], #b[0]['caption']
        'tokens': b[6],
        'lengths': b[5],
    } for b in batch]
    return collate(adapted_batch)



diffusion_latent/data_loaders/tensors.py
YoungSeng/UnifiedGesture
# This code is based on https://github.com/openai/guided-diffusion
"""
This code started out as a PyTorch port of Ho et al's diffusion models:
https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py

Docstrings have been added, as well as DDIM sampling and a new collection of beta schedules.
"""

import enum
import math
import pdb

import numpy as np
import torch
import torch as th
from copy import deepcopy
from diffusion.nn import mean_flat, sum_flat
from diffusion.losses import normal_kl, discretized_gaussian_log_likelihood
from data_loaders.humanml.scripts import motion_process

def get_named_beta_schedule(schedule_name, num_diffusion_timesteps, scale_betas=1.):
    """
    Get a pre-defined beta schedule for the given name.

    The beta schedule library consists of beta schedules which remain similar
    in the limit of num_diffusion_timesteps.
    Beta schedules may be added, but should not be removed or changed once
    they are committed to maintain backwards compatibility.
    """
    if schedule_name == "linear":
        # Linear schedule from Ho et al, extended to work for any number of
        # diffusion steps.
        scale = scale_betas * 1000 / num_diffusion_timesteps
        beta_start = scale * 0.0001
        beta_end = scale * 0.02
        return np.linspace(
            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
        )
    elif schedule_name == "cosine":
        return betas_for_alpha_bar(
            num_diffusion_timesteps,
            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
        )
    else:
        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")


def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
    """
    Create a beta schedule that discretizes the given alpha_t_bar function,
    which defines the cumulative product of (1-beta) over time from t = [0,1].

    :param num_diffusion_timesteps: the number of betas to produce.
    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
                      produces the cumulative product of (1-beta) up to that
                      part of the diffusion process.
    :param max_beta: the maximum beta to use; use values lower than 1 to
                     prevent singularities.
    """
    betas = []
    for i in range(num_diffusion_timesteps):
        t1 = i / num_diffusion_timesteps
        t2 = (i + 1) / num_diffusion_timesteps
        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
    return np.array(betas)


class ModelMeanType(enum.Enum):
    """
    Which type of output the model predicts.
    """

    PREVIOUS_X = enum.auto()  # the model predicts x_{t-1}
    START_X = enum.auto()  # the model predicts x_0
    EPSILON = enum.auto()  # the model predicts epsilon


class ModelVarType(enum.Enum):
    """
    What is used as the model's output variance.

    The LEARNED_RANGE option has been added to allow the model to predict
    values between FIXED_SMALL and FIXED_LARGE, making its job easier.
    """

    LEARNED = enum.auto()
    FIXED_SMALL = enum.auto()
    FIXED_LARGE = enum.auto()
    LEARNED_RANGE = enum.auto()


class LossType(enum.Enum):
    MSE = enum.auto()  # use raw MSE loss (and KL when learning variances)
    RESCALED_MSE = (
        enum.auto()
    )  # use raw MSE loss (with RESCALED_KL when learning variances)
    KL = enum.auto()  # use the variational lower-bound
    RESCALED_KL = enum.auto()  # like KL, but rescale to estimate the full VLB

    def is_vb(self):
        return self == LossType.KL or self == LossType.RESCALED_KL


class GaussianDiffusion:
    """
    Utilities for training and sampling diffusion models.

    Ported directly from here, and then adapted over time to further experimentation.
    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42

    :param betas: a 1-D numpy array of betas for each diffusion timestep,
                  starting at T and going to 1.
    :param model_mean_type: a ModelMeanType determining what the model outputs.
    :param model_var_type: a ModelVarType determining how variance is output.
    :param loss_type: a LossType determining the loss function to use.
    :param rescale_timesteps: if True, pass floating point timesteps into the
                              model so that they are always scaled like in the
                              original paper (0 to 1000).
    """

    def __init__(
        self,
        *,
        betas,
        model_mean_type,
        model_var_type,
        loss_type,
        rescale_timesteps=False,
        lambda_rcxyz=0.,
        lambda_vel=0.,
        lambda_pose=1.,
        lambda_orient=1.,
        lambda_loc=1.,
        data_rep='rot6d',
        lambda_root_vel=0.,
        lambda_vel_rcxyz=0.,
        lambda_fc=0.,
    ):
        self.model_mean_type = model_mean_type
        self.model_var_type = model_var_type
        self.loss_type = loss_type
        self.rescale_timesteps = rescale_timesteps
        self.data_rep = data_rep

        if data_rep != 'rot_vel' and lambda_pose != 1.:
            raise ValueError('lambda_pose is relevant only when training on velocities!')
        self.lambda_pose = lambda_pose
        self.lambda_orient = lambda_orient
        self.lambda_loc = lambda_loc

        self.lambda_rcxyz = lambda_rcxyz
        self.lambda_vel = lambda_vel
        self.lambda_root_vel = lambda_root_vel
        self.lambda_vel_rcxyz = lambda_vel_rcxyz
        self.lambda_fc = lambda_fc

        if self.lambda_rcxyz > 0. or self.lambda_vel > 0. or self.lambda_root_vel > 0. or \
                self.lambda_vel_rcxyz > 0. or self.lambda_fc > 0.:
            assert self.loss_type == LossType.MSE, 'Geometric losses are supported by MSE loss type only!'

        # Use float64 for accuracy.
        betas = np.array(betas, dtype=np.float64)
        self.betas = betas
        assert len(betas.shape) == 1, "betas must be 1-D"
        assert (betas > 0).all() and (betas <= 1).all()

        self.num_timesteps = int(betas.shape[0])

        alphas = 1.0 - betas
        self.alphas_cumprod = np.cumprod(alphas, axis=0)
        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)

        # calculations for diffusion q(x_t | x_{t-1}) and others
        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)

        # calculations for posterior q(x_{t-1} | x_t, x_0)
        self.posterior_variance = (
            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        # log calculation clipped because the posterior variance is 0 at the
        # beginning of the diffusion chain.
        self.posterior_log_variance_clipped = np.log(
            np.append(self.posterior_variance[1], self.posterior_variance[1:])
        )
        self.posterior_mean_coef1 = (
            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        self.posterior_mean_coef2 = (
            (1.0 - self.alphas_cumprod_prev)
            * np.sqrt(alphas)
            / (1.0 - self.alphas_cumprod)
        )

        self.l2_loss = lambda a, b: (a - b) ** 2  # th.nn.MSELoss(reduction='none')  # must be None for handling mask later on.
        self.smooth_l1_loss = th.nn.SmoothL1Loss(reduction='none')

    def masked_l2(self, a, b, mask):
        # assuming a.shape == b.shape == bs, J, Jdim, seqlen
        # assuming mask.shape == bs, 1, 1, seqlen
        # loss = self.l2_loss(a, b)     # 20221217
        loss = self.smooth_l1_loss(a, b)
        loss = sum_flat(loss * mask.float())  # gives \sigma_euclidean over unmasked elements
        n_entries = a.shape[1] * a.shape[2]
        non_zero_elements = sum_flat(mask) * n_entries
        # print('mask', mask.shape)
        # print('non_zero_elements', non_zero_elements)
        # print('loss', loss)
        mse_loss_val = loss / non_zero_elements
        # print('mse_loss_val', mse_loss_val)
        return mse_loss_val


    def q_mean_variance(self, x_start, t):
        """
        Get the distribution q(x_t | x_0).

        :param x_start: the [N x C x ...] tensor of noiseless inputs.
        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
        """
        mean = (
            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
        )
        variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
        log_variance = _extract_into_tensor(
            self.log_one_minus_alphas_cumprod, t, x_start.shape
        )
        return mean, variance, log_variance

    def q_sample(self, x_start, t, noise=None):
        """
        Diffuse the dataset for a given number of diffusion steps.

        In other words, sample from q(x_t | x_0).

        :param x_start: the initial dataset batch.
        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
        :param noise: if specified, the split-out normal noise.
        :return: A noisy version of x_start.
        """
        if noise is None:
            noise = th.randn_like(x_start)
        assert noise.shape == x_start.shape
        return (
            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
            + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
            * noise
        )

    def q_posterior_mean_variance(self, x_start, x_t, t):
        """
        Compute the mean and variance of the diffusion posterior:

            q(x_{t-1} | x_t, x_0)

        """
        assert x_start.shape == x_t.shape
        posterior_mean = (
            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
        )
        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
        posterior_log_variance_clipped = _extract_into_tensor(
            self.posterior_log_variance_clipped, t, x_t.shape
        )
        assert (
            posterior_mean.shape[0]
            == posterior_variance.shape[0]
            == posterior_log_variance_clipped.shape[0]
            == x_start.shape[0]
        )
        return posterior_mean, posterior_variance, posterior_log_variance_clipped

    def p_mean_variance(
        self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None
    ):
        """
        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
        the initial x, x_0.

        :param model: the model, which takes a signal and a batch of timesteps
                      as input.
        :param x: the [N x C x ...] tensor at time t.
        :param t: a 1-D Tensor of timesteps.
        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample. Applies before
            clip_denoised.
        :param model_kwargs: if not None, a dict of extra keyword arguments to
            pass to the model. This can be used for conditioning.
        :return: a dict with the following keys:
                 - 'mean': the model mean output.
                 - 'variance': the model variance output.
                 - 'log_variance': the log of 'variance'.
                 - 'pred_xstart': the prediction for x_0.
        """
        if model_kwargs is None:
            model_kwargs = {}

        B, C = x.shape[:2]
        assert t.shape == (B,)
        model_output = model(x, self._scale_timesteps(t), **model_kwargs)

        if 'inpainting_mask' in model_kwargs['y'].keys() and 'inpainted_motion' in model_kwargs['y'].keys():
            inpainting_mask, inpainted_motion = model_kwargs['y']['inpainting_mask'], model_kwargs['y']['inpainted_motion']
            assert self.model_mean_type == ModelMeanType.START_X, 'This feature supports only X_start pred for mow!'
            assert model_output.shape == inpainting_mask.shape == inpainted_motion.shape
            model_output = (model_output * ~inpainting_mask) + (inpainted_motion * inpainting_mask)
            # print('model_output', model_output.shape, model_output)
            # print('inpainting_mask', inpainting_mask.shape, inpainting_mask[0,0,0,:])
            # print('inpainted_motion', inpainted_motion.shape, inpainted_motion)

        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
            assert model_output.shape == (B, C * 2, *x.shape[2:])
            model_output, model_var_values = th.split(model_output, C, dim=1)
            if self.model_var_type == ModelVarType.LEARNED:
                model_log_variance = model_var_values
                model_variance = th.exp(model_log_variance)
            else:
                min_log = _extract_into_tensor(
                    self.posterior_log_variance_clipped, t, x.shape
                )
                max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
                # The model_var_values is [-1, 1] for [min_var, max_var].
                frac = (model_var_values + 1) / 2
                model_log_variance = frac * max_log + (1 - frac) * min_log
                model_variance = th.exp(model_log_variance)
        else:
            model_variance, model_log_variance = {
                # for fixedlarge, we set the initial (log-)variance like so
                # to get a better decoder log likelihood.
                ModelVarType.FIXED_LARGE: (
                    np.append(self.posterior_variance[1], self.betas[1:]),
                    np.log(np.append(self.posterior_variance[1], self.betas[1:])),
                ),
                ModelVarType.FIXED_SMALL: (
                    self.posterior_variance,
                    self.posterior_log_variance_clipped,
                ),
            }[self.model_var_type]
            # print('model_variance', model_variance)
            # print('model_log_variance',model_log_variance)
            # print('self.posterior_variance', self.posterior_variance)
            # print('self.posterior_log_variance_clipped', self.posterior_log_variance_clipped)
            # print('self.model_var_type', self.model_var_type)


            model_variance = _extract_into_tensor(model_variance, t, x.shape)
            model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)

        def process_xstart(x):
            if denoised_fn is not None:
                x = denoised_fn(x)
            if clip_denoised:
                # print('clip_denoised', clip_denoised)
                return x.clamp(-1, 1)
            return x

        if self.model_mean_type == ModelMeanType.PREVIOUS_X:
            pred_xstart = process_xstart(
                self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)
            )
            model_mean = model_output
        elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]:  # THIS IS US!
            if self.model_mean_type == ModelMeanType.START_X:
                pred_xstart = process_xstart(model_output)
            else:
                pred_xstart = process_xstart(
                    self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
                )
            model_mean, _, _ = self.q_posterior_mean_variance(
                x_start=pred_xstart, x_t=x, t=t
            )
        else:
            raise NotImplementedError(self.model_mean_type)

        assert (
            model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
        )
        return {
            "mean": model_mean,
            "variance": model_variance,
            "log_variance": model_log_variance,
            "pred_xstart": pred_xstart,
        }

    def _predict_xstart_from_eps(self, x_t, t, eps):
        assert x_t.shape == eps.shape
        return (
            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
        )

    def _predict_xstart_from_xprev(self, x_t, t, xprev):
        assert x_t.shape == xprev.shape
        return (  # (xprev - coef2*x_t) / coef1
            _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
            - _extract_into_tensor(
                self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
            )
            * x_t
        )

    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
        return (
            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
            - pred_xstart
        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)

    def _scale_timesteps(self, t):
        if self.rescale_timesteps:
            return t.float() * (1000.0 / self.num_timesteps)
        return t

    def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
        """
        Compute the mean for the previous step, given a function cond_fn that
        computes the gradient of a conditional log probability with respect to
        x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
        condition on y.

        This uses the conditioning strategy from Sohl-Dickstein et al. (2015).
        """
        gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs)
        new_mean = (
            p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float()
        )
        return new_mean

    def condition_mean_with_grad(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
        """
        Compute the mean for the previous step, given a function cond_fn that
        computes the gradient of a conditional log probability with respect to
        x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
        condition on y.

        This uses the conditioning strategy from Sohl-Dickstein et al. (2015).
        """
        gradient = cond_fn(x, t, p_mean_var, **model_kwargs)
        new_mean = (
            p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float()
        )
        return new_mean

    def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
        """
        Compute what the p_mean_variance output would have been, should the
        model's score function be conditioned by cond_fn.

        See condition_mean() for details on cond_fn.

        Unlike condition_mean(), this instead uses the conditioning strategy
        from Song et al (2020).
        """
        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)

        eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])
        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(
            x, self._scale_timesteps(t), **model_kwargs
        )

        out = p_mean_var.copy()
        out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)
        out["mean"], _, _ = self.q_posterior_mean_variance(
            x_start=out["pred_xstart"], x_t=x, t=t
        )
        return out

    def condition_score_with_grad(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
        """
        Compute what the p_mean_variance output would have been, should the
        model's score function be conditioned by cond_fn.

        See condition_mean() for details on cond_fn.

        Unlike condition_mean(), this instead uses the conditioning strategy
        from Song et al (2020).
        """
        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)

        eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])
        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(
            x, t, p_mean_var, **model_kwargs
        )

        out = p_mean_var.copy()
        out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)
        out["mean"], _, _ = self.q_posterior_mean_variance(
            x_start=out["pred_xstart"], x_t=x, t=t
        )
        return out

    def p_sample(
        self,
        model,
        x,
        t,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
        const_noise=False,
    ):
        """
        Sample x_{t-1} from the model at the given timestep.

        :param model: the model to sample from.
        :param x: the current tensor at x_{t-1}.
        :param t: the value of t, starting at 0 for the first diffusion step.
        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample.
        :param cond_fn: if not None, this is a gradient function that acts
                        similarly to the model.
        :param model_kwargs: if not None, a dict of extra keyword arguments to
            pass to the model. This can be used for conditioning.
        :return: a dict containing the following keys:
                 - 'sample': a random sample from the model.
                 - 'pred_xstart': a prediction of x_0.
        """
        out = self.p_mean_variance(
            model,
            x,
            t,
            clip_denoised=clip_denoised,
            denoised_fn=denoised_fn,
            model_kwargs=model_kwargs,
        )       # 'mean' （1, 135, 1, 240）, 'variance', 'log_variance', 'pred_xstart'
        noise = th.randn_like(x)
        # print('const_noise', const_noise)
        if const_noise:
            noise = noise[[0]].repeat(x.shape[0], 1, 1, 1)

        nonzero_mask = (
            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
        )  # no noise when t == 0
        if cond_fn is not None:
            out["mean"] = self.condition_mean(
                cond_fn, out, x, t, model_kwargs=model_kwargs
            )
        # print('mean', out["mean"].shape, out["mean"])
        # print('log_variance', out["log_variance"].shape, out["log_variance"])
        # print('nonzero_mask', nonzero_mask.shape, nonzero_mask)
        sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
        return {"sample": sample, "pred_xstart": out["pred_xstart"]}

    def p_sample_with_grad(
        self,
        model,
        x,
        t,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
    ):
        """
        Sample x_{t-1} from the model at the given timestep.

        :param model: the model to sample from.
        :param x: the current tensor at x_{t-1}.
        :param t: the value of t, starting at 0 for the first diffusion step.
        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample.
        :param cond_fn: if not None, this is a gradient function that acts
                        similarly to the model.
        :param model_kwargs: if not None, a dict of extra keyword arguments to
            pass to the model. This can be used for conditioning.
        :return: a dict containing the following keys:
                 - 'sample': a random sample from the model.
                 - 'pred_xstart': a prediction of x_0.
        """
        with th.enable_grad():
            x = x.detach().requires_grad_()
            out = self.p_mean_variance(
                model,
                x,
                t,
                clip_denoised=clip_denoised,
                denoised_fn=denoised_fn,
                model_kwargs=model_kwargs,
            )
            noise = th.randn_like(x)
            nonzero_mask = (
                (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
            )  # no noise when t == 0
            if cond_fn is not None:
                out["mean"] = self.condition_mean_with_grad(
                    cond_fn, out, x, t, model_kwargs=model_kwargs
                )
        sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
        # return {"sample": sample, "pred_xstart": out["pred_xstart"].detach()}
        return {"sample": sample, "pred_xstart": out["pred_xstart"]}        # 20230428

    def p_sample_loop(
        self,
        model,
        shape,
        noise=None,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
        device=None,
        progress=False,
        skip_timesteps=0,
        init_image=None,
        randomize_class=False,
        cond_fn_with_grad=False,
        dump_steps=None,
        const_noise=False,
        num_timesteps=1000,
    ):
        """
        Generate samples from the model.

        :param model: the model module.
        :param shape: the shape of the samples, (N, C, H, W).
        :param noise: if specified, the noise from the encoder to sample.
                      Should be of the same shape as `shape`.
        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample.
        :param cond_fn: if not None, this is a gradient function that acts
                        similarly to the model.
        :param model_kwargs: if not None, a dict of extra keyword arguments to
            pass to the model. This can be used for conditioning.
        :param device: if specified, the device to create the samples on.
                       If not specified, use a model parameter's device.
        :param progress: if True, show a tqdm progress bar.
        :param const_noise: If True, will noise all samples with the same noise throughout sampling
        :return: a non-differentiable batch of samples.
        """
        final = None
        if dump_steps is not None:
            dump = []

        for i, sample in enumerate(self.p_sample_loop_progressive(
            model,
            shape,
            noise=noise,
            clip_denoised=clip_denoised,
            denoised_fn=denoised_fn,
            cond_fn=cond_fn,
            model_kwargs=model_kwargs,
            device=device,
            progress=progress,
            skip_timesteps=skip_timesteps,
            init_image=init_image,
            randomize_class=randomize_class,
            cond_fn_with_grad=cond_fn_with_grad,
            const_noise=const_noise,
            num_timesteps=num_timesteps,
        )):
            if dump_steps is not None and i in dump_steps:
                dump.append(deepcopy(sample["sample"]))
            final = sample
        if dump_steps is not None:
            return dump
        return final["sample"]

    def p_sample_loop_progressive(
        self,
        model,
        shape,
        noise=None,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
        device=None,
        progress=False,
        skip_timesteps=0,
        init_image=None,
        randomize_class=False,
        cond_fn_with_grad=False,
        const_noise=False,
        num_timesteps=1000,
    ):
        """
        Generate samples from the model and yield intermediate samples from
        each timestep of diffusion.

        Arguments are the same as p_sample_loop().
        Returns a generator over dicts, where each dict is the return value of
        p_sample().
        """
        if device is None:
            device = next(model.parameters()).device
        assert isinstance(shape, (tuple, list))
        if noise is not None:
            img = noise
        else:
            img = th.randn(*shape, device=device)

        if skip_timesteps and init_image is None:
            init_image = th.zeros_like(img)

        indices = list(range(num_timesteps - skip_timesteps))[::-1]

        if init_image is not None:
            my_t = th.ones([shape[0]], device=device, dtype=th.long) * indices[0]
            img = self.q_sample(init_image, my_t, img)

        if progress:
            # Lazy import so that we don't depend on tqdm.
            from tqdm.auto import tqdm

            indices = tqdm(indices)

        for i in indices:
            t = th.tensor([i] * shape[0], device=device)
            if randomize_class and 'y' in model_kwargs:
                model_kwargs['y'] = th.randint(low=0, high=model.num_classes,
                                               size=model_kwargs['y'].shape,
                                               device=model_kwargs['y'].device)
            if not cond_fn_with_grad:
                with th.no_grad():
                    # sample_fn = self.p_sample_with_grad if cond_fn_with_grad else self.p_sample       # 20230429
                    sample_fn = self.p_sample
                    out = sample_fn(
                        model,
                        img,
                        t,
                        clip_denoised=clip_denoised,
                        denoised_fn=denoised_fn,
                        cond_fn=cond_fn,
                        model_kwargs=model_kwargs,
                        # const_noise=const_noise,
                    )
                    yield out
                    img = out["sample"]
            else:           # 20230428
                # sample_fn = self.p_sample_with_grad if cond_fn_with_grad else self.p_sample       # 20230429
                sample_fn = self.p_sample
                out = sample_fn(
                    model,
                    img,
                    t,
                    clip_denoised=clip_denoised,
                    denoised_fn=denoised_fn,
                    cond_fn=cond_fn,
                    model_kwargs=model_kwargs,
                    # const_noise=const_noise,
                )
                yield out
                img = out["sample"]

    def ddim_sample(
        self,
        model,
        x,
        t,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
        eta=0.0,
    ):
        """
        Sample x_{t-1} from the model using DDIM.

        Same usage as p_sample().
        """
        out_orig = self.p_mean_variance(
            model,
            x,
            t,
            clip_denoised=clip_denoised,
            denoised_fn=denoised_fn,
            model_kwargs=model_kwargs,
        )
        if cond_fn is not None:
            out = self.condition_score(cond_fn, out_orig, x, t, model_kwargs=model_kwargs)
        else:
            out = out_orig

        # Usually our model outputs epsilon, but we re-derive it
        # in case we used x_start or x_prev prediction.
        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])

        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
        sigma = (
            eta
            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
        )
        # Equation 12.
        noise = th.randn_like(x)
        mean_pred = (
            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
            + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps
        )
        nonzero_mask = (
            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
        )  # no noise when t == 0
        sample = mean_pred + nonzero_mask * sigma * noise
        return {"sample": sample, "pred_xstart": out_orig["pred_xstart"]}

    def ddim_sample_with_grad(
        self,
        model,
        x,
        t,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
        eta=0.0,
    ):
        """
        Sample x_{t-1} from the model using DDIM.

        Same usage as p_sample().
        """
        with th.enable_grad():
            x = x.detach().requires_grad_()
            out_orig = self.p_mean_variance(
                model,
                x,
                t,
                clip_denoised=clip_denoised,
                denoised_fn=denoised_fn,
                model_kwargs=model_kwargs,
            )
            if cond_fn is not None:
                out = self.condition_score_with_grad(cond_fn, out_orig, x, t,
                                                     model_kwargs=model_kwargs)
            else:
                out = out_orig

        out["pred_xstart"] = out["pred_xstart"].detach()

        # Usually our model outputs epsilon, but we re-derive it
        # in case we used x_start or x_prev prediction.
        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])

        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
        sigma = (
            eta
            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
        )
        # Equation 12.
        noise = th.randn_like(x)
        mean_pred = (
            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
            + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps
        )
        nonzero_mask = (
            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
        )  # no noise when t == 0
        sample = mean_pred + nonzero_mask * sigma * noise
        return {"sample": sample, "pred_xstart": out_orig["pred_xstart"].detach()}

    def ddim_reverse_sample(
        self,
        model,
        x,
        t,
        clip_denoised=True,
        denoised_fn=None,
        model_kwargs=None,
        eta=0.0,
    ):
        """
        Sample x_{t+1} from the model using DDIM reverse ODE.
        """
        assert eta == 0.0, "Reverse ODE only for deterministic path"
        out = self.p_mean_variance(
            model,
            x,
            t,
            clip_denoised=clip_denoised,
            denoised_fn=denoised_fn,
            model_kwargs=model_kwargs,
        )
        # Usually our model outputs epsilon, but we re-derive it
        # in case we used x_start or x_prev prediction.
        eps = (
            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
            - out["pred_xstart"]
        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
        alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)

        # Equation 12. reversed
        mean_pred = (
            out["pred_xstart"] * th.sqrt(alpha_bar_next)
            + th.sqrt(1 - alpha_bar_next) * eps
        )

        return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}

    def ddim_sample_loop(
        self,
        model,
        shape,
        noise=None,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
        device=None,
        progress=False,
        eta=0.0,
        skip_timesteps=0,
        init_image=None,
        randomize_class=False,
        cond_fn_with_grad=False,
        dump_steps=None,
        const_noise=False,
    ):
        """
        Generate samples from the model using DDIM.

        Same usage as p_sample_loop().
        """
        if dump_steps is not None:
            raise NotImplementedError()
        if const_noise == True:
            raise NotImplementedError()

        final = None
        for sample in self.ddim_sample_loop_progressive(
            model,
            shape,
            noise=noise,
            clip_denoised=clip_denoised,
            denoised_fn=denoised_fn,
            cond_fn=cond_fn,
            model_kwargs=model_kwargs,
            device=device,
            progress=progress,
            eta=eta,
            skip_timesteps=skip_timesteps,
            init_image=init_image,
            randomize_class=randomize_class,
            cond_fn_with_grad=cond_fn_with_grad,
        ):
            final = sample
        return final["sample"]

    def ddim_sample_loop_progressive(
        self,
        model,
        shape,
        noise=None,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
        device=None,
        progress=False,
        eta=0.0,
        skip_timesteps=0,
        init_image=None,
        randomize_class=False,
        cond_fn_with_grad=False,
    ):
        """
        Use DDIM to sample from the model and yield intermediate samples from
        each timestep of DDIM.

        Same usage as p_sample_loop_progressive().
        """
        if device is None:
            device = next(model.parameters()).device
        assert isinstance(shape, (tuple, list))
        if noise is not None:
            img = noise
        else:
            img = th.randn(*shape, device=device)

        if skip_timesteps and init_image is None:
            init_image = th.zeros_like(img)

        indices = list(range(self.num_timesteps - skip_timesteps))[::-1]

        if init_image is not None:
            my_t = th.ones([shape[0]], device=device, dtype=th.long) * indices[0]
            img = self.q_sample(init_image, my_t, img)

        if progress:
            # Lazy import so that we don't depend on tqdm.
            from tqdm.auto import tqdm

            indices = tqdm(indices)

        for i in indices:
            t = th.tensor([i] * shape[0], device=device)
            if randomize_class and 'y' in model_kwargs:
                model_kwargs['y'] = th.randint(low=0, high=model.num_classes,
                                               size=model_kwargs['y'].shape,
                                               device=model_kwargs['y'].device)
            with th.no_grad():
                sample_fn = self.ddim_sample_with_grad if cond_fn_with_grad else self.ddim_sample
                out = sample_fn(
                    model,
                    img,
                    t,
                    clip_denoised=clip_denoised,
                    denoised_fn=denoised_fn,
                    cond_fn=cond_fn,
                    model_kwargs=model_kwargs,
                    eta=eta,
                )
                yield out
                img = out["sample"]

    def plms_sample(
        self,
        model,
        x,
        t,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
        cond_fn_with_grad=False,
        order=2,
        old_out=None,
    ):
        """
        Sample x_{t-1} from the model using Pseudo Linear Multistep.

        Same usage as p_sample().
        """
        if not int(order) or not 1 <= order <= 4:
            raise ValueError('order is invalid (should be int from 1-4).')

        def get_model_output(x, t):
            with th.set_grad_enabled(cond_fn_with_grad and cond_fn is not None):
                x = x.detach().requires_grad_() if cond_fn_with_grad else x
                out_orig = self.p_mean_variance(
                    model,
                    x,
                    t,
                    clip_denoised=clip_denoised,
                    denoised_fn=denoised_fn,
                    model_kwargs=model_kwargs,
                )
                if cond_fn is not None:
                    if cond_fn_with_grad:
                        out = self.condition_score_with_grad(cond_fn, out_orig, x, t, model_kwargs=model_kwargs)
                        x = x.detach()
                    else:
                        out = self.condition_score(cond_fn, out_orig, x, t, model_kwargs=model_kwargs)
                else:
                    out = out_orig

            # Usually our model outputs epsilon, but we re-derive it
            # in case we used x_start or x_prev prediction.
            eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
            return eps, out, out_orig

        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
        eps, out, out_orig = get_model_output(x, t)

        if order > 1 and old_out is None:
            # Pseudo Improved Euler
            old_eps = [eps]
            mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_prev) + th.sqrt(1 - alpha_bar_prev) * eps
            eps_2, _, _ = get_model_output(mean_pred, t - 1)
            eps_prime = (eps + eps_2) / 2
            pred_prime = self._predict_xstart_from_eps(x, t, eps_prime)
            mean_pred = pred_prime * th.sqrt(alpha_bar_prev) + th.sqrt(1 - alpha_bar_prev) * eps_prime
        else:
            # Pseudo Linear Multistep (Adams-Bashforth)
            old_eps = old_out["old_eps"]
            old_eps.append(eps)
            cur_order = min(order, len(old_eps))
            if cur_order == 1:
                eps_prime = old_eps[-1]
            elif cur_order == 2:
                eps_prime = (3 * old_eps[-1] - old_eps[-2]) / 2
            elif cur_order == 3:
                eps_prime = (23 * old_eps[-1] - 16 * old_eps[-2] + 5 * old_eps[-3]) / 12
            elif cur_order == 4:
                eps_prime = (55 * old_eps[-1] - 59 * old_eps[-2] + 37 * old_eps[-3] - 9 * old_eps[-4]) / 24
            else:
                raise RuntimeError('cur_order is invalid.')
            pred_prime = self._predict_xstart_from_eps(x, t, eps_prime)
            mean_pred = pred_prime * th.sqrt(alpha_bar_prev) + th.sqrt(1 - alpha_bar_prev) * eps_prime

        if len(old_eps) >= order:
            old_eps.pop(0)

        nonzero_mask = (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
        sample = mean_pred * nonzero_mask + out["pred_xstart"] * (1 - nonzero_mask)

        return {"sample": sample, "pred_xstart": out_orig["pred_xstart"], "old_eps": old_eps}

    def plms_sample_loop(
        self,
        model,
        shape,
        noise=None,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
        device=None,
        progress=False,
        skip_timesteps=0,
        init_image=None,
        randomize_class=False,
        cond_fn_with_grad=False,
        order=2,
    ):
        """
        Generate samples from the model using Pseudo Linear Multistep.

        Same usage as p_sample_loop().
        """
        final = None
        for sample in self.plms_sample_loop_progressive(
            model,
            shape,
            noise=noise,
            clip_denoised=clip_denoised,
            denoised_fn=denoised_fn,
            cond_fn=cond_fn,
            model_kwargs=model_kwargs,
            device=device,
            progress=progress,
            skip_timesteps=skip_timesteps,
            init_image=init_image,
            randomize_class=randomize_class,
            cond_fn_with_grad=cond_fn_with_grad,
            order=order,
        ):
            final = sample
        return final["sample"]

    def plms_sample_loop_progressive(
        self,
        model,
        shape,
        noise=None,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
        device=None,
        progress=False,
        skip_timesteps=0,
        init_image=None,
        randomize_class=False,
        cond_fn_with_grad=False,
        order=2,
    ):
        """
        Use PLMS to sample from the model and yield intermediate samples from each
        timestep of PLMS.

        Same usage as p_sample_loop_progressive().
        """
        if device is None:
            device = next(model.parameters()).device
        assert isinstance(shape, (tuple, list))
        if noise is not None:
            img = noise
        else:
            img = th.randn(*shape, device=device)

        if skip_timesteps and init_image is None:
            init_image = th.zeros_like(img)

        indices = list(range(self.num_timesteps - skip_timesteps))[::-1]

        if init_image is not None:
            my_t = th.ones([shape[0]], device=device, dtype=th.long) * indices[0]
            img = self.q_sample(init_image, my_t, img)

        if progress:
            # Lazy import so that we don't depend on tqdm.
            from tqdm.auto import tqdm

            indices = tqdm(indices)

        old_out = None

        for i in indices:
            t = th.tensor([i] * shape[0], device=device)
            if randomize_class and 'y' in model_kwargs:
                model_kwargs['y'] = th.randint(low=0, high=model.num_classes,
                                               size=model_kwargs['y'].shape,
                                               device=model_kwargs['y'].device)
            with th.no_grad():
                out = self.plms_sample(
                    model,
                    img,
                    t,
                    clip_denoised=clip_denoised,
                    denoised_fn=denoised_fn,
                    cond_fn=cond_fn,
                    model_kwargs=model_kwargs,
                    cond_fn_with_grad=cond_fn_with_grad,
                    order=order,
                    old_out=old_out,
                )
                yield out
                old_out = out
                img = out["sample"]

    def _vb_terms_bpd(
        self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None
    ):
        """
        Get a term for the variational lower-bound.

        The resulting units are bits (rather than nats, as one might expect).
        This allows for comparison to other papers.

        :return: a dict with the following keys:
                 - 'output': a shape [N] tensor of NLLs or KLs.
                 - 'pred_xstart': the x_0 predictions.
        """
        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
            x_start=x_start, x_t=x_t, t=t
        )
        out = self.p_mean_variance(
            model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
        )
        kl = normal_kl(
            true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
        )
        kl = mean_flat(kl) / np.log(2.0)

        decoder_nll = -discretized_gaussian_log_likelihood(
            x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
        )
        assert decoder_nll.shape == x_start.shape
        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)

        # At the first timestep return the decoder NLL,
        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
        output = th.where((t == 0), decoder_nll, kl)
        return {"output": output, "pred_xstart": out["pred_xstart"]}

    def training_losses(self, model, x_start, t, model_kwargs=None, noise=None, dataset=None):
        """
        Compute training losses for a single timestep.

        :param model: the model to evaluate loss on.
        :param x_start: the [N x C x ...] tensor of inputs.
        :param t: a batch of timestep indices.
        :param model_kwargs: if not None, a dict of extra keyword arguments to
            pass to the model. This can be used for conditioning.
        :param noise: if specified, the specific Gaussian noise to try to remove.
        :return: a dict with the key "loss" containing a tensor of shape [N].
                 Some mean or variance settings may also have other keys.
        """

        # enc = model.model._modules['module']
        enc = model.model
        mask = model_kwargs['y']['mask']
        # get_xyz = lambda sample: enc.rot2xyz(sample, mask=None, pose_rep=enc.pose_rep, translation=enc.translation,
        #                                      glob=enc.glob,
        #                                      # jointstype='vertices',  # 3.4 iter/sec # USED ALSO IN MotionCLIP
        #                                      jointstype='smpl',  # 3.4 iter/sec
        #                                      vertstrans=False)

        if model_kwargs is None:
            model_kwargs = {}
        if noise is None:
            noise = th.randn_like(x_start)
        x_t = self.q_sample(x_start, t, noise=noise)        # torch.Size([64, 251, 1, 196]), add noisy

        terms = {}

        if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:     # LossType.MSE
            terms["loss"] = self._vb_terms_bpd(
                model=model,
                x_start=x_start,
                x_t=x_t,
                t=t,
                clip_denoised=False,
                model_kwargs=model_kwargs,
            )["output"]
            if self.loss_type == LossType.RESCALED_KL:
                terms["loss"] *= self.num_timesteps
        elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
            model_output = model(x_t, self._scale_timesteps(t), **model_kwargs)

            if self.model_var_type in [     # ModelVarType.FIXED_SMALL: 2
                ModelVarType.LEARNED,
                ModelVarType.LEARNED_RANGE,
            ]:
                B, C = x_t.shape[:2]
                assert model_output.shape == (B, C * 2, *x_t.shape[2:])
                model_output, model_var_values = th.split(model_output, C, dim=1)
                # Learn the variance using the variational bound, but don't let
                # it affect our mean prediction.
                frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
                terms["vb"] = self._vb_terms_bpd(
                    model=lambda *args, r=frozen_out: r,
                    x_start=x_start,
                    x_t=x_t,
                    t=t,
                    clip_denoised=False,
                )["output"]
                if self.loss_type == LossType.RESCALED_MSE:
                    # Divide by 1000 for equivalence with initial implementation.
                    # Without a factor of 1/1000, the VB term hurts the MSE term.
                    terms["vb"] *= self.num_timesteps / 1000.0

            target = {
                ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(
                    x_start=x_start, x_t=x_t, t=t
                )[0],
                ModelMeanType.START_X: x_start,
                ModelMeanType.EPSILON: noise,
            }[self.model_mean_type]     # ModelMeanType.START_X: 2
            assert model_output.shape == target.shape == x_start.shape  # [bs, njoints, nfeats, nframes]

            # pdb.set_trace()     # target (2, 135, 1, 240)

            terms["rot_mse"] = self.masked_l2(target, model_output, mask) # mean_flat(rot_mse)      # [64, 251, 1, 196], -, [64, 1, 1, 196]

            target_xyz, model_output_xyz = None, None

            if self.lambda_rcxyz > 0.:      # 0.0
                target_xyz = get_xyz(target)  # [bs, nvertices(vertices)/njoints(smpl), 3, nframes]
                model_output_xyz = get_xyz(model_output)  # [bs, nvertices, 3, nframes]
                terms["rcxyz_mse"] = self.masked_l2(target_xyz, model_output_xyz, mask)  # mean_flat((target_xyz - model_output_xyz) ** 2)

            if self.lambda_vel_rcxyz > 0.:      # 0.0
                if self.data_rep == 'rot6d' and dataset.dataname in ['humanact12', 'uestc']:
                    target_xyz = get_xyz(target) if target_xyz is None else target_xyz
                    model_output_xyz = get_xyz(model_output) if model_output_xyz is None else model_output_xyz
                    target_xyz_vel = (target_xyz[:, :, :, 1:] - target_xyz[:, :, :, :-1])
                    model_output_xyz_vel = (model_output_xyz[:, :, :, 1:] - model_output_xyz[:, :, :, :-1])
                    terms["vel_xyz_mse"] = self.masked_l2(target_xyz_vel, model_output_xyz_vel, mask[:, :, :, 1:])

            if self.lambda_fc > 0.:     # 0.0
                torch.autograd.set_detect_anomaly(True)
                if self.data_rep == 'rot6d' and dataset.dataname in ['humanact12', 'uestc']:
                    target_xyz = get_xyz(target) if target_xyz is None else target_xyz
                    model_output_xyz = get_xyz(model_output) if model_output_xyz is None else model_output_xyz
                    # 'L_Ankle',  # 7, 'R_Ankle',  # 8 , 'L_Foot',  # 10, 'R_Foot',  # 11
                    l_ankle_idx, r_ankle_idx, l_foot_idx, r_foot_idx = 7, 8, 10, 11
                    relevant_joints = [l_ankle_idx, l_foot_idx, r_ankle_idx, r_foot_idx]
                    gt_joint_xyz = target_xyz[:, relevant_joints, :, :]  # [BatchSize, 4, 3, Frames]
                    gt_joint_vel = torch.linalg.norm(gt_joint_xyz[:, :, :, 1:] - gt_joint_xyz[:, :, :, :-1], axis=2)  # [BatchSize, 4, Frames]
                    fc_mask = torch.unsqueeze((gt_joint_vel <= 0.01), dim=2).repeat(1, 1, 3, 1)
                    pred_joint_xyz = model_output_xyz[:, relevant_joints, :, :]  # [BatchSize, 4, 3, Frames]
                    pred_vel = pred_joint_xyz[:, :, :, 1:] - pred_joint_xyz[:, :, :, :-1]
                    pred_vel[~fc_mask] = 0
                    terms["fc"] = self.masked_l2(pred_vel,
                                                 torch.zeros(pred_vel.shape, device=pred_vel.device),
                                                 mask[:, :, :, 1:])
            if self.lambda_vel > 0.:        # 0.0
                target_vel = (target[..., 1:] - target[..., :-1])
                model_output_vel = (model_output[..., 1:] - model_output[..., :-1])
                terms["vel_mse"] = self.masked_l2(target_vel[:, :-1, :, :], # Remove last joint, is the root location!
                                                  model_output_vel[:, :-1, :, :],
                                                  mask[:, :, :, 1:])  # mean_flat((target_vel - model_output_vel) ** 2)

            terms["loss"] = terms["rot_mse"] + terms.get('vb', 0.) +\
                            (self.lambda_vel * terms.get('vel_mse', 0.)) +\
                            (self.lambda_rcxyz * terms.get('rcxyz_mse', 0.)) + \
                            (self.lambda_fc * terms.get('fc', 0.))

        else:
            raise NotImplementedError(self.loss_type)

        return terms

    def fc_loss_rot_repr(self, gt_xyz, pred_xyz, mask):
        def to_np_cpu(x):
            return x.detach().cpu().numpy()
        """
        pose_xyz: SMPL batch tensor of shape: [BatchSize, 24, 3, Frames]
        """
        # 'L_Ankle',  # 7, 'R_Ankle',  # 8 , 'L_Foot',  # 10, 'R_Foot',  # 11

        l_ankle_idx, r_ankle_idx = 7, 8
        l_foot_idx, r_foot_idx = 10, 11
        """ Contact calculated by 'Kfir Method' Commented code)"""
        # contact_signal = torch.zeros((pose_xyz.shape[0], pose_xyz.shape[3], 2), device=pose_xyz.device) # [BatchSize, Frames, 2]
        # left_xyz = 0.5 * (pose_xyz[:, l_ankle_idx, :, :] + pose_xyz[:, l_foot_idx, :, :]) # [BatchSize, 3, Frames]
        # right_xyz = 0.5 * (pose_xyz[:, r_ankle_idx, :, :] + pose_xyz[:, r_foot_idx, :, :])
        # left_z, right_z = left_xyz[:, 2, :], right_xyz[:, 2, :] # [BatchSize, Frames]
        # left_velocity = torch.linalg.norm(left_xyz[:, :, 2:] - left_xyz[:, :, :-2], axis=1)  # [BatchSize, Frames]
        # right_velocity = torch.linalg.norm(left_xyz[:, :, 2:] - left_xyz[:, :, :-2], axis=1)
        #
        # left_z_mask = left_z <= torch.mean(torch.sort(left_z)[0][:, :left_z.shape[1] // 5], axis=-1)
        # left_z_mask = torch.stack([left_z_mask, left_z_mask], dim=-1) # [BatchSize, Frames, 2]
        # left_z_mask[:, :, 1] = False  # Blank right side
        # contact_signal[left_z_mask] = 0.4
        #
        # right_z_mask = right_z <= torch.mean(torch.sort(right_z)[0][:, :right_z.shape[1] // 5], axis=-1)
        # right_z_mask = torch.stack([right_z_mask, right_z_mask], dim=-1) # [BatchSize, Frames, 2]
        # right_z_mask[:, :, 0] = False  # Blank left side
        # contact_signal[right_z_mask] = 0.4
        # contact_signal[left_z <= (torch.mean(torch.sort(left_z)[:left_z.shape[0] // 5]) + 20), 0] = 1
        # contact_signal[right_z <= (torch.mean(torch.sort(right_z)[:right_z.shape[0] // 5]) + 20), 1] = 1

        # plt.plot(to_np_cpu(left_z[0]), label='left_z')
        # plt.plot(to_np_cpu(left_velocity[0]), label='left_velocity')
        # plt.plot(to_np_cpu(contact_signal[0, :, 0]), label='left_fc')
        # plt.grid()
        # plt.legend()
        # plt.show()
        # plt.plot(to_np_cpu(right_z[0]), label='right_z')
        # plt.plot(to_np_cpu(right_velocity[0]), label='right_velocity')
        # plt.plot(to_np_cpu(contact_signal[0, :, 1]), label='right_fc')
        # plt.grid()
        # plt.legend()
        # plt.show()

        gt_joint_xyz = gt_xyz[:, [l_ankle_idx, l_foot_idx, r_ankle_idx, r_foot_idx], :, :]  # [BatchSize, 4, 3, Frames]
        gt_joint_vel = torch.linalg.norm(gt_joint_xyz[:, :, :, 1:] - gt_joint_xyz[:, :, :, :-1], axis=2)  # [BatchSize, 4, Frames]
        fc_mask = (gt_joint_vel <= 0.01)
        pred_joint_xyz = pred_xyz[:, [l_ankle_idx, l_foot_idx, r_ankle_idx, r_foot_idx], :, :]  # [BatchSize, 4, 3, Frames]
        pred_joint_vel = torch.linalg.norm(pred_joint_xyz[:, :, :, 1:] - pred_joint_xyz[:, :, :, :-1], axis=2)  # [BatchSize, 4, Frames]
        pred_joint_vel[~fc_mask] = 0  # Blank non-contact velocities frames. [BS,4,FRAMES]
        pred_joint_vel = torch.unsqueeze(pred_joint_vel, dim=2)

        """DEBUG CODE"""
        # print(f'mask: {mask.shape}')
        # print(f'pred_joint_vel: {pred_joint_vel.shape}')
        # plt.title(f'Joint: {joint_idx}')
        # plt.plot(to_np_cpu(gt_joint_vel[0]), label='velocity')
        # plt.plot(to_np_cpu(fc_mask[0]), label='fc')
        # plt.grid()
        # plt.legend()
        # plt.show()
        return self.masked_l2(pred_joint_vel, torch.zeros(pred_joint_vel.shape, device=pred_joint_vel.device),
                              mask[:, :, :, 1:])
    # TODO - NOT USED YET, JUST COMMITING TO NOT DELETE THIS AND KEEP INITIAL IMPLEMENTATION, NOT DONE!
    def foot_contact_loss_humanml3d(self, target, model_output):
        # root_rot_velocity (B, seq_len, 1)
        # root_linear_velocity (B, seq_len, 2)
        # root_y (B, seq_len, 1)
        # ric_data (B, seq_len, (joint_num - 1)*3) , XYZ
        # rot_data (B, seq_len, (joint_num - 1)*6) , 6D
        # local_velocity (B, seq_len, joint_num*3) , XYZ
        # foot contact (B, seq_len, 4) ,

        target_fc = target[:, -4:, :, :]
        root_rot_velocity = target[:, :1, :, :]
        root_linear_velocity = target[:, 1:3, :, :]
        root_y = target[:, 3:4, :, :]
        ric_data = target[:, 4:67, :, :]  # 4+(3*21)=67
        rot_data = target[:, 67:193, :, :]  # 67+(6*21)=193
        local_velocity = target[:, 193:259, :, :]  # 193+(3*22)=259
        contact = target[:, 259:, :, :]  # 193+(3*22)=259
        contact_mask_gt = contact > 0.5  # contact mask order for indexes are fid_l [7, 10], fid_r [8, 11]
        vel_lf_7 = local_velocity[:, 7 * 3:8 * 3, :, :]
        vel_rf_8 = local_velocity[:, 8 * 3:9 * 3, :, :]
        vel_lf_10 = local_velocity[:, 10 * 3:11 * 3, :, :]
        vel_rf_11 = local_velocity[:, 11 * 3:12 * 3, :, :]

        calc_vel_lf_7 = ric_data[:, 6 * 3:7 * 3, :, 1:] - ric_data[:, 6 * 3:7 * 3, :, :-1]
        calc_vel_rf_8 = ric_data[:, 7 * 3:8 * 3, :, 1:] - ric_data[:, 7 * 3:8 * 3, :, :-1]
        calc_vel_lf_10 = ric_data[:, 9 * 3:10 * 3, :, 1:] - ric_data[:, 9 * 3:10 * 3, :, :-1]
        calc_vel_rf_11 = ric_data[:, 10 * 3:11 * 3, :, 1:] - ric_data[:, 10 * 3:11 * 3, :, :-1]

        # vel_foots = torch.stack([vel_lf_7, vel_lf_10, vel_rf_8, vel_rf_11], dim=1)
        for chosen_vel_foot_calc, chosen_vel_foot, joint_idx, contact_mask_idx in zip(
                [calc_vel_lf_7, calc_vel_rf_8, calc_vel_lf_10, calc_vel_rf_11],
                [vel_lf_7, vel_lf_10, vel_rf_8, vel_rf_11],
                [7, 10, 8, 11],
                [0, 1, 2, 3]):
            tmp_mask_gt = contact_mask_gt[:, contact_mask_idx, :, :].cpu().detach().numpy().reshape(-1).astype(int)
            chosen_vel_norm = np.linalg.norm(chosen_vel_foot.cpu().detach().numpy().reshape((3, -1)), axis=0)
            chosen_vel_calc_norm = np.linalg.norm(chosen_vel_foot_calc.cpu().detach().numpy().reshape((3, -1)),
                                                  axis=0)

            print(tmp_mask_gt.shape)
            print(chosen_vel_foot.shape)
            print(chosen_vel_calc_norm.shape)
            import matplotlib.pyplot as plt
            plt.plot(tmp_mask_gt, label='FC mask')
            plt.plot(chosen_vel_norm, label='Vel. XYZ norm (from vector)')
            plt.plot(chosen_vel_calc_norm, label='Vel. XYZ norm (calculated diff XYZ)')

            plt.title(f'FC idx {contact_mask_idx}, Joint Index {joint_idx}')
            plt.legend()
            plt.show()
        # print(vel_foots.shape)
        return 0
    # TODO - NOT USED YET, JUST COMMITING TO NOT DELETE THIS AND KEEP INITIAL IMPLEMENTATION, NOT DONE!
    def velocity_consistency_loss_humanml3d(self, target, model_output):
        # root_rot_velocity (B, seq_len, 1)
        # root_linear_velocity (B, seq_len, 2)
        # root_y (B, seq_len, 1)
        # ric_data (B, seq_len, (joint_num - 1)*3) , XYZ
        # rot_data (B, seq_len, (joint_num - 1)*6) , 6D
        # local_velocity (B, seq_len, joint_num*3) , XYZ
        # foot contact (B, seq_len, 4) ,

        target_fc = target[:, -4:, :, :]
        root_rot_velocity = target[:, :1, :, :]
        root_linear_velocity = target[:, 1:3, :, :]
        root_y = target[:, 3:4, :, :]
        ric_data = target[:, 4:67, :, :]  # 4+(3*21)=67
        rot_data = target[:, 67:193, :, :]  # 67+(6*21)=193
        local_velocity = target[:, 193:259, :, :]  # 193+(3*22)=259
        contact = target[:, 259:, :, :]  # 193+(3*22)=259

        calc_vel_from_xyz = ric_data[:, :, :, 1:] - ric_data[:, :, :, :-1]
        velocity_from_vector = local_velocity[:, 3:, :, 1:]  # Slicing out root
        r_rot_quat, r_pos = motion_process.recover_root_rot_pos(target.permute(0, 2, 3, 1).type(th.FloatTensor))
        print(f'r_rot_quat: {r_rot_quat.shape}')
        print(f'calc_vel_from_xyz: {calc_vel_from_xyz.shape}')
        calc_vel_from_xyz = calc_vel_from_xyz.permute(0, 2, 3, 1)
        calc_vel_from_xyz = calc_vel_from_xyz.reshape((1, 1, -1, 21, 3)).type(th.FloatTensor)
        r_rot_quat_adapted = r_rot_quat[..., :-1, None, :].repeat((1,1,1,21,1)).to(calc_vel_from_xyz.device)
        print(f'calc_vel_from_xyz: {calc_vel_from_xyz.shape} , {calc_vel_from_xyz.device}')
        print(f'r_rot_quat_adapted: {r_rot_quat_adapted.shape}, {r_rot_quat_adapted.device}')

        calc_vel_from_xyz = motion_process.qrot(r_rot_quat_adapted, calc_vel_from_xyz)
        calc_vel_from_xyz = calc_vel_from_xyz.reshape((1, 1, -1, 21 * 3))
        calc_vel_from_xyz = calc_vel_from_xyz.permute(0, 3, 1, 2)
        print(f'calc_vel_from_xyz: {calc_vel_from_xyz.shape} , {calc_vel_from_xyz.device}')

        import matplotlib.pyplot as plt
        for i in range(21):
            plt.plot(np.linalg.norm(calc_vel_from_xyz[:,i*3:(i+1)*3,:,:].cpu().detach().numpy().reshape((3, -1)), axis=0), label='Calc Vel')
            plt.plot(np.linalg.norm(velocity_from_vector[:,i*3:(i+1)*3,:,:].cpu().detach().numpy().reshape((3, -1)), axis=0), label='Vector Vel')
            plt.title(f'Joint idx: {i}')
            plt.legend()
            plt.show()
        print(calc_vel_from_xyz.shape)
        print(velocity_from_vector.shape)
        diff = calc_vel_from_xyz-velocity_from_vector
        print(np.linalg.norm(diff.cpu().detach().numpy().reshape((63, -1)), axis=0))

        return 0


    def _prior_bpd(self, x_start):
        """
        Get the prior KL term for the variational lower-bound, measured in
        bits-per-dim.

        This term can't be optimized, as it only depends on the encoder.

        :param x_start: the [N x C x ...] tensor of inputs.
        :return: a batch of [N] KL values (in bits), one per batch element.
        """
        batch_size = x_start.shape[0]
        t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
        kl_prior = normal_kl(
            mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0
        )
        return mean_flat(kl_prior) / np.log(2.0)

    def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None):
        """
        Compute the entire variational lower-bound, measured in bits-per-dim,
        as well as other related quantities.

        :param model: the model to evaluate loss on.
        :param x_start: the [N x C x ...] tensor of inputs.
        :param clip_denoised: if True, clip denoised samples.
        :param model_kwargs: if not None, a dict of extra keyword arguments to
            pass to the model. This can be used for conditioning.

        :return: a dict containing the following keys:
                 - total_bpd: the total variational lower-bound, per batch element.
                 - prior_bpd: the prior term in the lower-bound.
                 - vb: an [N x T] tensor of terms in the lower-bound.
                 - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
                 - mse: an [N x T] tensor of epsilon MSEs for each timestep.
        """
        device = x_start.device
        batch_size = x_start.shape[0]

        vb = []
        xstart_mse = []
        mse = []
        for t in list(range(self.num_timesteps))[::-1]:
            t_batch = th.tensor([t] * batch_size, device=device)
            noise = th.randn_like(x_start)
            x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
            # Calculate VLB term at the current timestep
            with th.no_grad():
                out = self._vb_terms_bpd(
                    model,
                    x_start=x_start,
                    x_t=x_t,
                    t=t_batch,
                    clip_denoised=clip_denoised,
                    model_kwargs=model_kwargs,
                )
            vb.append(out["output"])
            xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2))
            eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"])
            mse.append(mean_flat((eps - noise) ** 2))

        vb = th.stack(vb, dim=1)
        xstart_mse = th.stack(xstart_mse, dim=1)
        mse = th.stack(mse, dim=1)

        prior_bpd = self._prior_bpd(x_start)
        total_bpd = vb.sum(dim=1) + prior_bpd
        return {
            "total_bpd": total_bpd,
            "prior_bpd": prior_bpd,
            "vb": vb,
            "xstart_mse": xstart_mse,
            "mse": mse,
        }


def _extract_into_tensor(arr, timesteps, broadcast_shape):
    """
    Extract values from a 1-D numpy array for a batch of indices.

    :param arr: the 1-D numpy array.
    :param timesteps: a tensor of indices into the array to extract.
    :param broadcast_shape: a larger shape of K dimensions with the batch
                            dimension equal to the length of timesteps.
    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
    """
    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
    while len(res.shape) < len(broadcast_shape):
        res = res[..., None]
    return res.expand(broadcast_shape)

diffusion_latent/diffusion/gaussian_diffusion.py
YoungSeng/UnifiedGesture
import pdb

import logging
logging.getLogger().setLevel(logging.INFO)

from torch.utils.data import DataLoader
from data_loader.lmdb_data_loader import TrinityDataset
import torch
import yaml
from pprint import pprint
from easydict import EasyDict
from configs.parse_args import parse_args
import os
import numpy as np
# import sys
# [sys.path.append(i) for i in ['.', '../diffusion_latent']]
from utils_.model_util import create_gaussian_diffusion
from train.training_loop import TrainLoop
from model.mdm import MDM


args = parse_args()
mydevice = torch.device('cuda:' + args.gpu)
torch.cuda.set_device(int(args.gpu))


def create_model_and_diffusion(args):
    model = MDM(modeltype='', njoints=16 * 7 * 3 + 27, nfeats=1, translation=True, pose_rep='rot6d', glob=True,
                glob_rot=True, cond_mode='cross_local_attention3_style1', action_emb='tensor', audio_feat='wavlm',
                arch='trans_enc', latent_dim=256, n_seed=args.n_seed, cond_mask_prob=0.1)
    diffusion = create_gaussian_diffusion()
    return model, diffusion


def main(args):
    # dataset
    train_dataset = TrinityDataset(args.train_data_path,
                                   n_poses=args.n_poses,
                                   subdivision_stride=args.subdivision_stride,
                                   pose_resampling_fps=args.motion_resampling_framerate, model='WavLM_36_aux')        # , model='Long_1200'
    train_loader = DataLoader(dataset=train_dataset, batch_size=args.batch_size,
                              shuffle=True, drop_last=True, num_workers=args.loader_workers, pin_memory=True)

    val_dataset = TrinityDataset(args.val_data_path,
                                       n_poses=args.n_poses,
                                       subdivision_stride=args.subdivision_stride,
                                       pose_resampling_fps=args.motion_resampling_framerate, model='WavLM_36_aux')         # , model='Long_1200'
    test_loader = DataLoader(dataset=val_dataset, batch_size=args.batch_size,
                             shuffle=False, drop_last=True, num_workers=args.loader_workers, pin_memory=False)

    logging.info('len of train loader:{}, len of test loader:{}'.format(len(train_loader), len(test_loader)))

    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)

    model, diffusion = create_model_and_diffusion(args)
    model.to(mydevice)
    TrainLoop(args, model, diffusion, mydevice, data=train_loader).run_loop()


if __name__ == '__main__':
    '''
    cd /ceph/hdd/yangsc21/Python/My_3/deep-motion-editing/diffusion_latent/
    cd diffusion_latent/
    python end2end.py --config=./configs/all_data.yml --gpu 1 --save_dir "./result/my_diffusion"
    modelarts:
cd /nfs7/y50021900/My_2/denoising-diffusion-pytorch-main/mydiffusion/ && pip install --upgrade pip && pip install -r requirements.txt && python end2end.py --config=./configs/codebook.yml --train --no_cuda 0 --gpu 0
cd /nfs7/y50021900/My_2/mydiffwave && pip install . && cd /nfs7/y50021900/My_2/denoising-diffusion-pytorch-main/mydiffusion/ && pip install --upgrade pip && pip install -r requirements.txt && python end2end.py --config=./configs/codebook.yml --train --no_cuda 0 --gpu 0
cd /nfs7/y50021900/My_2/mydiffwave && pip install . && cd /nfs7/y50021900/My_2/denoising-diffusion-pytorch-main/mydiffusion/ && pip install --upgrade pip && pip install blobfile spacy && pip install -r requirements.txt && pip install setuptools==59.5.0 && cd /nfs7/y50021900/My_2/mymdm/mydiffusion/ && python end2end.py --config=./configs/codebook.yml --train --no_cuda 0 --gpu 0
    '''

    with open(args.config) as f:
        config = yaml.safe_load(f)

    for k, v in vars(args).items():
        config[k] = v
    pprint(config)

    config = EasyDict(config)
    config.no_cuda = config.gpu
    main(config)

diffusion_latent/end2end.py
YoungSeng/UnifiedGesture
import pdb

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.autograd import Variable
import math
import sys
[sys.path.append(i) for i in ['.', '..', '../../..']]
from generate.generate import WavEncoder

from mydiffwave.src.diffwave.model import DiffWave
from mydiffwave.src.diffwave.params import params
import numpy as np


class diffwav_model(nn.Module):
    def __init__(self):
        super().__init__()
        self.WavEncoder = WavEncoder()
        self.diffwav_model = DiffWave(params)
        self.criterion = nn.SmoothL1Loss()

    def sample(self, batch_size, tmp_audio, beta, T):
        wav_feature = self.WavEncoder(tmp_audio).transpose(1, 2)  # (b, 240, 32)
        noisy_pose = torch.randn(batch_size, 240, 135).transpose(1, 2).to(tmp_audio.device)
        alpha = 1 - beta
        alpha_cum = np.cumprod(alpha)
        for n in range(len(alpha) - 1, -1, -1):
            c1 = 1 / alpha[n] ** 0.5
            c2 = beta[n] / (1 - alpha_cum[n]) ** 0.5
            noisy_pose = c1 * (noisy_pose - c2 * self.diffwav_model(noisy_pose, torch.tensor([T[n]], device=noisy_pose.device), wav_feature).squeeze(1))
            if n > 0:
                noise = torch.randn_like(noisy_pose)
                sigma = ((1.0 - alpha_cum[n - 1]) / (1.0 - alpha_cum[n]) * beta[n]) ** 0.5
                noisy_pose += sigma * noise
            noisy_pose = torch.clamp(noisy_pose, -1.0, 1.0)
        return noisy_pose.transpose(1, 2)

    def forward(self, noisy_pose, t, audio, noise):      # (b, len, 13)
        wav_feature = self.WavEncoder(audio).transpose(1, 2)        # (b, 240, 32)
        predicted = self.diffwav_model(noisy_pose, t, wav_feature)
        loss = self.criterion(predicted, noise)
        return loss


if __name__ == '__main__':
    '''
    cd mydiffusion/generate/
    python diffwav.py
    '''
    # z = torch.arange(0, 60).reshape(2, 30)

    device = torch.device('cuda:2')
    audio = torch.rand(2, 64000).to(device)
    pose = torch.rand(2, 240, 135).transpose(1, 2).to(device)
    model = diffwav_model().to(device)

    n_frames = 240
    n_pose_dims = 135
    n_audio_dim = 32
    hop_samples = 1

    N = pose.shape[0]  # 1, 15872
    device = pose.device

    beta = np.linspace(1e-4, 0.05, 50)
    noise_level = np.cumprod(1 - beta)
    noise_level = torch.tensor(noise_level.astype(np.float32)).to(device)

    t = torch.randint(0, len(beta), [N], device=pose.device)  # (batch)
    noise_scale = noise_level[t].unsqueeze(1).unsqueeze(1)  # (batch, 1)
    noise_scale_sqrt = noise_scale ** 0.5  # (batch, 1)
    noise = torch.randn_like(pose)  # (batch, 15872)
    noisy_pose = noise_scale_sqrt * pose + (1.0 - noise_scale) ** 0.5 * noise  # (batch, 15872)

    loss = model(noisy_pose, t, audio, pose)  # (batch, 1, 15872)
    print(loss)

    talpha = 1 - beta
    talpha_cum = np.cumprod(talpha)
    alpha = 1 - beta
    alpha_cum = np.cumprod(alpha)

    T = []
    for s in range(len(beta)):
      for t in range(len(beta) - 1):
        if talpha_cum[t+1] <= alpha_cum[s] <= talpha_cum[t]:
          twiddle = (talpha_cum[t]**0.5 - alpha_cum[s]**0.5) / (talpha_cum[t]**0.5 - talpha_cum[t+1]**0.5)
          T.append(t + twiddle)
          break
    T = np.array(T, dtype=np.float32)

    tmp_audio = torch.rand(1, 64000).to(device)
    sampled_seq = model.sample(batch_size=1, tmp_audio=tmp_audio, beta=beta, T=T)
    print(sampled_seq.shape)      # (4, 32, 128)

diffusion_latent/generate/diffwav.py
YoungSeng/UnifiedGesture
import pdb

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.autograd import Variable
import math
import sys
[sys.path.append(i) for i in ['.', '..', '../..']]
# from denoising_diffusion_pytorch import myUnet1D, myGaussianDiffusion1D

# class WavEncoder(nn.Module):        # (b, 64000) -> (b, 30, 32)
#     def __init__(self):
#         super().__init__()
#         self.feat_extractor = nn.Sequential(
#             nn.Conv1d(1, 8, 16, stride=3),
#             nn.BatchNorm1d(8),
#             nn.LeakyReLU(0.3, inplace=True),
#             nn.Conv1d(8, 16, 16, stride=3),
#             nn.BatchNorm1d(16),
#             nn.LeakyReLU(0.3, inplace=True),
#             nn.Conv1d(16, 32, 16, stride=6),
#             nn.BatchNorm1d(32),
#             nn.LeakyReLU(0.3, inplace=True),
#             nn.Conv1d(32, 64, 16, stride=6),
#             nn.BatchNorm1d(64),
#             nn.LeakyReLU(0.3, inplace=True),
#             nn.Conv1d(64, 32, 16, stride=6),
#         )
#
#     def forward(self, wav_data):
#         wav_data = wav_data.unsqueeze(1)  # add channel dim
#         out = self.feat_extractor(wav_data)
#         return out.transpose(1, 2)  # to (batch x seq x dim)

class WavEncoder(nn.Module):        # (b, 64000) -> (b, 240, 32)
    def __init__(self):
        super().__init__()
        self.feat_extractor = nn.Sequential(
            nn.Conv1d(1, 16, 15, stride=3, padding=800),
            nn.BatchNorm1d(16),
            nn.LeakyReLU(0.3, inplace=True),
            nn.Conv1d(16, 32, 15, stride=3),
            nn.BatchNorm1d(32),
            nn.LeakyReLU(0.3, inplace=True),
            nn.Conv1d(32, 64, 15, stride=5),
            nn.BatchNorm1d(64),
            nn.LeakyReLU(0.3, inplace=True),
            nn.Conv1d(64, 32, 15, stride=6),
        )

    def forward(self, wav_data):
        wav_data = wav_data.unsqueeze(1)  # add channel dim
        out = self.feat_extractor(wav_data)
        return out.transpose(1, 2)  # to (batch x seq x dim)


class Generator_linear(nn.Module):
    def __init__(self):
        super().__init__()
        self.WavEncoder = WavEncoder()
        self.project = nn.Linear(32, 512, bias=False)
        self.norm = nn.LayerNorm(32)

    def sample(self, x):
        wav_feature = self.WavEncoder(x)
        wav_feature = self.norm(wav_feature)          # (1, 30, 512)
        codebook_embedding = self.project(wav_feature).squeeze()
        # codebook_embedding = self.norm(codebook_embedding).squeeze()      # (1, 30, 512)
        code = torch.tensor([]).to(x.device)
        for k in codebook_embedding:
            probs = F.softmax(k, dim=-1)
            _, ix = torch.topk(probs, k=1, dim=-1)
            code = torch.cat((code, ix))
        return [code.unsqueeze(0).int()]

    def forward(self, x, target=None):
        wav_feature = self.WavEncoder(x)
        wav_feature = self.norm(wav_feature)        # norm before linear
        codebook_embedding = self.project(wav_feature)
        loss = None
        if target is not None:
            loss = F.cross_entropy(codebook_embedding.view(-1, codebook_embedding.size(-1)), target.view(-1))
        return codebook_embedding, loss

'''
Based on the following Se2Seq implementations:
- https://github.com/AuCson/PyTorch-Batch-Attention-Seq2seq
- https://github.com/spro/practical-pytorch/blob/master/seq2seq-translation/seq2seq-translation-batched.ipynb
'''


class EncoderRNN(nn.Module):
    def __init__(self, input_size, embed_size, hidden_size, n_layers=1, dropout=0.5, pre_trained_embedding=None):
        super(EncoderRNN, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.embed_size = embed_size
        self.n_layers = n_layers
        self.dropout = dropout

        if pre_trained_embedding is not None:  # use pre-trained embedding (e.g., word2vec, glove)
            assert pre_trained_embedding.shape[0] == input_size
            assert pre_trained_embedding.shape[1] == embed_size
            self.embedding = nn.Embedding.from_pretrained(torch.FloatTensor(pre_trained_embedding), freeze=False)
        else:
            self.embedding = nn.Embedding(input_size, embed_size)

        self.gru = nn.GRU(embed_size, hidden_size, n_layers, dropout=self.dropout, bidirectional=True)

        self.do_flatten_parameters = False
        if torch.cuda.device_count() > 1:
            self.do_flatten_parameters = True

    def forward(self, input_seqs, input_lengths, hidden=None):
        '''
        :param input_seqs:
            Variable of shape (num_step(T),batch_size(B)), sorted decreasingly by lengths(for packing)
        :param input_lengths:
            list of sequence length
        :param hidden:
            initial state of GRU
        :returns:
            GRU outputs in shape (T,B,hidden_size(H))
            last hidden stat of RNN(i.e. last output for GRU)
        '''
        if self.do_flatten_parameters:
            self.gru.flatten_parameters()

        embedded = self.embedding(input_seqs)
        packed = torch.nn.utils.rnn.pack_padded_sequence(embedded, input_lengths)
        outputs, hidden = self.gru(packed, hidden)
        outputs, output_lengths = torch.nn.utils.rnn.pad_packed_sequence(outputs)  # unpack (back to padded)
        outputs = outputs[:, :, :self.hidden_size] + outputs[:, :, self.hidden_size:]  # Sum bidirectional outputs
        return outputs, hidden


class Attn(nn.Module):
    def __init__(self, hidden_size):
        super(Attn, self).__init__()
        self.hidden_size = hidden_size
        self.attn = nn.Linear(self.hidden_size * 2, hidden_size)
        self.v = nn.Parameter(torch.rand(hidden_size))
        stdv = 1. / math.sqrt(self.v.size(0))
        self.v.data.normal_(mean=0, std=stdv)

    def forward(self, hidden, encoder_outputs):
        '''
        :param hidden:
            previous hidden state of the decoder, in shape (layers*directions,B,H)
        :param encoder_outputs:
            encoder outputs from Encoder, in shape (T,B,H)
        :return
            attention energies in shape (B,T)
        '''
        max_len = encoder_outputs.size(0)
        this_batch_size = encoder_outputs.size(1)
        H = hidden.repeat(max_len, 1, 1).transpose(0, 1)
        encoder_outputs = encoder_outputs.transpose(0, 1)  # [B*T*H]
        attn_energies = self.score(H, encoder_outputs)  # compute attention score
        return F.softmax(attn_energies, dim=1).unsqueeze(1)  # normalize with softmax

    def score(self, hidden, encoder_outputs):
        energy = torch.tanh(self.attn(torch.cat([hidden, encoder_outputs], 2)))  # [B*T*2H]->[B*T*H]
        energy = energy.transpose(2, 1)  # [B*H*T]
        v = self.v.repeat(encoder_outputs.data.shape[0], 1).unsqueeze(1)  # [B*1*H]
        energy = torch.bmm(v, energy)  # [B*1*T]
        return energy.squeeze(1)  # [B*T]


class BahdanauAttnDecoderRNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, n_layers=1, dropout_p=0.1,
                 discrete_representation=False, speaker_model=None):
        super(BahdanauAttnDecoderRNN, self).__init__()

        # define parameters
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.n_layers = n_layers
        self.dropout_p = dropout_p
        self.discrete_representation = discrete_representation
        self.speaker_model = speaker_model

        # define embedding layer
        if self.discrete_representation:
            self.embedding = nn.Embedding(output_size, hidden_size)
            self.dropout = nn.Dropout(dropout_p)

        if self.speaker_model:
            self.speaker_embedding = nn.Embedding(speaker_model.n_words, 8)

        # calc input size
        if self.discrete_representation:
            input_size = hidden_size  # embedding size
        linear_input_size = input_size + hidden_size
        if self.speaker_model:
            linear_input_size += 8

        # define layers
        self.pre_linear = nn.Sequential(
            nn.Linear(linear_input_size, hidden_size),
            nn.BatchNorm1d(hidden_size),
            nn.ReLU(inplace=True)
        )
        self.attn = Attn(hidden_size)
        self.gru = nn.GRU(hidden_size, hidden_size, n_layers, dropout=dropout_p)

        # self.out = nn.Linear(hidden_size * 2, output_size)
        self.out = nn.Linear(hidden_size, output_size)

        self.do_flatten_parameters = False
        if torch.cuda.device_count() > 1:
            self.do_flatten_parameters = True

    def freeze_attn(self):
        for param in self.attn.parameters():
            param.requires_grad = False

    def forward(self, motion_input, last_hidden, encoder_outputs, vid_indices=None):
        '''
        :param motion_input:
            motion input for current time step, in shape [batch x dim]
        :param last_hidden:
            last hidden state of the decoder, in shape [layers x batch x hidden_size]
        :param encoder_outputs:
            encoder outputs in shape [steps x batch x hidden_size]
        :param vid_indices:
        :return
            decoder output
        Note: we run this one step at a time i.e. you should use a outer loop
            to process the whole sequence
        '''

        if self.do_flatten_parameters:
            self.gru.flatten_parameters()

        if self.discrete_representation:
            word_embedded = self.embedding(motion_input).view(1, motion_input.size(0), -1)  # [1 x B x embedding_dim]
            motion_input = self.dropout(word_embedded)
        else:
            motion_input = motion_input.view(1, motion_input.size(0), -1)  # [1 x batch x dim]

        # attention
        attn_weights = self.attn(last_hidden[-1], encoder_outputs)  # [batch x 1 x T]
        context = attn_weights.bmm(encoder_outputs.transpose(0, 1))  # [batch x 1 x attn_size]
        context = context.transpose(0, 1)  # [1 x batch x attn_size]

        # make input vec
        rnn_input = torch.cat((motion_input, context), 2)  # [1 x batch x (dim + attn_size)]

        if self.speaker_model:
            assert vid_indices is not None
            speaker_context = self.speaker_embedding(vid_indices).unsqueeze(0)
            rnn_input = torch.cat((rnn_input, speaker_context), 2)  # [1 x batch x (dim + attn_size + embed_size)]

        rnn_input = self.pre_linear(rnn_input.squeeze(0))
        rnn_input = rnn_input.unsqueeze(0)

        # rnn
        output, hidden = self.gru(rnn_input, last_hidden)

        # post-fc
        output = output.squeeze(0)  # [1 x batch x hidden_size] -> [batch x hidden_size]
        output = self.out(output)

        return output, hidden, attn_weights


class Generator(nn.Module):
    def __init__(self, args, motion_dim, discrete_representation=False, speaker_model=None):
        super(Generator, self).__init__()
        self.output_size = motion_dim
        self.n_layers = args.n_layers
        self.discrete_representation = discrete_representation
        self.decoder = BahdanauAttnDecoderRNN(input_size=motion_dim,
                                              hidden_size=args.hidden_size,
                                              output_size=self.output_size,
                                              n_layers=self.n_layers,
                                              dropout_p=args.dropout_prob,
                                              discrete_representation=discrete_representation,
                                              speaker_model=speaker_model)

    def freeze_attn(self):
        self.decoder.freeze_attn()

    def forward(self, z, motion_input, last_hidden, encoder_output, vid_indices=None):
        if z is None:
            input_with_noise_vec = motion_input
        else:
            assert not self.discrete_representation  # not valid for discrete representation
            input_with_noise_vec = torch.cat([motion_input, z], dim=1)  # [bs x (10+z_size)]

        return self.decoder(input_with_noise_vec, last_hidden, encoder_output, vid_indices)


class Seq2SeqNet(nn.Module):
    def __init__(self, args, pose_dim, n_frames, n_words, word_embed_size, word_embeddings, speaker_model=None):
        super().__init__()
        self.encoder = EncoderRNN(
            n_words, word_embed_size, args.hidden_size, args.n_layers,
            dropout=args.dropout_prob, pre_trained_embedding=word_embeddings)
        self.decoder = Generator(args, pose_dim, speaker_model=speaker_model)

        self.n_frames = n_frames
        self.n_pre_poses = args.n_pre_poses
        self.pose_dim = pose_dim

    def forward(self, in_text, in_lengths, poses, vid_indices):
        # reshape to (seq x batch x dim)
        in_text = in_text.transpose(0, 1)
        poses = poses.transpose(0, 1)

        outputs = torch.zeros(self.n_frames, poses.size(1), self.decoder.output_size).to(poses.device)

        # run words through encoder
        encoder_outputs, encoder_hidden = self.encoder(in_text, in_lengths, None)
        decoder_hidden = encoder_hidden[:self.decoder.n_layers]  # use last hidden state from encoder

        # run through decoder one time step at a time
        decoder_input = poses[0]  # initial pose from the dataset
        outputs[0] = decoder_input

        for t in range(1, self.n_frames):
            decoder_output, decoder_hidden, _ = self.decoder(None, decoder_input, decoder_hidden, encoder_outputs,
                                                             vid_indices)
            outputs[t] = decoder_output

            if t < self.n_pre_poses:
                decoder_input = poses[t]  # next input is current target
            else:
                decoder_input = decoder_output  # next input is current prediction

        return outputs.transpose(0, 1)


class Generator_gru(nn.Module):
    def __init__(self):
        super().__init__()
        self.WavEncoder = WavEncoder()
        self.hidden_size = 200
        self.output_size = 512

        self.project = nn.GRU(input_size=32, hidden_size=self.hidden_size, num_layers=2, dropout=0.1, bidirectional=True, batch_first=True)
        self.norm = nn.LayerNorm(self.hidden_size)
        self.out = nn.Linear(self.hidden_size, self.output_size)

    def sample(self, x):
        wav_feature = self.WavEncoder(x)
        hidden = None
        outputs, hidden = self.project(wav_feature, hidden)
        outputs = outputs[:, :, :self.hidden_size] + outputs[:, :, self.hidden_size:]  # Sum bidirectional outputs, ((batch, seq_len=240, )
        outputs = self.norm(outputs)
        codebook_embedding = self.out(outputs)
        code = torch.tensor([]).to(x.device)
        for k in codebook_embedding:
            probs = F.softmax(k, dim=-1)
            _, ix = torch.topk(probs, k=1, dim=-1)
            code = torch.cat((code, ix.squeeze(-1)))
        return [code.unsqueeze(0).long()]

    def forward(self, x, target=None):      # (b, len, 13)
        wav_feature = self.WavEncoder(x)        # (b, 30, 32)
        hidden = None

        outputs, hidden = self.project(wav_feature, hidden)       # (batch, seq_len=240, num_directions * hidden_size)
        outputs = outputs[:, :, :self.hidden_size] + outputs[:, :, self.hidden_size:]  # Sum bidirectional outputs, ((batch, seq_len=240, )

        outputs = self.norm(outputs)
        codebook_embedding = self.out(outputs)

        loss = None
        if target is not None:
            loss = F.cross_entropy(codebook_embedding.view(-1, codebook_embedding.size(-1)), target.view(-1))
        return codebook_embedding, loss


class Generator_diff(nn.Module):
    def __init__(self):
        super().__init__()
        self.WavEncoder = WavEncoder()
        seq_len = 240
        joints = 15
        n_dim = 9
        n_channels = joints * n_dim
        audio_dim = 32

        model = myUnet1D(
            dim=64,
            dim_mults=(1, 2, 4, 8),
            channels=n_channels,
            self_condition=True,
            audio_dim=audio_dim
        )

        self.diffusion = myGaussianDiffusion1D(
            model,
            seq_length=seq_len,
            timesteps=250,
            objective='pred_v',
            audio_dim=audio_dim,
            loss_type='huber'
        )

    def sample(self, batch_size, tmp_audio):
        wav_feature = self.WavEncoder(tmp_audio).transpose(1, 2)  # (b, 240, 32)
        sampled_seq = self.diffusion.sample(batch_size=batch_size, tmp_audio_feat=wav_feature)
        return sampled_seq

    def forward(self, target, x):      # (b, len, 13)
        wav_feature = self.WavEncoder(x).transpose(1, 2)        # (b, 240, 32)
        loss = self.diffusion(target, wav_feature)
        return loss


if __name__ == '__main__':
    '''
    cd mydiffusion/generate/
    python generate.py
    '''
    audio = torch.rand(2, 64000)
    pose = torch.rand(2, 240, 135).transpose(1, 2)

    # z = torch.arange(0, 60).reshape(2, 30)
    # model = Generator_gru()
    model = Generator_diff()
    loss = model(pose, audio)        # (b, 30, 32)
    pdb.set_trace()
    sampled_seq = model.sample(batch_size=1)
    print(sampled_seq.shape)      # (4, 32, 128)




diffusion_latent/generate/generate.py
YoungSeng/UnifiedGesture
# Copyright (c) 2006 Carnegie Mellon University
#
# You may copy and modify this freely under the same terms as
# Sphinx-III

"""Compute MFCC coefficients.

This module provides functions for computing MFCC (mel-frequency
cepstral coefficients) as used in the Sphinx speech recognition
system.
"""

__author__ = "David Huggins-Daines <dhuggins@cs.cmu.edu>"
__version__ = "$Revision: 6390 $"

import pdb

import numpy, numpy.fft
import math
import librosa
import os


def mel(f):
    return 2595. * numpy.log10(1. + f / 700.)


def melinv(m):
    return 700. * (numpy.power(10., m / 2595.) - 1.)


class MFCC(object):
    def __init__(self, nfilt=40, ncep=13,
                 lowerf=133.3333, upperf=6855.4976, alpha=0.97,
                 samprate=16000, frate=100, wlen=0.0256,
                 nfft=512):
        # Store parameters
        self.samprate = samprate
        self.lowerf = lowerf
        self.upperf = upperf
        self.nfft = nfft
        self.ncep = ncep
        self.nfilt = nfilt
        self.frate = frate
        self.fshift = float(samprate) / frate

        # Build Hamming window
        self.wlen = int(wlen * samprate)
        self.win = numpy.hamming(self.wlen)

        # Prior sample for pre-emphasis
        self.prior = 0
        self.alpha = alpha

        # Build mel filter matrix
        self.filters = numpy.zeros((nfft // 2 + 1, nfilt), 'd')
        dfreq = float(samprate) / nfft
        if upperf > samprate / 2:
            raise (Exception,
                   "Upper frequency %f exceeds Nyquist %f" % (upperf, samprate / 2))
        melmax = mel(upperf)
        melmin = mel(lowerf)
        dmelbw = (melmax - melmin) / (nfilt + 1)
        # Filter edges, in Hz
        filt_edge = melinv(melmin + dmelbw * numpy.arange(nfilt + 2, dtype='d'))

        for whichfilt in range(0, nfilt):
            # Filter triangles, in DFT points
            leftfr = round(filt_edge[whichfilt] / dfreq)
            centerfr = round(filt_edge[whichfilt + 1] / dfreq)
            rightfr = round(filt_edge[whichfilt + 2] / dfreq)
            # For some reason this is calculated in Hz, though I think
            # it doesn't really matter
            fwidth = (rightfr - leftfr) * dfreq
            height = 2. / fwidth

            if centerfr != leftfr:
                leftslope = height / (centerfr - leftfr)
            else:
                leftslope = 0
            freq = leftfr + 1
            while freq < centerfr:
                self.filters[freq, whichfilt] = (freq - leftfr) * leftslope
                freq = freq + 1
            if freq == centerfr:  # This is always true
                self.filters[freq, whichfilt] = height
                freq = freq + 1
            if centerfr != rightfr:
                rightslope = height / (centerfr - rightfr)
            while freq < rightfr:
                self.filters[freq, whichfilt] = (freq - rightfr) * rightslope
                freq = freq + 1
        #             print("Filter %d: left %d=%f center %d=%f right %d=%f width %d" %
        #                   (whichfilt,
        #                   leftfr, leftfr*dfreq,
        #                   centerfr, centerfr*dfreq,
        #                   rightfr, rightfr*dfreq,
        #                   freq - leftfr))
        #             print self.filters[leftfr:rightfr,whichfilt]

        # Build DCT matrix
        self.s2dct = s2dctmat(nfilt, ncep, 1. / nfilt)
        self.dct = dctmat(nfilt, ncep, numpy.pi / nfilt)

    def sig2s2mfc(self, sig):
        nfr = int(len(sig) / self.fshift + 1)
        mfcc = numpy.zeros((nfr, self.ncep), 'd')
        fr = 0
        while fr < nfr:
            start = round(fr * self.fshift)
            end = min(len(sig), start + self.wlen)
            frame = sig[start:end]
            if len(frame) < self.wlen:
                frame = numpy.resize(frame, self.wlen)
                frame[self.wlen:] = 0
            mfcc[fr] = self.frame2s2mfc(frame)
            fr = fr + 1
        return mfcc

    def sig2logspec(self, sig):
        nfr = int(len(sig) / self.fshift + 1)
        mfcc = numpy.zeros((nfr, self.nfilt), 'd')
        fr = 0
        while fr < nfr:
            start = round(fr * self.fshift)
            end = min(len(sig), start + self.wlen)
            frame = sig[start:end]
            if len(frame) < self.wlen:
                frame = numpy.resize(frame, self.wlen)
                frame[self.wlen:] = 0
            mfcc[fr] = self.frame2logspec(frame)
            fr = fr + 1
        return mfcc

    def pre_emphasis(self, frame):
        # FIXME: Do this with matrix multiplication
        outfr = numpy.empty(len(frame), 'd')
        outfr[0] = frame[0] - self.alpha * self.prior
        for i in range(1, len(frame)):
            outfr[i] = frame[i] - self.alpha * frame[i - 1]
        self.prior = frame[-1]
        return outfr

    def frame2logspec(self, frame):
        frame = self.pre_emphasis(frame) * self.win
        fft = numpy.fft.rfft(frame, self.nfft)
        # Square of absolute value
        power = fft.real * fft.real + fft.imag * fft.imag
        return numpy.log(numpy.dot(power, self.filters).clip(1e-5, numpy.inf))

    def frame2s2mfc(self, frame):
        logspec = self.frame2logspec(frame)
        return numpy.dot(logspec, self.s2dct.T) / self.nfilt

    def sig2s2mfc_energy(self, sig, dn):
        nfr = int(len(sig) / self.fshift + 1)

        mfcc = numpy.zeros((nfr, self.ncep + 2), 'd')
        fr = 0
        while fr < nfr:
            start = int(round(fr * self.fshift))
            end = min(len(sig), start + self.wlen)
            frame = sig[start:end]
            if len(frame) < self.wlen:
                frame = numpy.resize(frame, self.wlen)
                frame[self.wlen:] = 0
            mfcc[fr, :-2] = self.frame2s2mfc(frame)
            mfcc[fr, -2] = math.log(1 + numpy.mean(numpy.power(frame.astype(float), 2)))
            mid = 0.5 * (start + end - 1)
            mfcc[fr, -1] = mid / self.samprate

            fr = fr + 1
        return mfcc


def s2dctmat(nfilt, ncep, freqstep):
    """Return the 'legacy' not-quite-DCT matrix used by Sphinx"""
    melcos = numpy.empty((ncep, nfilt), 'double')
    for i in range(0, ncep):
        freq = numpy.pi * float(i) / nfilt
        melcos[i] = numpy.cos(freq * numpy.arange(0.5, float(nfilt) + 0.5, 1.0, 'double'))
    melcos[:, 0] = melcos[:, 0] * 0.5
    return melcos


def logspec2s2mfc(logspec, ncep=13):
    """Convert log-power-spectrum bins to MFCC using the 'legacy'
    Sphinx transform"""
    nframes, nfilt = logspec.shape
    melcos = s2dctmat(nfilt, ncep, 1. / nfilt)
    return numpy.dot(logspec, melcos.T) / nfilt


def dctmat(N, K, freqstep, orthogonalize=True):
    """Return the orthogonal DCT-II/DCT-III matrix of size NxK.
    For computing or inverting MFCCs, N is the number of
    log-power-spectrum bins while K is the number of cepstra."""
    cosmat = numpy.zeros((N, K), 'double')
    for n in range(0, N):
        for k in range(0, K):
            cosmat[n, k] = numpy.cos(freqstep * (n + 0.5) * k)
    if orthogonalize:
        cosmat[:, 0] = cosmat[:, 0] * 1. / numpy.sqrt(2)
    return cosmat


def dct(input, K=13):
    """Convert log-power-spectrum to MFCC using the orthogonal DCT-II"""
    nframes, N = input.shape
    freqstep = numpy.pi / N
    cosmat = dctmat(N, K, freqstep)
    return numpy.dot(input, cosmat) * numpy.sqrt(2.0 / N)


def dct2(input, K=13):
    """Convert log-power-spectrum to MFCC using the normalized DCT-II"""
    nframes, N = input.shape
    freqstep = numpy.pi / N
    cosmat = dctmat(N, K, freqstep, False)
    return numpy.dot(input, cosmat) * (2.0 / N)


def idct(input, K=40):
    """Convert MFCC to log-power-spectrum using the orthogonal DCT-III"""
    nframes, N = input.shape
    freqstep = numpy.pi / K
    cosmat = dctmat(K, N, freqstep).T
    return numpy.dot(input, cosmat) * numpy.sqrt(2.0 / K)


def dct3(input, K=40):
    """Convert MFCC to log-power-spectrum using the unnormalized DCT-III"""
    nframes, N = input.shape
    freqstep = numpy.pi / K
    cosmat = dctmat(K, N, freqstep, False)
    cosmat[:, 0] = cosmat[:, 0] * 0.5
    return numpy.dot(input, cosmat.T)


if __name__ == '__main__':
    '''
    cd codebook/Speech2GestureMatching/
    python mfcc.py
    '''
    obj = MFCC(frate=20)        # 60 -> 20
    root = '/mnt/nfs7/y50021900/My/data/BEAT0909/'
    target_path = 'MFCC_20'
    if not os.path.exists(os.path.join(root, target_path)):
        os.mkdir(os.path.join(root, target_path))

    for item in os.listdir(os.path.join(root, 'Audio_normalized')):
        print(item)
        wav_path = os.path.join(root, 'Audio_normalized', item)
        wav, fs = librosa.load(wav_path, sr=16000)  # 2117849 / 16000 = 132.3655625
        mfcc = obj.sig2s2mfc_energy(wav, None)  # (13237, 15) with -1:1.275、2.275、3.375...
        print(mfcc[:, :-2].shape)       # -1 -> -2
        numpy.savez_compressed(os.path.join(root, target_path, item[:-4] + '.npz'), mfcc=mfcc[:, :-2])

diffusion_latent/mfcc.py
YoungSeng/UnifiedGesture
import math

import torch
from torch import nn, einsum
import torch.nn.functional as F

from einops import rearrange, repeat, pack, unpack

from local_attention.rotary import SinusoidalEmbeddings, apply_rotary_pos_emb

# constant

TOKEN_SELF_ATTN_VALUE = -5e4

# helper functions

def exists(val):
    return val is not None

def default(value, d):
    return d if not exists(value) else value

def to(t):
    return {'device': t.device, 'dtype': t.dtype}

def max_neg_value(tensor):
    return -torch.finfo(tensor.dtype).max

def l2norm(tensor):
    dtype = tensor.dtype
    normed = F.normalize(tensor, dim = -1)
    return normed.type(dtype)

def pad_to_multiple(tensor, multiple, dim=-1, value=0):
    seqlen = tensor.shape[dim]
    m = seqlen / multiple
    if m.is_integer():
        return False, tensor
    remainder = math.ceil(m) * multiple - seqlen
    pad_offset = (0,) * (-1 - dim) * 2
    return True, F.pad(tensor, (*pad_offset, 0, remainder), value = value)

def look_around(x, backward = 1, forward = 0, pad_value = -1, dim = 2):
    t = x.shape[1]
    dims = (len(x.shape) - dim) * (0, 0)
    padded_x = F.pad(x, (*dims, backward, forward), value = pad_value)
    tensors = [padded_x[:, ind:(ind + t), ...] for ind in range(forward + backward + 1)]
    return torch.cat(tensors, dim = dim)

# main class

class LocalAttention(nn.Module):
    def __init__(
        self,
        window_size,
        causal = False,
        look_backward = 1,
        look_forward = None,
        dropout = 0.,
        shared_qk = False,
        rel_pos_emb_config = None,
        dim = None,
        autopad = False,
        exact_windowsize = False
    ):
        super().__init__()
        look_forward = default(look_forward, 0 if causal else 1)
        assert not (causal and look_forward > 0), 'you cannot look forward if causal'

        self.window_size = window_size
        self.autopad = autopad
        self.exact_windowsize = exact_windowsize

        self.causal = causal

        self.look_backward = look_backward
        self.look_forward = look_forward

        self.dropout = nn.Dropout(dropout)

        self.shared_qk = shared_qk

        # relative positions

        self.rel_pos = None
        if exists(rel_pos_emb_config) or exists(dim):  # backwards compatible with old `rel_pos_emb_config` deprecated argument
            if exists(rel_pos_emb_config):
                dim = rel_pos_emb_config[0]
            self.rel_pos = SinusoidalEmbeddings(dim)

    def forward(self, q, k, v, mask = None, input_mask = None):
        mask = default(mask, input_mask)

        shape, autopad, pad_value, window_size, causal, look_backward, look_forward, shared_qk = q.shape, self.autopad, -1, self.window_size, self.causal, self.look_backward, self.look_forward, self.shared_qk

        # https://github.com/arogozhnikov/einops/blob/master/docs/4-pack-and-unpack.ipynb
        (q, packed_shape), (k, _), (v, _) = map(lambda t: pack([t], '* n d'), (q, k, v))

        # rotary embeddings

        if exists(self.rel_pos):
            pos_emb = self.rel_pos(q)
            q, k = apply_rotary_pos_emb(q, k, pos_emb)

        # auto padding

        if autopad:
            orig_seq_len = q.shape[1]
            (needed_pad, q), (_, k), (_, v) = map(lambda t: pad_to_multiple(t, self.window_size, dim = -2), (q, k, v))

        b, n, dim_head, device, dtype = *q.shape, q.device, q.dtype
        scale = dim_head ** -0.5

        assert (n % window_size) == 0, f'sequence length {t} must be divisible by window size {window_size} for local attention'

        windows = n // window_size

        if shared_qk:
            k = l2norm(k)

        seq = torch.arange(n, device = device)
        b_t = rearrange(seq, '(w n) -> 1 w n', w = windows, n = window_size)

        bq, bk, bv = map(lambda t: rearrange(t, 'b (w n) d -> b w n d', w = windows), (q, k, v))

        look_around_kwargs = dict(
            backward =  look_backward,
            forward =  look_forward,
            pad_value = pad_value
        )

        bk = look_around(bk, **look_around_kwargs)
        bv = look_around(bv, **look_around_kwargs)

        bq_t = b_t
        bq_k = look_around(b_t, **look_around_kwargs)

        bq_t = rearrange(bq_t, '... i -> ... i 1')
        bq_k = rearrange(bq_k, '... j -> ... 1 j')

        sim = einsum('b h i e, b h j e -> b h i j', bq, bk) * scale

        mask_value = max_neg_value(sim)

        if shared_qk:
            self_mask = bq_t == bq_k
            sim = sim.masked_fill(self_mask, TOKEN_SELF_ATTN_VALUE)
            del self_mask

        if causal:
            causal_mask = bq_t < bq_k

            if self.exact_windowsize:
                max_causal_window_size = (self.window_size * self.look_backward)
                causal_mask = causal_mask | (bq_t > (bq_k + max_causal_window_size))

            sim = sim.masked_fill(causal_mask, mask_value)
            del causal_mask

        # mask out padding value

        if autopad and needed_pad:
            pad_mask = bq_k == pad_value
            sim = sim.masked_fill(pad_mask, mask_value)
            del pad_mask

        if exists(mask):
            batch = mask.shape[0]
            assert (b % batch) == 0

            h = b // mask.shape[0]

            if autopad:
                _, mask = pad_to_multiple(mask, window_size, dim = -1, value = False)

            mask = rearrange(mask, '... (w n) -> (...) w n', w = windows, n = window_size)
            mask = look_around(mask, **{**look_around_kwargs, 'pad_value': False})
            mask = rearrange(mask, '... j -> ... 1 j')
            mask = repeat(mask, 'b ... -> (b h) ...', h = h)
            sim = sim.masked_fill(~mask, mask_value)
            del mask

        # attention

        attn = sim.softmax(dim = -1)
        attn = self.dropout(attn)

        # aggregation

        out = einsum('b h i j, b h j e -> b h i e', attn, bv)
        out = rearrange(out, 'b w n d -> b (w n) d')

        if autopad:
            out = out[:, :orig_seq_len, :]

        out, *_ = unpack(out, packed_shape, '* n d')
        return out

diffusion_latent/model/local_attention.py
YoungSeng/UnifiedGesture
from local_attention.local_attention import LocalAttention
from local_attention.transformer import LocalTransformer, LocalMHA

diffusion_latent/model/local_attention/__init__.py
YoungSeng/UnifiedGesture
import math

import torch
from torch import nn, einsum
import torch.nn.functional as F

from einops import rearrange, repeat, pack, unpack

from local_attention.rotary import SinusoidalEmbeddings, apply_rotary_pos_emb

# constant

TOKEN_SELF_ATTN_VALUE = -5e4

# helper functions

def exists(val):
    return val is not None

def default(value, d):
    return d if not exists(value) else value

def to(t):
    return {'device': t.device, 'dtype': t.dtype}

def max_neg_value(tensor):
    return -torch.finfo(tensor.dtype).max

def l2norm(tensor):
    dtype = tensor.dtype
    normed = F.normalize(tensor, dim = -1)
    return normed.type(dtype)

def pad_to_multiple(tensor, multiple, dim=-1, value=0):
    seqlen = tensor.shape[dim]
    m = seqlen / multiple
    if m.is_integer():
        return False, tensor
    remainder = math.ceil(m) * multiple - seqlen
    pad_offset = (0,) * (-1 - dim) * 2
    return True, F.pad(tensor, (*pad_offset, 0, remainder), value = value)

def look_around(x, backward = 1, forward = 0, pad_value = -1, dim = 2):
    t = x.shape[1]
    dims = (len(x.shape) - dim) * (0, 0)
    padded_x = F.pad(x, (*dims, backward, forward), value = pad_value)
    tensors = [padded_x[:, ind:(ind + t), ...] for ind in range(forward + backward + 1)]
    return torch.cat(tensors, dim = dim)

# main class

class LocalAttention(nn.Module):
    def __init__(
        self,
        window_size,
        causal = False,
        look_backward = 1,
        look_forward = None,
        dropout = 0.,
        shared_qk = False,
        rel_pos_emb_config = None,
        dim = None,
        autopad = False,
        exact_windowsize = False
    ):
        super().__init__()
        look_forward = default(look_forward, 0 if causal else 1)
        assert not (causal and look_forward > 0), 'you cannot look forward if causal'

        self.window_size = window_size
        self.autopad = autopad
        self.exact_windowsize = exact_windowsize

        self.causal = causal

        self.look_backward = look_backward
        self.look_forward = look_forward

        self.dropout = nn.Dropout(dropout)

        self.shared_qk = shared_qk

        # relative positions

        # self.rel_pos = None
        # if exists(rel_pos_emb_config) or exists(dim):  # backwards compatible with old `rel_pos_emb_config` deprecated argument
        #     if exists(rel_pos_emb_config):
        #         dim = rel_pos_emb_config[0]
        #     self.rel_pos = SinusoidalEmbeddings(dim)

    def forward(self, q, k, v, packed_shape, mask = None, input_mask = None):
        mask = default(mask, input_mask)

        autopad, pad_value, window_size, causal, look_backward, look_forward, shared_qk = self.autopad, -1, self.window_size, self.causal, self.look_backward, self.look_forward, self.shared_qk

        # https://github.com/arogozhnikov/einops/blob/master/docs/4-pack-and-unpack.ipynb
        # (q, packed_shape), (k, _), (v, _) = map(lambda t: pack([t], '* n d'), (q, k, v))        # (2, 8, 2048, 64) -> (16, 2048, 64)

        # rotary embeddings

        # if exists(self.rel_pos):
        #     pos_emb = self.rel_pos(q)     # (16, 2048, 64)
        #     q, k = apply_rotary_pos_emb(q, k, pos_emb)

        # auto padding

        # if autopad:
        #     orig_seq_len = q.shape[1]
        #     (needed_pad, q), (_, k), (_, v) = map(lambda t: pad_to_multiple(t, self.window_size, dim = -2), (q, k, v))

        b, n, dim_head, device, dtype = *q.shape, q.device, q.dtype
        scale = dim_head ** -0.5

        if n % window_size != 0:
            print('sequence length must be divisible by window size for local attention', n, window_size)
            assert True

        windows = n // window_size

        if shared_qk:
            k = l2norm(k)

        seq = torch.arange(n, device = device)
        b_t = rearrange(seq, '(w n) -> 1 w n', w = windows, n = window_size)

        bq, bk, bv = map(lambda t: rearrange(t, 'b (w n) d -> b w n d', w = windows), (q, k, v))

        look_around_kwargs = dict(
            backward =  look_backward,
            forward =  look_forward,
            pad_value = pad_value
        )

        bk = look_around(bk, **look_around_kwargs)
        bv = look_around(bv, **look_around_kwargs)

        bq_t = b_t
        bq_k = look_around(b_t, **look_around_kwargs)

        bq_t = rearrange(bq_t, '... i -> ... i 1')
        bq_k = rearrange(bq_k, '... j -> ... 1 j')

        sim = einsum('b h i e, b h j e -> b h i j', bq, bk) * scale

        mask_value = max_neg_value(sim)

        if shared_qk:
            self_mask = bq_t == bq_k
            sim = sim.masked_fill(self_mask, TOKEN_SELF_ATTN_VALUE)
            del self_mask

        if causal:
            causal_mask = bq_t < bq_k

            if self.exact_windowsize:
                max_causal_window_size = (self.window_size * self.look_backward)
                causal_mask = causal_mask | (bq_t > (bq_k + max_causal_window_size))

            sim = sim.masked_fill(causal_mask, mask_value)
            del causal_mask

        # mask out padding value

        # if autopad and needed_pad:
        #     pad_mask = bq_k == pad_value
        #     sim = sim.masked_fill(pad_mask, mask_value)
        #     del pad_mask

        if exists(mask):
            batch = mask.shape[0]
            assert (b % batch) == 0

            h = b // mask.shape[0]

            # if autopad:
            #     _, mask = pad_to_multiple(mask, window_size, dim = -1, value = False)

            mask = rearrange(mask, '... (w n) -> (...) w n', w = windows, n = window_size)
            mask = look_around(mask, **{**look_around_kwargs, 'pad_value': False})
            mask = rearrange(mask, '... j -> ... 1 j')
            mask = repeat(mask, 'b ... -> (b h) ...', h = h)
            sim = sim.masked_fill(~mask, mask_value)
            del mask

        # attention

        attn = sim.softmax(dim = -1)
        attn = self.dropout(attn)

        # aggregation

        out = einsum('b h i j, b h j e -> b h i e', attn, bv)
        out = rearrange(out, 'b w n d -> b (w n) d')

        # if autopad:
        #     out = out[:, :orig_seq_len, :]

        out, *_ = unpack(out, packed_shape, '* n d')
        return out

diffusion_latent/model/local_attention/local_attention.py
YoungSeng/UnifiedGesture
import torch
from torch import nn, einsum

from einops import rearrange

class SinusoidalEmbeddings(nn.Module):
    def __init__(self, dim):
        super().__init__()
        inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, x):
        n = x.shape[-2]
        t = torch.arange(n, device = x.device).type_as(self.inv_freq)
        freqs = torch.einsum('i , j -> i j', t, self.inv_freq)
        return torch.cat((freqs, freqs), dim=-1)

def rotate_half(x):
    x = rearrange(x, 'b ... (r d) -> b (...) r d', r = 2)
    x1, x2 = x.unbind(dim = -2)
    return torch.cat((-x2, x1), dim = -1)

def apply_rotary_pos_emb(q, k, freqs):
    q, k = map(lambda t: (t * freqs.cos()) + (rotate_half(t) * freqs.sin()), (q, k))
    return q, k

diffusion_latent/model/local_attention/rotary.py
YoungSeng/UnifiedGesture
import torch
from torch import nn
import torch.nn.functional as F

from einops import rearrange

from local_attention.local_attention import LocalAttention

# helper function

def exists(val):
    return val is not None

def eval_decorator(fn):
    def inner(model, *args, **kwargs):
        was_training = model.training
        model.eval()
        out = fn(model, *args, **kwargs)
        model.train(was_training)
        return out
    return inner

# sampling functions

def top_k(logits, thres = 0.9):
    k = int((1 - thres) * logits.shape[-1])
    val, ind = torch.topk(logits, k)
    probs = torch.full_like(logits, float('-inf'))
    probs.scatter_(1, ind, val)
    return probs

# multi-head attention

class LocalMHA(nn.Module):
    def __init__(
        self,
        *,
        dim,
        window_size,
        dim_head = 64,
        heads = 8,
        dropout = 0.,
        causal = False,
        prenorm = False,
        **kwargs
    ):
        super().__init__()        
        inner_dim = dim_head * heads

        self.norm = nn.LayerNorm(dim) if prenorm else None

        self.heads = heads
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)

        self.attn_fn = LocalAttention(
            dim = dim_head,
            window_size = window_size,
            causal = causal,
            autopad = True,
            exact_windowsize = True,
            **kwargs
        )

        self.to_out = nn.Linear(inner_dim, dim, bias = False)

    def forward(self, x, mask = None):
        if exists(self.norm):
            x = self.norm(x)

        q, k, v = self.to_qkv(x).chunk(3, dim = -1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), (q, k, v)) 

        out = self.attn_fn(q, k, v, mask = mask)

        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)

# feedforward

class GEGLU(nn.Module):
    def forward(self, x):
        x, gate = x.chunk(2, dim = -1)
        return x * F.gelu(gate)

def FeedForward(dim, mult = 4, dropout = 0.):
    inner_dim = int(dim * mult * 2 / 3)

    return nn.Sequential(
        nn.LayerNorm(dim),
        nn.Linear(dim, inner_dim * 2, bias = False),
        GEGLU(),
        nn.Dropout(dropout),
        nn.Linear(inner_dim, dim, bias = False)
    )

# main transformer class

class LocalTransformer(nn.Module):
    def __init__(
        self,
        *,
        num_tokens,
        max_seq_len,
        dim,
        depth,
        causal = True,
        local_attn_window_size = 512,
        dim_head = 64,
        heads = 8,
        ff_mult = 4,
        attn_dropout = 0.,
        ff_dropout = 0.,
        ignore_index = -1,
        **kwargs
    ):
        super().__init__()
        self.token_emb = nn.Embedding(num_tokens, dim)
        self.pos_emb = nn.Embedding(max_seq_len, dim)

        self.max_seq_len = max_seq_len
        self.layers = nn.ModuleList([])

        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                LocalMHA(dim = dim, dim_head = dim_head, heads = heads, dropout = attn_dropout, causal = causal, window_size = local_attn_window_size, prenorm = True, **kwargs),
                FeedForward(dim = dim, mult = ff_mult, dropout = ff_dropout)
            ]))

        self.ignore_index = ignore_index
        self.to_logits = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_tokens, bias = False)
        )

    @torch.no_grad()
    @eval_decorator
    def generate(
        self,
        prime,
        seq_len,
        temperature = 1.,
        filter_thres = 0.9,
        **kwargs
    ):
        n, device = prime.shape[1], prime.device

        out = prime

        for _ in range(seq_len):
            logits = self.forward(out[:, -self.max_seq_len:], **kwargs)
            filtered_logits = top_k(logits[:, -1], thres = filter_thres)
            probs = F.softmax(filtered_logits / temperature, dim = -1)
            sampled = torch.multinomial(probs, 1)
            out = torch.cat((out, sampled), dim = -1)

        return out[:, n:]

    def forward(self, x, mask = None, return_loss = False):
        if return_loss:
            x, labels = x[:, :-1], x[:, 1:]

        n, device = x.shape[1], x.device
        x = self.token_emb(x)

        assert n <= self.max_seq_len
        x = x + self.pos_emb(torch.arange(n, device = device))

        for attn, ff in self.layers:
            x = attn(x, mask = mask) + x
            x = ff(x) + x

        logits = self.to_logits(x)

        if not return_loss:
            return logits

        logits = rearrange(logits, 'b n c -> b c n')
        loss = F.cross_entropy(logits, labels, ignore_index = self.ignore_index)
        return loss

diffusion_latent/model/local_attention/transformer.py
YoungSeng/UnifiedGesture
import pdb

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
# import clip
import sys
[sys.path.append(i) for i in ['.', './model']]
from local_attention.rotary import SinusoidalEmbeddings, apply_rotary_pos_emb
from local_attention import LocalAttention

class MDM(nn.Module):
    def __init__(self, modeltype, njoints, nfeats, translation, pose_rep, glob, glob_rot,
                 latent_dim=256, ff_size=1024, num_layers=8, num_heads=4, dropout=0.1,
                 ablation=None, activation="gelu", legacy=False, data_rep='rot6d', dataset='amass', clip_dim=512,
                 arch='trans_enc', emb_trans_dec=False, audio_feat='', n_seed=1, cond_mode='', **kargs):
        super().__init__()

        self.legacy = legacy
        self.modeltype = modeltype
        self.njoints = njoints
        self.nfeats = nfeats
        self.data_rep = data_rep
        self.dataset = dataset

        self.pose_rep = pose_rep
        self.glob = glob
        self.glob_rot = glob_rot
        self.translation = translation

        self.latent_dim = latent_dim

        self.ff_size = ff_size
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.dropout = dropout

        self.ablation = ablation
        self.activation = activation
        self.clip_dim = clip_dim
        self.action_emb = kargs.get('action_emb', None)

        self.input_feats = self.njoints * self.nfeats

        self.normalize_output = kargs.get('normalize_encoder_output', False)

        self.cond_mask_prob = kargs.get('cond_mask_prob', 0.)
        self.arch = arch
        self.gru_emb_dim = self.latent_dim if self.arch == 'gru' else 0

        self.audio_feat = audio_feat
        if audio_feat == 'wav encoder':
            self.audio_feat_dim = 32
        elif audio_feat == 'mfcc':
            self.audio_feat_dim = 13
        elif self.audio_feat == 'wavlm':
            print('USE WAVLM')
            self.audio_feat_dim = 64        # Linear 1024 -> 64
            self.WavEncoder = WavEncoder()

        self.sequence_pos_encoder = PositionalEncoding(self.latent_dim, self.dropout)
        self.emb_trans_dec = emb_trans_dec

        self.cond_mode = cond_mode
        self.num_head = 8

        if 'style2' not in self.cond_mode:
            self.input_process = InputProcess(self.data_rep, self.input_feats + self.audio_feat_dim + self.gru_emb_dim, self.latent_dim)

        if self.arch == 'mytrans_enc':
            print("MY TRANS_ENC init")
            from mytransformer import TransformerEncoderLayer, TransformerEncoder

            self.embed_positions = RoFormerSinusoidalPositionalEmbedding(1536, self.latent_dim)

            seqTransEncoderLayer = TransformerEncoderLayer(d_model=self.latent_dim,
                                                              nhead=self.num_heads,
                                                              dim_feedforward=self.ff_size,
                                                              dropout=self.dropout,
                                                              activation=self.activation)
            self.seqTransEncoder = TransformerEncoder(seqTransEncoderLayer,
                                                         num_layers=self.num_layers)

        elif self.arch == 'trans_enc':
            print("TRANS_ENC init")
            seqTransEncoderLayer = nn.TransformerEncoderLayer(d_model=self.latent_dim,
                                                              nhead=self.num_heads,
                                                              dim_feedforward=self.ff_size,
                                                              dropout=self.dropout,
                                                              activation=self.activation)

            self.seqTransEncoder = nn.TransformerEncoder(seqTransEncoderLayer,
                                                         num_layers=self.num_layers)
        elif self.arch == 'trans_dec':
            print("TRANS_DEC init")
            seqTransDecoderLayer = nn.TransformerDecoderLayer(d_model=self.latent_dim,
                                                              nhead=self.num_heads,
                                                              dim_feedforward=self.ff_size,
                                                              dropout=self.dropout,
                                                              activation=activation)
            self.seqTransDecoder = nn.TransformerDecoder(seqTransDecoderLayer,
                                                         num_layers=self.num_layers)
        elif self.arch == 'gru':
            print("GRU init")
            self.gru = nn.GRU(self.latent_dim, self.latent_dim, num_layers=self.num_layers, batch_first=False)
        else:
            raise ValueError('Please choose correct architecture [trans_enc, trans_dec, gru]')

        self.embed_timestep = TimestepEmbedder(self.latent_dim, self.sequence_pos_encoder)
        self.n_seed = n_seed
        if 'style1' in self.cond_mode:
            print('EMBED STYLE BEGIN TOKEN')
            if self.n_seed != 0:
                self.style_dim = 64
                self.embed_style = nn.Linear(7, self.style_dim)
                self.embed_text = nn.Linear(self.njoints * n_seed, self.latent_dim - self.style_dim)
            else:
                self.style_dim = self.latent_dim
                self.embed_style = nn.Linear(6, self.style_dim)

        elif 'style2' in self.cond_mode:
            print('EMBED STYLE ALL FRAMES')
            self.style_dim = 64
            self.embed_style = nn.Linear(7, self.style_dim)
            self.input_process = InputProcess(self.data_rep, self.input_feats + self.audio_feat_dim + self.gru_emb_dim + self.style_dim,
                                              self.latent_dim)
            if self.n_seed != 0:
                self.embed_text = nn.Linear(self.njoints * n_seed, self.latent_dim)
        elif self.n_seed != 0:
            self.embed_text = nn.Linear(self.njoints * n_seed, self.latent_dim)

        self.output_process = OutputProcess(self.data_rep, self.input_feats, self.latent_dim, self.njoints,
                                            self.nfeats)

        if 'cross_local_attention' in self.cond_mode:
            self.rel_pos = SinusoidalEmbeddings(self.latent_dim // self.num_head)
            self.input_process = InputProcess(self.data_rep, self.input_feats + self.gru_emb_dim, self.latent_dim)
            self.cross_local_attention = LocalAttention(
                dim=32,  # dimension of each head (you need to pass this in for relative positional encoding)
                window_size=6,  # window size. 512 is optimal, but 256 or 128 yields good enough results        # 4 -> 6
                causal=True,  # auto-regressive or not
                look_backward=1,  # each window looks at the window before
                look_forward=0,     # for non-auto-regressive case, will default to 1, so each window looks at the window before and after it
                dropout=0.1,  # post-attention dropout
                exact_windowsize=False
                # if this is set to true, in the causal setting, each query will see at maximum the number of keys equal to the window size
            )
            self.input_process2 = nn.Linear(self.latent_dim * 2 + self.audio_feat_dim, self.latent_dim)

        if 'cross_local_attention2' in self.cond_mode:
            print('Cross Local Attention2')
            self.selfAttention = LinearTemporalCrossAttention(seq_len=0, latent_dim=256, text_latent_dim=256, num_head=8, dropout=0.1, time_embed_dim=0)

        if 'cross_local_attention3' in self.cond_mode:
            print('Cross Local Attention3')

        if 'cross_local_attention4' in self.cond_mode:
            print('Cross Local Attention4')

    def freeze_drop(self):
        for m in self.modules():
            if isinstance(m, nn.Dropout):
                m.eval()

    def parameters_wo_clip(self):
        return [p for name, p in self.named_parameters() if not name.startswith('clip_model.')]

    def mask_cond(self, cond, force_mask=False):
        bs, d = cond.shape
        if force_mask:
            return torch.zeros_like(cond)
        elif self.training and self.cond_mask_prob > 0.:
            mask = torch.bernoulli(torch.ones(bs, device=cond.device) * self.cond_mask_prob).view(bs, 1)  # 1-> use null_cond, 0-> use real cond
            return cond * (1. - mask)
        else:
            return cond

    def forward(self, x, timesteps, y=None):
        """
        x: [batch_size, njoints, nfeats, max_frames], denoted x_t in the paper
        timesteps: [batch_size] (int)
        seed: [batch_size, njoints, nfeats]
        """

        bs, njoints, nfeats, nframes = x.shape      # 64, 251, 1, 196
        emb_t = self.embed_timestep(timesteps)  # [1, bs, d], (1, 2, 256)

        force_mask = y.get('uncond', False)  # False
        if 'style1' in self.cond_mode:
            embed_style = self.mask_cond(self.embed_style(y['style']), force_mask=force_mask)       # (bs, 64)
            if self.n_seed != 0:

                embed_text = self.embed_text(self.mask_cond(y['seed'].squeeze(2).reshape(bs, -1), force_mask=force_mask))       # (bs, 256-64)
                emb_1 = torch.cat((embed_style, embed_text), dim=1)
            else:
                emb_1 = embed_style
        elif self.n_seed != 0:
            emb_1 = self.embed_text(self.mask_cond(y['seed'].squeeze(2).reshape(bs, -1), force_mask=force_mask))     # z_tk

        if self.audio_feat == 'wavlm':
            enc_text = self.WavEncoder(y['audio']).permute(1, 0, 2)
        else:
            enc_text = y['audio']

        if 'cross_local_attention' in self.cond_mode:
            if 'cross_local_attention4' in self.cond_mode:
                x = x.reshape(bs, njoints * nfeats, 1, nframes)  # [2, 135, 1, 240]
                # self-attention
                x_ = self.input_process(x)  # [2, 135, 1, 240] -> [240, 2, 256]
                xseq = torch.cat((emb_1 + emb_t, x_), axis=0)  # [seqlen+1, bs, d]     # [(1, 2, 256), (240, 2, 256)] -> (241, 2, 256)
                sinusoidal_pos = self.embed_positions(xseq.shape[0], 0)[None, None, :, :].chunk(2, dim=-1)
                xseq = self.apply_rotary(xseq.permute(1, 0, 2), sinusoidal_pos).squeeze(0).permute(1, 0, 2)

                if 'cross_local_attention2' in self.cond_mode:
                    xseq = (self.selfAttention(xseq).permute(1, 0, 2))[1:]
                else:
                    xseq = self.seqTransEncoder(xseq)[1:]

                # local-cross-attention
                packed_shape = [torch.Size([bs, self.num_head])]
                xseq = torch.cat((xseq, enc_text), axis=2)  # [bs, d+joints*feat, 1, #frames], (240, 2, 32)
                # all frames
                embed_style_2 = (emb_1 + emb_t).repeat(nframes, 1, 1)  # (bs, 64) -> (len, bs, 64)
                xseq = torch.cat((embed_style_2, xseq), axis=2)  # (seq, bs, dim)
                xseq = self.input_process2(xseq)
                sinusoidal_pos = self.embed_positions(xseq.shape[0], 0)[None, None, :, :].chunk(2, dim=-1)
                xseq = self.apply_rotary(xseq.permute(1, 0, 2), sinusoidal_pos).squeeze(0).permute(1, 0, 2)

                xseq = xseq.permute(1, 0, 2)  # (bs, len, dim)
                xseq = xseq.view(bs, nframes, self.num_head, -1)
                xseq = xseq.permute(0, 2, 1, 3)  # Need (2, 8, 2048, 64)
                xseq = xseq.reshape(bs * self.num_head, nframes, -1)
                xseq = self.cross_local_attention(xseq, xseq, xseq, packed_shape=packed_shape,
                                                  mask=y['mask_local'])  # (2, 8, 2048, 64)
                xseq = xseq.permute(0, 2, 1, 3)  # (bs, len, 8, 64)
                xseq = xseq.reshape(bs, nframes, -1)
                output = xseq.permute(1, 0, 2)

            elif 'cross_local_attention3' in self.cond_mode:
                x = x.reshape(bs, njoints * nfeats, 1, nframes)  # [2, 135, 1, 240]
                # self-attention
                x_ = self.input_process(x)  # [2, 135, 1, 240] -> [240, 2, 256]

                # local-cross-attention
                packed_shape = [torch.Size([bs, self.num_head])]
                xseq = torch.cat((x_, enc_text), axis=2)  # [bs, d+joints*feat, 1, #frames], (240, 2, 32)
                # all frames
                embed_style_2 = (emb_1 + emb_t).repeat(nframes, 1, 1)  # (bs, 64) -> (len, bs, 64)
                xseq = torch.cat((embed_style_2, xseq), axis=2)  # (seq, bs, dim)
                xseq = self.input_process2(xseq)
                xseq = xseq.permute(1, 0, 2)  # (bs, len, dim)
                xseq = xseq.view(bs, nframes, self.num_head, -1)
                xseq = xseq.permute(0, 2, 1, 3)  # Need (2, 8, 2048, 64)
                xseq = xseq.reshape(bs * self.num_head, nframes, -1)
                pos_emb = self.rel_pos(xseq)  # (89, 32)
                xseq, _ = apply_rotary_pos_emb(xseq, xseq, pos_emb)
                xseq = self.cross_local_attention(xseq, xseq, xseq, packed_shape=packed_shape,
                                                  mask=y['mask_local'])  # (2, 8, 2048, 64)
                xseq = xseq.permute(0, 2, 1, 3)  # (bs, len, 8, 64)
                xseq = xseq.reshape(bs, nframes, -1)
                xseq = xseq.permute(1, 0, 2)

                xseq = torch.cat((emb_1 + emb_t, xseq), axis=0)  # [seqlen+1, bs, d]     # [(1, 2, 256), (240, 2, 256)] -> (241, 2, 256)
                xseq = xseq.permute(1, 0, 2)  # (bs, len, dim)
                xseq = xseq.view(bs, nframes + 1, self.num_head, -1)
                xseq = xseq.permute(0, 2, 1, 3)  # Need (2, 8, 2048, 64)
                xseq = xseq.reshape(bs * self.num_head, nframes + 1, -1)
                pos_emb = self.rel_pos(xseq)  # (89, 32)
                xseq, _ = apply_rotary_pos_emb(xseq, xseq, pos_emb)
                xseq_rpe = xseq.reshape(bs, self.num_head, nframes + 1, -1)
                xseq = xseq_rpe.permute(0, 2, 1, 3)  # [seqlen+1, bs, d]
                xseq = xseq.view(bs, nframes + 1, -1)
                xseq = xseq.permute(1, 0, 2)
                if 'cross_local_attention2' in self.cond_mode:
                    xseq = (self.selfAttention(xseq).permute(1, 0, 2))[1:]
                else:
                    output = self.seqTransEncoder(xseq)[1:]

            elif 'cross_local_attention5' in self.cond_mode:
                x = x.reshape(bs, njoints * nfeats, 1, nframes)  # [2, 135, 1, 240]
                # self-attention
                x_ = self.input_process(x)  # [2, 135, 1, 240] -> [240, 2, 256]

                # local-cross-attention
                packed_shape = [torch.Size([bs, self.num_head])]
                xseq = torch.cat((x_, enc_text), axis=2)  # [bs, d+joints*feat, 1, #frames], (240, 2, 32)
                # all frames
                embed_style_2 = (emb_1 + emb_t).repeat(nframes, 1, 1)  # (bs, 64) -> (len, bs, 64)
                xseq = torch.cat((embed_style_2, xseq), axis=2)  # (seq, bs, dim)
                xseq = self.input_process2(xseq)
                xseq = xseq.permute(1, 0, 2)  # (bs, len, dim)
                xseq = xseq.view(bs, nframes, self.num_head, -1)
                xseq = xseq.permute(0, 2, 1, 3)  # Need (2, 8, 2048, 64)
                xseq = xseq.reshape(bs * self.num_head, nframes, -1)
                pos_emb = self.rel_pos(xseq)  # (89, 32)
                xseq, _ = apply_rotary_pos_emb(xseq, xseq, pos_emb)
                xseq = self.cross_local_attention(xseq, xseq, xseq, packed_shape=packed_shape,
                                                  mask=y['mask_local'])  # (2, 8, 2048, 64)
                xseq = xseq.permute(0, 2, 1, 3)  # (bs, len, 8, 64)
                xseq = xseq.reshape(bs, nframes, -1)
                output = xseq.permute(1, 0, 2)

            else:
                x = x.reshape(bs, njoints*nfeats, 1, nframes)      # [2, 135, 1, 240]
                # self-attention
                x_ = self.input_process(x)       # [2, 135, 1, 240] -> [240, 2, 256]
                xseq = torch.cat((emb_1 + emb_t, x_), axis=0)  # [seqlen+1, bs, d]     # [(1, 2, 256), (240, 2, 256)] -> (241, 2, 256)
                xseq = xseq.permute(1, 0, 2)        # (bs, len, dim)
                xseq = xseq.view(bs, nframes + 1, self.num_head, -1)
                xseq = xseq.permute(0, 2, 1, 3)     # Need (2, 8, 2048, 64)
                xseq = xseq.reshape(bs*self.num_head, nframes + 1, -1)
                pos_emb = self.rel_pos(xseq)        #  (89, 32)
                xseq, _ = apply_rotary_pos_emb(xseq, xseq, pos_emb)
                xseq_rpe = xseq.reshape(bs, self.num_head, nframes + 1, -1)
                xseq = xseq_rpe.permute(0, 2, 1, 3)       # [seqlen+1, bs, d]
                xseq = xseq.view(bs, nframes + 1, -1)
                xseq = xseq.permute(1, 0, 2)
                if 'cross_local_attention2' in self.cond_mode:
                    xseq = (self.selfAttention(xseq).permute(1, 0, 2))[1:]
                else:
                    xseq = self.seqTransEncoder(xseq)[1:]

                # local-cross-attention
                packed_shape = [torch.Size([bs, self.num_head])]
                xseq = torch.cat((xseq, enc_text), axis=2)  #[bs, d+joints*feat, 1, #frames], (240, 2, 32)
                # all frames
                embed_style_2 = (emb_1 + emb_t).repeat(nframes, 1, 1)       # (bs, 64) -> (len, bs, 64)
                xseq = torch.cat((embed_style_2, xseq), axis=2)     # (seq, bs, dim)
                xseq = self.input_process2(xseq)
                xseq = xseq.permute(1, 0, 2)  # (bs, len, dim)
                xseq = xseq.view(bs, nframes, self.num_head, -1)
                xseq = xseq.permute(0, 2, 1, 3)  # Need (2, 8, 2048, 64)
                xseq = xseq.reshape(bs * self.num_head, nframes, -1)
                pos_emb = self.rel_pos(xseq)  # (89, 32)
                xseq, _ = apply_rotary_pos_emb(xseq, xseq, pos_emb)
                xseq = self.cross_local_attention(xseq, xseq, xseq, packed_shape=packed_shape, mask=y['mask_local'])     # (2, 8, 2048, 64)
                xseq = xseq.permute(0, 2, 1, 3)     # (bs, len, 8, 64)
                xseq = xseq.reshape(bs, nframes, -1)
                output = xseq.permute(1, 0, 2)

        else:
            if self.arch == 'trans_enc' or self.arch == 'trans_dec' or self.arch == 'mytrans_enc':
                x_reshaped = x.reshape(bs, njoints*nfeats, 1, nframes)      # [2, 135, 1, 240]
                enc_text_gru = enc_text.permute(1, 2, 0)  # (240, 2, 32) -> (2, 32, 240)
                enc_text_gru = enc_text_gru.reshape(bs, self.audio_feat_dim, 1, nframes)
                x = torch.cat((x_reshaped, enc_text_gru), axis=1)  #[bs, d+joints*feat, 1, #frames]
                if 'style2' in self.cond_mode:
                    embed_style = self.mask_cond(self.embed_style(y['style']), force_mask=force_mask).repeat(nframes, 1, 1)  # (#frames, bs, 64)
                    embed_style = embed_style.unsqueeze(2)
                    embed_style = embed_style.permute(1, 3, 2, 0)
                    x = torch.cat((x, embed_style), axis=1)  # [bs, d+joints*feat, 1, #frames]

            if self.arch == 'gru':
                x_reshaped = x.reshape(bs, njoints*nfeats, 1, nframes)      # [2, 135, 1, 240]
                emb_gru = emb.repeat(nframes, 1, 1)     #[#frames, bs, d]

                enc_text_gru = enc_text.permute(1, 2, 0)        # (240, 2, 32) -> (2, 32, 240)
                enc_text_gru = enc_text_gru.reshape(bs, self.audio_feat_dim, 1, nframes)

                emb_gru = emb_gru.permute(1, 2, 0)      #[bs, d, #frames]
                emb_gru = emb_gru.reshape(bs, self.latent_dim, 1, nframes)  #[bs, d, 1, #frames]
                x = torch.cat((x_reshaped, emb_gru, enc_text_gru), axis=1)  #[bs, d+joints*feat, 1, #frames]

            x = self.input_process(x)       # [2, 135, 1, 240] -> [240, 2, 224]

            if self.arch == 'trans_enc':
                # adding the timestep embed
                # x = torch.cat((x, enc_text), axis=2)        # [[240, 2, 224], (240, 2, 32)] -> (240, 2, 256)
                xseq = torch.cat((emb, x), axis=0)  # [seqlen+1, bs, d]     # [(1, 2, 256), (240, 2, 256)] -> (241, 2, 256)

                xseq = self.sequence_pos_encoder(xseq)  # [seqlen+1, bs, d]
                output = self.seqTransEncoder(xseq)[1:]  # , src_key_padding_mask=~maskseq)  # [seqlen, bs, d]      # -> (240, 2, 256)

            elif self.arch == 'trans_dec':
                if self.emb_trans_dec:
                    xseq = torch.cat((emb, x), axis=0)
                else:
                    xseq = x
                xseq = self.sequence_pos_encoder(xseq)  # [seqlen+1, bs, d]
                if self.emb_trans_dec:
                    output = self.seqTransDecoder(tgt=xseq, memory=emb)[1:] # [seqlen, bs, d] # FIXME - maybe add a causal mask
                else:
                    output = self.seqTransDecoder(tgt=xseq, memory=emb)

            elif self.arch == 'gru':
                xseq = x
                xseq = self.sequence_pos_encoder(xseq)  # [seqlen, bs, d]
                # pdb.set_trace()
                output, _ = self.gru(xseq)

            elif self.arch == 'mytrans_enc':
                # adding the timestep embed
                # x = torch.cat((x, enc_text), axis=2)        # [[240, 2, 224], (240, 2, 32)] -> (240, 2, 256)
                xseq = torch.cat((emb, x), axis=0)  # [seqlen+1, bs, d]     # [(1, 2, 256), (240, 2, 256)] -> (241, 2, 256)

                sinusoidal_pos = self.embed_positions(xseq.shape[0], 0)[None, None, :, :].chunk(2, dim=-1)
                xseq = self.apply_rotary(xseq.permute(1, 0, 2), sinusoidal_pos).squeeze(0).permute(1, 0, 2)

                output = self.seqTransEncoder(xseq)[1:]  # , src_key_padding_mask=~maskseq)  # [seqlen, bs, d]      # -> (240, 2, 256)

        output = self.output_process(output)  # [bs, njoints, nfeats, nframes]
        return output


    @staticmethod
    def apply_rotary(x, sinusoidal_pos):
        sin, cos = sinusoidal_pos
        x1, x2 = x[..., 0::2], x[..., 1::2]
        # 如果是旋转query key的话，下面这个直接cat就行，因为要进行矩阵乘法，最终会在这个维度求和。（只要保持query和key的最后一个dim的每一个位置对应上就可以）
        # torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
        # 如果是旋转value的话，下面这个stack后再flatten才可以，因为训练好的模型最后一个dim是两两之间交替的。
        return torch.stack([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).flatten(-2, -1)


    # def train(self, *args, **kwargs):
    #     super().train(*args, **kwargs)
    #     self.rot2xyz.smpl_model.train(*args, **kwargs)


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)      # (5000, 128)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)     # (5000, 1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)

        self.register_buffer('pe', pe)

    def forward(self, x):
        # not used in the final model
        x = x + self.pe[:x.shape[0], :]
        return self.dropout(x)


# Copied from transformers.models.marian.modeling_marian.MarianSinusoidalPositionalEmbedding with Marian->RoFormer
class RoFormerSinusoidalPositionalEmbedding(nn.Embedding):
    """This module produces sinusoidal positional embeddings of any length."""

    def __init__(
        self, num_positions: int, embedding_dim: int
    ):
        super().__init__(num_positions, embedding_dim)
        self.weight = self._init_weight(self.weight)

    @staticmethod
    def _init_weight(out: nn.Parameter):
        """
        Identical to the XLM create_sinusoidal_embeddings except features are not interleaved. The cos features are in
        the 2nd half of the vector. [dim // 2:]
        """
        n_pos, dim = out.shape
        position_enc = np.array(
            [
                [pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)]
                for pos in range(n_pos)
            ]
        )
        out.requires_grad = False  # set early to avoid an error in pytorch-1.8+
        sentinel = dim // 2 if dim % 2 == 0 else (dim // 2) + 1
        out[:, 0:sentinel] = torch.FloatTensor(np.sin(position_enc[:, 0::2]))
        out[:, sentinel:] = torch.FloatTensor(np.cos(position_enc[:, 1::2]))
        out.detach_()
        return out

    @torch.no_grad()
    def forward(self, seq_len: int, past_key_values_length: int = 0):
        """`input_ids_shape` is expected to be [bsz x seqlen]."""
        positions = torch.arange(
            past_key_values_length,
            past_key_values_length + seq_len,
            dtype=torch.long,
            device=self.weight.device,
        )
        return super().forward(positions)


class TimestepEmbedder(nn.Module):
    def __init__(self, latent_dim, sequence_pos_encoder):
        super().__init__()
        self.latent_dim = latent_dim
        self.sequence_pos_encoder = sequence_pos_encoder

        time_embed_dim = self.latent_dim
        self.time_embed = nn.Sequential(
            nn.Linear(self.latent_dim, time_embed_dim),
            nn.SiLU(),
            nn.Linear(time_embed_dim, time_embed_dim),
        )

    def forward(self, timesteps):
        return self.time_embed(self.sequence_pos_encoder.pe[timesteps]).permute(1, 0, 2)


class InputProcess(nn.Module):
    def __init__(self, data_rep, input_feats, latent_dim):
        super().__init__()
        self.data_rep = data_rep
        self.input_feats = input_feats
        self.latent_dim = latent_dim
        self.poseEmbedding = nn.Linear(self.input_feats, self.latent_dim)
        if self.data_rep == 'rot_vel':
            self.velEmbedding = nn.Linear(self.input_feats, self.latent_dim)

    def forward(self, x):
        bs, njoints, nfeats, nframes = x.shape
        x = x.permute((3, 0, 1, 2)).reshape(nframes, bs, njoints*nfeats)

        if self.data_rep in ['rot6d', 'xyz', 'hml_vec']:
            x = self.poseEmbedding(x)  # [seqlen, bs, d]
            return x
        elif self.data_rep == 'rot_vel':
            first_pose = x[[0]]  # [1, bs, 150]
            first_pose = self.poseEmbedding(first_pose)  # [1, bs, d]
            vel = x[1:]  # [seqlen-1, bs, 150]
            vel = self.velEmbedding(vel)  # [seqlen-1, bs, d]
            return torch.cat((first_pose, vel), axis=0)  # [seqlen, bs, d]
        else:
            raise ValueError


class OutputProcess(nn.Module):
    def __init__(self, data_rep, input_feats, latent_dim, njoints, nfeats):
        super().__init__()
        self.data_rep = data_rep
        self.input_feats = input_feats
        self.latent_dim = latent_dim
        self.njoints = njoints
        self.nfeats = nfeats
        self.poseFinal = nn.Linear(self.latent_dim, self.input_feats)
        if self.data_rep == 'rot_vel':
            self.velFinal = nn.Linear(self.latent_dim, self.input_feats)

    def forward(self, output):
        nframes, bs, d = output.shape
        if self.data_rep in ['rot6d', 'xyz', 'hml_vec']:
            output = self.poseFinal(output)  # [seqlen, bs, 150]
        elif self.data_rep == 'rot_vel':
            first_pose = output[[0]]  # [1, bs, d]
            first_pose = self.poseFinal(first_pose)  # [1, bs, 150]
            vel = output[1:]  # [seqlen-1, bs, d]
            vel = self.velFinal(vel)  # [seqlen-1, bs, 150]
            output = torch.cat((first_pose, vel), axis=0)  # [seqlen, bs, 150]
        else:
            raise ValueError
        output = output.reshape(nframes, bs, self.njoints, self.nfeats)
        output = output.permute(1, 2, 3, 0)  # [bs, njoints, nfeats, nframes]
        return output


class LinearTemporalCrossAttention(nn.Module):

    def __init__(self, seq_len, latent_dim, text_latent_dim, num_head, dropout, time_embed_dim):
        super().__init__()
        self.num_head = num_head
        self.norm = nn.LayerNorm(latent_dim)
        self.text_norm = nn.LayerNorm(text_latent_dim)
        self.query = nn.Linear(latent_dim, latent_dim)
        self.key = nn.Linear(text_latent_dim, latent_dim)
        self.value = nn.Linear(text_latent_dim, latent_dim)
        self.dropout = nn.Dropout(dropout)
        self.proj_out = nn.Linear(latent_dim, latent_dim)

    def forward(self, x, xf=None, emb=None):
        """
        x: B, T, D      , [240, 2, 256]
        xf: B, N, L     , [1, 2, 256]
        """
        x = x.permute(1, 0, 2)
        # xf = xf.permute(1, 0, 2)
        B, T, D = x.shape
        # N = xf.shape[1]
        H = self.num_head
        # B, T, D
        query = self.query(self.norm(x))
        # B, N, D
        key = self.key(self.text_norm(x))
        query = F.softmax(query.view(B, T, H, -1), dim=-1)
        key = F.softmax(key.view(B, T, H, -1), dim=1)
        # B, N, H, HD
        value = self.value(self.text_norm(x)).view(B, T, H, -1)
        # B, H, HD, HD
        attention = torch.einsum('bnhd,bnhl->bhdl', key, value)
        y = torch.einsum('bnhd,bhdl->bnhl', query, attention).reshape(B, T, D)
        # y = x + self.proj_out(y, emb)
        return y


class WavEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        # import sys
        # [sys.path.append(i) for i in ['./WavLM', '../mydiffusion_zeggs/WavLM']]
        # from WavLM import WavLM, WavLMConfig
        # wavlm_model_path = '../../../GENEA/genea_challenge_2022/baselines/multimodal-deep-learning/My_MISA/wavlm_cache/WavLM-Large.pt'
        # # wavlm_model_path = '../../../My/process/WavLM-Base+.pt'
        # # load the pre-trained checkpoints
        # checkpoint = torch.load(wavlm_model_path, map_location=torch.device('cpu'))
        # cfg = WavLMConfig(checkpoint['cfg'])
        # self.model = WavLM(cfg)
        # # device = torch.device(device)
        # # model = model.to(device)
        # self.model.load_state_dict(checkpoint['model'])
        # self.model.eval()
        # for name, param in self.model.named_parameters():
        #     param.requires_grad = False
        self.audio_feature_map = nn.Linear(1024, 64)

    def forward(self, rep):
        # rep = self.model.extract_features(wav_input_16khz)[0]

        rep = self.audio_feature_map(rep)
        return rep


if __name__ == '__main__':
    '''
    cd ./diffusion_latent/model
    python mdm.py
    '''
    # device = torch.device('cuda:2')
    n_frames = 32
    # n_pose_dims = 135
    # n_audio_dim = 32
    #
    # '''
    # sequence_pos_encoder = PositionalEncoding(d_model=128, dropout=0.1)
    # xseq = torch.randn(n_frames + 1, 3, 128)
    # # sequence_pos_encoder.pe (5000, 1, 128)
    # pdb.set_trace()
    # xseq_ = sequence_pos_encoder(xseq)  # [seqlen+1, bs, d]
    # '''

    n_seed = 2

    model = MDM(modeltype='', njoints=16 * 7, nfeats=1, translation=True, pose_rep='rot6d', glob=True,
                glob_rot=True, cond_mode = 'cross_local_attention4_style1', action_emb = 'tensor', audio_feat='wavlm',     # text_style2
                arch='mytrans_enc', latent_dim=128, n_seed=n_seed, cond_mask_prob=0.1)       # trans_enc, trans_dec, gru

    x = torch.randn(2, 112, 1, 32)

    t = torch.tensor([12, 85])

    model_kwargs_ = {'y': {}}
    model_kwargs_['y']['mask'] = (torch.zeros([1, 1, 1, n_frames]) < 1)     # [..., n_seed:]
    model_kwargs_['y']['audio'] = torch.randn(2, 32, 1024)      # [n_seed:, ...]
    model_kwargs_['y']['style'] = torch.randn(2, 6)
    model_kwargs_['y']['mask_local'] = torch.ones(2, 32).bool()

    model_kwargs_['y']['seed'] = x[..., 0:n_seed]
    y = model(x, t, model_kwargs_['y'])
    pdb.set_trace()




    # model = LinearTemporalCrossAttention(seq_len=0, latent_dim=256, text_latent_dim=256, num_head=8, dropout=0.1, time_embed_dim=0)
    # x = torch.randn(240, 2, 256)
    # x2 = torch.randn(1, 2, 256)
    # y = model(x, x2)
    # pdb.set_trace()




    # import torch
    # from local_attention import LocalAttention
    #
    # q = torch.randn(2, 8, 2048, 64)
    # k = torch.randn(2, 8, 2048, 64)
    # v = torch.randn(2, 8, 2048, 64)
    #
    # attn = LocalAttention(
    #     dim=64,  # dimension of each head (you need to pass this in for relative positional encoding)
    #     window_size=512,  # window size. 512 is optimal, but 256 or 128 yields good enough results
    #     causal=True,  # auto-regressive or not
    #     look_backward=1,  # each window looks at the window before
    #     look_forward=0,     # for non-auto-regressive case, will default to 1, so each window looks at the window before and after it
    #     dropout=0.1,  # post-attention dropout
    #     exact_windowsize=False
    #     # if this is set to true, in the causal setting, each query will see at maximum the number of keys equal to the window size
    # )
    #
    # mask = torch.ones(2, 2048).bool()
    # out = attn(q, k, v, mask=mask)  # (2, 8, 2048, 64)
    # pdb.set_trace()



    # model = WavEncoder()
    # device = torch.device('cuda:2')
    # model = model.to(device)
    # x = torch.randn(2, 64000)
    # y = model(x)
    # pdb.set_trace()

diffusion_latent/model/mdm.py
YoungSeng/UnifiedGesture
import warnings
from typing import Tuple, Optional

import torch
from torch import Tensor
if float(torch.__version__.split('.')[0]) == 0 or (float(torch.__version__.split('.')[0]) == 1 and float(torch.__version__.split('.')[1])) < 9:
    from torch.nn.modules.linear import _LinearWithBias
else:
    from torch.nn.modules.linear import NonDynamicallyQuantizableLinear as _LinearWithBias
from torch.nn.init import xavier_uniform_
from torch.nn.init import constant_
from torch.nn.init import xavier_normal_
from torch.nn.parameter import Parameter
from torch.nn.modules.module import Module
from torch.nn import functional as F


class Threshold(Module):
    r"""Thresholds each element of the input Tensor.

    Threshold is defined as:

    .. math::
        y =
        \begin{cases}
        x, &\text{ if } x > \text{threshold} \\
        \text{value}, &\text{ otherwise }
        \end{cases}

    Args:
        threshold: The value to threshold at
        value: The value to replace with
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    Examples::

        >>> m = nn.Threshold(0.1, 20)
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['threshold', 'value', 'inplace']

    threshold: float
    value: float
    inplace: bool

    def __init__(self, threshold: float, value: float, inplace: bool = False) -> None:
        super(Threshold, self).__init__()
        self.threshold = threshold
        self.value = value
        self.inplace = inplace
        # TODO: check in THNN (if inplace == True, then assert value <= threshold)

    def forward(self, input: Tensor) -> Tensor:
        return F.threshold(input, self.threshold, self.value, self.inplace)

    def extra_repr(self):
        inplace_str = ', inplace=True' if self.inplace else ''
        return 'threshold={}, value={}{}'.format(
            self.threshold, self.value, inplace_str
        )


class ReLU(Module):
    r"""Applies the rectified linear unit function element-wise:

    :math:`\text{ReLU}(x) = (x)^+ = \max(0, x)`

    Args:
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/ReLU.png

    Examples::

        >>> m = nn.ReLU()
        >>> input = torch.randn(2)
        >>> output = m(input)


      An implementation of CReLU - https://arxiv.org/abs/1603.05201

        >>> m = nn.ReLU()
        >>> input = torch.randn(2).unsqueeze(0)
        >>> output = torch.cat((m(input),m(-input)))
    """
    __constants__ = ['inplace']
    inplace: bool

    def __init__(self, inplace: bool = False):
        super(ReLU, self).__init__()
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.relu(input, inplace=self.inplace)

    def extra_repr(self) -> str:
        inplace_str = 'inplace=True' if self.inplace else ''
        return inplace_str


class RReLU(Module):
    r"""Applies the randomized leaky rectified liner unit function, element-wise,
    as described in the paper:

    `Empirical Evaluation of Rectified Activations in Convolutional Network`_.

    The function is defined as:

    .. math::
        \text{RReLU}(x) =
        \begin{cases}
            x & \text{if } x \geq 0 \\
            ax & \text{ otherwise }
        \end{cases}

    where :math:`a` is randomly sampled from uniform distribution
    :math:`\mathcal{U}(\text{lower}, \text{upper})`.

     See: https://arxiv.org/pdf/1505.00853.pdf

    Args:
        lower: lower bound of the uniform distribution. Default: :math:`\frac{1}{8}`
        upper: upper bound of the uniform distribution. Default: :math:`\frac{1}{3}`
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    Examples::

        >>> m = nn.RReLU(0.1, 0.3)
        >>> input = torch.randn(2)
        >>> output = m(input)

    .. _`Empirical Evaluation of Rectified Activations in Convolutional Network`:
        https://arxiv.org/abs/1505.00853
    """
    __constants__ = ['lower', 'upper', 'inplace']

    lower: float
    upper: float
    inplace: bool

    def __init__(
        self,
        lower: float = 1. / 8,
        upper: float = 1. / 3,
        inplace: bool = False
    ):
        super(RReLU, self).__init__()
        self.lower = lower
        self.upper = upper
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.rrelu(input, self.lower, self.upper, self.training, self.inplace)

    def extra_repr(self):
        inplace_str = ', inplace=True' if self.inplace else ''
        return 'lower={}, upper={}{}'.format(self.lower, self.upper, inplace_str)


class Hardtanh(Module):
    r"""Applies the HardTanh function element-wise

    HardTanh is defined as:

    .. math::
        \text{HardTanh}(x) = \begin{cases}
            1 & \text{ if } x > 1 \\
            -1 & \text{ if } x < -1 \\
            x & \text{ otherwise } \\
        \end{cases}

    The range of the linear region :math:`[-1, 1]` can be adjusted using
    :attr:`min_val` and :attr:`max_val`.

    Args:
        min_val: minimum value of the linear region range. Default: -1
        max_val: maximum value of the linear region range. Default: 1
        inplace: can optionally do the operation in-place. Default: ``False``

    Keyword arguments :attr:`min_value` and :attr:`max_value`
    have been deprecated in favor of :attr:`min_val` and :attr:`max_val`.

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Hardtanh.png

    Examples::

        >>> m = nn.Hardtanh(-2, 2)
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['min_val', 'max_val', 'inplace']

    min_val: float
    max_val: float
    inplace: bool

    def __init__(
        self,
        min_val: float = -1.,
        max_val: float = 1.,
        inplace: bool = False,
        min_value: Optional[float] = None,
        max_value: Optional[float] = None
    ) -> None:
        super(Hardtanh, self).__init__()
        if min_value is not None:
            warnings.warn("keyword argument min_value is deprecated and rename to min_val")
            min_val = min_value
        if max_value is not None:
            warnings.warn("keyword argument max_value is deprecated and rename to max_val")
            max_val = max_value

        self.min_val = min_val
        self.max_val = max_val
        self.inplace = inplace
        assert self.max_val > self.min_val

    def forward(self, input: Tensor) -> Tensor:
        return F.hardtanh(input, self.min_val, self.max_val, self.inplace)

    def extra_repr(self) -> str:
        inplace_str = ', inplace=True' if self.inplace else ''
        return 'min_val={}, max_val={}{}'.format(
            self.min_val, self.max_val, inplace_str
        )


class ReLU6(Hardtanh):
    r"""Applies the element-wise function:

    .. math::
        \text{ReLU6}(x) = \min(\max(0,x), 6)

    Args:
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/ReLU6.png

    Examples::

        >>> m = nn.ReLU6()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """

    def __init__(self, inplace: bool = False):
        super(ReLU6, self).__init__(0., 6., inplace)

    def extra_repr(self) -> str:
        inplace_str = 'inplace=True' if self.inplace else ''
        return inplace_str


class Sigmoid(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{Sigmoid}(x) = \sigma(x) = \frac{1}{1 + \exp(-x)}


    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Sigmoid.png

    Examples::

        >>> m = nn.Sigmoid()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """

    def forward(self, input: Tensor) -> Tensor:
        return torch.sigmoid(input)


class Hardsigmoid(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{Hardsigmoid}(x) = \begin{cases}
            0 & \text{if~} x \le -3, \\
            1 & \text{if~} x \ge +3, \\
            x / 6 + 1 / 2 & \text{otherwise}
        \end{cases}

    Args:
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    Examples::

        >>> m = nn.Hardsigmoid()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['inplace']

    inplace: bool

    def __init__(self, inplace : bool = False) -> None:
        super(Hardsigmoid, self).__init__()
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.hardsigmoid(input, self.inplace)


class Tanh(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{Tanh}(x) = \tanh(x) = \frac{\exp(x) - \exp(-x)} {\exp(x) + \exp(-x)}

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Tanh.png

    Examples::

        >>> m = nn.Tanh()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """

    def forward(self, input: Tensor) -> Tensor:
        return torch.tanh(input)

class SiLU(Module):
    r"""Applies the silu function, element-wise.

    .. math::
        \text{silu}(x) = x * \sigma(x), \text{where } \sigma(x) \text{ is the logistic sigmoid.}

    .. note::
        See `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_
        where the SiLU (Sigmoid Linear Unit) was originally coined, and see
        `Sigmoid-Weighted Linear Units for Neural Network Function Approximation
        in Reinforcement Learning <https://arxiv.org/abs/1702.03118>`_ and `Swish:
        a Self-Gated Activation Function <https://arxiv.org/abs/1710.05941v1>`_
        where the SiLU was experimented with later.

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    Examples::

        >>> m = nn.SiLU()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['inplace']
    inplace: bool

    def __init__(self, inplace: bool = False):
        super(SiLU, self).__init__()
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.silu(input, inplace=self.inplace)

    def extra_repr(self) -> str:
        inplace_str = 'inplace=True' if self.inplace else ''
        return inplace_str

class Hardswish(Module):
    r"""Applies the hardswish function, element-wise, as described in the paper:

    `Searching for MobileNetV3`_.

    .. math::
        \text{Hardswish}(x) = \begin{cases}
            0 & \text{if~} x \le -3, \\
            x & \text{if~} x \ge +3, \\
            x \cdot (x + 3) /6 & \text{otherwise}
        \end{cases}

    Args:
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    Examples::

        >>> m = nn.Hardswish()
        >>> input = torch.randn(2)
        >>> output = m(input)

    .. _`Searching for MobileNetV3`:
        https://arxiv.org/abs/1905.02244
    """
    __constants__ = ['inplace']

    inplace: bool

    def __init__(self, inplace : bool = False) -> None:
        super(Hardswish, self).__init__()
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.hardswish(input, self.inplace)


class ELU(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{ELU}(x) = \begin{cases}
        x, & \text{ if } x > 0\\
        \alpha * (\exp(x) - 1), & \text{ if } x \leq 0
        \end{cases}

    Args:
        alpha: the :math:`\alpha` value for the ELU formulation. Default: 1.0
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/ELU.png

    Examples::

        >>> m = nn.ELU()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['alpha', 'inplace']
    alpha: float
    inplace: bool

    def __init__(self, alpha: float = 1., inplace: bool = False) -> None:
        super(ELU, self).__init__()
        self.alpha = alpha
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.elu(input, self.alpha, self.inplace)

    def extra_repr(self) -> str:
        inplace_str = ', inplace=True' if self.inplace else ''
        return 'alpha={}{}'.format(self.alpha, inplace_str)


class CELU(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x/\alpha) - 1))

    More details can be found in the paper `Continuously Differentiable Exponential Linear Units`_ .

    Args:
        alpha: the :math:`\alpha` value for the CELU formulation. Default: 1.0
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/CELU.png

    Examples::

        >>> m = nn.CELU()
        >>> input = torch.randn(2)
        >>> output = m(input)

    .. _`Continuously Differentiable Exponential Linear Units`:
        https://arxiv.org/abs/1704.07483
    """
    __constants__ = ['alpha', 'inplace']
    alpha: float
    inplace: bool

    def __init__(self, alpha: float = 1., inplace: bool = False) -> None:
        super(CELU, self).__init__()
        self.alpha = alpha
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.celu(input, self.alpha, self.inplace)

    def extra_repr(self) -> str:
        inplace_str = ', inplace=True' if self.inplace else ''
        return 'alpha={}{}'.format(self.alpha, inplace_str)


class SELU(Module):
    r"""Applied element-wise, as:

    .. math::
        \text{SELU}(x) = \text{scale} * (\max(0,x) + \min(0, \alpha * (\exp(x) - 1)))

    with :math:`\alpha = 1.6732632423543772848170429916717` and
    :math:`\text{scale} = 1.0507009873554804934193349852946`.

    More details can be found in the paper `Self-Normalizing Neural Networks`_ .

    Args:
        inplace (bool, optional): can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/SELU.png

    Examples::

        >>> m = nn.SELU()
        >>> input = torch.randn(2)
        >>> output = m(input)

    .. _Self-Normalizing Neural Networks: https://arxiv.org/abs/1706.02515
    """
    __constants__ = ['inplace']
    inplace: bool

    def __init__(self, inplace: bool = False) -> None:
        super(SELU, self).__init__()
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.selu(input, self.inplace)

    def extra_repr(self) -> str:
        inplace_str = 'inplace=True' if self.inplace else ''
        return inplace_str


class GLU(Module):
    r"""Applies the gated linear unit function
    :math:`{GLU}(a, b)= a \otimes \sigma(b)` where :math:`a` is the first half
    of the input matrices and :math:`b` is the second half.

    Args:
        dim (int): the dimension on which to split the input. Default: -1

    Shape:
        - Input: :math:`(\ast_1, N, \ast_2)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(\ast_1, M, \ast_2)` where :math:`M=N/2`

    Examples::

        >>> m = nn.GLU()
        >>> input = torch.randn(4, 2)
        >>> output = m(input)
    """
    __constants__ = ['dim']
    dim: int

    def __init__(self, dim: int = -1) -> None:
        super(GLU, self).__init__()
        self.dim = dim

    def forward(self, input: Tensor) -> Tensor:
        return F.glu(input, self.dim)

    def extra_repr(self) -> str:
        return 'dim={}'.format(self.dim)


class GELU(Module):
    r"""Applies the Gaussian Error Linear Units function:

    .. math:: \text{GELU}(x) = x * \Phi(x)

    where :math:`\Phi(x)` is the Cumulative Distribution Function for Gaussian Distribution.

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/GELU.png

    Examples::

        >>> m = nn.GELU()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    def forward(self, input: Tensor) -> Tensor:
        return F.gelu(input)


class Hardshrink(Module):
    r"""Applies the hard shrinkage function element-wise:

    .. math::
        \text{HardShrink}(x) =
        \begin{cases}
        x, & \text{ if } x > \lambda \\
        x, & \text{ if } x < -\lambda \\
        0, & \text{ otherwise }
        \end{cases}

    Args:
        lambd: the :math:`\lambda` value for the Hardshrink formulation. Default: 0.5

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Hardshrink.png

    Examples::

        >>> m = nn.Hardshrink()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['lambd']
    lambd: float

    def __init__(self, lambd: float = 0.5) -> None:
        super(Hardshrink, self).__init__()
        self.lambd = lambd

    def forward(self, input: Tensor) -> Tensor:
        return F.hardshrink(input, self.lambd)

    def extra_repr(self) -> str:
        return '{}'.format(self.lambd)


class LeakyReLU(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{LeakyReLU}(x) = \max(0, x) + \text{negative\_slope} * \min(0, x)


    or

    .. math::
        \text{LeakyRELU}(x) =
        \begin{cases}
        x, & \text{ if } x \geq 0 \\
        \text{negative\_slope} \times x, & \text{ otherwise }
        \end{cases}

    Args:
        negative_slope: Controls the angle of the negative slope. Default: 1e-2
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/LeakyReLU.png

    Examples::

        >>> m = nn.LeakyReLU(0.1)
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['inplace', 'negative_slope']
    inplace: bool
    negative_slope: float

    def __init__(self, negative_slope: float = 1e-2, inplace: bool = False) -> None:
        super(LeakyReLU, self).__init__()
        self.negative_slope = negative_slope
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.leaky_relu(input, self.negative_slope, self.inplace)

    def extra_repr(self) -> str:
        inplace_str = ', inplace=True' if self.inplace else ''
        return 'negative_slope={}{}'.format(self.negative_slope, inplace_str)


class LogSigmoid(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{LogSigmoid}(x) = \log\left(\frac{ 1 }{ 1 + \exp(-x)}\right)

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/LogSigmoid.png

    Examples::

        >>> m = nn.LogSigmoid()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """

    def forward(self, input: Tensor) -> Tensor:
        return F.logsigmoid(input)


class Softplus(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{Softplus}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x))

    SoftPlus is a smooth approximation to the ReLU function and can be used
    to constrain the output of a machine to always be positive.

    For numerical stability the implementation reverts to the linear function
    when :math:`input \times \beta > threshold`.

    Args:
        beta: the :math:`\beta` value for the Softplus formulation. Default: 1
        threshold: values above this revert to a linear function. Default: 20

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Softplus.png

    Examples::

        >>> m = nn.Softplus()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['beta', 'threshold']
    beta: int
    threshold: int

    def __init__(self, beta: int = 1, threshold: int = 20) -> None:
        super(Softplus, self).__init__()
        self.beta = beta
        self.threshold = threshold

    def forward(self, input: Tensor) -> Tensor:
        return F.softplus(input, self.beta, self.threshold)

    def extra_repr(self) -> str:
        return 'beta={}, threshold={}'.format(self.beta, self.threshold)


class Softshrink(Module):
    r"""Applies the soft shrinkage function elementwise:

    .. math::
        \text{SoftShrinkage}(x) =
        \begin{cases}
        x - \lambda, & \text{ if } x > \lambda \\
        x + \lambda, & \text{ if } x < -\lambda \\
        0, & \text{ otherwise }
        \end{cases}

    Args:
        lambd: the :math:`\lambda` (must be no less than zero) value for the Softshrink formulation. Default: 0.5

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Softshrink.png

    Examples::

        >>> m = nn.Softshrink()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['lambd']
    lambd: float

    def __init__(self, lambd: float = 0.5) -> None:
        super(Softshrink, self).__init__()
        self.lambd = lambd

    def forward(self, input: Tensor) -> Tensor:
        return F.softshrink(input, self.lambd)

    def extra_repr(self) -> str:
        return str(self.lambd)


class MultiheadAttention(Module):
    r"""Allows the model to jointly attend to information
    from different representation subspaces.
    See reference: Attention Is All You Need

    .. math::
        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
        \text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)

    Args:
        embed_dim: total dimension of the model.
        num_heads: parallel attention heads.
        dropout: a Dropout layer on attn_output_weights. Default: 0.0.
        bias: add bias as module parameter. Default: True.
        add_bias_kv: add bias to the key and value sequences at dim=0.
        add_zero_attn: add a new batch of zeros to the key and
                       value sequences at dim=1.
        kdim: total number of features in key. Default: None.
        vdim: total number of features in value. Default: None.

        Note: if kdim and vdim are None, they will be set to embed_dim such that
        query, key, and value have the same number of features.

    Examples::

        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
    """
    bias_k: Optional[torch.Tensor]
    bias_v: Optional[torch.Tensor]

    def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):
        super(MultiheadAttention, self).__init__()
        self.embed_dim = embed_dim
        self.kdim = kdim if kdim is not None else embed_dim
        self.vdim = vdim if vdim is not None else embed_dim
        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim

        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"

        if self._qkv_same_embed_dim is False:
            self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
            self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))
            self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))
            self.register_parameter('in_proj_weight', None)
        else:
            self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim))
            self.register_parameter('q_proj_weight', None)
            self.register_parameter('k_proj_weight', None)
            self.register_parameter('v_proj_weight', None)

        if bias:
            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))
        else:
            self.register_parameter('in_proj_bias', None)
        self.out_proj = _LinearWithBias(embed_dim, embed_dim)

        if add_bias_kv:
            self.bias_k = Parameter(torch.empty(1, 1, embed_dim))
            self.bias_v = Parameter(torch.empty(1, 1, embed_dim))
        else:
            self.bias_k = self.bias_v = None

        self.add_zero_attn = add_zero_attn

        self._reset_parameters()

    def _reset_parameters(self):
        if self._qkv_same_embed_dim:
            xavier_uniform_(self.in_proj_weight)
        else:
            xavier_uniform_(self.q_proj_weight)
            xavier_uniform_(self.k_proj_weight)
            xavier_uniform_(self.v_proj_weight)

        if self.in_proj_bias is not None:
            constant_(self.in_proj_bias, 0.)
            constant_(self.out_proj.bias, 0.)
        if self.bias_k is not None:
            xavier_normal_(self.bias_k)
        if self.bias_v is not None:
            xavier_normal_(self.bias_v)

    def __setstate__(self, state):
        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
        if '_qkv_same_embed_dim' not in state:
            state['_qkv_same_embed_dim'] = True

        super(MultiheadAttention, self).__setstate__(state)

    def forward(self, query, key, value, key_padding_mask=None,
                need_weights=True, attn_mask=None):
        # type: (Tensor, Tensor, Tensor, Optional[Tensor], bool, Optional[Tensor]) -> Tuple[Tensor, Optional[Tensor]]
        r"""
    Args:
        query, key, value: map a query and a set of key-value pairs to an output.
            See "Attention Is All You Need" for more details.
        key_padding_mask: if provided, specified padding elements in the key will
            be ignored by the attention. When given a binary mask and a value is True,
            the corresponding value on the attention layer will be ignored. When given
            a byte mask and a value is non-zero, the corresponding value on the attention
            layer will be ignored
        need_weights: output attn_output_weights.
        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
            the batches while a 3D mask allows to specify a different mask for the entries of each batch.

    Shape:
        - Inputs:
        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
          the embedding dimension.
        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
          the embedding dimension.
        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
          the embedding dimension.
        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
          If a ByteTensor is provided, the non-zero positions will be ignored while the position
          with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
          S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
          is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
          is provided, it will be added to the attention weight.

        - Outputs:
        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
          E is the embedding dimension.
        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
          L is the target sequence length, S is the source sequence length.
        """
        if not self._qkv_same_embed_dim:
            return F.multi_head_attention_forward(
                query, key, value, self.embed_dim, self.num_heads,
                self.in_proj_weight, self.in_proj_bias,
                self.bias_k, self.bias_v, self.add_zero_attn,
                self.dropout, self.out_proj.weight, self.out_proj.bias,
                training=self.training,
                key_padding_mask=key_padding_mask, need_weights=need_weights,
                attn_mask=attn_mask, use_separate_proj_weight=True,
                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
                v_proj_weight=self.v_proj_weight)
        else:
            return F.multi_head_attention_forward(
                query, key, value, self.embed_dim, self.num_heads,
                self.in_proj_weight, self.in_proj_bias,
                self.bias_k, self.bias_v, self.add_zero_attn,
                self.dropout, self.out_proj.weight, self.out_proj.bias,
                training=self.training,
                key_padding_mask=key_padding_mask, need_weights=need_weights,
                attn_mask=attn_mask)


class PReLU(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{PReLU}(x) = \max(0,x) + a * \min(0,x)

    or

    .. math::
        \text{PReLU}(x) =
        \begin{cases}
        x, & \text{ if } x \geq 0 \\
        ax, & \text{ otherwise }
        \end{cases}

    Here :math:`a` is a learnable parameter. When called without arguments, `nn.PReLU()` uses a single
    parameter :math:`a` across all input channels. If called with `nn.PReLU(nChannels)`,
    a separate :math:`a` is used for each input channel.


    .. note::
        weight decay should not be used when learning :math:`a` for good performance.

    .. note::
        Channel dim is the 2nd dim of input. When input has dims < 2, then there is
        no channel dim and the number of channels = 1.

    Args:
        num_parameters (int): number of :math:`a` to learn.
            Although it takes an int as input, there is only two values are legitimate:
            1, or the number of channels at input. Default: 1
        init (float): the initial value of :math:`a`. Default: 0.25

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    Attributes:
        weight (Tensor): the learnable weights of shape (:attr:`num_parameters`).

    .. image:: ../scripts/activation_images/PReLU.png

    Examples::

        >>> m = nn.PReLU()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['num_parameters']
    num_parameters: int

    def __init__(self, num_parameters: int = 1, init: float = 0.25) -> None:
        self.num_parameters = num_parameters
        super(PReLU, self).__init__()
        self.weight = Parameter(torch.Tensor(num_parameters).fill_(init))

    def forward(self, input: Tensor) -> Tensor:
        return F.prelu(input, self.weight)

    def extra_repr(self) -> str:
        return 'num_parameters={}'.format(self.num_parameters)


class Softsign(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{SoftSign}(x) = \frac{x}{ 1 + |x|}

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Softsign.png

    Examples::

        >>> m = nn.Softsign()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """

    def forward(self, input: Tensor) -> Tensor:
        return F.softsign(input)


class Tanhshrink(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{Tanhshrink}(x) = x - \tanh(x)

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Tanhshrink.png

    Examples::

        >>> m = nn.Tanhshrink()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """

    def forward(self, input: Tensor) -> Tensor:
        return F.tanhshrink(input)


class Softmin(Module):
    r"""Applies the Softmin function to an n-dimensional input Tensor
    rescaling them so that the elements of the n-dimensional output Tensor
    lie in the range `[0, 1]` and sum to 1.

    Softmin is defined as:

    .. math::
        \text{Softmin}(x_{i}) = \frac{\exp(-x_i)}{\sum_j \exp(-x_j)}

    Shape:
        - Input: :math:`(*)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(*)`, same shape as the input

    Arguments:
        dim (int): A dimension along which Softmin will be computed (so every slice
            along dim will sum to 1).

    Returns:
        a Tensor of the same dimension and shape as the input, with
        values in the range [0, 1]

    Examples::

        >>> m = nn.Softmin()
        >>> input = torch.randn(2, 3)
        >>> output = m(input)
    """
    __constants__ = ['dim']
    dim: Optional[int]

    def __init__(self, dim: Optional[int] = None) -> None:
        super(Softmin, self).__init__()
        self.dim = dim

    def __setstate__(self, state):
        self.__dict__.update(state)
        if not hasattr(self, 'dim'):
            self.dim = None

    def forward(self, input: Tensor) -> Tensor:
        return F.softmin(input, self.dim, _stacklevel=5)

    def extra_repr(self):
        return 'dim={dim}'.format(dim=self.dim)

class Softmax(Module):
    r"""Applies the Softmax function to an n-dimensional input Tensor
    rescaling them so that the elements of the n-dimensional output Tensor
    lie in the range [0,1] and sum to 1.

    Softmax is defined as:

    .. math::
        \text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}

    When the input Tensor is a sparse tensor then the unspecifed
    values are treated as ``-inf``.

    Shape:
        - Input: :math:`(*)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(*)`, same shape as the input

    Returns:
        a Tensor of the same dimension and shape as the input with
        values in the range [0, 1]

    Arguments:
        dim (int): A dimension along which Softmax will be computed (so every slice
            along dim will sum to 1).

    .. note::
        This module doesn't work directly with NLLLoss,
        which expects the Log to be computed between the Softmax and itself.
        Use `LogSoftmax` instead (it's faster and has better numerical properties).

    Examples::

        >>> m = nn.Softmax(dim=1)
        >>> input = torch.randn(2, 3)
        >>> output = m(input)

    """
    __constants__ = ['dim']
    dim: Optional[int]

    def __init__(self, dim: Optional[int] = None) -> None:
        super(Softmax, self).__init__()
        self.dim = dim

    def __setstate__(self, state):
        self.__dict__.update(state)
        if not hasattr(self, 'dim'):
            self.dim = None

    def forward(self, input: Tensor) -> Tensor:
        return F.softmax(input, self.dim, _stacklevel=5)

    def extra_repr(self) -> str:
        return 'dim={dim}'.format(dim=self.dim)


class Softmax2d(Module):
    r"""Applies SoftMax over features to each spatial location.

    When given an image of ``Channels x Height x Width``, it will
    apply `Softmax` to each location :math:`(Channels, h_i, w_j)`

    Shape:
        - Input: :math:`(N, C, H, W)`
        - Output: :math:`(N, C, H, W)` (same shape as input)

    Returns:
        a Tensor of the same dimension and shape as the input with
        values in the range [0, 1]

    Examples::

        >>> m = nn.Softmax2d()
        >>> # you softmax over the 2nd dimension
        >>> input = torch.randn(2, 3, 12, 13)
        >>> output = m(input)
    """

    def forward(self, input: Tensor) -> Tensor:
        assert input.dim() == 4, 'Softmax2d requires a 4D tensor as input'
        return F.softmax(input, 1, _stacklevel=5)


class LogSoftmax(Module):
    r"""Applies the :math:`\log(\text{Softmax}(x))` function to an n-dimensional
    input Tensor. The LogSoftmax formulation can be simplified as:

    .. math::
        \text{LogSoftmax}(x_{i}) = \log\left(\frac{\exp(x_i) }{ \sum_j \exp(x_j)} \right)

    Shape:
        - Input: :math:`(*)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(*)`, same shape as the input

    Arguments:
        dim (int): A dimension along which LogSoftmax will be computed.

    Returns:
        a Tensor of the same dimension and shape as the input with
        values in the range [-inf, 0)

    Examples::

        >>> m = nn.LogSoftmax()
        >>> input = torch.randn(2, 3)
        >>> output = m(input)
    """
    __constants__ = ['dim']
    dim: Optional[int]

    def __init__(self, dim: Optional[int] = None) -> None:
        super(LogSoftmax, self).__init__()
        self.dim = dim

    def __setstate__(self, state):
        self.__dict__.update(state)
        if not hasattr(self, 'dim'):
            self.dim = None

    def forward(self, input: Tensor) -> Tensor:
        return F.log_softmax(input, self.dim, _stacklevel=5)

    def extra_repr(self):
        return 'dim={dim}'.format(dim=self.dim)

diffusion_latent/model/myactivation.py
YoungSeng/UnifiedGesture
import copy
import pdb
from typing import Optional, Any

import torch
import torch.nn as nn
from torch import Tensor
from torch.nn import functional as F
from torch.nn.modules.module import Module
from myactivation import MultiheadAttention
from torch.nn.modules.container import ModuleList
from torch.nn.init import xavier_uniform_
from torch.nn.modules.dropout import Dropout
from torch.nn.modules.linear import Linear
from torch.nn.modules.normalization import LayerNorm


class Transformer(Module):
    r"""A transformer model. User is able to modify the attributes as needed. The architecture
    is based on the paper "Attention Is All You Need". Ashish Vaswani, Noam Shazeer,
    Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and
    Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information
    Processing Systems, pages 6000-6010. Users can build the BERT(https://arxiv.org/abs/1810.04805)
    model with corresponding parameters.

    Args:
        d_model: the number of expected features in the encoder/decoder inputs (default=512).
        nhead: the number of heads in the multiheadattention models (default=8).
        num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).
        num_decoder_layers: the number of sub-decoder-layers in the decoder (default=6).
        dim_feedforward: the dimension of the feedforward network model (default=2048).
        dropout: the dropout value (default=0.1).
        activation: the activation function of encoder/decoder intermediate layer, relu or gelu (default=relu).
        custom_encoder: custom encoder (default=None).
        custom_decoder: custom decoder (default=None).

    Examples::
        >>> transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
        >>> src = torch.rand((10, 32, 512))
        >>> tgt = torch.rand((20, 32, 512))
        >>> out = transformer_model(src, tgt)

    Note: A full example to apply nn.Transformer module for the word language model is available in
    https://github.com/pytorch/examples/tree/master/word_language_model
    """

    def __init__(self, d_model: int = 512, nhead: int = 8, num_encoder_layers: int = 6,
                 num_decoder_layers: int = 6, dim_feedforward: int = 2048, dropout: float = 0.1,
                 activation: str = "relu", custom_encoder: Optional[Any] = None, custom_decoder: Optional[Any] = None) -> None:
        super(Transformer, self).__init__()

        if custom_encoder is not None:
            self.encoder = custom_encoder
        else:
            encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
            encoder_norm = LayerNorm(d_model)
            self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)

        if custom_decoder is not None:
            self.decoder = custom_decoder
        else:
            decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
            decoder_norm = LayerNorm(d_model)
            self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)

        self._reset_parameters()

        self.d_model = d_model
        self.nhead = nhead

    def forward(self, src: Tensor, tgt: Tensor, src_mask: Optional[Tensor] = None, tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""Take in and process masked source/target sequences.

        Args:
            src: the sequence to the encoder (required).
            tgt: the sequence to the decoder (required).
            src_mask: the additive mask for the src sequence (optional).
            tgt_mask: the additive mask for the tgt sequence (optional).
            memory_mask: the additive mask for the encoder output (optional).
            src_key_padding_mask: the ByteTensor mask for src keys per batch (optional).
            tgt_key_padding_mask: the ByteTensor mask for tgt keys per batch (optional).
            memory_key_padding_mask: the ByteTensor mask for memory keys per batch (optional).

        Shape:
            - src: :math:`(S, N, E)`.
            - tgt: :math:`(T, N, E)`.
            - src_mask: :math:`(S, S)`.
            - tgt_mask: :math:`(T, T)`.
            - memory_mask: :math:`(T, S)`.
            - src_key_padding_mask: :math:`(N, S)`.
            - tgt_key_padding_mask: :math:`(N, T)`.
            - memory_key_padding_mask: :math:`(N, S)`.

            Note: [src/tgt/memory]_mask ensures that position i is allowed to attend the unmasked
            positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
            while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
            are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
            is provided, it will be added to the attention weight.
            [src/tgt/memory]_key_padding_mask provides specified elements in the key to be ignored by
            the attention. If a ByteTensor is provided, the non-zero positions will be ignored while the zero
            positions will be unchanged. If a BoolTensor is provided, the positions with the
            value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.

            - output: :math:`(T, N, E)`.

            Note: Due to the multi-head attention architecture in the transformer model,
            the output sequence length of a transformer is same as the input sequence
            (i.e. target) length of the decode.

            where S is the source sequence length, T is the target sequence length, N is the
            batch size, E is the feature number

        Examples:
            >>> output = transformer_model(src, tgt, src_mask=src_mask, tgt_mask=tgt_mask)
        """

        if src.size(1) != tgt.size(1):
            raise RuntimeError("the batch number of src and tgt must be equal")

        if src.size(2) != self.d_model or tgt.size(2) != self.d_model:
            raise RuntimeError("the feature number of src and tgt must be equal to d_model")

        memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
        output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
                              tgt_key_padding_mask=tgt_key_padding_mask,
                              memory_key_padding_mask=memory_key_padding_mask)
        return output

    def generate_square_subsequent_mask(self, sz: int) -> Tensor:
        r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
            Unmasked positions are filled with float(0.0).
        """
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask

    def _reset_parameters(self):
        r"""Initiate parameters in the transformer model."""

        for p in self.parameters():
            if p.dim() > 1:
                xavier_uniform_(p)


class TransformerEncoder(Module):
    r"""TransformerEncoder is a stack of N encoder layers

    Args:
        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
        num_layers: the number of sub-encoder-layers in the encoder (required).
        norm: the layer normalization component (optional).

    Examples::
        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
        >>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
        >>> src = torch.rand(10, 32, 512)
        >>> out = transformer_encoder(src)
    """
    __constants__ = ['norm']

    def __init__(self, encoder_layer, num_layers, norm=None):
        super(TransformerEncoder, self).__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""Pass the input through the encoder layers in turn.

        Args:
            src: the sequence to the encoder (required).
            mask: the mask for the src sequence (optional).
            src_key_padding_mask: the mask for the src keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """
        output = src

        for mod in self.layers:
            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)

        if self.norm is not None:
            output = self.norm(output)

        return output


class TransformerDecoder(Module):
    r"""TransformerDecoder is a stack of N decoder layers

    Args:
        decoder_layer: an instance of the TransformerDecoderLayer() class (required).
        num_layers: the number of sub-decoder-layers in the decoder (required).
        norm: the layer normalization component (optional).

    Examples::
        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
        >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
        >>> memory = torch.rand(10, 32, 512)
        >>> tgt = torch.rand(20, 32, 512)
        >>> out = transformer_decoder(tgt, memory)
    """
    __constants__ = ['norm']

    def __init__(self, decoder_layer, num_layers, norm=None):
        super(TransformerDecoder, self).__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""Pass the inputs (and mask) through the decoder layer in turn.

        Args:
            tgt: the sequence to the decoder (required).
            memory: the sequence from the last layer of the encoder (required).
            tgt_mask: the mask for the tgt sequence (optional).
            memory_mask: the mask for the memory sequence (optional).
            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
            memory_key_padding_mask: the mask for the memory keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """
        output = tgt

        for mod in self.layers:
            output = mod(output, memory, tgt_mask=tgt_mask,
                         memory_mask=memory_mask,
                         tgt_key_padding_mask=tgt_key_padding_mask,
                         memory_key_padding_mask=memory_key_padding_mask)

        if self.norm is not None:
            output = self.norm(output)

        return output

class TransformerEncoderLayer(Module):
    r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
    This standard encoder layer is based on the paper "Attention Is All You Need".
    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
    in a different way during application.

    Args:
        d_model: the number of expected features in the input (required).
        nhead: the number of heads in the multiheadattention models (required).
        dim_feedforward: the dimension of the feedforward network model (default=2048).
        dropout: the dropout value (default=0.1).
        activation: the activation function of intermediate layer, relu or gelu (default=relu).

    Examples::
        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
        >>> src = torch.rand(10, 32, 512)
        >>> out = encoder_layer(src)
    """

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu"):
        super(TransformerEncoderLayer, self).__init__()
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = Linear(d_model, dim_feedforward)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model)

        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        self.activation = _get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.relu
        super(TransformerEncoderLayer, self).__setstate__(state)

    def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""Pass the input through the encoder layer.

        Args:
            src: the sequence to the encoder layer (required).
            src_mask: the mask for the src sequence (optional).
            src_key_padding_mask: the mask for the src keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """
        # src: [81, 2, 128]
        src2 = self.self_attn(src, src, src, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src


class TransformerDecoderLayer(Module):
    r"""TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.
    This standard decoder layer is based on the paper "Attention Is All You Need".
    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
    in a different way during application.

    Args:
        d_model: the number of expected features in the input (required).
        nhead: the number of heads in the multiheadattention models (required).
        dim_feedforward: the dimension of the feedforward network model (default=2048).
        dropout: the dropout value (default=0.1).
        activation: the activation function of intermediate layer, relu or gelu (default=relu).

    Examples::
        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
        >>> memory = torch.rand(10, 32, 512)
        >>> tgt = torch.rand(20, 32, 512)
        >>> out = decoder_layer(tgt, memory)
    """

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu"):
        super(TransformerDecoderLayer, self).__init__()
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = Linear(d_model, dim_feedforward)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model)

        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.norm3 = LayerNorm(d_model)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)
        self.dropout3 = Dropout(dropout)

        self.activation = _get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.relu
        super(TransformerDecoderLayer, self).__setstate__(state)

    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""Pass the inputs (and mask) through the decoder layer.

        Args:
            tgt: the sequence to the decoder layer (required).
            memory: the sequence from the last layer of the encoder (required).
            tgt_mask: the mask for the tgt sequence (optional).
            memory_mask: the mask for the memory sequence (optional).
            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
            memory_key_padding_mask: the mask for the memory keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """
        tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt


def _get_clones(module, N):
    return ModuleList([copy.deepcopy(module) for i in range(N)])


def _get_activation_fn(activation):
    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu

    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))

diffusion_latent/model/mytransformer.py
YoungSeng/UnifiedGesture
import torch
from torch import nn, einsum

from einops import rearrange

class SinusoidalEmbeddings(nn.Module):
    def __init__(self, dim):
        super().__init__()
        inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, x):
        n = x.shape[-2]
        t = torch.arange(n, device = x.device).type_as(self.inv_freq)
        freqs = torch.einsum('i , j -> i j', t, self.inv_freq)
        return torch.cat((freqs, freqs), dim=-1)

def rotate_half(x):
    x = rearrange(x, 'b ... (r d) -> b (...) r d', r = 2)
    x1, x2 = x.unbind(dim = -2)
    return torch.cat((-x2, x1), dim = -1)

def apply_rotary_pos_emb(q, k, freqs):
    q, k = map(lambda t: (t * freqs.cos()) + (rotate_half(t) * freqs.sin()), (q, k))
    return q, k

diffusion_latent/model/rotary.py
YoungSeng/UnifiedGesture
import pdb

import torch

from torch import nn


class Tisa(nn.Module):
    def __init__(self, num_attention_heads: int = 12, num_kernels: int = 5):
        super().__init__()
        self.num_attention_heads = num_attention_heads
        self.num_kernels = num_kernels

        self.kernel_offsets = nn.Parameter(
            torch.Tensor(self.num_kernels, self.num_attention_heads)
        )
        self.kernel_amplitudes = nn.Parameter(
            torch.Tensor(self.num_kernels, self.num_attention_heads)
        )
        self.kernel_sharpness = nn.Parameter(
            torch.Tensor(self.num_kernels, self.num_attention_heads)
        )
        self._init_weights()

    def create_relative_offsets(self, seq_len: int):
        """Creates offsets for all the relative distances between
        -seq_len + 1 to seq_len - 1."""
        return torch.arange(-seq_len, seq_len + 1)

    def compute_positional_scores(self, relative_offsets):
        """Takes seq_len and outputs position scores for each relative distance.
        This implementation uses radial basis functions. Override this function to
        use other scoring functions than the example in the paper."""
        rbf_scores = (
            self.kernel_amplitudes.unsqueeze(-1)
            * torch.exp(
                -torch.abs(self.kernel_sharpness.unsqueeze(-1))
                * ((self.kernel_offsets.unsqueeze(-1) - relative_offsets) ** 2)
            )
        ).sum(axis=0)
        return rbf_scores

    def scores_to_toeplitz_matrix(self, positional_scores, seq_len: int):
        """Converts the TISA positional scores into the final matrix for the
        self-attention equation. PRs with memory and/or speed optimizations are
        welcome."""
        deformed_toeplitz = (
            (
                (torch.arange(0, -(seq_len ** 2), step=-1) + (seq_len - 1)).view(
                    seq_len, seq_len
                )
                + (seq_len + 1) * torch.arange(seq_len).view(-1, 1)
            )
            .view(-1)
            .long()
            .to(device=positional_scores.device)
        )
        expanded_positional_scores = torch.take_along_dim(
            positional_scores, deformed_toeplitz.view(1, -1), 1
        ).view(self.num_attention_heads, seq_len, seq_len)
        return expanded_positional_scores

    def forward(self, seq_len: int):
        """Computes the translation-invariant positional contribution to the
        attention matrix in the self-attention module of transformer models."""
        if not self.num_kernels:
            return torch.zeros((self.num_attention_heads, seq_len, seq_len))
        positional_scores_vector = self.compute_positional_scores(
            self.create_relative_offsets(seq_len)
        )
        positional_scores_matrix = self.scores_to_toeplitz_matrix(
            positional_scores_vector, seq_len
        )
        return positional_scores_matrix

    def visualize(self, seq_len: int = 10, attention_heads=None):
        """Visualizes the TISA interpretability by plotting position scores as
        a function of relative distance for each attention head."""
        if attention_heads is None:
            attention_heads = list(range(self.num_attention_heads))
        import matplotlib.pyplot as plt

        x = self.create_relative_offsets(seq_len).detach().numpy()
        y = (
            self.compute_positional_scores(self.create_relative_offsets(seq_len))
            .detach()
            .numpy()
        )
        for i in attention_heads:
            plt.plot(x, y[i])
        plt.savefig('./pic-tisa.png')
        plt.show()

    def _init_weights(self):
        """Initialize the weights"""
        ampl_init_mean = 0.1
        sharpness_init_mean = 0.1
        torch.nn.init.normal_(self.kernel_offsets, mean=0.0, std=5.0)
        torch.nn.init.normal_(
            self.kernel_amplitudes, mean=ampl_init_mean, std=0.1 * ampl_init_mean
        )
        torch.nn.init.normal_(
            self.kernel_sharpness,
            mean=sharpness_init_mean,
            std=0.1 * sharpness_init_mean,
        )


def main():
    tisa = Tisa()
    positional_scores = tisa(20)
    pdb.set_trace()
    tisa.visualize(seq_len=20)


if __name__ == "__main__":
    main()


diffusion_latent/model/tisa.py
YoungSeng/UnifiedGesture
import pdb
import sys
from utils_.model_util import create_gaussian_diffusion, load_model_wo_clip
[sys.path.append(i) for i in ['.', '..', '../model', '../codebook']]
from model.mdm import MDM
import os
from datetime import datetime
import librosa
import numpy as np
import yaml
from pprint import pprint
from configs.parse_args import parse_args
import torch
import torch.nn.functional as F
from easydict import EasyDict
import math
import glob


args = parse_args()
mydevice = torch.device('cuda:' + args.gpu)
torch.cuda.set_device(int(args.gpu))

batch_size = 1


style2onehot = {
'Happy':[1, 0, 0, 0, 0, 0, 0],
'Sad':[0, 1, 0, 0, 0, 0, 0],
'Neutral':[0, 0, 1, 0, 0, 0, 0],
'Old':[0, 0, 0, 1, 0, 0, 0],
'Angry':[0, 0, 0, 0, 1, 0, 0],
'Relaxed':[0, 0, 0, 0, 0, 1, 0],
'Still':[0, 0, 0, 0, 0, 0, 1]
}

feature_dim = 16 * 7 * 3 + 27        #  + 27
seed = 6        # 4
n_poses = 36
dim = 256           # 256
MAX_LEN = 0       # 0 for generate all frames, 1 * 32 used debug


def vqvae_init(device=torch.device('cuda:2')):
    import torch.nn as nn
    from codebook.models.vqvae import VQVAE_ulr
    model = VQVAE_ulr(config.VQVAE, 7 * 16)  # n_joints * n_chanels
    model_path = "../codebook/result/train_codebook_upper_lower_root_downsample2/codebook_checkpoint_best.bin"
    # model_path = "/ceph/home/wangzl21/Projects/My_3/deep-motion-editing/diffusion_latent/experiments/reinforce_diffusion_seed0/ckpt/vqvae_epoch_1.bin"
    model = nn.DataParallel(model, device_ids=[eval(i) for i in config.no_cuda])
    model = model.to(device)
    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
    model.load_state_dict(checkpoint['model_dict'])
    # model.load_state_dict(checkpoint['model'])       # zerlin 20230428
    model = model.eval()
    return model


def latent2code(model, source_upper, save_path="../codebook/result/inference/Trinity", prefix = 'train_codebook_bvh2upper_lower_root_again', w_grad=False, code=None):
    from VisualizeCodebook import main_2
    # print(source_upper.shape)       # (len, 64)
    recon, code, out_distance = main_2(source_upper, model, save_path, prefix, w_grad, code=code)
    return recon, code, out_distance


def wavlm_init(device=torch.device('cuda:2')):
    import sys
    [sys.path.append(i) for i in ['./WavLM']]
    from WavLM import WavLM, WavLMConfig
    wavlm_model_path = './wavlm_cache/WavLM-Large.pt'
    # load the pre-trained checkpoints
    checkpoint = torch.load(wavlm_model_path, map_location=torch.device('cpu'))
    cfg = WavLMConfig(checkpoint['cfg'])
    model = WavLM(cfg)
    model = model.to(device)
    model.load_state_dict(checkpoint['model'])
    # model.eval()
    return model


def wav2wavlm(model, wav_input_16khz, device=torch.device('cuda:2')):
    with torch.no_grad():
        wav_input_16khz = wav_input_16khz.to(device)
        rep = model.extract_features(wav_input_16khz)[0]
        rep = F.interpolate(rep.transpose(1, 2), size=n_poses, align_corners=True, mode='linear').transpose(1, 2)
    return rep


def create_model_and_diffusion():
    model = MDM(modeltype='', njoints=feature_dim, nfeats=1, translation=True, pose_rep='rot6d', glob=True,
                glob_rot=True, cond_mode='cross_local_attention3_style1', action_emb='tensor', audio_feat='wavlm',
                arch='trans_enc', latent_dim=dim, n_seed=seed, cond_mask_prob=0.1)
    diffusion = create_gaussian_diffusion()
    return model, diffusion


def inference(args, save_dir, save_name, wavlm_model, audio, sample_fn, model, n_frames=0, smoothing=False, SG_filter=False, minibatch=False, skip_timesteps=False, n_seed=seed, style=None, seed=123456, time_step=False, time=None, return_seed=False, w_grad=False, save_result=False):

    torch.manual_seed(seed)

    if n_frames == 0:
        n_frames = int(audio.shape[0] * 7.5 // 16000)
    if minibatch:
        stride_poses = args.n_poses - n_seed        # 32
        if n_frames < stride_poses:
            num_subdivision = 1
        else:
            num_subdivision = math.floor(n_frames / stride_poses)
            n_frames = num_subdivision * stride_poses       # 8
            print(
                '{}, {}, {}'.format(num_subdivision, stride_poses, n_frames))

    audio = audio[:num_subdivision * math.floor(stride_poses * 16000 / 7.5)]             # (1, 76800)

    model_kwargs_ = {'y': {}}
    model_kwargs_['y']['mask'] = (torch.zeros([1, 1, 1, n_frames]) < 1).to(mydevice)
    model_kwargs_['y']['style'] = torch.as_tensor([style]).float().to(mydevice)
    model_kwargs_['y']['mask_local'] = torch.ones(1, args.n_poses).bool().to(mydevice)

    if minibatch:
        audio_reshape = torch.from_numpy(audio).to(torch.float32).reshape(num_subdivision, math.floor(stride_poses * 16000 / 7.5)).to(mydevice).transpose(0, 1)     #        # mfcc[:, :-2]
        shape_ = (1, model.njoints, model.nfeats, args.n_poses)
        out_list = []
        wavlm_result = []
        time_step_list = []
        for i in range(0, num_subdivision):
            print(i, num_subdivision)
            model_kwargs_['y']['audio'] = audio_reshape[:, i:i + 1]
            if time_step:
                if type(time) is np.ndarray:
                    t = time[i]
                else:
                    t = np.random.randint(1, 1000 + 1)
                skip_t = t - 1
                time_step_list.append(t)
            else:
                t = 1000
                skip_t = 0
            if i == 0:
                if n_seed != 0:
                    pad_zeros = torch.zeros([int(n_seed * 16000 / 7.5), 1]).to(mydevice)        # wavlm dims are 1024
                    model_kwargs_['y']['audio'] = torch.cat((pad_zeros, model_kwargs_['y']['audio']), 0)
                    model_kwargs_['y']['seed'] = torch.zeros([1, feature_dim, 1, n_seed]).to(mydevice)
            else:
                if n_seed != 0:
                    pad_audio = audio_reshape[-int(n_seed * 16000 / 7.5):, i - 1:i]
                    model_kwargs_['y']['audio'] = torch.cat((pad_audio, model_kwargs_['y']['audio']), 0)
                    # out_list[-1].shape, [1, 224, 1, 32]
                    # if feature_dim == 1:
                    #     model_kwargs_['y']['seed'] = out_list[-1][..., -n_seed:].to(mydevice)
                    # elif feature_dim >= 2:
                    #     model_kwargs_['y']['seed'] = torch.cat((out_list[-1][:, :16 * 7][..., -n_seed:].to(mydevice), torch.zeros([1, 16 * 7 * (feature_dim - 1), 1, n_seed]).to(mydevice)), dim=1)
                    model_kwargs_['y']['seed'] = out_list[-1][..., -n_seed:].to(mydevice)

            wavlm_feature = wav2wavlm(wavlm_model, model_kwargs_['y']['audio'].transpose(0, 1), device=mydevice)
            if w_grad:
                wavlm_feature.requires_grad = False
            model_kwargs_['y']['audio'] = wavlm_feature

            wavlm_feature_downsample = F.interpolate(wavlm_feature.transpose(1, 2), size=n_poses//2, align_corners=True, mode='linear').transpose(1, 2)
            wavlm_result.append(wavlm_feature_downsample.transpose(0, 1).detach().data.cpu().numpy())
            sample = sample_fn(
                model,
                shape_,
                clip_denoised=False,
                model_kwargs=model_kwargs_,
                skip_timesteps=skip_t,  # 0 is the default value - i.e. don't skip any step
                init_image=None,
                progress=True,
                dump_steps=None,
                noise=None,  # None, torch.randn(*shape_, device=mydevice)
                const_noise=False,
                cond_fn_with_grad=w_grad,
                num_timesteps=t,
            )
            # smoothing motion transition
            if len(out_list) > 0 and n_seed != 0:
                last_poses = out_list[-1][..., -n_seed:]        # [1, 336, 1, 6]
                if not return_seed:
                    out_list[-1] = out_list[-1][..., :-n_seed]  # delete last 4 frames
                    wavlm_result[-1] = wavlm_result[-1][:-n_seed//2]
                if smoothing:
                    # Extract predictions
                    last_poses_root_pos = last_poses[:, :16 * 7]        # (1, 3, 1, 8)

                    next_poses_root_pos = sample[:, :16 * 7]        # (1, 3, 1, 88)

                    root_pos = last_poses_root_pos[..., 1]      # (1, 3, 1)
                    predict_pos = next_poses_root_pos[..., 0]
                    delta_pos = (predict_pos - root_pos).unsqueeze(-1)      # # (1, 3, 1, 1)
                    sample[:, :16 * 7] = sample[:, :16 * 7] - delta_pos

                for j in range(len(last_poses)):
                    n = len(last_poses)
                    prev = last_poses[..., j]
                    next = sample[..., j]
                    sample[..., j] = prev * (n - j) / (n + 1) + next * (j + 1) / (n + 1)
            out_list.append(sample)     # [1, 112, 1, 32]


        if not return_seed:
            out_list[-1] = out_list[-1][..., :-n_seed]
        if not w_grad:
            out_list = [i.detach().data.cpu().numpy() for i in out_list]
            out_dir_vec = np.vstack(out_list)
        else:
            out_dir_vec = torch.cat(out_list, dim=0)
        if not return_seed:
            if not w_grad:
                sampled_seq = out_dir_vec.squeeze(2).transpose(0, 2, 1).reshape(batch_size, n_frames, model.njoints)
            else:
                sampled_seq = out_dir_vec.squeeze(2).permute(0, 2, 1).reshape(batch_size, n_frames, model.njoints)
        else:
            if not w_grad:
                sampled_seq = out_dir_vec.squeeze(2).transpose(0, 2, 1).reshape(1, num_subdivision * args.n_poses, model.njoints)     # (batch, 36, 336)
            else:
                sampled_seq = out_dir_vec.squeeze(2).permute(0, 2, 1).reshape(1, num_subdivision * args.n_poses, model.njoints)

    if not w_grad:
        out_poses = sampled_seq[..., :16 * 7].transpose(0, 2, 1)  # (1, 238, 112) -> (1, 112, 238)
    else:
        out_poses = sampled_seq[..., :16 * 7].permute(0, 2, 1)
    print(out_poses.shape)
    # prefix = str(datetime.now().strftime('%Y%m%d_%H%M%S'))
    prefix = save_name
    if smoothing: prefix += '_smoothing'
    if SG_filter: prefix += '_SG'
    if minibatch: prefix += '_minibatch'
    prefix += '_%s' % (n_frames)
    prefix += '_' + str(style)
    prefix += '_' + str(seed)

    output_path = save_dir + '/' + prefix
    if save_result:
        np.save(output_path + '.npy', out_poses)

    if time_step:
        return out_poses, output_path, np.vstack(wavlm_result), np.vstack(time_step_list)        # (1, 112, len)
    else:
        return out_poses, output_path, np.vstack(wavlm_result), None        # (1, 112, len)

def main(save_dir, model_path):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    # sample
    print("Creating model and diffusion...")
    model, diffusion = create_model_and_diffusion()
    print(f"Loading checkpoints from [{model_path}]...")
    state_dict = torch.load(model_path, map_location='cpu')
    load_model_wo_clip(model, state_dict)
    model.to(mydevice)
    # model.eval()

    sample_fn = diffusion.p_sample_loop     # predict x_start

    # style = mfcc_path.split('/')[-1].split('_')[1]
    wavlm_model = wavlm_init(device=mydevice)
    vqvae_model = vqvae_init(mydevice)

    return sample_fn, wavlm_model, vqvae_model, model


def inference_2(args, sample_fn, wavlm_model, vqvae_model, model, audionpy_path=None, audio_path=None, audio_wav=None, time_step=False, time=None, return_seed=False, w_grad=False, save_result=False):
    if type(audio_wav) is np.ndarray and audio_path != None:
        save_name = os.path.split(audio_path)[1][:-4]
        audio_feat = audio_wav
    elif audio_path != None:
        audio_feat, fs = librosa.load(audio_path, sr=16000)
        save_name = os.path.split(audio_path)[1][:-4]
    elif audionpy_path != None:
        audio_feat = np.load(audionpy_path)[0]
        save_name = os.path.split(audionpy_path)[1][:-4]

    if '_' in save_name:
        style = save_name.split('_')[1]
        print(save_name, style)
        if style in style2onehot:
            style = np.array(style2onehot[style])     # 20230430        #  * 1.0
    else:
        style = style2onehot['Neutral']

    x, output_path, wavlm_feat, t_list = inference(args, save_dir, save_name, wavlm_model, audio_feat, sample_fn, model,
                                                   n_frames=MAX_LEN, smoothing=False, SG_filter=False, minibatch=True,
                                                   skip_timesteps=time_step, style=[0, 0, 0, 0, 0, 3, 0], seed=123456, time_step=time_step,
                                                   time=time,
                                                   return_seed=return_seed, w_grad=w_grad, save_result=save_result)      # style2onehot['Happy']

    # x = np.load("./result/Trinity/Recording_006_minibatch_96_[0, 0, 0, 1, 0, 0, 0]_123456.npy")  # (1, 112, len)
    x_ = x[0].transpose(1, 0)[..., 16 * 2:-16 * 1]  # (len, 112)
    # my_code = np.array([[101], [140], [140], [140], [140], [140], [368], [239], [239], [239], [ 16], [ 16], [ 16], [ 16], [ 16]])
    # my_code = np.array([[480, 480, 140, 140, 140, 140, 368, 239, 239, 239, 16, 16, 16, 16, 16]])

    # my_code = np.array([[480, 480, 480, 182, 272, 272, 272, 272, 272, 272, 272, 272, 16, 16, 16]])

    recon, code, distance = latent2code(vqvae_model, x_, None, None, w_grad, code=None)
    if w_grad:
        x = x.detach().data.cpu().numpy()
    y = np.concatenate((x[:, :16 * 2], recon, x[:, -16 * 1:]), axis=1)  # (1, 112, len)
    np.save(output_path + "_recon.npy", y)
    print(output_path, code)
    np.save(output_path + "_code.npy", code)
    # print(code.shape, x.shape, y.shape)
    return y, x, code, wavlm_feat, distance, t_list


def single_inference(audio, config, sample_fn, wavlm_model, vqvae_model, diffusion_model, audio_wav=None, time_step=False, time=None, return_seed=False, w_grad=False, save_result=False):
    if not w_grad:
        vqvae_model = vqvae_model.eval()
        diffusion_model = diffusion_model.eval()
    else:
        vqvae_model = vqvae_model.train()
        vqvae_model.module.freeze_drop()
        diffusion_model = diffusion_model.train()
        diffusion_model.freeze_drop()

    if type(audio_wav) is np.ndarray:
        motion_vq, motion_diff, code, wavlm_feat, distance, t_list = inference_2(config, sample_fn, wavlm_model, vqvae_model, diffusion_model,
                audio_path=audio, audio_wav=audio_wav, time_step=time_step, time=time, return_seed=return_seed, w_grad=w_grad, save_result=save_result)
    else:
        if audio.endswith('.wav'):
            motion_vq, motion_diff, code, wavlm_feat, distance, t_list = inference_2(config, sample_fn, wavlm_model, vqvae_model, diffusion_model,
                                       audio_path=audio, time_step=time_step, time=time, return_seed=return_seed, w_grad=w_grad, save_result=save_result)
        elif audio.endswith('.npy'):
            motion_vq, motion_diff, code, wavlm_feat, distance, t_list = inference_2(config, sample_fn, wavlm_model, vqvae_model, diffusion_model,
                                       audionpy_path=audio, time_step=time_step, time=time, return_seed=return_seed, w_grad=w_grad, save_result=save_result)

    return motion_vq, motion_diff, code, wavlm_feat, distance, t_list


def all_inference(ZEGGS_audio_fold, Trinity_audio_fold, ZEGGS_audio_fold_valid, Trinity_audio_fold_valid, config, sample_fn, wavlm_model, vqvae_model, diffusion_model, return_GT=False):
    train_audio_list = sorted(glob.glob(ZEGGS_audio_fold + '/*.npy') + glob.glob(Trinity_audio_fold + '/*.npy'))
    valid_audio_list = sorted(glob.glob(ZEGGS_audio_fold_valid + '/*.npy') + glob.glob(Trinity_audio_fold_valid + '/*.npy'))

    GT_list = []

    # motion_list = []
    code_list = []

    for audio in train_audio_list:
        name = os.path.split(audio)[1][:-4]
        audio_wav = np.load(audio)[0]
        print(name)
        num_subdivision = math.floor(int(audio_wav.shape[0] * 7.5 // 16000 / (config.n_poses - seed)))
        time = np.array([500] * num_subdivision)        # [1, 1000]

        motion_vq, motion_diff, code, wavlm_feat, distance, t_list = single_inference(audio, config, sample_fn, wavlm_model, vqvae_model, diffusion_model, audio_wav=audio_wav, time_step=True, time=time, return_seed=True, w_grad=False)

        # motion_list.append(motion)
        code = code.squeeze(1)
        wavlm_feat = wavlm_feat.squeeze(1)[:code.shape[0]]
        code_list.append(code)
        if return_GT:
            if name.split('_')[0] == 'Recording':
                GT_code = np.load("../dataset/Trinity/all_gesture_aux/" + name + '_upper.npy')
            else:
                GT_code = np.load("../dataset/ZEGGS/all_gesture_aux/" + name + '_upper.npy')
            GT_code = GT_code[:code.shape[0]]
            GT_list.append(GT_code)
        # break

        '''
        # np.argmin(distance, axis=1) == code
        # TODO: save wavlm_feat, code(distance) and t
        '''
        code = code.reshape(-1, config.n_poses//2)
        distance = distance.reshape(-1, config.n_poses//2, 512)
        wavlm_feat = wavlm_feat.reshape(-1, config.n_poses//2, 1024)
        pdb.set_trace()

    return code_list, GT_list


def generate_result(ZEGGS_audio_fold_valid, Trinity_audio_fold_valid, config, sample_fn, wavlm_model, vqvae_model, diffusion_model, return_GT=False):
    valid_audio_list = sorted(glob.glob(ZEGGS_audio_fold_valid + '/*.npy') + glob.glob(Trinity_audio_fold_valid + '/*.npy'))

    GT_list = []

    # motion_list = []
    code_list = []

    for audio in valid_audio_list:
        name = os.path.split(audio)[1][:-4]
        audio_wav = np.load(audio)[0]
        print(name)

        motion_vq, motion_diff, code, wavlm_feat, distance2, t_list = single_inference(audio, config, sample_fn, wavlm_model, vqvae_model, diffusion_model, audio_wav=audio_wav, time_step=False, save_result=False)

        # motion_list.append(motion)
        code = code.squeeze(1)
        wavlm_feat = wavlm_feat.squeeze(1)[:code.shape[0]]
        code_list.append(code)
        # break

        '''
        # np.argmin(distance, axis=1) == code
        # TODO: save wavlm_feat, code(distance) and t
        '''
        # code = code.reshape(-1, config.n_poses//2)
        # distance = distance.reshape(-1, config.n_poses//2, 512)
        # wavlm_feat = wavlm_feat.reshape(-1, config.n_poses//2, 1024)

        # [sys.path.append(i) for i in ['../retargeting']]
        # from retargeting.demo import latent2bvh
        np.save(save_dir + '/' + name + '_final' + '.npy', motion_vq)
        np.save(save_dir + '/' + name + '_wovq' + '.npy', motion_diff)


if __name__ == '__main__':
    '''
    /ceph/hdd/yangsc21/Python/My_3/deep-motion-editing/diffusion_latent/
    # pip install ffmpeg-normalize, see https://github.com/slhck/ffmpeg-normalize
    '''

    with open(args.config) as f:
        config = yaml.safe_load(f)

    for k, v in vars(args).items():
        config[k] = v
    # pprint(config)

    config = EasyDict(config)

    # model_path = './result/inference/new_DiffuseStyleGesture_model3_256/model001100000.pt'
    # model_path = './result/inference/new_DiffuseStyleGesture_model3_128/model001150000.pt'
    # model_path = "/ceph/home/wangzl21/Projects/My_3/deep-motion-editing/diffusion_latent/experiments/reinforce_diffusion_onlydiff_gradnorm0.1_lr1e-6_max0_seed0/ckpt/diffusion_epoch_1.pt"
    # model_path = '/ceph/home/wangzl21/Projects/My_3/deep-motion-editing/diffusion_latent/experiments/256_seed_4_aux_model001700000_reinforce_diffusion_onlydiff_gradnorm0.1_lr1e-6_max0_seed0/ckpt/diffusion_epoch_1.pt'

    # model_path = "/ceph/home/wangzl21/Projects/My_3/deep-motion-editing/diffusion_latent/experiments/256_seed_6_aux_model001700000_reinforce_diffusion_onlydiff_gradnorm0.1_lr1e-6_max0_seed0/ckpt/diffusion_epoch_1.pt"
    # model_path = "/ceph/home/wangzl21/Projects/My_3/deep-motion-editing/diffusion_latent/experiments/reinforce_diffusion_baseline_onlydiff_withoutreturnnorm_gradnorm0.2_lr1e-5_seed0/ckpt/diffusion_epoch_1.pt"

    save_dir = config.save_dir
    model_path = config.model_path
    sample_fn, wavlm_model, vqvae_model, diffusion_model = main(save_dir, model_path)

    audio = config.audio_path       # "../dataset/ZEGGS/all_speech/067_Speech_2_x_1_0.npy"
    _, _, code, _, distance, _ = single_inference(audio, config, sample_fn, wavlm_model, vqvae_model, diffusion_model, time_step=False, w_grad=False, save_result=True)      # for debug, motion (1, 112, len)  code (len//2, 1)

    # generate_result('../dataset/ZEGGS/valid_speech/', '../dataset/Trinity/valid_speech/', config, sample_fn, wavlm_model, vqvae_model, diffusion_model, return_GT=False)

    # all_inference('../dataset/ZEGGS/all_speech/', '../dataset/Trinity/all_speech/', '../dataset/ZEGGS/valid_speech/', '../dataset/Trinity/valid_speech/',
    #               config, sample_fn, wavlm_model, vqvae_model, diffusion_model, return_GT=False)       # you may change return_GT to True to get GT code (y)


diffusion_latent/sample.py
YoungSeng/UnifiedGesture
import pdb
import sys
[sys.path.append(i) for i in ['.', '..']]

from model.mdm import MDM
from utils.model_util import create_gaussian_diffusion

# from data_loaders.get_data import get_dataset_loader
from train.training_loop import TrainLoop

import torch
import os
import json


device = torch.device('cuda:2')
n_frames = 240
n_pose_dims = 135
n_audio_dim = 32

# n_frames = 240
# n_pose_dims = 251


def create_model_and_diffusion():
    model = MDM(modeltype='', njoints=n_pose_dims, nfeats=1, translation=True, pose_rep='rot6d', glob=True,
                glob_rot=True, cond_mode = 'text', clip_version = 'ViT-B/32', action_emb = 'tensor')
    diffusion = create_gaussian_diffusion()
    return model, diffusion


if __name__ == '__main__':
    '''
    python train/mytrain.py --overwrite --save_dir save/mydebug --dataset kit --device 1
    '''
    # modify data/dataset.py

    model, diffusion = create_model_and_diffusion()
    model.to(device)
    # model.rot2xyz.smpl_model.eval()

    print('Total params: %.2fM' % (sum(p.numel() for p in model.parameters_wo_clip()) / 1000000.0))


    from utils.parser_util import train_args
    from utils.fixseed import fixseed

    args = train_args()
    fixseed(args.seed)

    if args.save_dir is None:
        raise FileNotFoundError('save_dir was not specified.')
    elif os.path.exists(args.save_dir) and not args.overwrite:
        raise FileExistsError('save_dir [{}] already exists.'.format(args.save_dir))
    elif not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)
    args_path = os.path.join(args.save_dir, 'args.json')
    with open(args_path, 'w') as fw:
        json.dump(vars(args), fw, indent=4, sort_keys=True)

    # print("creating data loader...")
    # data = get_dataset_loader(name=args.dataset, batch_size=args.batch_size, num_frames=args.num_frames)

    # print(iter(data).next()[1]['y'].keys())
    print("Training...")
    TrainLoop(args, model, diffusion, device).run_loop()


diffusion_latent/train/mytrain.py
YoungSeng/UnifiedGesture
import copy
import functools
import os
import pdb

import torch.nn.functional as F

import blobfile as bf
import torch
from torch.optim import AdamW

from diffusion import logger
from diffusion.fp16_util import MixedPrecisionTrainer
from diffusion.resample import LossAwareSampler, UniformSampler
from tqdm import tqdm
from diffusion.resample import create_named_schedule_sampler
# from eval import eval_humanml, eval_humanact12_uestc

from datetime import datetime

import sys
[sys.path.append(i) for i in ['./mydiffusion', '../../My', '../../My/process', '../process', '../../ubisoft-laforge-ZeroEGGS-main', '../../ubisoft-laforge-ZeroEGGS-main/ZEGGS']]
from generate.generate import WavEncoder
# from process.process_bvh import make_bvh_GENEA2020_BT
# from process_zeggs_bvh import pose2bvh

# For ImageNet experiments, this was a good default value.
# We found that the lg_loss_scale quickly climbed to
# 20-21 within the first ~1K steps of training.
INITIAL_LOG_LOSS_SCALE = 20.0


class TrainLoop:
    def __init__(self, args, model, diffusion, device, data=None):
        self.args = args
        self.data = data
        self.model = model
        self.diffusion = diffusion
        self.cond_mode = model.cond_mode
        self.batch_size = args.batch_size
        self.microbatch = args.batch_size  # deprecating this option
        self.lr = args.lr
        self.log_interval = args.log_interval
        # self.save_interval = args.save_interval
        # self.resume_checkpoint = args.resume_checkpoint
        self.use_fp16 = False  # deprecating this option
        self.fp16_scale_growth = 1e-3  # deprecating this option
        self.weight_decay = args.weight_decay
        self.lr_anneal_steps = args.lr_anneal_steps

        self.step = 0
        self.resume_step = 0
        self.global_batch = self.batch_size # * dist.get_world_size()
        # self.num_steps = args.num_steps
        self.num_epochs = 40000
        self.n_seed = args.n_seed

        self.sync_cuda = torch.cuda.is_available()

        # self._load_and_sync_parameters()
        self.mp_trainer = MixedPrecisionTrainer(
            model=self.model,
            use_fp16=self.use_fp16,
            fp16_scale_growth=self.fp16_scale_growth,
        )

        self.save_dir = args.save_dir

        self.device = device
        if args.audio_feat == "wav encoder":
            self.WavEncoder = WavEncoder().to(self.device)
            self.opt = AdamW([
                {'params': self.mp_trainer.master_params, 'lr':self.lr, 'weight_decay':self.weight_decay},
                {'params': self.WavEncoder.parameters(), 'lr':self.lr}
            ])
        elif args.audio_feat == "mfcc" or args.audio_feat == 'wavlm':
            self.opt = AdamW([
                {'params': self.mp_trainer.master_params, 'lr':self.lr, 'weight_decay':self.weight_decay}
            ])

        # if self.resume_step:
        #     self._load_optimizer_state()
            # Model was resumed, either due to a restart or a checkpoint
            # being specified at the command line.

        self.schedule_sampler_type = 'uniform'
        self.schedule_sampler = create_named_schedule_sampler(self.schedule_sampler_type, diffusion)
        self.eval_wrapper, self.eval_data, self.eval_gt_data = None, None, None

        self.use_ddp = False
        self.ddp_model = self.model
        self.mask_train = (torch.zeros([self.batch_size, 1, 1, args.n_poses]) < 1).to(self.device)
        self.mask_test = (torch.zeros([1, 1, 1, args.n_poses]) < 1).to(self.device)
        self.mask_local_train = torch.ones(self.batch_size, args.n_poses).bool().to(self.device)
        self.mask_local_test = torch.ones(1, args.n_poses).bool().to(self.device)

    def run_loop(self):

        for epoch in range(self.num_epochs):
            print(f'Starting epoch {epoch}')
            # for _ in tqdm(range(10)):     # 4 steps, batch size, chmod 777
            for batch in tqdm(self.data):
                if not (not self.lr_anneal_steps or self.step + self.resume_step < self.lr_anneal_steps):
                    break

                cond_ = {'y':{}}

                pose_seq,  _, style, wavlm = batch  # (batch, 240, 135), (batch, 30), (batch, 64000)
                # pose_vel = F.pad((pose_seq[:, 1:] - pose_seq[:, :-1]), (0, 0, 1, 0))        # 20230424
                # pose_seq = torch.cat([pose_seq, pose_vel], dim=-1)

                # vel = pose_seq[:, 1:] - pose_seq[:, :-1]
                # acc = vel[:, 1:] - vel[:, :-1]
                # vel = F.pad(vel, (0, 0, 1, 0))
                # acc = F.pad(acc, (0, 0, 2, 0))  # 20230424
                # pose_seq = torch.cat([pose_seq, vel, acc], dim=-1)

                motion = pose_seq.permute(0, 2, 1).unsqueeze(2).to(self.device)
                cond_['y']['seed'] = motion[..., 0:self.n_seed]
                cond_['y']['style'] = style.to(self.device)
                cond_['y']['mask_local'] = self.mask_local_train

                if self.args.audio_feat == 'wav encoder':
                    # cond_['y']['audio'] = torch.rand(240, 2, 32).to(self.device)
                    cond_['y']['audio'] = self.WavEncoder(audio.to(self.device)).permute(1, 0, 2)       # (batch, 240, 32)
                elif self.args.audio_feat == 'mfcc':
                    # cond_['y']['audio'] = torch.rand(80, 2, 13).to(self.device)
                    cond_['y']['audio'] = mfcc.to(torch.float32).to(self.device).permute(1, 0, 2)       # [self.n_seed:, ...]      # (batch, 80, 13)
                elif self.args.audio_feat == 'wavlm':
                    cond_['y']['audio'] = wavlm.to(torch.float32).to(self.device)

                cond_['y']['mask'] = self.mask_train        # [..., self.n_seed:]

                self.run_step(motion, cond_)
                if self.step % self.log_interval == 0:
                    for k,v in logger.get_current().name2val.items():
                        if k == 'loss':
                            print('step[{}]: loss[{:0.5f}]'.format(self.step+self.resume_step, v))

                if self.step % 50000 == 0:
                    self.save()

                    # Run for a finite amount of time in integration tests.
                    if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0:
                        return
                self.step += 1
            if not (not self.lr_anneal_steps or self.step + self.resume_step < self.lr_anneal_steps):
                break


    def run_step(self, batch, cond):
        self.forward_backward(batch, cond)      # torch.Size([64, 251, 1, 196]) cond['y'].keys() dict_keys(['mask', 'lengths', 'text', 'tokens'])
        self.mp_trainer.optimize(self.opt)
        self._anneal_lr()
        self.log_step()

    def forward_backward(self, batch, cond):
        self.mp_trainer.zero_grad()
        for i in range(0, batch.shape[0], self.microbatch):
            # Eliminates the microbatch feature
            assert i == 0
            assert self.microbatch == self.batch_size
            micro = batch
            micro_cond = cond
            last_batch = (i + self.microbatch) >= batch.shape[0]
            t, weights = self.schedule_sampler.sample(micro.shape[0], self.device)

            compute_losses = functools.partial(
                self.diffusion.training_losses,
                self.ddp_model,
                micro,  # [bs, ch, image_size, image_size]      # x_start, (2, 135, 1, 240)
                t,  # [bs](int) sampled timesteps
                model_kwargs=micro_cond,
                dataset='kit'
            )

            if last_batch or not self.use_ddp:
                losses = compute_losses()
            else:
                with self.ddp_model.no_sync():
                    losses = compute_losses()

            if isinstance(self.schedule_sampler, LossAwareSampler):
                self.schedule_sampler.update_with_local_losses(
                    t, losses["loss"].detach()
                )

            loss = (losses["loss"] * weights).mean()
            log_loss_dict(
                self.diffusion, t, {k: v * weights for k, v in losses.items()}
            )
            self.mp_trainer.backward(loss)

    def _anneal_lr(self):
        if not self.lr_anneal_steps:
            return
        frac_done = (self.step + self.resume_step) / self.lr_anneal_steps
        lr = self.lr * (1 - frac_done)
        for param_group in self.opt.param_groups:
            param_group["lr"] = lr

    def log_step(self):
        logger.logkv("step", self.step + self.resume_step)
        logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch)


    def ckpt_file_name(self):
        return f"model{(self.step+self.resume_step):09d}.pt"


    def save(self):
        def save_checkpoint(params):
            state_dict = self.mp_trainer.master_params_to_state_dict(params)

            # Do not save CLIP weights
            clip_weights = [e for e in state_dict.keys() if e.startswith('clip_model.')]
            for e in clip_weights:
                del state_dict[e]

            logger.log(f"saving model...")
            filename = self.ckpt_file_name()
            with bf.BlobFile(bf.join(self.save_dir, filename), "wb") as f:
                torch.save(state_dict, f)

        save_checkpoint(self.mp_trainer.master_params)

        with bf.BlobFile(
            bf.join(self.save_dir, f"opt{(self.step+self.resume_step):09d}.pt"),
            "wb",
        ) as f:
            torch.save(self.opt.state_dict(), f)


def parse_resume_step_from_filename(filename):
    """
    Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the
    checkpoint's number of steps.
    """
    split = filename.split("model")
    if len(split) < 2:
        return 0
    split1 = split[-1].split(".")[0]
    try:
        return int(split1)
    except ValueError:
        return 0


def get_blob_logdir():
    # You can change this to be a separate path to save checkpoints to
    # a blobstore or some external drive.
    return logger.get_dir()


def find_resume_checkpoint():
    # On your infrastructure, you may want to override this to automatically
    # discover the latest checkpoint on your blob storage, etc.
    return None


def log_loss_dict(diffusion, ts, losses):
    for key, values in losses.items():
        logger.logkv_mean(key, values.mean().item())
        # Log the quantiles (four quartiles, in particular).
        for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()):
            quartile = int(4 * sub_t / diffusion.num_timesteps)
            logger.logkv_mean(f"{key}_q{quartile}", sub_loss)

diffusion_latent/train/training_loop.py
YoungSeng/UnifiedGesture
import pdb

from model.mdm import MDM
from diffusion import gaussian_diffusion as gd
from diffusion.respace import SpacedDiffusion, space_timesteps


def load_model_wo_clip(model, state_dict):
    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
    print(missing_keys, unexpected_keys)
    assert len(unexpected_keys) == 0
    assert all([k.startswith('clip_model.') for k in missing_keys])


def create_model_and_diffusion(args, data):
    model = MDM(**get_model_args(args, data))
    diffusion = create_gaussian_diffusion(args)
    return model, diffusion


def get_model_args(args, data):

    # default args
    clip_version = 'ViT-B/32'
    action_emb = 'tensor'
    if args.unconstrained:
        cond_mode = 'no_cond'
    elif args.dataset in ['kit', 'humanml']:
        cond_mode = 'text'
    else:
        cond_mode = 'action'
    if hasattr(data.dataset, 'num_actions'):
        num_actions = data.dataset.num_actions
    else:
        num_actions = 1

    # SMPL defaults
    data_rep = 'rot6d'
    njoints = 25
    nfeats = 6

    if args.dataset == 'humanml':
        data_rep = 'hml_vec'
        njoints = 263
        nfeats = 1
    elif args.dataset == 'kit':
        data_rep = 'hml_vec'
        njoints = 251
        nfeats = 1

    return {'modeltype': '', 'njoints': njoints, 'nfeats': nfeats, 'num_actions': num_actions,
            'translation': True, 'pose_rep': 'rot6d', 'glob': True, 'glob_rot': True,
            'latent_dim': args.latent_dim, 'ff_size': 1024, 'num_layers': args.layers, 'num_heads': 4,
            'dropout': 0.1, 'activation': "gelu", 'data_rep': data_rep, 'cond_mode': cond_mode,
            'cond_mask_prob': args.cond_mask_prob, 'action_emb': action_emb, 'arch': args.arch,
            'emb_trans_dec': args.emb_trans_dec, 'clip_version': clip_version, 'dataset': args.dataset}


def create_gaussian_diffusion():
    noise_schedule = 'cosine'
    sigma_small = True
    lambda_vel = 0.0
    lambda_rcxyz = 0.0
    lambda_fc = 0.0

    # default params
    predict_xstart = True  # we always predict x_start (a.k.a. x0), that's our deal!
    steps = 1000
    scale_beta = 1.  # no scaling
    timestep_respacing = ''  # can be used for ddim sampling, we don't use it.
    learn_sigma = False
    rescale_timesteps = False

    betas = gd.get_named_beta_schedule(noise_schedule, steps, scale_beta)
    loss_type = gd.LossType.MSE

    if not timestep_respacing:
        timestep_respacing = [steps]

    return SpacedDiffusion(
        use_timesteps=space_timesteps(steps, timestep_respacing),
        betas=betas,
        model_mean_type=(
            gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X
        ),
        model_var_type=(
            (
                gd.ModelVarType.FIXED_LARGE
                if not sigma_small
                else gd.ModelVarType.FIXED_SMALL
            )
            if not learn_sigma
            else gd.ModelVarType.LEARNED_RANGE
        ),
        loss_type=loss_type,
        rescale_timesteps=rescale_timesteps,
        lambda_vel=lambda_vel,
        lambda_rcxyz=lambda_rcxyz,
        lambda_fc=lambda_fc,
    )
diffusion_latent/utils_/model_util.py
YoungSeng/UnifiedGesture
from argparse import ArgumentParser
import argparse
import os
import json


def parse_and_load_from_model(parser):
    # args according to the loaded model
    # do not try to specify them from cmd line since they will be overwritten
    add_data_options(parser)
    add_model_options(parser)
    add_diffusion_options(parser)
    args = parser.parse_args()
    args_to_overwrite = []
    for group_name in ['dataset', 'model', 'diffusion']:
        args_to_overwrite += get_args_per_group_name(parser, args, group_name)

    # load args from model
    model_path = get_model_path_from_args()
    args_path = os.path.join(os.path.dirname(model_path), 'args.json')
    assert os.path.exists(args_path), 'Arguments json file was not found!'
    with open(args_path, 'r') as fr:
        model_args = json.load(fr)

    for a in args_to_overwrite:
        if a in model_args.keys():
            setattr(args, a, model_args[a])

        elif 'cond_mode' in model_args: # backward compitability
            unconstrained = (model_args['cond_mode'] == 'no_cond')
            setattr(args, 'unconstrained', unconstrained)

        else:
            print('Warning: was not able to load [{}], using default value [{}] instead.'.format(a, args.__dict__[a]))

    if args.cond_mask_prob == 0:
        args.guidance_param = 1
    return args


def get_args_per_group_name(parser, args, group_name):
    for group in parser._action_groups:
        if group.title == group_name:
            group_dict = {a.dest: getattr(args, a.dest, None) for a in group._group_actions}
            return list(argparse.Namespace(**group_dict).__dict__.keys())
    return ValueError('group_name was not found.')

def get_model_path_from_args():
    try:
        dummy_parser = ArgumentParser()
        dummy_parser.add_argument('model_path')
        dummy_args, _ = dummy_parser.parse_known_args()
        return dummy_args.model_path
    except:
        raise ValueError('model_path argument must be specified.')


def add_base_options(parser):
    group = parser.add_argument_group('base')
    group.add_argument("--cuda", default=True, type=bool, help="Use cuda device, otherwise use CPU.")
    group.add_argument("--device", default=0, type=int, help="Device id to use.")
    group.add_argument("--seed", default=10, type=int, help="For fixing random seed.")
    group.add_argument("--batch_size", default=64, type=int, help="Batch size during training.")


def add_diffusion_options(parser):
    group = parser.add_argument_group('diffusion')
    group.add_argument("--noise_schedule", default='cosine', choices=['linear', 'cosine'], type=str,
                       help="Noise schedule type")
    group.add_argument("--diffusion_steps", default=1000, type=int,
                       help="Number of diffusion steps (denoted T in the paper)")
    group.add_argument("--sigma_small", default=True, type=bool, help="Use smaller sigma values.")


def add_model_options(parser):
    group = parser.add_argument_group('model')
    group.add_argument("--arch", default='trans_enc',
                       choices=['trans_enc', 'trans_dec', 'gru'], type=str,
                       help="Architecture types as reported in the paper.")
    group.add_argument("--emb_trans_dec", default=False, type=bool,
                       help="For trans_dec architecture only, if true, will inject condition as a class token"
                            " (in addition to cross-attention).")
    group.add_argument("--layers", default=8, type=int,
                       help="Number of layers.")
    group.add_argument("--latent_dim", default=512, type=int,
                       help="Transformer/GRU width.")
    group.add_argument("--cond_mask_prob", default=.1, type=float,
                       help="The probability of masking the condition during training."
                            " For classifier-free guidance learning.")
    group.add_argument("--lambda_rcxyz", default=0.0, type=float, help="Joint positions loss.")
    group.add_argument("--lambda_vel", default=0.0, type=float, help="Joint velocity loss.")
    group.add_argument("--lambda_fc", default=0.0, type=float, help="Foot contact loss.")
    group.add_argument("--unconstrained", action='store_true',
                       help="Model is trained unconditionally. That is, it is constrained by neither text nor action. "
                            "Currently tested on HumanAct12 only.")



def add_data_options(parser):
    group = parser.add_argument_group('dataset')
    group.add_argument("--dataset", default='humanml', choices=['humanml', 'kit', 'humanact12', 'uestc'], type=str,
                       help="Dataset name (choose from list).")
    group.add_argument("--data_dir", default="", type=str,
                       help="If empty, will use defaults according to the specified dataset.")


def add_training_options(parser):
    group = parser.add_argument_group('training')
    group.add_argument("--save_dir", required=True, type=str,
                       help="Path to save checkpoints and results.")
    group.add_argument("--overwrite", action='store_true',
                       help="If True, will enable to use an already existing save_dir.")
    group.add_argument("--train_platform_type", default='NoPlatform', choices=['NoPlatform', 'ClearmlPlatform', 'TensorboardPlatform'], type=str,
                       help="Choose platform to log results. NoPlatform means no logging.")
    group.add_argument("--lr", default=1e-4, type=float, help="Learning rate.")
    group.add_argument("--weight_decay", default=0.0, type=float, help="Optimizer weight decay.")
    group.add_argument("--lr_anneal_steps", default=0, type=int, help="Number of learning rate anneal steps.")
    group.add_argument("--eval_batch_size", default=32, type=int,
                       help="Batch size during evaluation loop. Do not change this unless you know what you are doing. "
                            "T2m precision calculation is based on fixed batch size 32.")
    group.add_argument("--eval_split", default='test', choices=['val', 'test'], type=str,
                       help="Which split to evaluate on during training.")
    group.add_argument("--eval_during_training", action='store_true',
                       help="If True, will run evaluation during training.")
    group.add_argument("--eval_rep_times", default=3, type=int,
                       help="Number of repetitions for evaluation loop during training.")
    group.add_argument("--eval_num_samples", default=1_000, type=int,
                       help="If -1, will use all samples in the specified split.")
    group.add_argument("--log_interval", default=1_000, type=int,
                       help="Log losses each N steps")
    group.add_argument("--save_interval", default=50_000, type=int,
                       help="Save checkpoints and run evaluation each N steps")
    group.add_argument("--num_steps", default=600_000, type=int,
                       help="Training will stop after the specified number of steps.")
    group.add_argument("--num_frames", default=60, type=int,
                       help="Limit for the maximal number of frames. In HumanML3D and KIT this field is ignored.")
    group.add_argument("--resume_checkpoint", default="", type=str,
                       help="If not empty, will start from the specified checkpoint (path to model###.pt file).")


def add_sampling_options(parser):
    group = parser.add_argument_group('sampling')
    group.add_argument("--model_path", required=True, type=str,
                       help="Path to model####.pt file to be sampled.")
    group.add_argument("--output_dir", default='', type=str,
                       help="Path to results dir (auto created by the script). "
                            "If empty, will create dir in parallel to checkpoint.")
    group.add_argument("--num_samples", default=10, type=int,
                       help="Maximal number of prompts to sample, "
                            "if loading dataset from file, this field will be ignored.")
    group.add_argument("--num_repetitions", default=3, type=int,
                       help="Number of repetitions, per sample (text prompt/action)")
    group.add_argument("--guidance_param", default=2.5, type=float,
                       help="For classifier-free sampling - specifies the s parameter, as defined in the paper.")


def add_generate_options(parser):
    group = parser.add_argument_group('generate')
    group.add_argument("--motion_length", default=6.0, type=float,
                       help="The length of the sampled motion [in seconds]. "
                            "Maximum is 9.8 for HumanML3D (text-to-motion), and 2.0 for HumanAct12 (action-to-motion)")
    group.add_argument("--input_text", default='', type=str,
                       help="Path to a text file lists text prompts to be synthesized. If empty, will take text prompts from dataset.")
    group.add_argument("--action_file", default='', type=str,
                       help="Path to a text file that lists names of actions to be synthesized. Names must be a subset of dataset/uestc/info/action_classes.txt if sampling from uestc, "
                            "or a subset of [warm_up,walk,run,jump,drink,lift_dumbbell,sit,eat,turn steering wheel,phone,boxing,throw] if sampling from humanact12. "
                            "If no file is specified, will take action names from dataset.")
    group.add_argument("--text_prompt", default='', type=str,
                       help="A text prompt to be generated. If empty, will take text prompts from dataset.")
    group.add_argument("--action_name", default='', type=str,
                       help="An action name to be generated. If empty, will take text prompts from dataset.")


def add_edit_options(parser):
    group = parser.add_argument_group('edit')
    group.add_argument("--edit_mode", default='in_between', choices=['in_between', 'upper_body'], type=str,
                       help="Defines which parts of the input motion will be edited.\n"
                            "(1) in_between - suffix and prefix motion taken from input motion, "
                            "middle motion is generated.\n"
                            "(2) upper_body - lower body joints taken from input motion, "
                            "upper body is generated.")
    group.add_argument("--text_condition", default='', type=str,
                       help="Editing will be conditioned on this text prompt. "
                            "If empty, will perform unconditioned editing.")
    group.add_argument("--prefix_end", default=0.25, type=float,
                       help="For in_between editing - Defines the end of input prefix (ratio from all frames).")
    group.add_argument("--suffix_start", default=0.75, type=float,
                       help="For in_between editing - Defines the start of input suffix (ratio from all frames).")


def add_evaluation_options(parser):
    group = parser.add_argument_group('eval')
    group.add_argument("--model_path", required=True, type=str,
                       help="Path to model####.pt file to be sampled.")
    group.add_argument("--eval_mode", default='wo_mm', choices=['wo_mm', 'mm_short', 'debug', 'full'], type=str,
                       help="wo_mm (t2m only) - 20 repetitions without multi-modality metric; "
                            "mm_short (t2m only) - 5 repetitions with multi-modality metric; "
                            "debug - short run, less accurate results."
                            "full (a2m only) - 20 repetitions.")
    group.add_argument("--guidance_param", default=2.5, type=float,
                       help="For classifier-free sampling - specifies the s parameter, as defined in the paper.")


def train_args():
    parser = ArgumentParser()
    add_base_options(parser)
    add_data_options(parser)
    add_model_options(parser)
    add_diffusion_options(parser)
    add_training_options(parser)
    return parser.parse_args()


def generate_args():
    parser = ArgumentParser()
    # args specified by the user: (all other will be loaded from the model)
    add_base_options(parser)
    add_sampling_options(parser)
    add_generate_options(parser)
    return parse_and_load_from_model(parser)


def edit_args():
    parser = ArgumentParser()
    # args specified by the user: (all other will be loaded from the model)
    add_base_options(parser)
    add_sampling_options(parser)
    add_edit_options(parser)
    return parse_and_load_from_model(parser)


def evaluation_parser():
    parser = ArgumentParser()
    # args specified by the user: (all other will be loaded from the model)
    add_base_options(parser)
    add_evaluation_options(parser)
    return parse_and_load_from_model(parser)
diffusion_latent/utils_/parser_util.py
YoungSeng/UnifiedGesture
import pdb

import numpy as np
import os

def extract_kinetic_features(positions):
    assert len(positions.shape) == 3  # (seq_len, n_joints, 3)
    features = KineticFeatures(positions)
    kinetic_feature_vector = []
    for i in range(positions.shape[1]):
        feature_vector = np.hstack(
            [
                features.average_kinetic_energy_horizontal(i),
                features.average_kinetic_energy_vertical(i),
                features.average_energy_expenditure(i),
            ]
        )
        kinetic_feature_vector.extend(feature_vector)
    kinetic_feature_vector = np.array(kinetic_feature_vector, dtype=np.float32)
    return kinetic_feature_vector


def calc_average_velocity_vertical(
    positions, i, joint_idx, sliding_window, frame_time, up_vec
):
    current_window = 0
    average_velocity = np.zeros(len(positions[0][joint_idx]))
    for j in range(-sliding_window, sliding_window + 1):
        if i + j - 1 < 0 or i + j >= len(positions):
            continue
        average_velocity += (
            positions[i + j][joint_idx] - positions[i + j - 1][joint_idx]
        )
        current_window += 1
    if up_vec == "y":
        average_velocity = np.array([average_velocity[1]]) / (
            current_window * frame_time
        )
    elif up_vec == "z":
        average_velocity = np.array([average_velocity[2]]) / (
            current_window * frame_time
        )
    else:
        raise NotImplementedError
    return np.linalg.norm(average_velocity)


def calc_average_acceleration(
    positions, i, joint_idx, sliding_window, frame_time
):
    current_window = 0
    average_acceleration = np.zeros(len(positions[0][joint_idx]))
    for j in range(-sliding_window, sliding_window + 1):
        if i + j - 1 < 0 or i + j + 1 >= len(positions):
            continue
        v2 = (
            positions[i + j + 1][joint_idx] - positions[i + j][joint_idx]
        ) / frame_time
        v1 = (
            positions[i + j][joint_idx]
            - positions[i + j - 1][joint_idx] / frame_time
        )
        average_acceleration += (v2 - v1) / frame_time
        current_window += 1
    return np.linalg.norm(average_acceleration / current_window)


def calc_average_velocity_horizontal(
    positions, i, joint_idx, sliding_window, frame_time, up_vec="z"
):
    current_window = 0
    average_velocity = np.zeros(len(positions[0][joint_idx]))
    for j in range(-sliding_window, sliding_window + 1):
        if i + j - 1 < 0 or i + j >= len(positions):
            continue
        average_velocity += (
            positions[i + j][joint_idx] - positions[i + j - 1][joint_idx]
        )
        current_window += 1
    if up_vec == "y":
        average_velocity = np.array(
            [average_velocity[0], average_velocity[2]]
        ) / (current_window * frame_time)
    elif up_vec == "z":
        average_velocity = np.array(
            [average_velocity[0], average_velocity[1]]
        ) / (current_window * frame_time)
    else:
        raise NotImplementedError
    return np.linalg.norm(average_velocity)


def calc_average_velocity(positions, i, joint_idx, sliding_window, frame_time):
    current_window = 0
    average_velocity = np.zeros(len(positions[0][joint_idx]))
    for j in range(-sliding_window, sliding_window + 1):
        if i + j - 1 < 0 or i + j >= len(positions):
            continue
        average_velocity += (
            positions[i + j][joint_idx] - positions[i + j - 1][joint_idx]
        )
        current_window += 1
    return np.linalg.norm(average_velocity / (current_window * frame_time))


class KineticFeatures:
    def __init__(
        self, positions, frame_time=1./30, up_vec="y", sliding_window=2
    ):
        self.positions = positions
        self.frame_time = frame_time
        self.up_vec = up_vec
        self.sliding_window = sliding_window

    def average_kinetic_energy(self, joint):
        average_kinetic_energy = 0
        for i in range(1, len(self.positions)):
            average_velocity = calc_average_velocity(
                self.positions, i, joint, self.sliding_window, self.frame_time
            )
            average_kinetic_energy += average_velocity ** 2
        average_kinetic_energy = average_kinetic_energy / (
            len(self.positions) - 1.0
        )
        return average_kinetic_energy

    def average_kinetic_energy_horizontal(self, joint):
        val = 0
        for i in range(1, len(self.positions)):
            average_velocity = calc_average_velocity_horizontal(
                self.positions,
                i,
                joint,
                self.sliding_window,
                self.frame_time,
                self.up_vec,
            )
            val += average_velocity ** 2
        val = val / (len(self.positions) - 1.0)
        return val

    def average_kinetic_energy_vertical(self, joint):
        val = 0
        for i in range(1, len(self.positions)):
            average_velocity = calc_average_velocity_vertical(
                self.positions,
                i,
                joint,
                self.sliding_window,
                self.frame_time,
                self.up_vec,
            )
            val += average_velocity ** 2
        val = val / (len(self.positions) - 1.0)
        return val

    def average_energy_expenditure(self, joint):
        val = 0.0
        for i in range(1, len(self.positions)):
            val += calc_average_acceleration(
                self.positions, i, joint, self.sliding_window, self.frame_time
            )
        val = val / (len(self.positions) - 1.0)
        return val


def calculate_avg_distance(feature_list, mean=None, std=None):
    feature_list = np.stack(feature_list)
    n = feature_list.shape[0]
    # normalize the scale
    if (mean is not None) and (std is not None):
        feature_list = (feature_list - mean) / std
    dist = 0
    print(n)
    for i in range(n):
        print(i, end='\r')
        for j in range(i + 1, n):
            dist += np.linalg.norm(feature_list[i] - feature_list[j])
    dist /= (n * n - n) / 2
    return dist


def normalize(feat, feat2):
    mean = feat.mean(axis=0)
    std = feat.std(axis=0)

    return (feat - mean) / (std + 1e-10), (feat2 - mean) / (std + 1e-10)


def calc_diversity(feats):
    feat_array = np.array(feats)
    n, c = feat_array.shape
    diff = np.array([feat_array] * n) - feat_array.reshape(n, 1, c)
    return np.sqrt(np.sum(diff**2, axis=2)).sum() / n / (n-1)


def cal_div(GT_path, source_path):
    pred_features = []
    for item in os.listdir(source_path):
        if item.endswith('.npy'):
            feat_source = np.load(os.path.join(source_path, item))
            pred_features.append(feat_source)
    GT_features = []
    for item in os.listdir(GT_path):
        if item.endswith('.npy'):
            GT_feat_source = np.load(os.path.join(GT_path, item))
            GT_features.append(GT_feat_source)

    pred_features_k = np.stack(pred_features)
    GT_features_k = np.stack(GT_features)
    GT_features_k, pred_features_k = normalize(GT_features_k, pred_features_k)
    # return calc_diversity(pred_features_k)
    return calculate_avg_distance(pred_features_k)


def extract_feat(source_path, save_path):
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    for item in os.listdir(source_path):
        if item.endswith('.npy'):
            joint3d = np.load(os.path.join(source_path, item)).reshape(-1, 33 * 3)
            roott = joint3d[:1, :3]
            joint3d = joint3d - np.tile(roott, (1, 33))
            feat = extract_kinetic_features(joint3d.reshape(-1, 33, 3))
            print(item, feat.shape)
            np.save(os.path.join(save_path, item), feat)


if __name__ == '__main__':
    '''
    python Diversity.py
    '''
    source_path = '/ceph/hdd/yangsc21/Python/My_3/worl_npy'
    GT_path = '/ceph/hdd/yangsc21/Python/My_3/GT_Gesture_npy'
    # extract_feat(source_path, source_path + '_feat')
    print(cal_div(GT_path + '_feat', source_path + '_feat'))

genea_numerical_evaluations/Diversity.py
YoungSeng/UnifiedGesture
import pdb

import torch
import torch.nn as nn


def ConvNormRelu(in_channels, out_channels, downsample=False, padding=0, batchnorm=True):
    if not downsample:
        k = 3
        s = 1
    else:
        k = 4
        s = 2

    conv_block = nn.Conv1d(in_channels, out_channels, kernel_size=k, stride=s, padding=padding)
    norm_block = nn.BatchNorm1d(out_channels)

    if batchnorm:
        net = nn.Sequential(
            conv_block,
            norm_block,
            nn.LeakyReLU(0.2, True)
        )
    else:
        net = nn.Sequential(
            conv_block,
            nn.LeakyReLU(0.2, True)
        )

    return net


class PoseEncoderConv(nn.Module):
    def __init__(self, dim, length):
        super().__init__()

        self.net = nn.Sequential(
            ConvNormRelu(dim, 32, batchnorm=True),
            ConvNormRelu(32, 64, batchnorm=True),
            ConvNormRelu(64, 64, True, batchnorm=True),
            nn.Conv1d(64, 32, 3)
        )


        if length == 120:
            in_channels = 1760
        else:
            assert False

        self.out_net = nn.Sequential(
            nn.Linear(in_channels, 256),
            nn.BatchNorm1d(256),
            nn.LeakyReLU(True),
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.LeakyReLU(True),
            nn.Linear(128, 32),
        )

    def forward(self, poses):
        poses = poses.transpose(1, 2)  # to (bs, dim, seq)
        out = self.net(poses)
        out = out.flatten(1)
        z = self.out_net(out)

        return z


class PoseDecoderConv(nn.Module):
    def __init__(self, dim, length):
        super().__init__()

        if length == 30:
            out_channels = 120
        elif length == 60:
            out_channels = 240
        elif length == 90:
            out_channels = 360
        elif length == 120:
            out_channels = 120*4
        else:
            assert False

        self.pre_net = nn.Sequential(
            nn.Linear(32, 64),
            nn.BatchNorm1d(64),
            nn.LeakyReLU(True),
            nn.Linear(64, out_channels),
        )

        self.net = nn.Sequential(
            nn.ConvTranspose1d(4, 32, 3),
            nn.BatchNorm1d(32),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose1d(32, 32, 3),
            nn.BatchNorm1d(32),
            nn.LeakyReLU(0.2, True),
            nn.Conv1d(32, 32, 3),
            nn.Conv1d(32, dim, 3),
        )

    def forward(self, feat):
        out = self.pre_net(feat)
        out = out.view(feat.shape[0], 4, -1)
        out = self.net(out)
        out = out.transpose(1, 2)
        return out


class EmbeddingNet(nn.Module):
    def __init__(self, pose_dim, n_frames):
        super().__init__()
        self.pose_encoder = PoseEncoderConv(pose_dim, n_frames)
        self.decoder = PoseDecoderConv(pose_dim, n_frames)

    def forward(self, poses):
        poses_feat = self.pose_encoder(poses)
        out_poses = self.decoder(poses_feat)
        return poses_feat, out_poses


if __name__ == '__main__':  # model test
    '''
    python embedding_net.py
    '''
    n_frames = 120
    pose_dim = 33*3
    encoder = PoseEncoderConv(pose_dim, n_frames)
    decoder = PoseDecoderConv(pose_dim, n_frames)

    poses = torch.randn((4, n_frames, pose_dim))
    feat = encoder(poses)
    recon_poses = decoder(feat)

    print('input', poses.shape)
    print('feat', feat.shape)
    print('output', recon_poses.shape)

genea_numerical_evaluations/FGD/embedding_net.py
YoungSeng/UnifiedGesture
import os

import numpy as np
import torch

from embedding_space_evaluator import EmbeddingSpaceEvaluator
from train_AE import make_tensor

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def run_fgd(fgd_evaluator, gt_data, test_data):
    fgd_evaluator.reset()

    fgd_evaluator.push_real_samples(gt_data)
    fgd_evaluator.push_generated_samples(test_data)
    fgd_on_feat = fgd_evaluator.get_fgd(use_feat_space=True)
    # fdg_on_raw = fgd_evaluator.get_fgd(use_feat_space=False)
    fdg_on_raw = 0
    return fgd_on_feat, fdg_on_raw


def exp_base(chunk_len):
    # AE model
    ae_path = f'output/model_checkpoint_{chunk_len}.bin'
    fgd_evaluator = EmbeddingSpaceEvaluator(ae_path, chunk_len, device)

    # load GT data
    gt_data = make_tensor("/ceph/hdd/yangsc21/Python/My_3/GT_Gesture_npy/", chunk_len).to(device)

    # load generated data
    generated_data_path = 'data'
    # folders = sorted([f.path for f in os.scandir(generated_data_path) if f.is_dir() and 'Cond_' in f.path])

    folders = ["/ceph/hdd/yangsc21/Python/My_3/wTrinity_npy/"]

    print(f'----- Experiment (motion chunk length: {chunk_len}) -----')
    print('FGDs on feature space and raw data space')
    for folder in folders:
        test_data = make_tensor(folder, chunk_len).to(device)
        fgd_on_feat, fgd_on_raw = run_fgd(fgd_evaluator, gt_data, test_data)
        print(f'{os.path.basename(folder)}: {fgd_on_feat:8.3f}, {fgd_on_raw:8.3f}')
    print()


# def exp_per_seq(chunk_len, stride=1):
#     n_test_seq = 10
#
#     # AE model
#     ae_path = f'output/model_checkpoint_{chunk_len}.bin'
#     fgd_evaluator = EmbeddingSpaceEvaluator(ae_path, chunk_len, device)
#
#     # run
#     print(f'----- Experiment (motion chunk length: {chunk_len}, stride: {stride}) -----')
#     print('FGDs on feature space and raw data space for each system and each test speech sequence')
#
#     results = []
#     for i in range(n_test_seq):
#         name = f'TestSeq{i+1:03d}'
#
#         # load GT data
#         gt_data = make_tensor(os.path.join('data/GroundTruth', name + '.npz'), chunk_len, stride=stride)
#         gt_data = gt_data.to(device)
#         # print(gt_data.shape)
#
#         # load generated data
#         test_data = make_tensor(os.path.join(f'data/Cond_{system_name}', name + '.npz'), chunk_len, stride=stride)
#         test_data = test_data.to(device)
#         # print(test_data.shape)
#
#         # calculate fgd
#         fgd_on_feat, fgd_on_raw = run_fgd(fgd_evaluator, gt_data, test_data)
#         print(f' {name}: {fgd_on_feat:8.3f}, {fgd_on_raw:8.3f}')
#         results.append([fgd_on_feat, fgd_on_raw])
#     results = np.array(results)
#     print(f' M=({np.mean(results[:, 0])}, {np.mean(results[:, 1])}), SD=({np.std(results[:, 0])}, {np.std(results[:, 1])})')


if __name__ == '__main__':
    '''
    python evaluate_FGD.py
    '''
    # calculate fgd per system
    exp_base(120)

    # calculate fgd per system per test speech sequence
    # exp_per_seq(120, stride=1)

genea_numerical_evaluations/FGD/evaluate_FGD.py
YoungSeng/UnifiedGesture
import glob
import os

import numpy as np
import torch
import torch.nn.functional as F
from torch import optim
from torch.utils.data import TensorDataset, DataLoader

from embedding_net import EmbeddingNet

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self, name, fmt=':f'):
        self.name = name
        self.fmt = fmt
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

    def __str__(self):
        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
        return fmtstr.format(**self.__dict__)


def train_iter(target_data, net, optim):
    # zero gradients
    optim.zero_grad()

    # reconstruction loss
    feat, recon_data = net(target_data)
    recon_loss = F.l1_loss(recon_data, target_data, reduction='none')
    recon_loss = torch.mean(recon_loss, dim=(1, 2))

    if True:  # use pose diff
        target_diff = target_data[:, 1:] - target_data[:, :-1]
        recon_diff = recon_data[:, 1:] - recon_data[:, :-1]
        recon_loss += torch.mean(F.l1_loss(recon_diff, target_diff, reduction='none'), dim=(1, 2))

    recon_loss = torch.sum(recon_loss)

    recon_loss.backward()
    optim.step()

    ret_dict = {'loss': recon_loss.item()}
    return ret_dict


def make_tensor(path, n_frames, stride=None):

    if os.path.isdir(path):
        files = glob.glob(os.path.join(path, '*.npy'))
    else:
        files = [path]

    samples = []
    stride = n_frames // 2 if stride is None else stride
    for file in files:
        data = np.load(file).reshape(-1, 33 * 3)
        for i in range(0, len(data) - n_frames, stride):
            sample = data[i:i+n_frames]
            samples.append(sample)

    return torch.Tensor(samples)


def main(n_frames):
    # dataset
    train_dataset = TensorDataset(make_tensor('/ceph/hdd/yangsc21/Python/My_3/GT_Gesture_npy/', n_frames, stride=10))
    train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True, drop_last=False)

    # train
    loss_meters = [AverageMeter('loss')]

    # interval params
    print_interval = int(len(train_loader) / 5)

    # init model and optimizer
    pose_dim = 33 * 3
    generator = EmbeddingNet(pose_dim, n_frames).to(device)
    gen_optimizer = optim.Adam(generator.parameters(), lr=0.0004, betas=(0.5, 0.999))

    # training
    for epoch in range(100):
        for iter_idx, target in enumerate(train_loader, 0):
            target = target[0]
            batch_size = target.size(0)
            target_vec = target.to(device)
            loss = train_iter(target_vec, generator, gen_optimizer)

            # loss values
            for loss_meter in loss_meters:
                name = loss_meter.name
                if name in loss:
                    loss_meter.update(loss[name], batch_size)

            # print training status
            if (iter_idx + 1) % print_interval == 0:
                print_summary = 'EP {} ({:3d}) | '.format(epoch, iter_idx + 1)
                for loss_meter in loss_meters:
                    if loss_meter.count > 0:
                        print_summary += '{}: {:.3f}, '.format(loss_meter.name, loss_meter.avg)
                        loss_meter.reset()
                print(print_summary)

    # save model
    gen_state_dict = generator.state_dict()
    save_name = f'output/model_checkpoint_{n_frames}.bin'
    torch.save({'pose_dim': pose_dim, 'n_frames': n_frames, 'gen_dict': gen_state_dict}, save_name)


if __name__ == '__main__':
    '''
    cd ./genea_numerical_evaluations/FGD
    python train_AE.py
    '''
    n_frames = 120  # 30, 60, 90
    main(n_frames)

genea_numerical_evaluations/FGD/train_AE.py
YoungSeng/UnifiedGesture
import librosa
import numpy as np
import pdb
import os
from  scipy.ndimage import gaussian_filter as G
from scipy.signal import argrelextrema

def calc_db(keypoints, name=''):
    keypoints = np.array(keypoints).reshape(-1, 33, 3)
    kinetic_vel = np.mean(np.sqrt(np.sum((keypoints[1:] - keypoints[:-1]) ** 2, axis=2)), axis=1)
    kinetic_vel = G(kinetic_vel, 5)
    motion_beats = argrelextrema(kinetic_vel, np.less)
    return motion_beats, len(kinetic_vel)


def BA(music_beats, motion_beats):
    ba = 0
    for bb in music_beats:
        ba +=  np.exp(-np.min((motion_beats[0] - bb)**2) / 2 / 9)
    return (ba / len(music_beats))


def cal_beat_alien(audio_path, motion_path):
    ba_scores = []
    for audio in os.listdir(audio_path):
        audio_wav = np.load(os.path.join(audio_path, audio))[0]
        audio_beat_time = librosa.onset.onset_detect(y=audio_wav, sr=16000, units='frames')
        for item in os.listdir(motion_path):
            if audio.split('.')[0] in item.split('.')[0]:
                joint3d = np.load(os.path.join(motion_path, item))
                break
        gesture_beats, length = calc_db(joint3d)
        ba_scores.append(BA(audio_beat_time, gesture_beats))
    return np.mean(ba_scores), len(ba_scores)


if __name__ == '__main__':
    '''
    python beat.py
    '''
    audio_path = '/ceph/hdd/yangsc21/Python/My_3/10_npy/'
    motion_path = '/ceph/hdd/yangsc21/Python/My_3/wZEGGS_npy'
    print(cal_beat_alien(audio_path, motion_path))

genea_numerical_evaluations/beat.py
YoungSeng/UnifiedGesture
import pdb

import numpy as np
import os
import sys
sys.path.append('../utils')
import BVH as BVH
import Animation as Animation

def bvh2positon(source_path, save_path):
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    for file in os.listdir(source_path):
        print(file)
        if file.endswith('.bvh'):
            bvh_file = os.path.join(source_path, file)
            anim, names, frametime = BVH.load(bvh_file)
            glb = Animation.positions_global(anim)
            # root_position = glb[:, names.index('Hips')].copy()
            np.save(os.path.join(save_path, file[:-4] + '.npy'), glb)


if __name__ == '__main__':
    '''
    cd genea_numerical_evaluations
    python bvh2position.py
    '''
    source_path = "/ceph/hdd/yangsc21/Python/My_3/proposed"
    savepath = "/ceph/hdd/yangsc21/Python/My_3/proposed" + '_npy'
    bvh2positon(source_path, savepath)


genea_numerical_evaluations/bvh2position.py
YoungSeng/UnifiedGesture
# -*- coding: utf-8 -*-
"""
Calculating CCA (Canonical Correlation Analysis) over the produced and ground truth gestures

@author: Taras Kucherenko
"""

import argparse
import glob
import os
import numpy as np

from cca import calculate_CCA_score, find_CCA_scaling_vectors


def shorten(arr_one, arr_two):
    """
    Make sure that two arrays have the same length
    Args:
        arr_one:  array one
        arr_two:  array two

    Returns:
        shortened versions of both arrays

    """
    min_len = min(arr_one.shape[0], arr_two.shape[0])
    arr_one = arr_one[:min_len]
    arr_two = arr_two[:min_len]

    return arr_one, arr_two


def save_result(lines, out_dir):
    """Write computed measure to CSV

      Args:
          lines:        list of strings to be written
          out_dir:      output directory
          measure:      used measure
    """

    # Make output directory
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    outname = os.path.join(out_dir, 'cca.csv')

    with open(outname, 'w') as out_file:
        out_file.writelines(lines)


# def evaluate_folder(cond_name, coords_dir):
def evaluate_folder(cond_dir, gt_dir):
    """
    Calculate numerical measure for the coordinates in the given folder
    Args:
        cond_name:   name of the condition / folder to evaluate
        coords_dir:  folder where all the data for the current model is stored

    Returns:
        nothing, prints out the metrics results

    """

    # cond_dir = os.path.join(coords_dir, cond_name)
    # gt_dir = os.path.join(coords_dir, "GT")

    generated_files = sorted(glob.glob(os.path.join(cond_dir, '*.npy')))

    gt_files = sorted(glob.glob(os.path.join(gt_dir, '*.npy')))

    # First - find the CCA scaling vectors using all the data for the given model
    all_predicted_frames = []
    all_ground_tr_frames = []
    for predicted_file, gt_file in zip(generated_files, gt_files):

        # read and flatten the predicted values
        predicted_coords = np.load(predicted_file)
        predicted_coords = np.reshape(predicted_coords, (predicted_coords.shape[0], -1))

        # read and flatten the ground truth values
        original_coords = np.load(gt_file)
        original_coords = np.reshape(original_coords, (original_coords.shape[0], -1))

        # make sure sequences have the same length
        predicted_coords, original_coords = shorten(predicted_coords, original_coords)

        if len(all_predicted_frames) != 0:
            all_predicted_frames = np.concatenate((all_predicted_frames, predicted_coords), axis=0)
            all_ground_tr_frames = np.concatenate((all_ground_tr_frames, original_coords), axis=0)
        else:
            all_predicted_frames = predicted_coords
            all_ground_tr_frames = original_coords

    # find CCA models
    cca_model = find_CCA_scaling_vectors(all_predicted_frames, all_ground_tr_frames)

    # calculate Global CCA value
    global_cca_value = calculate_CCA_score(all_predicted_frames, all_ground_tr_frames, cca_model)
    print('Global CCA value: {:.5f}'.format(global_cca_value))


    # calculate CCA value for each sequence
    predicted_out_lines = [','.join(['file']) + '\n']

    predicted_errors = []
    for predicted_file, gt_file in zip(generated_files, gt_files):
        # read and flatten the predicted values
        predicted_coords = np.load(predicted_file)
        predicted_coords = np.reshape(predicted_coords, (predicted_coords.shape[0], -1))

        # read and flatten the ground truth values
        original_coords = np.load(gt_file)
        original_coords = np.reshape(original_coords, (original_coords.shape[0], -1))

        # make sure sequences have the same length
        predicted_coords, original_coords = shorten(predicted_coords, original_coords)

        # calculate CCA value
        current_cca_value = calculate_CCA_score(original_coords, predicted_coords, cca_model)

        predicted_errors.append(current_cca_value)

        basename = os.path.basename(predicted_file)
        predicted_line = basename

        predicted_line += ',' + str(current_cca_value) + '\n'

        predicted_out_lines.append(predicted_line)

    predicted_average_line = 'Average'
    error_avgs = np.mean(predicted_errors, axis=0)
    error_stds = np.std(predicted_errors, axis=0)

    predicted_average_line += ',' + str(error_avgs)

    predicted_out_lines.append(predicted_average_line)

    # predicted_out_dir = os.path.join("result", cond_name)
    #
    # save_result(predicted_out_lines, predicted_out_dir)
    #
    # print('{:s}: {:.2f} +- {:.2F}'.format(cond_name, np.mean(predicted_errors), error_stds))
    print('{:.2f} +- {:.2F}'.format(np.mean(predicted_errors), error_stds))


if __name__ == '__main__':
    '''
    python calc_cca.py --cond_dir "<..your path/GENEA/genea_challenge_2022/baselines/Tri/output/infer_sample/output_2_new_wavlm/npy/>" --gt_dir "<..your path/GENEA/genea_challenge_2022/dataset/v1_18_1/val/npy/>"
    '''
    parser = argparse.ArgumentParser(
        description='Calculate prediction errors')
    # parser.add_argument('--coords_dir', '-c', default='data',
    #                     help='Predicted gesture directory')
    parser.add_argument('--cond_dir', '-c', default='data',
                        help='Predicted gesture directory')
    parser.add_argument('--gt_dir', '-g', default='data',
                        help='Predicted gesture directory')
    args = parser.parse_args()

    # Make sure that data is stored in the correct folder
    # if not os.listdir(args.coords_dir):
    #     print("--coords_dir argument is wrong. there is no data at the folder '", args.coords_dir, "'")
    #     exit(-1)

    print('CCA:')

    # for cond_name in os.listdir(args.coords_dir):
    #     evaluate_folder(cond_name, args.coords_dir)

    evaluate_folder(args.cond_dir, args.gt_dir)

    '''
    CCA 
    output_2_new_wavlm  0.97 +- 0.01
    '''

genea_numerical_evaluations/calc_cca.py
YoungSeng/UnifiedGesture
# -*- coding: utf-8 -*-
"""
Calculating statistics over the produced and ground truth gestures

@author: Naoshi Kaneko and Taras Kucherenko
"""

import argparse
import glob
import os
import pdb
import warnings

import matplotlib.pyplot as plt
import numpy as np



def compute_velocity(data, dim=3):
    """Compute velocity between adjacent frames

      Args:
          data:         array containing joint positions of gesture
          dim:          gesture dimensionality

      Returns:
          vel_norms:    velocities of each joint between each adjacent frame
    """

    # Flatten the array of 3d coords
    coords = data.reshape(data.shape[0], -1)

    # First derivative of position is velocity
    vels = np.diff(coords, n=1, axis=0)

    num_vels = vels.shape[0]
    num_joints = vels.shape[1] // dim

    vel_norms = np.zeros((num_vels, num_joints))

    # calculate vector norms
    for i in range(num_vels):
        for j in range(num_joints):
            x1 = j * dim + 0
            x2 = j * dim + dim
            vel_norms[i, j] = np.linalg.norm(vels[i, x1:x2])

    # multiply on 20 to compensate for the time-frame being  0.05s
    return vel_norms * 20


def compute_acceleration(data, dim=3):
    """Compute acceleration between adjacent frames

      Args:
          data:         array containing joint positions of gesture
          dim:          gesture dimensionality

      Returns:
          acc_norms:    accelerations of each joint between each adjacent frame
    """

    # Second derivative of position is acceleration
    accs = np.diff(data, n=2, axis=0)

    num_accs = accs.shape[0]
    num_joints = accs.shape[1] // dim

    acc_norms = np.zeros((num_accs, num_joints))

    # calculate vector norms
    for i in range(num_accs):
        for j in range(num_joints):
            x1 = j * dim + 0
            x2 = j * dim + dim
            acc_norms[i, j] = np.linalg.norm(accs[i, x1:x2])

    # multiply on 20 to compensate for the time-frame being  0.05s
    return acc_norms * 20 * 20


def save_result(lines, out_dir, width, measure):
    """Write computed histogram to CSV

      Args:
          lines:        list of strings to be written
          out_dir:      output directory
          width:        bin width of the histogram
          measure:      used measure for histogram calculation
    """

    # Make output directory
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    hist_type = measure[:3]  # 'vel' or 'acc'
    filename = 'hmd_{}_{}.csv'.format(hist_type, width)
    outname = os.path.join(out_dir, filename)

    with open(outname, 'w') as out_file:
        out_file.writelines(lines)


def make_histogram(cond_name, measure, coord_dir_, out_dir, visualize=True, width=1):
    """
    Calculate frequencies histogram for a given measure
    Args:
        cond_name:   name of the folder to consider
        measure:     measure to be used
        coord_dir:   folder where all the data for the current model is stored
        out_dir:     folder where the results should be stored
        visualize:   flag if we want to visualize the result
        width:       width of the histogram bins

    Returns:
        Saved in the "result" folder

    """

    # coord_dir = "<..your path/GENEA/genea_challenge_2022/dataset/v1_18_1/val/>"
    # cond_name = 'npy'

    cond_name = coord_dir_.split('/')[-2]
    coord_dir = '/'.join(coord_dir_.split('/')[:-2])
    predicted_dir = os.path.join(coord_dir, cond_name)

    measures = {
        'velocity': compute_velocity,
        'acceleration': compute_acceleration,
    }

    # Check if error measure was correct
    if measure not in measures:
        raise ValueError('Unknown measure: \'{}\'. Choose from {}'
                         ''.format(measure, list(measures.keys())))

    predicted_files = sorted(glob.glob(os.path.join(predicted_dir, '*.npy')))

    predicted_out_lines = [','.join([''] + ['Total']) + '\n']

    predicted_distances = []
    for predicted_file in predicted_files:
        # read the values and remove hips which are fixed
        predicted_coords = np.load(predicted_file)[:, 2:]

        predicted_distance = measures[measure](predicted_coords)  # [:, selected_joints]

        predicted_distances.append(predicted_distance)

    predicted_distances = np.concatenate(predicted_distances)

    # Compute histogram for each joint
    bins = np.arange(0, 59 + width, width)
    num_joints = predicted_distances.shape[1]
    predicted_hists = []
    for i in range(num_joints):
        predicted_hist, _ = np.histogram(predicted_distances[:, i], bins=bins)

        predicted_hists.append(predicted_hist)

    # Sum over all joints
    predicted_total = np.sum(predicted_hists, axis=0)

    # Append total number of bin counts to the last
    predicted_hists = np.stack(predicted_hists + [predicted_total], axis=1)

    num_bins = bins.size - 1
    for i in range(num_bins):
        predicted_line = str(bins[i])
        for j in range(num_joints + 1):
            predicted_line += ',' + str(predicted_hists[i, j])
        predicted_line += '\n'
        predicted_out_lines.append(predicted_line)

    out_dir = os.path.join('/ceph/hdd/yangsc21/Python/My_3/', 'result')
    predicted_out_dir = os.path.join(out_dir, cond_name)

    # predicted_out_dir = out_dir

    if visualize:
        import matplotlib
        matplotlib.use('AGG')

        sum = np.sum(predicted_total)
        actual_freq = predicted_total / sum
        plt.plot(bins[:-1], actual_freq, label=cond_name)
        plt.legend()
        plt.xlabel('Velocity (cm/s)')
        plt.ylabel('Frequency')
        plt.title('Frequencies of Moving Distance ({})'.format(measure))
        plt.tight_layout()
        plt.savefig('2.jpg')
        plt.show()

    save_result(predicted_out_lines, predicted_out_dir, width, measure)

    print('HMD ({}):'.format(measure))
    print('bins: {}'.format(bins))
    print('population: {}'.format(predicted_total))

def main():
    '''
    python calc_histogram.py
    python calc_histogram.py -p your_prediction_dir -m velocity # You can change the bin width of the histogram
    '''
    parser = argparse.ArgumentParser(
        description='Calculate histograms of moving distances')
    parser.add_argument('--coords_dir', '-p', default='data',
                        help='Predicted gestures directory')
    parser.add_argument('--width', '-w', type=float, default=1,
                        help='Bin width of the histogram')
    parser.add_argument('--measure', '-m', default='acceleration',      # acceleration/velocity
                        help='Measure to calculate (velocity or acceleration)')
    parser.add_argument('--select', '-s', nargs='+',
                        help='Joint subset to compute (if omitted, use all)')
    parser.add_argument('--visualize', '-v', action='store_true',
                        help='Visualize histograms')
    parser.add_argument('--out_dir', default='resulting',
                        help='Directory to output the result')
    args = parser.parse_args()

    # Make sure that data is stored in the correct folder
    # if not os.listdir(args.coords_dir):
    #     print("--coords_dir argument is wrong. there is no data at the folder '", args.coords_dir, "'")
    #     exit(-1)
    #
    # for cond_name in os.listdir(args.coords_dir):
    #     print("\nConsider", cond_name)
    #     make_histogram(cond_name, args.measure, args.coords_dir, args.out_dir, args.visualize, args.width)

    make_histogram(None, args.measure, args.coords_dir, args.out_dir, True, args.width)

    print('\nMore detailed result was writen to the files in the ',  args.out_dir, ' folder ')
    print('')

if __name__ == '__main__':
    main()
    '''
    HMD (velocity):
bins: [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47
 48 49 50 51 52 53 54 55 56 57 58 59]
population: [625273 320410 205760 155278 124341 102980  86938  70758  65529  53843
  48623  45588  35938  37620  28903  31020  25701  25250  21911  20780
  19008  18130  16544  16517  13187  14551  11563  12837  10702  11117
   9425   9779   8670   9162   7326   8170   6646   7215   6046   6602
   5633   6001   5100   5238   4769   5030   4246   4717   3826   4341
   3946   3797   3502   3520   3311   3234   3046   3110   2797]
    
    HMD (acceleration):
bins: [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47
 48 49 50 51 52 53 54 55 56 57 58 59]
population: [227618  45027  37958  34907  25682  19830  16445  54049  27785  13353
  10919  10299  16592  20684  10678  19987  13479   7893   7355  10320
  13200   7920   7246  14707  16362  13671   8736   7863   6880   7136
   7503  26425   7367  17333   8120   9997  11934  31960  19732 123117
  38035  17649  13405  10394   9272   8852  18737  16649   8916  10373
  18598   9610   8088   7064  11053  11297  24689   6829   6914]

    '''
genea_numerical_evaluations/calc_histogram.py
YoungSeng/UnifiedGesture
# -*- coding: utf-8 -*-
"""
Calculating average jerk over the produced and ground truth gestures

@author: Taras Kucherenko and Naoshi Kaneko
"""

import argparse
import glob
import os
import pdb
import warnings

import numpy as np


def compute_jerks(data, dim=3):
    """Compute jerk between adjacent frames

      Args:
          data:         array containing joint positions of gesture
          dim:          gesture dimensionality

      Returns:
          np.ndarray:   jerks of each joint averaged over all frames
    """

    # Third derivative of position is jerk
    jerks = np.diff(data, n=3, axis=0)

    num_jerks = jerks.shape[0]
    num_joints = jerks.shape[1] // dim

    jerk_norms = np.zeros((num_jerks, num_joints))

    for i in range(num_jerks):
        for j in range(num_joints):
            x1 = j * dim + 0
            x2 = j * dim + dim
            jerk_norms[i, j] = np.linalg.norm(jerks[i, x1:x2])

    average = np.mean(jerk_norms, axis=0)

    # Take into account that frame rate was 20 fps
    scaled_av = average * 20 * 20 * 20

    return scaled_av


def compute_acceleration(data, dim=3):
    """Compute acceleration between adjacent frames

      Args:
          data:         array containing joint positions of gesture
          dim:          gesture dimensionality

      Returns:
          np.ndarray:   accelerations of each joint averaged over all frames
    """

    # Second derivative of position is acceleration
    accs = np.diff(data, n=2, axis=0)

    num_accs = accs.shape[0]
    num_joints = accs.shape[1] // dim

    acc_norms = np.zeros((num_accs, num_joints))

    for i in range(num_accs):
        for j in range(num_joints):
            x1 = j * dim + 0
            x2 = j * dim + dim
            acc_norms[i, j] = np.linalg.norm(accs[i, x1:x2])

    average = np.mean(acc_norms, axis=0)

    # Take into account that frame rate was 20 fps
    scaled_av = average * 20 * 20

    return scaled_av


def save_result(lines, out_dir, measure):
    """Write computed measure to CSV

      Args:
          lines:        list of strings to be written
          out_dir:      output directory
          measure:      used measure
    """

    # Make output directory
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    if measure == "jerk":
        outname = os.path.join(out_dir, 'aj.csv')
    elif measure == "acceleration":
        outname = os.path.join(out_dir, 'aa.csv')

    with open(outname, 'w') as out_file:
        out_file.writelines(lines)


def evaluate_folder(cond_name, coord_dir, measure):
    """
    Calculate numerical measure for the coordinates in the given folder
    Args:
        cond_name:   name of the condition / folder to evaluate
        coord_dir:   folder where all the data for the current model is stored
        measure:     measure to be used

    Returns:
        nothing, prints out the metrics results

    """

    # cond_dir = os.path.join(coord_dir, cond_name)

    cond_dir = coord_dir

    cond_files = sorted(glob.glob(os.path.join(cond_dir, '*.npy')))

    # define possible measures
    measures = {
        'jerk': compute_jerks,
        'acceleration': compute_acceleration,
    }

    # Check if error measure was correct
    if measure not in measures:
        raise ValueError('Unknown measure: \'{}\'. Choose from {}'
                         ''.format(measure, list(measures.keys())))

    predicted_out_lines = [','.join(['file']) + '\n']

    all_motion_stats = []
    for predicted_file in cond_files:
        predicted_coords = np.load(predicted_file)

        # flatten the values
        predicted_coords = np.reshape(predicted_coords, (predicted_coords.shape[0], -1))

        current_motion_stats = measures[measure](predicted_coords)

        all_motion_stats.append(current_motion_stats)

        basename = os.path.basename(predicted_file)
        predicted_line = basename

        for ov in current_motion_stats:
            predicted_line += ',' + str(ov)

        predicted_line += '\n'

        predicted_out_lines.append(predicted_line)

    predicted_average_line = 'Average'

    avgs_for_each_joint = np.mean(all_motion_stats, axis=0)

    avgs_for_each_sequence = np.mean(all_motion_stats, axis=1)

    stds_over_sequences = np.std(avgs_for_each_sequence)
    mean_over_sequences = np.mean(avgs_for_each_sequence)

    for oa in avgs_for_each_joint:
        predicted_average_line += ',' + str(oa)

    predicted_out_lines.append(predicted_average_line)

    # predicted_out_dir = os.path.join("result", cond_name)

    # save_result(predicted_out_lines, predicted_out_dir, measure)

    print('{:.2f} +- {:.2F}'.format(mean_over_sequences, stds_over_sequences))


if __name__ == '__main__':
    '''
    python calk_jerk_or_acceleration.py -m jerk --coords_dir "<..your path/GENEA/genea_challenge_2022/baselines/Tri/output/infer_sample/output_2_new_wavlm/npy/>"
    '''
    parser = argparse.ArgumentParser(
        description='Calculate prediction errors')
    parser.add_argument('--coords_dir', '-c', default='data',
                        help='Predicted gesture directory')
    parser.add_argument('--measure', '-m', default='acceleration',
                        help='Measure to calculate (jerk or acceleration)')
    args = parser.parse_args()

    # Make sure that data is stored in the correct folder
    if not os.listdir(args.coords_dir):
        print("--coords_dir argument is wrong. there is no data at the folder '", args.coords_dir, "'")
        exit(-1)

    if args.measure == 'jerk':
        print('AJ:')
    elif args.measure == 'acceleration':
        print('AA:')

    # for cond_name in os.listdir(args.coords_dir):
    #     if cond_name == "GT":
    #         continue
    evaluate_folder(None, args.coords_dir, args.measure)

    print('More detailed result was writen to the files in the "result" folder ')
    print('')

    '''
    GT  AJ  6184.28 +- 2235.72 cm/s3    AA  198.10 +- 69.12 cm/s2
    output_2_new_wavlm  AJ 550.74 +- 297.20      AA  38.76 +- 14.76
    
    '''

genea_numerical_evaluations/calk_jerk_or_acceleration.py
YoungSeng/UnifiedGesture
# -*- coding: utf-8 -*-
"""
Created on Mon Mar  2 11:30:40 2020

@authors: Naoshi Kaneko, Taras Kucherenko
"""

import argparse
import glob
import os
import pdb
import re

import matplotlib
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from matplotlib.patches import Rectangle
import numpy as np
import pandas as pd
import seaborn as sns


def normalize(hist):
    return hist / np.sum(hist)


def hellinger(hist1, hist2):
    """Compute Hellinger distance between two histograms

      Args:
          hist1:        first histogram
          hist2:        second histogram of the same size as hist1

      Returns:
          float:        Hellinger distance between hist1 and hist2
    """

    return np.sqrt(1.0 - np.sum(np.sqrt(normalize(hist1) * normalize(hist2))))


# https://stackoverflow.com/questions/4836710/does-python-have-a-built-in-function-for-string-natural-sort  # NOQA
def natural_sort_key(s, _nsre=re.compile('([0-9]+)')):
    return [int(text) if text.isdigit() else text.lower()
            for text in _nsre.split(s)]


def natural_sort(l, key=natural_sort_key):
    return sorted(l, key=key)


def main():
    '''
    python calc_histogram.py --original result/GT --predicted result/ --file hmd_vel_1.csv
    :return:
    '''
    parser = argparse.ArgumentParser(
        description='Calculate histograms of moving distances')
    parser.add_argument('--original', default="/ceph/hdd/yangsc21/Python/My_3/result/GT_Gesture_npy/",
                        help='Original gesture directory')
    parser.add_argument('--predicted', '-p', default="/ceph/hdd/yangsc21/Python/My_3/result/Audio2Gesture-20fps_npy/",
                        help='Predicted gesture directory')
    parser.add_argument('--file', '-f', default='hmd_vel_0.05.csv',
                        help='File name to load')
    parser.add_argument('--select', '-s', nargs='+',
                        help='Joint subset to compute (if omitted, use all)')
    parser.add_argument('--visualize', '-v', action='store_true',
                        help='Visualize histograms')
    parser.add_argument('--out', '-o', default='results',
                        help='Directory to output the result')
    args = parser.parse_args()


    def get_directories(directory):
        return sorted(filter(lambda x: os.path.isdir(x), glob.glob(directory)))

    def get_histograms(data_dir, hist_file):

        # Read original gesture's distribution
        hist_path = os.path.join(data_dir, hist_file)
        original_val = pd.read_csv(hist_path, header=None, skiprows=1)
        original_array = np.array(original_val)

        # Calculate histograms for wrists and normalize it
        actual_hist = (original_array[:, -2] + original_array[:, -5]) / original_array[:, -1]

        return actual_hist

    original_hist = get_histograms(args.original, args.file)

    # List of predicted gesture directories
    # predicted_dirs = get_directories(os.path.join(args.predicted, '*'))
    predicted_dirs = get_directories(args.predicted)

    results = {os.path.basename(d): None for d in predicted_dirs}

    # Iterate over the list of directories
    for predicted_dir in predicted_dirs:
        # Does this directory have a target file?
        try:
            predicted_hist = get_histograms(predicted_dir, args.file)
        except FileNotFoundError:
            # Are there any subdirectories which have integer names?
            sub_dirs = sorted(
                filter(lambda x: os.path.basename(x).isdecimal(),
                get_directories(os.path.join(predicted_dir, '*'))))

            # If no, raise an exception
            if not sub_dirs:
                raise FileNotFoundError(
                    'There is neither ' + args.file
                    + ' nor subdirectories in ' + predicted_dir)

            predicted = None
            for sub_dir in sub_dirs:
                predicted_file = os.path.join(sub_dir, args.file)
                tmp = pd.read_csv(predicted_file, index_col=0)

                if predicted is None:
                    predicted = tmp
                else:
                    predicted = predicted + tmp
                
            predicted_hist = predicted / float(len(sub_dirs))

        assert len(original_hist) == len(predicted_hist)

        # Hellinger distance between two histograms
        dist = hellinger(original_hist, predicted_hist)

        # Store results
        key = os.path.basename(predicted_dir)
        results[key] = {'dist': dist, 'hist': predicted_hist}

    # Print and save results
    keys = natural_sort(results.keys())

    result_str = ['Hellinger distances:']
    for key in keys:
        result_str.append('\t{}: {}'.format(key, results[key]['dist']))
    
    result_str = '\n'.join(result_str)
    
    print(result_str)
    print('')

    # Make output directory
    out = os.path.join(args.out, os.path.basename(args.predicted))
    if not os.path.isdir(out):
        os.makedirs(out)
    
    with open(os.path.join(out, 'distances.txt'), 'w') as f:
        f.write(result_str)

    if args.visualize:
        # Set color and style
        mpl_default = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728',
                       '#9467bd', '#8c564b', '#e377c2', '#7f7f7f',
                       '#bcbd22', '#17becf']
        sns.set(context='poster', palette=sns.color_palette(mpl_default), font_scale=1.05)
        sns.set_style('white', {'legend.frameon':True})

        index = original_hist
        bins = [format(i, '.2f') for i in list(index)]

        fig = plt.figure(figsize=(8, 5))
        ax = fig.add_subplot(111)

        # Convert frequency to percentage
        gt_handle, = ax.plot(bins, normalize(original_hist) * 100, color='C4')

        # Awesome way to create a tabular-style legend
        # https://stackoverflow.com/questions/25830780/tabular-legend-layout-for-matplotlib
        # Create a blank rectangle
        blank = Rectangle((0, 0), 1, 1, fc="w", fill=False, edgecolor='none', linewidth=0)

        # Correspond to each columns of the tabular
        legend_handles = [blank, gt_handle]
        legend_names = ['Name', 'Ground Truth']
        legend_dists = ['Hell. Dist.', '0'.center(16)]

        colors = ['C1', 'C3', 'C0', 'C2'] if len(keys) <= 4 else \
                 ['C1', 'C0', 'C6', 'C7', 'C8', 'C9', 'C5', 'C2', 'C3', 'C4']
        
        assert len(keys) <= len(colors)

        for color, key in zip(colors, keys):
            predicted_hist = results[key]['hist'][:-4]
            label = key

            if 'Aud2Pose' in label:
                label += ' [14]'

            handle, = ax.plot(bins, normalize(predicted_hist) * 100, color=color)

            legend_handles.append(handle)
            legend_names.append(label)
            legend_dists.append('{:.3f}'.format(results[key]['dist']).center(12))

        # Legend will have a tabular of (rows x 3)
        rows = len(legend_handles)
        empty_label = ['']

        legend_handles = legend_handles + [blank] * (rows * 2)
        legend_labels = np.concatenate([empty_label * rows, legend_names, legend_dists])

        ax.legend(legend_handles, legend_labels,
                  ncol=3, handletextpad=0.5, columnspacing=-2.15,
                  labelspacing=0.35)
        ax.set_xlabel('Velocity (cm/s)')
        ax.set_ylabel('Frequency (%)')
        ax.set_xticks(np.arange(16))
        ax.tick_params(pad=6)
        ax.yaxis.set_major_locator(
            MaxNLocator(nbins='auto', steps=[1, 2, 2.5, 5, 10], integer=True))

        plt.subplots_adjust(left=0.09, right=0.98, top=0.98, bottom=0.12)
        plt.savefig(os.path.join(out, 'velocity_histogram.pdf'))
        plt.show()
    
    print('Results were writen in ' + out)
    print('')


if __name__ == '__main__':
    main()

genea_numerical_evaluations/hellinger_distance.py
YoungSeng/UnifiedGesture
import os
import pdb

import numpy as np
import yaml
from pprint import pprint
from easydict import EasyDict
import torch
from configs.parse_args import parse_args
import glob
import sys
[sys.path.append(i) for i in ['.', './conditional_gpt']]
from Hierarchical_XTransformer import Hierarchical_XTransformer, cascade_XTransformer, XTransformer_GRU, XTransformer_GPT, GPT_GRU
import torch.nn as nn
import math

args = parse_args()
mydevice = torch.device('cuda:' + args.gpu)

with open(args.config) as f:
    config = yaml.safe_load(f)

for k, v in vars(args).items():
    config[k] = v
# pprint(config)

config = EasyDict(config)


def main(args, wav_code, model, save_path, prefix=None, max_codes=None):
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    clip_length = wav_code.shape[0]
    unit_time = int(75 * args.n_poses / args.motion_resampling_framerate)        # 4 * 75
    if clip_length < unit_time:
        num_subdivision = 1
    else:
        num_subdivision = math.ceil((clip_length - unit_time) / unit_time) + 1

    if max_codes is not None and num_subdivision >= int(max_codes / args.n_poses):
        num_subdivision = int(max_codes / args.n_poses)

    print('num_subdivision: {}, unit_time: {}, clip_length: {}'.format(num_subdivision, unit_time, clip_length))

    begin_token_1 = torch.tensor(torch.randint(0, 512, (1, 1))).to(mydevice)
    begin_token_2 = torch.tensor(torch.randint(0, 512, (1, 1))).to(mydevice)

    with torch.no_grad():
        code = []
        motion = []
        for i in range(0, num_subdivision):
            if i == 0:
                xs_1 = torch.randint(1, 512, (1, 1)).to(mydevice)
                hidden = None
            else:
                xs_1 = out_zs[:, -1:]
            start_time = i * unit_time

            # prepare pose input
            audio_start = math.floor(start_time)
            audio_end = audio_start + unit_time
            in_audio = wav_code[audio_start:audio_end]
            if len(in_audio) < unit_time:
                in_audio = np.pad(in_audio, [(0, unit_time - len(in_audio)), (0, 0)], mode='constant')

            in_audio = torch.tensor(in_audio).unsqueeze(0).to(mydevice)
            # in_audio = in_audio.reshape(in_audio.shape[0], -1)
            # out_zs = model.module.generate(seq_in=in_audio, seq_out_start=begin_token, seq_len=args.n_poses * 3)

            # out_zs = model.module.generate(in_audio, begin_token_1, begin_token_2)
            # out_zs = model.module.generate(in_audio)

            out_zs, out_lower, hidden = model.module.generate(in_audio, xs_1, hidden)
            code.append(out_zs[0].squeeze(0).data.cpu().numpy())
            motion.append(out_lower[0].squeeze(0).data.cpu().numpy())

    out_code = np.vstack(code).flatten()
    lower_motion = np.vstack(motion)
    lower_motion = np.expand_dims(lower_motion, 0).transpose((0, 2, 1))
    print(out_code.shape, lower_motion.shape)
    np.save(os.path.join(save_path, 'code_upper_' + prefix + '.npy'), out_code)
    np.save(os.path.join(save_path, 'motion_lower_' + prefix + '.npy'), lower_motion)


if __name__ == '__main__':
    '''
    python inference.py --config=./configs/pretrain.yml --train --no_cuda 3 --gpu 3
    '''
    wav_code_path = './dataset/Trinity/audio_clips/Recording_006.npy'
    model_path = './result/GPT_GRU_block_size_again/XTransformer_checkpoint_250.bin'
    save_path = './result/inference'

    # wav_code = np.random.randint(0, 1024, size=(580, 2))
    wav_code = np.load(wav_code_path)
    print('wav code shape: ', wav_code.shape)

    # model = Hierarchical_XTransformer(batch_size=1, n_poses=15, mydevice=mydevice)
    # model = cascade_XTransformer(batch_size=1, n_poses=15, mydevice=mydevice)
    # model = XTransformer_GRU(batch_size=1, n_poses=15, mydevice=mydevice)
    # model = XTransformer_GPT(batch_size=1, n_poses=15, mydevice=mydevice)

    model = GPT_GRU(batch_size=1, n_poses=15, mydevice=mydevice)
    model = nn.DataParallel(model, device_ids=[eval(i) for i in config.no_cuda])
    model = model.to(mydevice)
    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
    model.load_state_dict(checkpoint['model_dict'])
    model = model.eval()

    name = os.path.split(wav_code_path)[1][:-4]
    main(config, wav_code, model, save_path, prefix=name, max_codes=15/4*32)     # 32s

inference.py
YoungSeng/UnifiedGesture
'''
Download from  https://pan.baidu.com/s/1z1cQiqLUgjfxlWoajIPr0g (ye1q)

cd ./deep-motion-editing/retargeting/datasets/
tar -jxvf test_set.tar.bz2
cd ..
python demo.py

python test.py
Collecting test error...
Intra-retargeting error: 0.00047083793936272735
Cross-retargeting error: 0.002252287446935629
Evaluation finished!

Download from https://pan.baidu.com/s/1PM0maLNCJxBsR9RIv18cTA (4rgv)
cd ./datasets/
tar -jxvf training_set.tar.bz2

# git filter-branch --tree-filter 'rm -f ./deep-motion-editing/retargeting/datasets/training_set.tar.bz2' --tag-name-filter cat -- --all


Download from https://www.mixamo.com/ Silly Dancing and Hip Hop Dancing
# 确保blender可执行程序所在目录添加到系统环境变量path, https://zhuanlan.zhihu.com/p/525475118
# windows
blender -b -P fbx2bvh.py

cd ..
python ./datasets/preprocess.py
((135, 75), (116, 75))
python train.py --save_dir=./training/
# windows
pip install tensorboard
tensorboard --logdir=./logs/
'''

import pdb

import torch
from vector_quantize_pytorch import ResidualVQ

# python main.py
residual_vq = ResidualVQ(
    dim = 7 * 16,
    codebook_size = 512,
    num_quantizers = 1,
    kmeans_init = True,   # set to True
    kmeans_iters = 10     # number of kmeans iterations to calculate the centroids for the codebook on init
)

x = torch.randn(2, 11, 7 * 16)
quantized, indices, commit_loss = residual_vq(x)

pdb.set_trace()

main.py
YoungSeng/UnifiedGesture

import argparse
import os
import glob
import pdb
from pathlib import Path

# import librosa
import numpy as np
import lmdb
import pyarrow


style2onehot = {
'Happy':[1, 0, 0, 0, 0, 0, 0],
'Sad':[0, 1, 0, 0, 0, 0, 0],
'Neutral':[0, 0, 1, 0, 0, 0, 0],
'Old':[0, 0, 0, 1, 0, 0, 0],
'Angry':[0, 0, 0, 0, 1, 0, 0],
'Relaxed':[0, 0, 0, 0, 0, 1, 0],
'Still':[0, 0, 0, 0, 0, 0, 1]
}

def make_lmdb_latent_dataset(base_path):
    Trinity_path = os.path.join(base_path, 'Trinity')
    ZEGGS_path = os.path.join(base_path, 'ZEGGS')
    out_path = os.path.join(base_path, 'all_lmdb_aux')
    if not os.path.exists(out_path):
        os.makedirs(out_path)

    map_size = 1024 * 20  # in MB
    map_size <<= 20  # in B
    db = [lmdb.open(os.path.join(out_path, 'lmdb_train'), map_size=map_size),
          lmdb.open(os.path.join(out_path, 'lmdb_test'), map_size=map_size)]

    # delete existing files
    for i in range(2):
        with db[i].begin(write=True) as txn:
            txn.drop(db[i].open_db())

    total_files = sorted(glob.glob(Trinity_path + "/all_gesture_aux" + "/*_upper.npy") + glob.glob(ZEGGS_path + "/all_gesture_aux" + "/*_upper.npy"))

    save_idx = 0
    for v_i, code_file in enumerate(total_files):
        path = os.path.split(code_file)[0]
        name = os.path.split(code_file)[1][:-4].replace('_upper', '')
        style = name.split('_')[1]

        print(name)
        if style in style2onehot:
            print(style)
            style = style2onehot[style]
        else:
            style = style2onehot['Neutral']

        # load
        upper_code = np.load(code_file)
        full_body = np.load(os.path.join(path, name + '_all_body.npy'))
        audio_wav = np.load(os.path.join(path.replace('all_gesture_aux', 'all_speech'), name) + '.npy')

        MIN_LEN = min(full_body.shape[0], int(audio_wav.shape[-1]*30/16000/4))
        full_body = full_body[:MIN_LEN]
        audio_wav = audio_wav[:int(MIN_LEN*16000*4/30)]

        # process
        print(upper_code.shape, full_body.shape, audio_wav.shape)

        # process
        clips = [{'vid': name, 'clips': []},  # train
                 {'vid': name, 'clips': []}]  # validation

        # split
        if save_idx % 10 == 0:
            dataset_idx = 1  # validation
        else:
            dataset_idx = 0  # train

        clips[dataset_idx]['clips'].append(
            {
              'audio_wav': audio_wav,
              'full_body': full_body,
              'upper_code': upper_code,
              'style': np.array(style),
             })

        # write to db
        for i in range(2):
            with db[i].begin(write=True) as txn:
                if len(clips[i]['clips']) > 0:
                    k = '{:010}'.format(v_i).encode('ascii')
                    v = pyarrow.serialize(clips[i]).to_buffer()
                    txn.put(k, v)
        save_idx += 1

    # close db
    for i in range(2):
        db[i].sync()
        db[i].close()


if __name__ == '__main__':
    '''
    python ./make_lmdb.py --base_path ./dataset/
    '''
    parser = argparse.ArgumentParser()
    parser.add_argument("--base_path", type=Path)
    args = parser.parse_args()

    make_lmdb_latent_dataset(args.base_path)

make_lmdb.py
YoungSeng/UnifiedGesture
from local_attention.local_attention import LocalAttention
from local_attention.transformer import LocalTransformer, LocalMHA

model/local_attention/__init__.py
YoungSeng/UnifiedGesture
import math

import torch
from torch import nn, einsum
import torch.nn.functional as F

from einops import rearrange, repeat, pack, unpack

from local_attention.rotary import SinusoidalEmbeddings, apply_rotary_pos_emb

# constant

TOKEN_SELF_ATTN_VALUE = -5e4

# helper functions

def exists(val):
    return val is not None

def default(value, d):
    return d if not exists(value) else value

def to(t):
    return {'device': t.device, 'dtype': t.dtype}

def max_neg_value(tensor):
    return -torch.finfo(tensor.dtype).max

def l2norm(tensor):
    dtype = tensor.dtype
    normed = F.normalize(tensor, dim = -1)
    return normed.type(dtype)

def pad_to_multiple(tensor, multiple, dim=-1, value=0):
    seqlen = tensor.shape[dim]
    m = seqlen / multiple
    if m.is_integer():
        return False, tensor
    remainder = math.ceil(m) * multiple - seqlen
    pad_offset = (0,) * (-1 - dim) * 2
    return True, F.pad(tensor, (*pad_offset, 0, remainder), value = value)

def look_around(x, backward = 1, forward = 0, pad_value = -1, dim = 2):
    t = x.shape[1]
    dims = (len(x.shape) - dim) * (0, 0)
    padded_x = F.pad(x, (*dims, backward, forward), value = pad_value)
    tensors = [padded_x[:, ind:(ind + t), ...] for ind in range(forward + backward + 1)]
    return torch.cat(tensors, dim = dim)

# main class

class LocalAttention(nn.Module):
    def __init__(
        self,
        window_size,
        causal = False,
        look_backward = 1,
        look_forward = None,
        dropout = 0.,
        shared_qk = False,
        rel_pos_emb_config = None,
        dim = None,
        autopad = False,
        exact_windowsize = False
    ):
        super().__init__()
        look_forward = default(look_forward, 0 if causal else 1)
        assert not (causal and look_forward > 0), 'you cannot look forward if causal'

        self.window_size = window_size
        self.autopad = autopad
        self.exact_windowsize = exact_windowsize

        self.causal = causal

        self.look_backward = look_backward
        self.look_forward = look_forward

        self.dropout = nn.Dropout(dropout)

        self.shared_qk = shared_qk

        # relative positions

        # self.rel_pos = None
        # if exists(rel_pos_emb_config) or exists(dim):  # backwards compatible with old `rel_pos_emb_config` deprecated argument
        #     if exists(rel_pos_emb_config):
        #         dim = rel_pos_emb_config[0]
        #     self.rel_pos = SinusoidalEmbeddings(dim)

    def forward(self, q, k, v, packed_shape, mask = None, input_mask = None):
        mask = default(mask, input_mask)

        autopad, pad_value, window_size, causal, look_backward, look_forward, shared_qk = self.autopad, -1, self.window_size, self.causal, self.look_backward, self.look_forward, self.shared_qk

        # https://github.com/arogozhnikov/einops/blob/master/docs/4-pack-and-unpack.ipynb
        # (q, packed_shape), (k, _), (v, _) = map(lambda t: pack([t], '* n d'), (q, k, v))        # (2, 8, 2048, 64) -> (16, 2048, 64)

        # rotary embeddings

        # if exists(self.rel_pos):
        #     pos_emb = self.rel_pos(q)     # (16, 2048, 64)
        #     q, k = apply_rotary_pos_emb(q, k, pos_emb)

        # auto padding

        # if autopad:
        #     orig_seq_len = q.shape[1]
        #     (needed_pad, q), (_, k), (_, v) = map(lambda t: pad_to_multiple(t, self.window_size, dim = -2), (q, k, v))

        b, n, dim_head, device, dtype = *q.shape, q.device, q.dtype
        scale = dim_head ** -0.5

        if n % window_size != 0:
            print('sequence length must be divisible by window size for local attention', n, window_size)
            assert True

        windows = n // window_size

        if shared_qk:
            k = l2norm(k)

        seq = torch.arange(n, device = device)
        b_t = rearrange(seq, '(w n) -> 1 w n', w = windows, n = window_size)

        bq, bk, bv = map(lambda t: rearrange(t, 'b (w n) d -> b w n d', w = windows), (q, k, v))

        look_around_kwargs = dict(
            backward =  look_backward,
            forward =  look_forward,
            pad_value = pad_value
        )

        bk = look_around(bk, **look_around_kwargs)
        bv = look_around(bv, **look_around_kwargs)

        bq_t = b_t
        bq_k = look_around(b_t, **look_around_kwargs)

        bq_t = rearrange(bq_t, '... i -> ... i 1')
        bq_k = rearrange(bq_k, '... j -> ... 1 j')

        sim = einsum('b h i e, b h j e -> b h i j', bq, bk) * scale

        mask_value = max_neg_value(sim)

        if shared_qk:
            self_mask = bq_t == bq_k
            sim = sim.masked_fill(self_mask, TOKEN_SELF_ATTN_VALUE)
            del self_mask

        if causal:
            causal_mask = bq_t < bq_k

            if self.exact_windowsize:
                max_causal_window_size = (self.window_size * self.look_backward)
                causal_mask = causal_mask | (bq_t > (bq_k + max_causal_window_size))

            sim = sim.masked_fill(causal_mask, mask_value)
            del causal_mask

        # mask out padding value

        # if autopad and needed_pad:
        #     pad_mask = bq_k == pad_value
        #     sim = sim.masked_fill(pad_mask, mask_value)
        #     del pad_mask

        if exists(mask):
            batch = mask.shape[0]
            assert (b % batch) == 0

            h = b // mask.shape[0]

            # if autopad:
            #     _, mask = pad_to_multiple(mask, window_size, dim = -1, value = False)

            mask = rearrange(mask, '... (w n) -> (...) w n', w = windows, n = window_size)
            mask = look_around(mask, **{**look_around_kwargs, 'pad_value': False})
            mask = rearrange(mask, '... j -> ... 1 j')
            mask = repeat(mask, 'b ... -> (b h) ...', h = h)
            sim = sim.masked_fill(~mask, mask_value)
            del mask

        # attention

        attn = sim.softmax(dim = -1)
        attn = self.dropout(attn)

        # aggregation

        out = einsum('b h i j, b h j e -> b h i e', attn, bv)
        out = rearrange(out, 'b w n d -> b (w n) d')

        # if autopad:
        #     out = out[:, :orig_seq_len, :]

        out, *_ = unpack(out, packed_shape, '* n d')
        return out

model/local_attention/local_attention.py
YoungSeng/UnifiedGesture
import torch
from torch import nn, einsum

from einops import rearrange

class SinusoidalEmbeddings(nn.Module):
    def __init__(self, dim):
        super().__init__()
        inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, x):
        n = x.shape[-2]
        t = torch.arange(n, device = x.device).type_as(self.inv_freq)
        freqs = torch.einsum('i , j -> i j', t, self.inv_freq)
        return torch.cat((freqs, freqs), dim=-1)

def rotate_half(x):
    x = rearrange(x, 'b ... (r d) -> b (...) r d', r = 2)
    x1, x2 = x.unbind(dim = -2)
    return torch.cat((-x2, x1), dim = -1)

def apply_rotary_pos_emb(q, k, freqs):
    q, k = map(lambda t: (t * freqs.cos()) + (rotate_half(t) * freqs.sin()), (q, k))
    return q, k

model/local_attention/rotary.py
YoungSeng/UnifiedGesture
import torch
from torch import nn
import torch.nn.functional as F

from einops import rearrange

from local_attention.local_attention import LocalAttention

# helper function

def exists(val):
    return val is not None

def eval_decorator(fn):
    def inner(model, *args, **kwargs):
        was_training = model.training
        model.eval()
        out = fn(model, *args, **kwargs)
        model.train(was_training)
        return out
    return inner

# sampling functions

def top_k(logits, thres = 0.9):
    k = int((1 - thres) * logits.shape[-1])
    val, ind = torch.topk(logits, k)
    probs = torch.full_like(logits, float('-inf'))
    probs.scatter_(1, ind, val)
    return probs

# multi-head attention

class LocalMHA(nn.Module):
    def __init__(
        self,
        *,
        dim,
        window_size,
        dim_head = 64,
        heads = 8,
        dropout = 0.,
        causal = False,
        prenorm = False,
        **kwargs
    ):
        super().__init__()        
        inner_dim = dim_head * heads

        self.norm = nn.LayerNorm(dim) if prenorm else None

        self.heads = heads
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)

        self.attn_fn = LocalAttention(
            dim = dim_head,
            window_size = window_size,
            causal = causal,
            autopad = True,
            exact_windowsize = True,
            **kwargs
        )

        self.to_out = nn.Linear(inner_dim, dim, bias = False)

    def forward(self, x, mask = None):
        if exists(self.norm):
            x = self.norm(x)

        q, k, v = self.to_qkv(x).chunk(3, dim = -1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), (q, k, v)) 

        out = self.attn_fn(q, k, v, mask = mask)

        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)

# feedforward

class GEGLU(nn.Module):
    def forward(self, x):
        x, gate = x.chunk(2, dim = -1)
        return x * F.gelu(gate)

def FeedForward(dim, mult = 4, dropout = 0.):
    inner_dim = int(dim * mult * 2 / 3)

    return nn.Sequential(
        nn.LayerNorm(dim),
        nn.Linear(dim, inner_dim * 2, bias = False),
        GEGLU(),
        nn.Dropout(dropout),
        nn.Linear(inner_dim, dim, bias = False)
    )

# main transformer class

class LocalTransformer(nn.Module):
    def __init__(
        self,
        *,
        num_tokens,
        max_seq_len,
        dim,
        depth,
        causal = True,
        local_attn_window_size = 512,
        dim_head = 64,
        heads = 8,
        ff_mult = 4,
        attn_dropout = 0.,
        ff_dropout = 0.,
        ignore_index = -1,
        **kwargs
    ):
        super().__init__()
        self.token_emb = nn.Embedding(num_tokens, dim)
        self.pos_emb = nn.Embedding(max_seq_len, dim)

        self.max_seq_len = max_seq_len
        self.layers = nn.ModuleList([])

        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                LocalMHA(dim = dim, dim_head = dim_head, heads = heads, dropout = attn_dropout, causal = causal, window_size = local_attn_window_size, prenorm = True, **kwargs),
                FeedForward(dim = dim, mult = ff_mult, dropout = ff_dropout)
            ]))

        self.ignore_index = ignore_index
        self.to_logits = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_tokens, bias = False)
        )

    @torch.no_grad()
    @eval_decorator
    def generate(
        self,
        prime,
        seq_len,
        temperature = 1.,
        filter_thres = 0.9,
        **kwargs
    ):
        n, device = prime.shape[1], prime.device

        out = prime

        for _ in range(seq_len):
            logits = self.forward(out[:, -self.max_seq_len:], **kwargs)
            filtered_logits = top_k(logits[:, -1], thres = filter_thres)
            probs = F.softmax(filtered_logits / temperature, dim = -1)
            sampled = torch.multinomial(probs, 1)
            out = torch.cat((out, sampled), dim = -1)

        return out[:, n:]

    def forward(self, x, mask = None, return_loss = False):
        if return_loss:
            x, labels = x[:, :-1], x[:, 1:]

        n, device = x.shape[1], x.device
        x = self.token_emb(x)

        assert n <= self.max_seq_len
        x = x + self.pos_emb(torch.arange(n, device = device))

        for attn, ff in self.layers:
            x = attn(x, mask = mask) + x
            x = ff(x) + x

        logits = self.to_logits(x)

        if not return_loss:
            return logits

        logits = rearrange(logits, 'b n c -> b c n')
        loss = F.cross_entropy(logits, labels, ignore_index = self.ignore_index)
        return loss

model/local_attention/transformer.py
YoungSeng/UnifiedGesture
import warnings
from typing import Tuple, Optional

import torch
from torch import Tensor
if float(torch.__version__.split('.')[0]) == 0 or (float(torch.__version__.split('.')[0]) == 1 and float(torch.__version__.split('.')[1])) < 9:
    from torch.nn.modules.linear import _LinearWithBias
else:
    from torch.nn.modules.linear import NonDynamicallyQuantizableLinear as _LinearWithBias
from torch.nn.init import xavier_uniform_
from torch.nn.init import constant_
from torch.nn.init import xavier_normal_
from torch.nn.parameter import Parameter
from torch.nn.modules.module import Module
from torch.nn import functional as F


class Threshold(Module):
    r"""Thresholds each element of the input Tensor.

    Threshold is defined as:

    .. math::
        y =
        \begin{cases}
        x, &\text{ if } x > \text{threshold} \\
        \text{value}, &\text{ otherwise }
        \end{cases}

    Args:
        threshold: The value to threshold at
        value: The value to replace with
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    Examples::

        >>> m = nn.Threshold(0.1, 20)
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['threshold', 'value', 'inplace']

    threshold: float
    value: float
    inplace: bool

    def __init__(self, threshold: float, value: float, inplace: bool = False) -> None:
        super(Threshold, self).__init__()
        self.threshold = threshold
        self.value = value
        self.inplace = inplace
        # TODO: check in THNN (if inplace == True, then assert value <= threshold)

    def forward(self, input: Tensor) -> Tensor:
        return F.threshold(input, self.threshold, self.value, self.inplace)

    def extra_repr(self):
        inplace_str = ', inplace=True' if self.inplace else ''
        return 'threshold={}, value={}{}'.format(
            self.threshold, self.value, inplace_str
        )


class ReLU(Module):
    r"""Applies the rectified linear unit function element-wise:

    :math:`\text{ReLU}(x) = (x)^+ = \max(0, x)`

    Args:
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/ReLU.png

    Examples::

        >>> m = nn.ReLU()
        >>> input = torch.randn(2)
        >>> output = m(input)


      An implementation of CReLU - https://arxiv.org/abs/1603.05201

        >>> m = nn.ReLU()
        >>> input = torch.randn(2).unsqueeze(0)
        >>> output = torch.cat((m(input),m(-input)))
    """
    __constants__ = ['inplace']
    inplace: bool

    def __init__(self, inplace: bool = False):
        super(ReLU, self).__init__()
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.relu(input, inplace=self.inplace)

    def extra_repr(self) -> str:
        inplace_str = 'inplace=True' if self.inplace else ''
        return inplace_str


class RReLU(Module):
    r"""Applies the randomized leaky rectified liner unit function, element-wise,
    as described in the paper:

    `Empirical Evaluation of Rectified Activations in Convolutional Network`_.

    The function is defined as:

    .. math::
        \text{RReLU}(x) =
        \begin{cases}
            x & \text{if } x \geq 0 \\
            ax & \text{ otherwise }
        \end{cases}

    where :math:`a` is randomly sampled from uniform distribution
    :math:`\mathcal{U}(\text{lower}, \text{upper})`.

     See: https://arxiv.org/pdf/1505.00853.pdf

    Args:
        lower: lower bound of the uniform distribution. Default: :math:`\frac{1}{8}`
        upper: upper bound of the uniform distribution. Default: :math:`\frac{1}{3}`
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    Examples::

        >>> m = nn.RReLU(0.1, 0.3)
        >>> input = torch.randn(2)
        >>> output = m(input)

    .. _`Empirical Evaluation of Rectified Activations in Convolutional Network`:
        https://arxiv.org/abs/1505.00853
    """
    __constants__ = ['lower', 'upper', 'inplace']

    lower: float
    upper: float
    inplace: bool

    def __init__(
        self,
        lower: float = 1. / 8,
        upper: float = 1. / 3,
        inplace: bool = False
    ):
        super(RReLU, self).__init__()
        self.lower = lower
        self.upper = upper
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.rrelu(input, self.lower, self.upper, self.training, self.inplace)

    def extra_repr(self):
        inplace_str = ', inplace=True' if self.inplace else ''
        return 'lower={}, upper={}{}'.format(self.lower, self.upper, inplace_str)


class Hardtanh(Module):
    r"""Applies the HardTanh function element-wise

    HardTanh is defined as:

    .. math::
        \text{HardTanh}(x) = \begin{cases}
            1 & \text{ if } x > 1 \\
            -1 & \text{ if } x < -1 \\
            x & \text{ otherwise } \\
        \end{cases}

    The range of the linear region :math:`[-1, 1]` can be adjusted using
    :attr:`min_val` and :attr:`max_val`.

    Args:
        min_val: minimum value of the linear region range. Default: -1
        max_val: maximum value of the linear region range. Default: 1
        inplace: can optionally do the operation in-place. Default: ``False``

    Keyword arguments :attr:`min_value` and :attr:`max_value`
    have been deprecated in favor of :attr:`min_val` and :attr:`max_val`.

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Hardtanh.png

    Examples::

        >>> m = nn.Hardtanh(-2, 2)
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['min_val', 'max_val', 'inplace']

    min_val: float
    max_val: float
    inplace: bool

    def __init__(
        self,
        min_val: float = -1.,
        max_val: float = 1.,
        inplace: bool = False,
        min_value: Optional[float] = None,
        max_value: Optional[float] = None
    ) -> None:
        super(Hardtanh, self).__init__()
        if min_value is not None:
            warnings.warn("keyword argument min_value is deprecated and rename to min_val")
            min_val = min_value
        if max_value is not None:
            warnings.warn("keyword argument max_value is deprecated and rename to max_val")
            max_val = max_value

        self.min_val = min_val
        self.max_val = max_val
        self.inplace = inplace
        assert self.max_val > self.min_val

    def forward(self, input: Tensor) -> Tensor:
        return F.hardtanh(input, self.min_val, self.max_val, self.inplace)

    def extra_repr(self) -> str:
        inplace_str = ', inplace=True' if self.inplace else ''
        return 'min_val={}, max_val={}{}'.format(
            self.min_val, self.max_val, inplace_str
        )


class ReLU6(Hardtanh):
    r"""Applies the element-wise function:

    .. math::
        \text{ReLU6}(x) = \min(\max(0,x), 6)

    Args:
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/ReLU6.png

    Examples::

        >>> m = nn.ReLU6()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """

    def __init__(self, inplace: bool = False):
        super(ReLU6, self).__init__(0., 6., inplace)

    def extra_repr(self) -> str:
        inplace_str = 'inplace=True' if self.inplace else ''
        return inplace_str


class Sigmoid(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{Sigmoid}(x) = \sigma(x) = \frac{1}{1 + \exp(-x)}


    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Sigmoid.png

    Examples::

        >>> m = nn.Sigmoid()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """

    def forward(self, input: Tensor) -> Tensor:
        return torch.sigmoid(input)


class Hardsigmoid(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{Hardsigmoid}(x) = \begin{cases}
            0 & \text{if~} x \le -3, \\
            1 & \text{if~} x \ge +3, \\
            x / 6 + 1 / 2 & \text{otherwise}
        \end{cases}

    Args:
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    Examples::

        >>> m = nn.Hardsigmoid()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['inplace']

    inplace: bool

    def __init__(self, inplace : bool = False) -> None:
        super(Hardsigmoid, self).__init__()
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.hardsigmoid(input, self.inplace)


class Tanh(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{Tanh}(x) = \tanh(x) = \frac{\exp(x) - \exp(-x)} {\exp(x) + \exp(-x)}

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Tanh.png

    Examples::

        >>> m = nn.Tanh()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """

    def forward(self, input: Tensor) -> Tensor:
        return torch.tanh(input)

class SiLU(Module):
    r"""Applies the silu function, element-wise.

    .. math::
        \text{silu}(x) = x * \sigma(x), \text{where } \sigma(x) \text{ is the logistic sigmoid.}

    .. note::
        See `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_
        where the SiLU (Sigmoid Linear Unit) was originally coined, and see
        `Sigmoid-Weighted Linear Units for Neural Network Function Approximation
        in Reinforcement Learning <https://arxiv.org/abs/1702.03118>`_ and `Swish:
        a Self-Gated Activation Function <https://arxiv.org/abs/1710.05941v1>`_
        where the SiLU was experimented with later.

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    Examples::

        >>> m = nn.SiLU()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['inplace']
    inplace: bool

    def __init__(self, inplace: bool = False):
        super(SiLU, self).__init__()
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.silu(input, inplace=self.inplace)

    def extra_repr(self) -> str:
        inplace_str = 'inplace=True' if self.inplace else ''
        return inplace_str

class Hardswish(Module):
    r"""Applies the hardswish function, element-wise, as described in the paper:

    `Searching for MobileNetV3`_.

    .. math::
        \text{Hardswish}(x) = \begin{cases}
            0 & \text{if~} x \le -3, \\
            x & \text{if~} x \ge +3, \\
            x \cdot (x + 3) /6 & \text{otherwise}
        \end{cases}

    Args:
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    Examples::

        >>> m = nn.Hardswish()
        >>> input = torch.randn(2)
        >>> output = m(input)

    .. _`Searching for MobileNetV3`:
        https://arxiv.org/abs/1905.02244
    """
    __constants__ = ['inplace']

    inplace: bool

    def __init__(self, inplace : bool = False) -> None:
        super(Hardswish, self).__init__()
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.hardswish(input, self.inplace)


class ELU(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{ELU}(x) = \begin{cases}
        x, & \text{ if } x > 0\\
        \alpha * (\exp(x) - 1), & \text{ if } x \leq 0
        \end{cases}

    Args:
        alpha: the :math:`\alpha` value for the ELU formulation. Default: 1.0
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/ELU.png

    Examples::

        >>> m = nn.ELU()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['alpha', 'inplace']
    alpha: float
    inplace: bool

    def __init__(self, alpha: float = 1., inplace: bool = False) -> None:
        super(ELU, self).__init__()
        self.alpha = alpha
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.elu(input, self.alpha, self.inplace)

    def extra_repr(self) -> str:
        inplace_str = ', inplace=True' if self.inplace else ''
        return 'alpha={}{}'.format(self.alpha, inplace_str)


class CELU(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x/\alpha) - 1))

    More details can be found in the paper `Continuously Differentiable Exponential Linear Units`_ .

    Args:
        alpha: the :math:`\alpha` value for the CELU formulation. Default: 1.0
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/CELU.png

    Examples::

        >>> m = nn.CELU()
        >>> input = torch.randn(2)
        >>> output = m(input)

    .. _`Continuously Differentiable Exponential Linear Units`:
        https://arxiv.org/abs/1704.07483
    """
    __constants__ = ['alpha', 'inplace']
    alpha: float
    inplace: bool

    def __init__(self, alpha: float = 1., inplace: bool = False) -> None:
        super(CELU, self).__init__()
        self.alpha = alpha
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.celu(input, self.alpha, self.inplace)

    def extra_repr(self) -> str:
        inplace_str = ', inplace=True' if self.inplace else ''
        return 'alpha={}{}'.format(self.alpha, inplace_str)


class SELU(Module):
    r"""Applied element-wise, as:

    .. math::
        \text{SELU}(x) = \text{scale} * (\max(0,x) + \min(0, \alpha * (\exp(x) - 1)))

    with :math:`\alpha = 1.6732632423543772848170429916717` and
    :math:`\text{scale} = 1.0507009873554804934193349852946`.

    More details can be found in the paper `Self-Normalizing Neural Networks`_ .

    Args:
        inplace (bool, optional): can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/SELU.png

    Examples::

        >>> m = nn.SELU()
        >>> input = torch.randn(2)
        >>> output = m(input)

    .. _Self-Normalizing Neural Networks: https://arxiv.org/abs/1706.02515
    """
    __constants__ = ['inplace']
    inplace: bool

    def __init__(self, inplace: bool = False) -> None:
        super(SELU, self).__init__()
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.selu(input, self.inplace)

    def extra_repr(self) -> str:
        inplace_str = 'inplace=True' if self.inplace else ''
        return inplace_str


class GLU(Module):
    r"""Applies the gated linear unit function
    :math:`{GLU}(a, b)= a \otimes \sigma(b)` where :math:`a` is the first half
    of the input matrices and :math:`b` is the second half.

    Args:
        dim (int): the dimension on which to split the input. Default: -1

    Shape:
        - Input: :math:`(\ast_1, N, \ast_2)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(\ast_1, M, \ast_2)` where :math:`M=N/2`

    Examples::

        >>> m = nn.GLU()
        >>> input = torch.randn(4, 2)
        >>> output = m(input)
    """
    __constants__ = ['dim']
    dim: int

    def __init__(self, dim: int = -1) -> None:
        super(GLU, self).__init__()
        self.dim = dim

    def forward(self, input: Tensor) -> Tensor:
        return F.glu(input, self.dim)

    def extra_repr(self) -> str:
        return 'dim={}'.format(self.dim)


class GELU(Module):
    r"""Applies the Gaussian Error Linear Units function:

    .. math:: \text{GELU}(x) = x * \Phi(x)

    where :math:`\Phi(x)` is the Cumulative Distribution Function for Gaussian Distribution.

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/GELU.png

    Examples::

        >>> m = nn.GELU()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    def forward(self, input: Tensor) -> Tensor:
        return F.gelu(input)


class Hardshrink(Module):
    r"""Applies the hard shrinkage function element-wise:

    .. math::
        \text{HardShrink}(x) =
        \begin{cases}
        x, & \text{ if } x > \lambda \\
        x, & \text{ if } x < -\lambda \\
        0, & \text{ otherwise }
        \end{cases}

    Args:
        lambd: the :math:`\lambda` value for the Hardshrink formulation. Default: 0.5

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Hardshrink.png

    Examples::

        >>> m = nn.Hardshrink()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['lambd']
    lambd: float

    def __init__(self, lambd: float = 0.5) -> None:
        super(Hardshrink, self).__init__()
        self.lambd = lambd

    def forward(self, input: Tensor) -> Tensor:
        return F.hardshrink(input, self.lambd)

    def extra_repr(self) -> str:
        return '{}'.format(self.lambd)


class LeakyReLU(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{LeakyReLU}(x) = \max(0, x) + \text{negative\_slope} * \min(0, x)


    or

    .. math::
        \text{LeakyRELU}(x) =
        \begin{cases}
        x, & \text{ if } x \geq 0 \\
        \text{negative\_slope} \times x, & \text{ otherwise }
        \end{cases}

    Args:
        negative_slope: Controls the angle of the negative slope. Default: 1e-2
        inplace: can optionally do the operation in-place. Default: ``False``

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/LeakyReLU.png

    Examples::

        >>> m = nn.LeakyReLU(0.1)
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['inplace', 'negative_slope']
    inplace: bool
    negative_slope: float

    def __init__(self, negative_slope: float = 1e-2, inplace: bool = False) -> None:
        super(LeakyReLU, self).__init__()
        self.negative_slope = negative_slope
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.leaky_relu(input, self.negative_slope, self.inplace)

    def extra_repr(self) -> str:
        inplace_str = ', inplace=True' if self.inplace else ''
        return 'negative_slope={}{}'.format(self.negative_slope, inplace_str)


class LogSigmoid(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{LogSigmoid}(x) = \log\left(\frac{ 1 }{ 1 + \exp(-x)}\right)

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/LogSigmoid.png

    Examples::

        >>> m = nn.LogSigmoid()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """

    def forward(self, input: Tensor) -> Tensor:
        return F.logsigmoid(input)


class Softplus(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{Softplus}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x))

    SoftPlus is a smooth approximation to the ReLU function and can be used
    to constrain the output of a machine to always be positive.

    For numerical stability the implementation reverts to the linear function
    when :math:`input \times \beta > threshold`.

    Args:
        beta: the :math:`\beta` value for the Softplus formulation. Default: 1
        threshold: values above this revert to a linear function. Default: 20

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Softplus.png

    Examples::

        >>> m = nn.Softplus()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['beta', 'threshold']
    beta: int
    threshold: int

    def __init__(self, beta: int = 1, threshold: int = 20) -> None:
        super(Softplus, self).__init__()
        self.beta = beta
        self.threshold = threshold

    def forward(self, input: Tensor) -> Tensor:
        return F.softplus(input, self.beta, self.threshold)

    def extra_repr(self) -> str:
        return 'beta={}, threshold={}'.format(self.beta, self.threshold)


class Softshrink(Module):
    r"""Applies the soft shrinkage function elementwise:

    .. math::
        \text{SoftShrinkage}(x) =
        \begin{cases}
        x - \lambda, & \text{ if } x > \lambda \\
        x + \lambda, & \text{ if } x < -\lambda \\
        0, & \text{ otherwise }
        \end{cases}

    Args:
        lambd: the :math:`\lambda` (must be no less than zero) value for the Softshrink formulation. Default: 0.5

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Softshrink.png

    Examples::

        >>> m = nn.Softshrink()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['lambd']
    lambd: float

    def __init__(self, lambd: float = 0.5) -> None:
        super(Softshrink, self).__init__()
        self.lambd = lambd

    def forward(self, input: Tensor) -> Tensor:
        return F.softshrink(input, self.lambd)

    def extra_repr(self) -> str:
        return str(self.lambd)


class MultiheadAttention(Module):
    r"""Allows the model to jointly attend to information
    from different representation subspaces.
    See reference: Attention Is All You Need

    .. math::
        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
        \text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)

    Args:
        embed_dim: total dimension of the model.
        num_heads: parallel attention heads.
        dropout: a Dropout layer on attn_output_weights. Default: 0.0.
        bias: add bias as module parameter. Default: True.
        add_bias_kv: add bias to the key and value sequences at dim=0.
        add_zero_attn: add a new batch of zeros to the key and
                       value sequences at dim=1.
        kdim: total number of features in key. Default: None.
        vdim: total number of features in value. Default: None.

        Note: if kdim and vdim are None, they will be set to embed_dim such that
        query, key, and value have the same number of features.

    Examples::

        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
    """
    bias_k: Optional[torch.Tensor]
    bias_v: Optional[torch.Tensor]

    def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):
        super(MultiheadAttention, self).__init__()
        self.embed_dim = embed_dim
        self.kdim = kdim if kdim is not None else embed_dim
        self.vdim = vdim if vdim is not None else embed_dim
        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim

        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"

        if self._qkv_same_embed_dim is False:
            self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
            self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))
            self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))
            self.register_parameter('in_proj_weight', None)
        else:
            self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim))
            self.register_parameter('q_proj_weight', None)
            self.register_parameter('k_proj_weight', None)
            self.register_parameter('v_proj_weight', None)

        if bias:
            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))
        else:
            self.register_parameter('in_proj_bias', None)
        self.out_proj = _LinearWithBias(embed_dim, embed_dim)

        if add_bias_kv:
            self.bias_k = Parameter(torch.empty(1, 1, embed_dim))
            self.bias_v = Parameter(torch.empty(1, 1, embed_dim))
        else:
            self.bias_k = self.bias_v = None

        self.add_zero_attn = add_zero_attn

        self._reset_parameters()

    def _reset_parameters(self):
        if self._qkv_same_embed_dim:
            xavier_uniform_(self.in_proj_weight)
        else:
            xavier_uniform_(self.q_proj_weight)
            xavier_uniform_(self.k_proj_weight)
            xavier_uniform_(self.v_proj_weight)

        if self.in_proj_bias is not None:
            constant_(self.in_proj_bias, 0.)
            constant_(self.out_proj.bias, 0.)
        if self.bias_k is not None:
            xavier_normal_(self.bias_k)
        if self.bias_v is not None:
            xavier_normal_(self.bias_v)

    def __setstate__(self, state):
        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
        if '_qkv_same_embed_dim' not in state:
            state['_qkv_same_embed_dim'] = True

        super(MultiheadAttention, self).__setstate__(state)

    def forward(self, query, key, value, key_padding_mask=None,
                need_weights=True, attn_mask=None):
        # type: (Tensor, Tensor, Tensor, Optional[Tensor], bool, Optional[Tensor]) -> Tuple[Tensor, Optional[Tensor]]
        r"""
    Args:
        query, key, value: map a query and a set of key-value pairs to an output.
            See "Attention Is All You Need" for more details.
        key_padding_mask: if provided, specified padding elements in the key will
            be ignored by the attention. When given a binary mask and a value is True,
            the corresponding value on the attention layer will be ignored. When given
            a byte mask and a value is non-zero, the corresponding value on the attention
            layer will be ignored
        need_weights: output attn_output_weights.
        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
            the batches while a 3D mask allows to specify a different mask for the entries of each batch.

    Shape:
        - Inputs:
        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
          the embedding dimension.
        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
          the embedding dimension.
        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
          the embedding dimension.
        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
          If a ByteTensor is provided, the non-zero positions will be ignored while the position
          with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
          S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
          is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
          is provided, it will be added to the attention weight.

        - Outputs:
        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
          E is the embedding dimension.
        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
          L is the target sequence length, S is the source sequence length.
        """
        if not self._qkv_same_embed_dim:
            return F.multi_head_attention_forward(
                query, key, value, self.embed_dim, self.num_heads,
                self.in_proj_weight, self.in_proj_bias,
                self.bias_k, self.bias_v, self.add_zero_attn,
                self.dropout, self.out_proj.weight, self.out_proj.bias,
                training=self.training,
                key_padding_mask=key_padding_mask, need_weights=need_weights,
                attn_mask=attn_mask, use_separate_proj_weight=True,
                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
                v_proj_weight=self.v_proj_weight)
        else:
            return F.multi_head_attention_forward(
                query, key, value, self.embed_dim, self.num_heads,
                self.in_proj_weight, self.in_proj_bias,
                self.bias_k, self.bias_v, self.add_zero_attn,
                self.dropout, self.out_proj.weight, self.out_proj.bias,
                training=self.training,
                key_padding_mask=key_padding_mask, need_weights=need_weights,
                attn_mask=attn_mask)


class PReLU(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{PReLU}(x) = \max(0,x) + a * \min(0,x)

    or

    .. math::
        \text{PReLU}(x) =
        \begin{cases}
        x, & \text{ if } x \geq 0 \\
        ax, & \text{ otherwise }
        \end{cases}

    Here :math:`a` is a learnable parameter. When called without arguments, `nn.PReLU()` uses a single
    parameter :math:`a` across all input channels. If called with `nn.PReLU(nChannels)`,
    a separate :math:`a` is used for each input channel.


    .. note::
        weight decay should not be used when learning :math:`a` for good performance.

    .. note::
        Channel dim is the 2nd dim of input. When input has dims < 2, then there is
        no channel dim and the number of channels = 1.

    Args:
        num_parameters (int): number of :math:`a` to learn.
            Although it takes an int as input, there is only two values are legitimate:
            1, or the number of channels at input. Default: 1
        init (float): the initial value of :math:`a`. Default: 0.25

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    Attributes:
        weight (Tensor): the learnable weights of shape (:attr:`num_parameters`).

    .. image:: ../scripts/activation_images/PReLU.png

    Examples::

        >>> m = nn.PReLU()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """
    __constants__ = ['num_parameters']
    num_parameters: int

    def __init__(self, num_parameters: int = 1, init: float = 0.25) -> None:
        self.num_parameters = num_parameters
        super(PReLU, self).__init__()
        self.weight = Parameter(torch.Tensor(num_parameters).fill_(init))

    def forward(self, input: Tensor) -> Tensor:
        return F.prelu(input, self.weight)

    def extra_repr(self) -> str:
        return 'num_parameters={}'.format(self.num_parameters)


class Softsign(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{SoftSign}(x) = \frac{x}{ 1 + |x|}

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Softsign.png

    Examples::

        >>> m = nn.Softsign()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """

    def forward(self, input: Tensor) -> Tensor:
        return F.softsign(input)


class Tanhshrink(Module):
    r"""Applies the element-wise function:

    .. math::
        \text{Tanhshrink}(x) = x - \tanh(x)

    Shape:
        - Input: :math:`(N, *)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(N, *)`, same shape as the input

    .. image:: ../scripts/activation_images/Tanhshrink.png

    Examples::

        >>> m = nn.Tanhshrink()
        >>> input = torch.randn(2)
        >>> output = m(input)
    """

    def forward(self, input: Tensor) -> Tensor:
        return F.tanhshrink(input)


class Softmin(Module):
    r"""Applies the Softmin function to an n-dimensional input Tensor
    rescaling them so that the elements of the n-dimensional output Tensor
    lie in the range `[0, 1]` and sum to 1.

    Softmin is defined as:

    .. math::
        \text{Softmin}(x_{i}) = \frac{\exp(-x_i)}{\sum_j \exp(-x_j)}

    Shape:
        - Input: :math:`(*)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(*)`, same shape as the input

    Arguments:
        dim (int): A dimension along which Softmin will be computed (so every slice
            along dim will sum to 1).

    Returns:
        a Tensor of the same dimension and shape as the input, with
        values in the range [0, 1]

    Examples::

        >>> m = nn.Softmin()
        >>> input = torch.randn(2, 3)
        >>> output = m(input)
    """
    __constants__ = ['dim']
    dim: Optional[int]

    def __init__(self, dim: Optional[int] = None) -> None:
        super(Softmin, self).__init__()
        self.dim = dim

    def __setstate__(self, state):
        self.__dict__.update(state)
        if not hasattr(self, 'dim'):
            self.dim = None

    def forward(self, input: Tensor) -> Tensor:
        return F.softmin(input, self.dim, _stacklevel=5)

    def extra_repr(self):
        return 'dim={dim}'.format(dim=self.dim)

class Softmax(Module):
    r"""Applies the Softmax function to an n-dimensional input Tensor
    rescaling them so that the elements of the n-dimensional output Tensor
    lie in the range [0,1] and sum to 1.

    Softmax is defined as:

    .. math::
        \text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}

    When the input Tensor is a sparse tensor then the unspecifed
    values are treated as ``-inf``.

    Shape:
        - Input: :math:`(*)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(*)`, same shape as the input

    Returns:
        a Tensor of the same dimension and shape as the input with
        values in the range [0, 1]

    Arguments:
        dim (int): A dimension along which Softmax will be computed (so every slice
            along dim will sum to 1).

    .. note::
        This module doesn't work directly with NLLLoss,
        which expects the Log to be computed between the Softmax and itself.
        Use `LogSoftmax` instead (it's faster and has better numerical properties).

    Examples::

        >>> m = nn.Softmax(dim=1)
        >>> input = torch.randn(2, 3)
        >>> output = m(input)

    """
    __constants__ = ['dim']
    dim: Optional[int]

    def __init__(self, dim: Optional[int] = None) -> None:
        super(Softmax, self).__init__()
        self.dim = dim

    def __setstate__(self, state):
        self.__dict__.update(state)
        if not hasattr(self, 'dim'):
            self.dim = None

    def forward(self, input: Tensor) -> Tensor:
        return F.softmax(input, self.dim, _stacklevel=5)

    def extra_repr(self) -> str:
        return 'dim={dim}'.format(dim=self.dim)


class Softmax2d(Module):
    r"""Applies SoftMax over features to each spatial location.

    When given an image of ``Channels x Height x Width``, it will
    apply `Softmax` to each location :math:`(Channels, h_i, w_j)`

    Shape:
        - Input: :math:`(N, C, H, W)`
        - Output: :math:`(N, C, H, W)` (same shape as input)

    Returns:
        a Tensor of the same dimension and shape as the input with
        values in the range [0, 1]

    Examples::

        >>> m = nn.Softmax2d()
        >>> # you softmax over the 2nd dimension
        >>> input = torch.randn(2, 3, 12, 13)
        >>> output = m(input)
    """

    def forward(self, input: Tensor) -> Tensor:
        assert input.dim() == 4, 'Softmax2d requires a 4D tensor as input'
        return F.softmax(input, 1, _stacklevel=5)


class LogSoftmax(Module):
    r"""Applies the :math:`\log(\text{Softmax}(x))` function to an n-dimensional
    input Tensor. The LogSoftmax formulation can be simplified as:

    .. math::
        \text{LogSoftmax}(x_{i}) = \log\left(\frac{\exp(x_i) }{ \sum_j \exp(x_j)} \right)

    Shape:
        - Input: :math:`(*)` where `*` means, any number of additional
          dimensions
        - Output: :math:`(*)`, same shape as the input

    Arguments:
        dim (int): A dimension along which LogSoftmax will be computed.

    Returns:
        a Tensor of the same dimension and shape as the input with
        values in the range [-inf, 0)

    Examples::

        >>> m = nn.LogSoftmax()
        >>> input = torch.randn(2, 3)
        >>> output = m(input)
    """
    __constants__ = ['dim']
    dim: Optional[int]

    def __init__(self, dim: Optional[int] = None) -> None:
        super(LogSoftmax, self).__init__()
        self.dim = dim

    def __setstate__(self, state):
        self.__dict__.update(state)
        if not hasattr(self, 'dim'):
            self.dim = None

    def forward(self, input: Tensor) -> Tensor:
        return F.log_softmax(input, self.dim, _stacklevel=5)

    def extra_repr(self):
        return 'dim={dim}'.format(dim=self.dim)

model/myactivation.py
YoungSeng/UnifiedGesture
import copy
import pdb
from typing import Optional, Any

import torch
import torch.nn as nn
from torch import Tensor
from torch.nn import functional as F
from torch.nn.modules.module import Module
from myactivation import MultiheadAttention
from torch.nn.modules.container import ModuleList
from torch.nn.init import xavier_uniform_
from torch.nn.modules.dropout import Dropout
from torch.nn.modules.linear import Linear
from torch.nn.modules.normalization import LayerNorm


class Transformer(Module):
    r"""A transformer model. User is able to modify the attributes as needed. The architecture
    is based on the paper "Attention Is All You Need". Ashish Vaswani, Noam Shazeer,
    Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and
    Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information
    Processing Systems, pages 6000-6010. Users can build the BERT(https://arxiv.org/abs/1810.04805)
    model with corresponding parameters.

    Args:
        d_model: the number of expected features in the encoder/decoder inputs (default=512).
        nhead: the number of heads in the multiheadattention models (default=8).
        num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).
        num_decoder_layers: the number of sub-decoder-layers in the decoder (default=6).
        dim_feedforward: the dimension of the feedforward network model (default=2048).
        dropout: the dropout value (default=0.1).
        activation: the activation function of encoder/decoder intermediate layer, relu or gelu (default=relu).
        custom_encoder: custom encoder (default=None).
        custom_decoder: custom decoder (default=None).

    Examples::
        >>> transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
        >>> src = torch.rand((10, 32, 512))
        >>> tgt = torch.rand((20, 32, 512))
        >>> out = transformer_model(src, tgt)

    Note: A full example to apply nn.Transformer module for the word language model is available in
    https://github.com/pytorch/examples/tree/master/word_language_model
    """

    def __init__(self, d_model: int = 512, nhead: int = 8, num_encoder_layers: int = 6,
                 num_decoder_layers: int = 6, dim_feedforward: int = 2048, dropout: float = 0.1,
                 activation: str = "relu", custom_encoder: Optional[Any] = None, custom_decoder: Optional[Any] = None) -> None:
        super(Transformer, self).__init__()

        if custom_encoder is not None:
            self.encoder = custom_encoder
        else:
            encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
            encoder_norm = LayerNorm(d_model)
            self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)

        if custom_decoder is not None:
            self.decoder = custom_decoder
        else:
            decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
            decoder_norm = LayerNorm(d_model)
            self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)

        self._reset_parameters()

        self.d_model = d_model
        self.nhead = nhead

    def forward(self, src: Tensor, tgt: Tensor, src_mask: Optional[Tensor] = None, tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""Take in and process masked source/target sequences.

        Args:
            src: the sequence to the encoder (required).
            tgt: the sequence to the decoder (required).
            src_mask: the additive mask for the src sequence (optional).
            tgt_mask: the additive mask for the tgt sequence (optional).
            memory_mask: the additive mask for the encoder output (optional).
            src_key_padding_mask: the ByteTensor mask for src keys per batch (optional).
            tgt_key_padding_mask: the ByteTensor mask for tgt keys per batch (optional).
            memory_key_padding_mask: the ByteTensor mask for memory keys per batch (optional).

        Shape:
            - src: :math:`(S, N, E)`.
            - tgt: :math:`(T, N, E)`.
            - src_mask: :math:`(S, S)`.
            - tgt_mask: :math:`(T, T)`.
            - memory_mask: :math:`(T, S)`.
            - src_key_padding_mask: :math:`(N, S)`.
            - tgt_key_padding_mask: :math:`(N, T)`.
            - memory_key_padding_mask: :math:`(N, S)`.

            Note: [src/tgt/memory]_mask ensures that position i is allowed to attend the unmasked
            positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
            while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
            are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
            is provided, it will be added to the attention weight.
            [src/tgt/memory]_key_padding_mask provides specified elements in the key to be ignored by
            the attention. If a ByteTensor is provided, the non-zero positions will be ignored while the zero
            positions will be unchanged. If a BoolTensor is provided, the positions with the
            value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.

            - output: :math:`(T, N, E)`.

            Note: Due to the multi-head attention architecture in the transformer model,
            the output sequence length of a transformer is same as the input sequence
            (i.e. target) length of the decode.

            where S is the source sequence length, T is the target sequence length, N is the
            batch size, E is the feature number

        Examples:
            >>> output = transformer_model(src, tgt, src_mask=src_mask, tgt_mask=tgt_mask)
        """

        if src.size(1) != tgt.size(1):
            raise RuntimeError("the batch number of src and tgt must be equal")

        if src.size(2) != self.d_model or tgt.size(2) != self.d_model:
            raise RuntimeError("the feature number of src and tgt must be equal to d_model")

        memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
        output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
                              tgt_key_padding_mask=tgt_key_padding_mask,
                              memory_key_padding_mask=memory_key_padding_mask)
        return output

    def generate_square_subsequent_mask(self, sz: int) -> Tensor:
        r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
            Unmasked positions are filled with float(0.0).
        """
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask

    def _reset_parameters(self):
        r"""Initiate parameters in the transformer model."""

        for p in self.parameters():
            if p.dim() > 1:
                xavier_uniform_(p)


class TransformerEncoder(Module):
    r"""TransformerEncoder is a stack of N encoder layers

    Args:
        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
        num_layers: the number of sub-encoder-layers in the encoder (required).
        norm: the layer normalization component (optional).

    Examples::
        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
        >>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
        >>> src = torch.rand(10, 32, 512)
        >>> out = transformer_encoder(src)
    """
    __constants__ = ['norm']

    def __init__(self, encoder_layer, num_layers, norm=None):
        super(TransformerEncoder, self).__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""Pass the input through the encoder layers in turn.

        Args:
            src: the sequence to the encoder (required).
            mask: the mask for the src sequence (optional).
            src_key_padding_mask: the mask for the src keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """
        output = src

        for mod in self.layers:
            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)

        if self.norm is not None:
            output = self.norm(output)

        return output


class TransformerDecoder(Module):
    r"""TransformerDecoder is a stack of N decoder layers

    Args:
        decoder_layer: an instance of the TransformerDecoderLayer() class (required).
        num_layers: the number of sub-decoder-layers in the decoder (required).
        norm: the layer normalization component (optional).

    Examples::
        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
        >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
        >>> memory = torch.rand(10, 32, 512)
        >>> tgt = torch.rand(20, 32, 512)
        >>> out = transformer_decoder(tgt, memory)
    """
    __constants__ = ['norm']

    def __init__(self, decoder_layer, num_layers, norm=None):
        super(TransformerDecoder, self).__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""Pass the inputs (and mask) through the decoder layer in turn.

        Args:
            tgt: the sequence to the decoder (required).
            memory: the sequence from the last layer of the encoder (required).
            tgt_mask: the mask for the tgt sequence (optional).
            memory_mask: the mask for the memory sequence (optional).
            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
            memory_key_padding_mask: the mask for the memory keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """
        output = tgt

        for mod in self.layers:
            output = mod(output, memory, tgt_mask=tgt_mask,
                         memory_mask=memory_mask,
                         tgt_key_padding_mask=tgt_key_padding_mask,
                         memory_key_padding_mask=memory_key_padding_mask)

        if self.norm is not None:
            output = self.norm(output)

        return output

class TransformerEncoderLayer(Module):
    r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
    This standard encoder layer is based on the paper "Attention Is All You Need".
    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
    in a different way during application.

    Args:
        d_model: the number of expected features in the input (required).
        nhead: the number of heads in the multiheadattention models (required).
        dim_feedforward: the dimension of the feedforward network model (default=2048).
        dropout: the dropout value (default=0.1).
        activation: the activation function of intermediate layer, relu or gelu (default=relu).

    Examples::
        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
        >>> src = torch.rand(10, 32, 512)
        >>> out = encoder_layer(src)
    """

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu"):
        super(TransformerEncoderLayer, self).__init__()
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = Linear(d_model, dim_feedforward)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model)

        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        self.activation = _get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.relu
        super(TransformerEncoderLayer, self).__setstate__(state)

    def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""Pass the input through the encoder layer.

        Args:
            src: the sequence to the encoder layer (required).
            src_mask: the mask for the src sequence (optional).
            src_key_padding_mask: the mask for the src keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """
        # src: [81, 2, 128]
        src2 = self.self_attn(src, src, src, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src


class TransformerDecoderLayer(Module):
    r"""TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.
    This standard decoder layer is based on the paper "Attention Is All You Need".
    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
    in a different way during application.

    Args:
        d_model: the number of expected features in the input (required).
        nhead: the number of heads in the multiheadattention models (required).
        dim_feedforward: the dimension of the feedforward network model (default=2048).
        dropout: the dropout value (default=0.1).
        activation: the activation function of intermediate layer, relu or gelu (default=relu).

    Examples::
        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
        >>> memory = torch.rand(10, 32, 512)
        >>> tgt = torch.rand(20, 32, 512)
        >>> out = decoder_layer(tgt, memory)
    """

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu"):
        super(TransformerDecoderLayer, self).__init__()
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = Linear(d_model, dim_feedforward)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model)

        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.norm3 = LayerNorm(d_model)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)
        self.dropout3 = Dropout(dropout)

        self.activation = _get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.relu
        super(TransformerDecoderLayer, self).__setstate__(state)

    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""Pass the inputs (and mask) through the decoder layer.

        Args:
            tgt: the sequence to the decoder layer (required).
            memory: the sequence from the last layer of the encoder (required).
            tgt_mask: the mask for the tgt sequence (optional).
            memory_mask: the mask for the memory sequence (optional).
            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
            memory_key_padding_mask: the mask for the memory keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """
        tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt


def _get_clones(module, N):
    return ModuleList([copy.deepcopy(module) for i in range(N)])


def _get_activation_fn(activation):
    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu

    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))

model/mytransformer.py
YoungSeng/UnifiedGesture
import pdb

import torch

from torch import nn


class Tisa(nn.Module):
    def __init__(self, num_attention_heads: int = 12, num_kernels: int = 5):
        super().__init__()
        self.num_attention_heads = num_attention_heads
        self.num_kernels = num_kernels

        self.kernel_offsets = nn.Parameter(
            torch.Tensor(self.num_kernels, self.num_attention_heads)
        )
        self.kernel_amplitudes = nn.Parameter(
            torch.Tensor(self.num_kernels, self.num_attention_heads)
        )
        self.kernel_sharpness = nn.Parameter(
            torch.Tensor(self.num_kernels, self.num_attention_heads)
        )
        self._init_weights()

    def create_relative_offsets(self, seq_len: int):
        """Creates offsets for all the relative distances between
        -seq_len + 1 to seq_len - 1."""
        return torch.arange(-seq_len, seq_len + 1)

    def compute_positional_scores(self, relative_offsets):
        """Takes seq_len and outputs position scores for each relative distance.
        This implementation uses radial basis functions. Override this function to
        use other scoring functions than the example in the paper."""
        rbf_scores = (
            self.kernel_amplitudes.unsqueeze(-1)
            * torch.exp(
                -torch.abs(self.kernel_sharpness.unsqueeze(-1))
                * ((self.kernel_offsets.unsqueeze(-1) - relative_offsets) ** 2)
            )
        ).sum(axis=0)
        return rbf_scores

    def scores_to_toeplitz_matrix(self, positional_scores, seq_len: int):
        """Converts the TISA positional scores into the final matrix for the
        self-attention equation. PRs with memory and/or speed optimizations are
        welcome."""
        deformed_toeplitz = (
            (
                (torch.arange(0, -(seq_len ** 2), step=-1) + (seq_len - 1)).view(
                    seq_len, seq_len
                )
                + (seq_len + 1) * torch.arange(seq_len).view(-1, 1)
            )
            .view(-1)
            .long()
            .to(device=positional_scores.device)
        )
        expanded_positional_scores = torch.take_along_dim(
            positional_scores, deformed_toeplitz.view(1, -1), 1
        ).view(self.num_attention_heads, seq_len, seq_len)
        return expanded_positional_scores

    def forward(self, seq_len: int):
        """Computes the translation-invariant positional contribution to the
        attention matrix in the self-attention module of transformer models."""
        if not self.num_kernels:
            return torch.zeros((self.num_attention_heads, seq_len, seq_len))
        positional_scores_vector = self.compute_positional_scores(
            self.create_relative_offsets(seq_len)
        )
        positional_scores_matrix = self.scores_to_toeplitz_matrix(
            positional_scores_vector, seq_len
        )
        return positional_scores_matrix

    def visualize(self, seq_len: int = 10, attention_heads=None):
        """Visualizes the TISA interpretability by plotting position scores as
        a function of relative distance for each attention head."""
        if attention_heads is None:
            attention_heads = list(range(self.num_attention_heads))
        import matplotlib.pyplot as plt

        x = self.create_relative_offsets(seq_len).detach().numpy()
        y = (
            self.compute_positional_scores(self.create_relative_offsets(seq_len))
            .detach()
            .numpy()
        )
        for i in attention_heads:
            plt.plot(x, y[i])
        plt.savefig('./pic-tisa.png')
        plt.show()

    def _init_weights(self):
        """Initialize the weights"""
        ampl_init_mean = 0.1
        sharpness_init_mean = 0.1
        torch.nn.init.normal_(self.kernel_offsets, mean=0.0, std=5.0)
        torch.nn.init.normal_(
            self.kernel_amplitudes, mean=ampl_init_mean, std=0.1 * ampl_init_mean
        )
        torch.nn.init.normal_(
            self.kernel_sharpness,
            mean=sharpness_init_mean,
            std=0.1 * sharpness_init_mean,
        )


def main():
    tisa = Tisa()
    positional_scores = tisa(20)
    pdb.set_trace()
    tisa.visualize(seq_len=20)


if __name__ == "__main__":
    main()


model/tisa.py
YoungSeng/UnifiedGesture

import pdb

import logging
logging.getLogger().setLevel(logging.INFO)

from torch.utils.data import DataLoader
from data_loader.lmdb_data_loader import TrinityDataset
import time
import torch
import torch.nn as nn
import yaml
from pprint import pprint
from easydict import EasyDict
from configs.parse_args import parse_args
import os
from torch import optim
from Hierarchical_XTransformer import Hierarchical_XTransformer, cascade_XTransformer, XTransformer_GRU, XTransformer_GPT, GPT_GRU
from datetime import datetime
from torch.utils.tensorboard import SummaryWriter
from pathlib import Path


args = parse_args()
mydevice = torch.device('cuda:' + args.gpu)

def evaluate_testset(model, test_data_loader):
    start = time.time()
    model = model.eval()
    tot_euclidean_error_1 = 0
    tot_euclidean_error_2 = 0
    tot_eval_nums = 0

    with torch.no_grad():
        for iter_idx, data in enumerate(test_data_loader, 0):
            speech, gesture, lower, _ = data  # （batch, 40, 135）
            speech = speech.to(mydevice)
            gesture = gesture.to(mydevice)
            lower = lower.to(mydevice)
            _, _, loss1, loss2 = model(speech, gesture, lower)
            tot_euclidean_error_1 += loss1.item()
            tot_euclidean_error_2 += loss2.item()
            tot_eval_nums += 1

        print('generation took {:.2} s'.format(time.time() - start))
    model.train()
    return tot_euclidean_error_1 / (tot_eval_nums * 1.0), tot_euclidean_error_2 / (tot_eval_nums * 1.0)


def main(args):
    # dataset
    train_dataset = TrinityDataset(args.train_data_path,
                                   n_poses=args.n_poses,
                                   subdivision_stride=args.subdivision_stride,
                                   pose_resampling_fps=args.motion_resampling_framerate)
    train_loader = DataLoader(dataset=train_dataset, batch_size=args.batch_size,
                              shuffle=True, drop_last=True, num_workers=args.loader_workers, pin_memory=True)

    val_dataset = TrinityDataset(args.val_data_path,
                                       n_poses=args.n_poses,
                                       subdivision_stride=args.subdivision_stride,
                                       pose_resampling_fps=args.motion_resampling_framerate)
    test_loader = DataLoader(dataset=val_dataset, batch_size=args.batch_size,
                             shuffle=False, drop_last=True, num_workers=args.loader_workers, pin_memory=True)

    logging.info('len of train loader:{}, len of test loader:{}'.format(len(train_loader), len(test_loader)))

    model = GPT_GRU(args.batch_size, args.n_poses, mydevice)
    model = nn.DataParallel(model, device_ids=[eval(i) for i in config.no_cuda])
    model = model.to(mydevice)

    best_val_loss = (1e+2, 0)  # value, epoch

    if not os.path.exists(args.model_save_path):
        os.mkdir(args.model_save_path)

    # optimizer = optim.Adam(itertools.chain(model.module.parameters()), lr=args.lr, betas=args.betas)
    optimizer = optim.Adam(model.module.parameters(), lr=args.lr, betas=args.betas)
    schedular = optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.milestones, gamma=args.gamma)

    updates = 0
    total = len(train_loader)

    tb_path = args.name + '_' + str(datetime.now().strftime('%Y%m%d_%H%M%S'))
    tb_writer = SummaryWriter(log_dir=str(Path(args.model_save_path).parent / 'tensorboard_runs' / tb_path))

    for epoch in range(1, args.epochs + 1):
        logging.info(f'Epoch: {epoch}')
        i = 0
        total_loss1, total_loss2 = evaluate_testset(model, test_loader)
        logging.info('loss on validation: {:.3f}, {:.3f}'.format(total_loss1, total_loss2))
        # tb_writer.add_scalar('upper mean/validation', upper_mean, epoch)
        total_loss = total_loss1 + total_loss2
        is_best = total_loss < best_val_loss[0]
        if is_best:
            logging.info(' *** BEST VALIDATION LOSS : {:.3f}'.format(total_loss))
            best_val_loss = (total_loss, epoch)
        else:
            logging.info(' best validation loss so far: {:.3f} at EPOCH {}'.format(best_val_loss[0], best_val_loss[1]))

        if is_best or (epoch % args.save_per_epochs == 0):
            if is_best:
                save_name = '{}/{}_checkpoint_best.bin'.format(args.model_save_path, args.name)
            else:
                save_name = '{}/{}_checkpoint_{:03d}.bin'.format(args.model_save_path, args.name, epoch)

            torch.save({
                'args': args, "epoch": epoch, 'model_dict': model.state_dict()
            }, save_name)
            logging.info('Saved the checkpoint')

        # train model
        model = model.train()
        start = datetime.now()
        for batch_i, batch in enumerate(train_loader, 0):
            speech, gesture, lower, _ = batch
            optimizer.zero_grad()
            _, _, loss1, loss2 = model(speech, gesture, lower)

            # write to tensorboard
            # tb_writer.add_scalar('loss' + '/train', loss, updates)
            # for key in metrics.keys():
            #     tb_writer.add_scalar(key + '/train', metrics[key], updates)
            loss = loss1 + loss2
            loss.backward()
            optimizer.step()
            # log
            stats = {'updates': updates, 'loss1': loss1.item(), 'loss2': loss2.item()}
            stats_str = ' '.join(f'{key}[{val:.8f}]' for key, val in stats.items())
            i += 1
            remaining = str((datetime.now() - start) / i * (total - i))
            remaining = remaining.split('.')[0]
            logging.info(f'> epoch [{epoch}] updates[{i}] {stats_str} eta[{remaining}]')
            # if i == total:
            #     logging.debug('\n')
            #     logging.debug(f'elapsed time: {str(datetime.now() - start).split(".")[0]}')
            updates += 1
        schedular.step()
    tb_writer.close()
    # print best losses
    logging.info('--------- Final best loss values ---------')
    logging.info('diff mean: {:.3f} at EPOCH {}'.format(best_val_loss[0], best_val_loss[1]))


if __name__ == '__main__':
    '''
    pip install nni

    python pretrain.py --config=./configs/pretrain.yml --train --no_cuda 2 --gpu 2
    '''

    with open(args.config) as f:
        config = yaml.safe_load(f)

    for k, v in vars(args).items():
        config[k] = v
    pprint(config)

    config = EasyDict(config)

    main(config)

pretrain.py
YoungSeng/UnifiedGesture
import os
import pdb
import torchaudio
from conditional_gpt.encodec.utils import convert_audio
import numpy as np
import torch


def clip_audio(src_audio_path, src_motion_path, save_path, mode):

    from conditional_gpt.encodec import EncodecModel
    model = EncodecModel.encodec_model_24khz()
    model.set_target_bandwidth(1.5)

    if not os.path.exists(save_path):  # 判断是否存在文件夹如果不存在则创建为文件夹
        os.makedirs(save_path)  # makedirs 创建文件时如果路径不存在会创建这个路径
    if mode == 'ZEGGS':
        flag = 461
        audio_sample_rate = 16000
    elif mode == 'Trinity':
        flag = 343
        audio_sample_rate = 48000
    for bvh_file in os.listdir(src_motion_path):
        print(bvh_file)
        with open(os.path.join(src_motion_path, bvh_file), 'r') as f:
            content = f.readlines()
        motion_frames = eval(content[flag].strip().split(':')[-1])
        motion_frames = motion_frames // 4 * 4      # retargeting 4x
        length = motion_frames * audio_sample_rate // 30

        audio_file, sr = torchaudio.load(os.path.join(src_audio_path, bvh_file[:-4] + '.wav'))

        print(sr, motion_frames, length, audio_file.shape[1])
        audio_file = audio_file[:, :length]
        # assert audio_file.shape[1] == length

        wav = convert_audio(audio_file, sr, model.sample_rate, model.channels)
        wav = wav.unsqueeze(0)
        with torch.no_grad():
            encoded_frames = model.encode(wav)
        codes = torch.cat([encoded[0] for encoded in encoded_frames], dim=-1)  # [B, n_q, T]        # [1, 8, 1500]
        codes = codes.squeeze(0).transpose(0, 1).numpy()
        np.save(os.path.join(save_path, bvh_file[:-4] + '.npy'), codes)


def process_audio(src_audio_path, src_motion_path, save_path, mode):
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    if 'ZEGGS' in mode:
        flag = 461
        audio_sample_rate = 16000
    elif 'Trinity' in mode:
        flag = 343
        audio_sample_rate = 48000
    if 'valid' not in mode:
        for bvh_file in os.listdir(src_motion_path):
            print(bvh_file)
            with open(os.path.join(src_motion_path, bvh_file), 'r') as f:
                content = f.readlines()
            motion_frames = eval(content[flag].strip().split(':')[-1])
            motion_frames = motion_frames // 4 * 4      # retargeting 4x
            length = motion_frames * audio_sample_rate // 30

            audio_file, sr = torchaudio.load(os.path.join(src_audio_path, bvh_file[:-4] + '.wav'))

            print(sr, motion_frames, length, audio_file.shape[1])
            audio_file = audio_file[:, :length]
            # assert audio_file.shape[1] == length

            wav = convert_audio(audio_file, sr, 16000, 1)

            np.save(os.path.join(save_path, bvh_file[:-4] + '.npy'), wav.numpy())
    else:
        for wav_file in os.listdir(src_audio_path):
            print(wav_file)
            audio_file, sr = torchaudio.load(os.path.join(src_audio_path, wav_file))
            wav = convert_audio(audio_file, sr, 16000, 1)
            np.save(os.path.join(save_path, wav_file[:-4] + '.npy'), wav.numpy())


if __name__ == '__main__':
    '''
    python process_audio.py
    '''
    trinity_motion_path = './retargeting/datasets/Trinity_ZEGGS/Trinity/'
    zeggs_motion_path = './retargeting/datasets/Trinity_ZEGGS/ZEGGS/'
    trinity_audio_path = './dataset/Trinity/audio/'
    zeggs_audio_path = './dataset/ZEGGS/clean/'
    # trinity_audio_clip_path = './dataset/Trinity/audio_clips/'
    # zeggs_audio_clip_path = './dataset/ZEGGS/audio_clips/'
    trinity_audio_proceed_path = './dataset/Trinity/all_speech/'
    zeggs_audio_proceed_path = './dataset/ZEGGS/all_speech/'

    # 1. clip audio
    # clip_audio(zeggs_audio_path, zeggs_motion_path, zeggs_audio_clip_path, 'ZEGGS')
    # clip_audio(trinity_audio_path, trinity_motion_path, trinity_audio_clip_path, 'Trinity')

    # 2. process training audio
    process_audio(zeggs_audio_path, zeggs_motion_path, zeggs_audio_proceed_path, 'ZEGGS')
    process_audio(trinity_audio_path, trinity_motion_path, trinity_audio_proceed_path, 'Trinity')

    # 3. process valid audio
    # trinity_audio_proceed_path = './dataset/Trinity/valid_speech/'
    # zeggs_audio_proceed_path = './dataset/ZEGGS/valid_speech/'
    # process_audio('../ZEGGS/valid/', zeggs_motion_path, zeggs_audio_proceed_path, 'ZEGGS_valid')
    # process_audio('../trinity/test/', trinity_motion_path, trinity_audio_proceed_path, 'Trinity_valid')

process_audio.py
YoungSeng/UnifiedGesture
import os
import glob
import pdb

import numpy as np


def process_code(source_path, save_path):
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    upper_files = sorted(glob.glob(source_path + "/code_upper_*.npy"))
    for item in upper_files:
        name = os.path.split(item)[1].replace('code_upper_', '')
        print(name)
        upper = np.load(item)
        lower = np.load(item.replace('upper', 'lower'))
        # root = np.load(item.replace('upper', 'root'))
        # print(upper.shape, lower.shape, root.shape)
        # latent_code = np.concatenate([upper, lower, root], axis=1)

        print(upper.shape, lower.shape)
        latent_code = np.concatenate([upper, lower], axis=1)

        np.save(os.path.join(save_path, name), latent_code)


def process_motion(source_code_path, source_motion_path, source_aux_path, save_path):
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    root_motions = sorted(glob.glob(source_motion_path + "/*_root.npy"))
    for item in root_motions:
        name = os.path.split(item)[1].replace('_root.npy', '')
        print(name)
        root_motion = np.load(item)
        lower_motion = np.load(os.path.join(source_motion_path, name + '_lower' + '.npy'))
        upper_motion = np.load(os.path.join(source_motion_path, name + '_upper' + '.npy'))
        all_body_motion = np.concatenate([lower_motion, upper_motion, root_motion], axis=1)
        all_body_motion = all_body_motion[0].transpose(1, 0)      # (len, 112)

        all_body_vel = all_body_motion[:-1] - all_body_motion[1:]
        all_body_acc = all_body_vel[:-1] - all_body_vel[1:]
        clip_skeleton_vel = np.pad(all_body_vel, ((1, 0), (0, 0)), 'constant')
        clip_skeleton_acc = np.pad(all_body_acc, ((2, 0), (0, 0)), 'constant')

        aux = np.load(os.path.join(source_aux_path, name + '.bvh.npy'))
        aux = aux[::4][:aux.shape[0]//4]

        all_body = np.concatenate([all_body_motion, clip_skeleton_vel, clip_skeleton_acc, aux], axis=1)

        upper_code = np.load(os.path.join(source_code_path, 'code_upper_' + name + '.npy'))
        upper_code = upper_code[:, 0]      # (2694,)

        np.save(os.path.join(save_path, name + '_upper' + '.npy'), upper_code)
        np.save(os.path.join(save_path, name + '_all_body' + '.npy'), all_body)


if __name__ == '__main__':
    '''
    python process_code.py
    '''
    # Trinity_VQVAE_path = './retargeting/datasets/Trinity_ZEGGS/VQVAE_result/Trinity/'
    # ZEGGS_VQVAE_path = './retargeting/datasets/Trinity_ZEGGS/VQVAE_result/ZEGGS/'
    # Trinity_save_code_path = './retargeting/datasets/Trinity_ZEGGS/latent_code_2/Trinity/'
    # ZEGGS_save_code_path = './retargeting/datasets/Trinity_ZEGGS/latent_code_2/ZEGGS/'
    # process_code(Trinity_VQVAE_path, Trinity_save_code_path)
    # process_code(ZEGGS_VQVAE_path, ZEGGS_save_code_path)



    Trinity_save_code_path = './retargeting/datasets/Trinity_ZEGGS/VQVAE_result/Trinity/'
    ZEGGS_save_code_path = './retargeting/datasets/Trinity_ZEGGS/VQVAE_result/ZEGGS/'
    Trinity_source_motion_path = './retargeting/datasets/Trinity_ZEGGS/bvh2upper_lower_root/Trinity/'
    ZEGGS_source_motion_path = './retargeting/datasets/Trinity_ZEGGS/bvh2upper_lower_root/ZEGGS/'
    Trinity_source_aux_path = "./retargeting/datasets/Trinity_ZEGGS/Trinity_aux/"
    ZEGGS_source_aux_path = "./retargeting/datasets/Trinity_ZEGGS/ZEGGS_aux/"
    Trinity_save_path = './dataset/Trinity/all_gesture_aux/'
    ZEGGS_save_path = './dataset/datasets/ZEGGS/all_gesture_aux/'
    process_motion(Trinity_save_code_path, Trinity_source_motion_path, Trinity_source_aux_path, Trinity_save_path)
    process_motion(ZEGGS_save_code_path, ZEGGS_source_motion_path, ZEGGS_source_aux_path, ZEGGS_save_path)

process_code.py
YoungSeng/UnifiedGesture
import pdb


def get_character_names(args):
    if args.is_train:
        """
        Put the name of subdirectory in retargeting/datasets/Mixamo as [[names of group A], [names of group B]]
        """
        # characters = [['Aj', 'BigVegas', 'Kaya', 'SportyGranny'],
        #               ['Malcolm_m', 'Remy_m', 'Maria_m', 'Jasper_m', 'Knight_m',
        #                'Liam_m', 'ParasiteLStarkie_m', 'Pearl_m', 'Michelle_m', 'LolaB_m',
        #                'Pumpkinhulk_m', 'Ortiz_m', 'Paladin_m', 'James_m', 'Joe_m',
        #                'Olivia_m', 'Yaku_m', 'Timmy_m', 'Racer_m', 'Abe_m']]

        # characters = [['ch01'], ['ch02']]
        # characters = [['Talking_With_Hands'], ['Trinity']]
        characters = [['Trinity'], ['ZEGGS']]
        # characters = [['Talking_With_Hands'], ['ZEGGS']]

    else:
        """
        To run evaluation successfully, number of characters in both groups must be the same. Repeat is okay.
        """
        # characters = [['BigVegas', 'BigVegas', 'BigVegas', 'BigVegas'],  ['Mousey_m', 'Goblin_m', 'Mremireh_m', 'Vampire_m']]
        # tmp = characters[1][args.eval_seq]
        # characters[1][args.eval_seq] = characters[1][0]
        # characters[1][0] = tmp

        pass    # 20230320

    return characters


def create_dataset(args, character_names=None):
    from datasets.combined_motion import TestData, MixedData

    if args.is_train:       # 1
        return MixedData(args, character_names)
    else:
        return TestData(args, character_names)


def get_test_set():
    with open('./datasets/Mixamo/test_list.txt', 'r') as file:
        list = file.readlines()
        list = [f[:-1] for f in list]
        return list


# def get_train_list():
#     with open('./datasets/Mixamo/train_list.txt', 'r') as file:
#         list = file.readlines()
#         list = [f[:-1] for f in list]
#         return list


if __name__ == '__main__':
    '''
    python ./datasets/__init__.py
    '''
    import sys
    [sys.path.append(i) for i in ['./', '../', "../../utils"]]
    import option_parser
    args = option_parser.get_args()
    characters = get_character_names(args)      # [['ch01'], ['ch02']]
    # dataset = create_dataset(args, characters)
    # from torch.utils.data.dataloader import DataLoader
    # data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=2)
    # for step, motions in enumerate(data_loader):
    #     pass
        # print(motions.shape)        # [[[55, 91, 64], [55]], [[55, 91, 64], [55]]]

    # Test
    from datasets.combined_motion import TestData
    dataset = TestData(args, characters)
    from torch.utils.data.dataloader import DataLoader

    data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=2)
    # pdb.set_trace()
    for step, motions in enumerate(data_loader):
        # pdb.set_trace()
        print(motions.shape)

retargeting/datasets/__init__.py
YoungSeng/UnifiedGesture
import pdb
import torch
import sys
sys.path.append('../../utils')
sys.path.append("../utils")
sys.path.append('../')

import BVH_mod as BVH
import numpy as np
from Quaternions import Quaternions
from models.Kinematics import ForwardKinematics
from models.skeleton import build_edge_topology
from option_parser import get_std_bvh
from datasets.bvh_writer import write_bvh

"""
1.
Specify the joints that you want to use in training and test. Other joints will be discarded.
Please start with root joint, then left leg chain, right leg chain, head chain, left shoulder chain and right shoulder chain.
See the examples below.
"""
corps_name_1 = ['Pelvis', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'Hips', 'Spine', 'Spine1', 'Spine2', 'Neck', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']
corps_name_2 = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'LeftToe_End', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'RightToe_End', 'Spine', 'Spine1', 'Spine2', 'Neck', 'Head', 'HeadTop_End', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']
corps_name_3 = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'RightUpLeg', 'RightLeg', 'RightFoot', 'Spine', 'Spine1', 'Neck', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']
corps_name_boss = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'Spine', 'Spine1', 'Spine2', 'Neck', 'Neck1', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']
corps_name_boss2 = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'Left_End', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'Right_End', 'Spine', 'Spine1', 'Spine2', 'Neck', 'Neck1', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']
corps_name_cmu = ['Hips', 'LHipJoint', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'RHipJoint', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'LowerBack', 'Spine', 'Spine1', 'Neck', 'Neck1', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']
corps_name_monkey = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'Spine', 'Spine1', 'Neck', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']
corps_name_three_arms = ['Three_Arms_Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'Spine', 'Spine1', 'Neck', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']
corps_name_three_arms_split = ['Three_Arms_split_Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'Spine', 'Spine1', 'Neck', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'LeftHand_split', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand', 'RightHand_split']
corps_name_Prisoner = ['HipsPrisoner', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'LeftToe_End', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'RightToe_End', 'Spine', 'Spine1', 'Spine2', 'Neck', 'Head', 'HeadTop_End', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm']
corps_name_mixamo2_m = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'LeftToe_End', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'RightToe_End', 'Spine', 'Spine1', 'Spine1_split', 'Spine2', 'Neck', 'Head', 'HeadTop_End', 'LeftShoulder', 'LeftShoulder_split', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightShoulder_split', 'RightArm', 'RightForeArm', 'RightHand']
corps_name_GENEA_2020 = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftForeFoot', 'LeftToeBase', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightForeFoot', 'RightToeBase', 'Spine', 'Spine1', 'Spine2', 'Spine3', 'Neck', 'Neck1', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']      # Trinity
# corps_name_GENEA_2020 = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftForeFoot', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightForeFoot', 'Spine', 'Spine1', 'Spine2', 'Spine3', 'Neck', 'Neck1', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']        # model_8
# corps_name_GENEA_2020 = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'RightUpLeg', 'RightLeg', 'RightFoot', 'Spine', 'Spine1', 'Spine2', 'Spine3', 'Neck', 'Neck1', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']      # Trinity, model_9


# modify for test
# corps_name_GENEA_2022 = ['b_root', 'b_l_upleg', 'b_l_leg', 'b_l_foot_twist', 'b_l_foot', 'b_r_upleg', 'b_r_leg', 'b_r_foot_twist', 'b_r_foot', 'b_spine0', 'b_spine1', 'b_spine2', 'b_spine3', 'b_neck0', 'b_head', 'b_l_shoulder', 'p_l_scap', 'b_l_arm', 'b_l_arm_twist', 'b_l_forearm', 'b_l_wrist_twist', 'b_l_wrist', 'b_r_shoulder', 'p_r_scap', 'b_r_arm', 'b_r_arm_twist', 'b_r_forearm', 'b_r_wrist_twist', 'b_r_wrist']       # Talking_with_hand       model_3, model_4, model_5
corps_name_GENEA_2022 = ['b_root', 'b_l_upleg', 'b_l_leg', 'b_l_foot_twist', 'b_l_foot', 'b_r_upleg', 'b_r_leg', 'b_r_foot_twist', 'b_r_foot', 'b_spine0', 'b_spine1', 'b_spine2', 'b_spine3', 'b_neck0', 'b_head', 'b_l_shoulder', 'p_l_scap', 'b_l_arm', 'b_l_arm_twist', 'b_l_forearm', 'b_l_wrist_twist', 'b_r_shoulder', 'p_r_scap', 'b_r_arm', 'b_r_arm_twist', 'b_r_forearm', 'b_r_wrist_twist']       # Talking_with_hand     model_2, model_6


ZEGGS = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'LeftToeBaseEnd', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'RightToeBaseEnd', 'Spine', 'Spine1', 'Spine2', 'Spine3', 'Neck', 'Neck1', 'Head', 'HeadEnd', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']      # ZEGGS
# ZEGGS = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'LeftToeBaseEnd', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'RightToeBaseEnd', 'Spine', 'Spine1', 'Spine2', 'Spine3', 'Neck', 'Neck1', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']      # my_model_new_3


# corps_name_example = ['Root', 'LeftUpLeg', ..., 'LeftToe', 'RightUpLeg', ..., 'RightToe', 'Spine', ..., 'Head', 'LeftShoulder', ..., 'LeftHand', 'RightShoulder', ..., 'RightHand']

"""
2.
Specify five end effectors' name.
Please follow the same order as in 1.
"""
ee_name_1 = ['LeftToeBase', 'RightToeBase', 'Head', 'LeftHand', 'RightHand']
ee_name_2 = ['LeftToe_End', 'RightToe_End', 'HeadTop_End', 'LeftHand', 'RightHand']
ee_name_3 = ['LeftFoot', 'RightFoot', 'Head', 'LeftHand', 'RightHand']
ee_name_cmu = ['LeftToeBase', 'RightToeBase', 'Head', 'LeftHand', 'RightHand']
ee_name_monkey = ['LeftToeBase', 'RightToeBase', 'Head', 'LeftHand', 'RightHand']
ee_name_three_arms_split = ['LeftToeBase', 'RightToeBase', 'Head', 'LeftHand_split', 'RightHand_split']
ee_name_Prisoner = ['LeftToe_End', 'RightToe_End', 'HeadTop_End', 'LeftHand', 'RightForeArm']
ee_name_GENEA_2020 = ['LeftToeBase', 'RightToeBase', 'Head', 'LeftHand', 'RightHand']       # new
# ee_name_GENEA_2020 = ['LeftForeFoot', 'RightForeFoot', 'Head', 'LeftHand', 'RightHand']       # model_8
# ee_name_GENEA_2020 = ['LeftFoot', 'RightFoot', 'Head', 'LeftHand', 'RightHand']       # model_9

# modify for test
# ee_name_GENEA_2022 = ['b_l_foot', 'b_r_foot', 'b_head', 'b_l_wrist', 'b_r_wrist']       # model_3, model_4, model_5
ee_name_GENEA_2022 = ['b_l_foot', 'b_r_foot', 'b_head', 'b_l_wrist_twist', 'b_r_wrist_twist']     # model_2, model_6

ee_name_zeggs = ['LeftToeBaseEnd', 'RightToeBaseEnd', 'Head', 'LeftHand', 'RightHand']

# ee_name_example = ['LeftToe', 'RightToe', 'Head', 'LeftHand', 'RightHand']


# corps_names = [corps_name_1, corps_name_2, corps_name_3, corps_name_cmu, corps_name_monkey, corps_name_boss,
#                corps_name_boss, corps_name_three_arms, corps_name_three_arms_split, corps_name_Prisoner, corps_name_mixamo2_m,
#                corps_name_GENEA_2022, corps_name_GENEA_2020, ZEGGS]
#
# ee_names = [ee_name_1, ee_name_2, ee_name_3, ee_name_cmu, ee_name_monkey, ee_name_1,
#             ee_name_1, ee_name_1, ee_name_three_arms_split, ee_name_Prisoner, ee_name_2,
#             ee_name_GENEA_2022, ee_name_GENEA_2020, ee_name_zeggs]

# corps_names = [corps_name_GENEA_2022, corps_name_GENEA_2020]        # model_4
# ee_names = [ee_name_GENEA_2022, ee_name_GENEA_2020]

corps_names = [corps_name_GENEA_2020, ZEGGS]      # new
ee_names = [ee_name_GENEA_2020, ee_name_zeggs]

# corps_names = [corps_name_GENEA_2022, ZEGGS]        # new_new
# ee_names = [ee_name_GENEA_2022, ee_name_zeggs]

# corps_names = [corps_name_GENEA_2022, corps_name_GENEA_2020, ZEGGS]        # new_new_new
# ee_names = [ee_name_GENEA_2022, ee_name_GENEA_2020, ee_name_zeggs]

"""
3.
Add previously added corps_name and ee_name at the end of the two above lists.
"""
# corps_names.append(corps_name_example)
# ee_names.append(ee_name_example)


class BVH_file:
    def __init__(self, file_path=None, args=None, dataset=None, new_root=None):
        if file_path is None:
            file_path = get_std_bvh(dataset=dataset)
        self.anim, self._names, self.frametime = BVH.load(file_path)
        if new_root is not None:
            self.set_new_root(new_root)
        self.skeleton_type = -1
        self.edges = []
        self.edge_mat = []
        self.edge_num = 0
        self._topology = None
        self.ee_length = []

        for i, name in enumerate(self._names):
            if ':' in name:
                name = name[name.find(':') + 1:]
                self._names[i] = name

        full_fill = [1] * len(corps_names)
        for i, ref_names in enumerate(corps_names):
            for ref_name in ref_names:
                if ref_name not in self._names:
                    full_fill[i] = 0
                    break

        print('full_fill', full_fill)
        # # '''
        # if full_fill[3]:
        #     self.skeleton_type = 3
        # else:
        #     for i, _ in enumerate(full_fill):
        #         if full_fill[i]:
        #             self.skeleton_type = i
        #             break
        #
        # if self.skeleton_type == 2 and full_fill[4]:
        #     self.skeleton_type = 4
        # # '''
        # if 'Neck1' in self._names:
        #     self.skeleton_type = 5
        # if 'Left_End' in self._names:
        #     self.skeleton_type = 6
        # if 'Three_Arms_Hips' in self._names:
        #     self.skeleton_type = 7
        # if 'Three_Arms_Hips_split' in self._names:
        #     self.skeleton_type = 8
        #
        # if 'LHipJoint' in self._names:
        #     self.skeleton_type = 3
        #
        # if 'HipsPrisoner' in self._names:
        #     self.skeleton_type = 9
        #
        # if 'Spine1_split' in self._names:
        #     self.skeleton_type = 10
        #
        # if 'LeftForeFoot' in self._names:
        #     self.skeleton_type = 12
        #
        # if 'LeftToeBaseEnd' in self._names:
        #     self.skeleton_type = 13

        # Talking with hands
        # if 'Hips' in self._names:
        #     self.skeleton_type = 1
        # if 'b_root' in self._names:
        #     self.skeleton_type = 0

        # ZEGGS
        if 'LeftToeBaseEnd' in self._names:
            self.skeleton_type = 1
        else:
            self.skeleton_type = 0

        # Three
        # if 'LeftToeBaseEnd' in self._names:
        #     self.skeleton_type = 2
        # elif 'b_root' in self._names:
        #     self.skeleton_type = 0
        # else:
        #     self.skeleton_type = 1


        """
        4. 
        Here, you need to assign self.skeleton_type the corresponding index of your own dataset in corps_names or ee_names list.
        You can use self._names, which contains the joints name in original bvh file, to write your own if statement.
        """
        # if ...:
        #     self.skeleton_type = 11

        if self.skeleton_type == -1:
            print(self._names)
            raise Exception('Unknown skeleton')

        # if self.skeleton_type == 0:
        #     self.set_new_root(1)

        self.details = [i for i, name in enumerate(self._names) if name not in corps_names[self.skeleton_type]]
        self.joint_num = self.anim.shape[1]
        self.corps = []
        self.simplified_name = []
        self.simplify_map = {}
        self.inverse_simplify_map = {}

        for name in corps_names[self.skeleton_type]:
            for j in range(self.anim.shape[1]):
                if name == self._names[j]:
                    self.corps.append(j)
                    break

        if len(self.corps) != len(corps_names[self.skeleton_type]):
            for i in self.corps: print(self._names[i], end=' ')
            print(self.corps, self.skeleton_type, len(self.corps), sep='\n')
            raise Exception('Problem in file', file_path)

        self.ee_id = []
        for i in ee_names[self.skeleton_type]:
            self.ee_id.append(corps_names[self.skeleton_type].index(i))

        self.joint_num_simplify = len(self.corps)
        for i, j in enumerate(self.corps):
            self.simplify_map[j] = i
            self.inverse_simplify_map[i] = j
            self.simplified_name.append(self._names[j])
        self.inverse_simplify_map[0] = -1
        for i in range(self.anim.shape[1]):
            if i in self.details:
                self.simplify_map[i] = -1

        self.edges = build_edge_topology(self.topology, self.offset)

    def scale(self, alpha):
        self.anim.offsets *= alpha
        global_position = self.anim.positions[:, 0, :]
        global_position[1:, :] *= alpha
        global_position[1:, :] += (1 - alpha) * global_position[0, :]

    def rotate(self, theta, axis):
        q = Quaternions(np.hstack((np.cos(theta/2), np.sin(theta/2) * axis)))
        position = self.anim.positions[:, 0, :].copy()
        rotation = self.anim.rotations[:, 0, :]
        position[1:, ...] -= position[0:-1, ...]
        q_position = Quaternions(np.hstack((np.zeros((position.shape[0], 1)), position)))
        q_rotation = Quaternions.from_euler(np.radians(rotation))
        q_rotation = q * q_rotation
        q_position = q * q_position * (-q)
        self.anim.rotations[:, 0, :] = np.degrees(q_rotation.euler())
        position = q_position.imaginaries
        for i in range(1, position.shape[0]):
            position[i] += position[i-1]
        self.anim.positions[:, 0, :] = position

    @property
    def topology(self):
        if self._topology is None:
            self._topology = self.anim.parents[self.corps].copy()
            for i in range(self._topology.shape[0]):
                if i >= 1: self._topology[i] = self.simplify_map[self._topology[i]]
            self._topology = tuple(self._topology)
        return self._topology

    def get_ee_id(self):
        return self.ee_id

    def to_numpy(self, quater=False, edge=True):
        rotations = self.anim.rotations[:, self.corps, :]
        if quater:
            rotations = Quaternions.from_euler(np.radians(rotations)).qs
            positions = self.anim.positions[:, 0, :]
        else:
            positions = self.anim.positions[:, 0, :]
        if edge:
            index = []
            for e in self.edges:
                index.append(e[0])
            rotations = rotations[:, index, :]

        rotations = rotations.reshape(rotations.shape[0], -1)
        return np.concatenate((rotations, positions), axis=1)       # ((135, 72), (135, 3))

    def to_tensor(self, quater=False, edge=True):
        res = self.to_numpy(quater, edge)
        res = torch.tensor(res, dtype=torch.float)
        res = res.permute(1, 0)
        res = res.reshape((-1, res.shape[-1]))
        return res

    def get_position(self):
        positions = self.anim.positions
        positions = positions[:, self.corps, :]
        return positions

    @property
    def offset(self):
        return self.anim.offsets[self.corps]

    @property
    def names(self):
        return self.simplified_name

    def get_height(self):
        offset = self.offset
        topo = self.topology

        res = 0
        p = self.ee_id[0]
        while p != 0:
            res += np.dot(offset[p], offset[p]) ** 0.5
            p = topo[p]

        p = self.ee_id[2]
        while p != 0:
            res += np.dot(offset[p], offset[p]) ** 0.5
            p = topo[p]

        return res

    def write(self, file_path):
        motion = self.to_numpy(quater=False, edge=False)
        # print('bvh_parser', motion.shape)       # (1800, 90)
        rotations = motion[..., :-3].reshape(motion.shape[0], -1, 3)
        positions = motion[..., -3:]
        write_bvh(self.topology, self.offset, rotations, positions, self.names, 1.0/30, 'xyz', file_path)       # xyz

    def get_ee_length(self):
        if len(self.ee_length): return self.ee_length
        degree = [0] * len(self.topology)
        for i in self.topology:
            if i < 0: continue
            degree[i] += 1

        for j in self.ee_id:
            length = 0
            while degree[j] <= 1:
                t = self.offset[j]
                length += np.dot(t, t) ** 0.5
                j = self.topology[j]

            self.ee_length.append(length)

        height = self.get_height()
        ee_group = [[0, 1], [2], [3, 4]]
        for group in ee_group:
            maxv = 0
            for j in group:
                maxv = max(maxv, self.ee_length[j])
            for j in group:
                self.ee_length[j] *= height / maxv

        return self.ee_length

    def set_new_root(self, new_root):
        euler = torch.tensor(self.anim.rotations[:, 0, :], dtype=torch.float)
        transform = ForwardKinematics.transform_from_euler(euler, 'xyz')
        offset = torch.tensor(self.anim.offsets[new_root], dtype=torch.float)
        new_pos = torch.matmul(transform, offset)
        new_pos = new_pos.numpy() + self.anim.positions[:, 0, :]
        self.anim.offsets[0] = -self.anim.offsets[new_root]
        self.anim.offsets[new_root] = np.zeros((3, ))
        self.anim.positions[:, new_root, :] = new_pos
        rot0 = Quaternions.from_euler(np.radians(self.anim.rotations[:, 0, :]), order='xyz')
        rot1 = Quaternions.from_euler(np.radians(self.anim.rotations[:, new_root, :]), order='xyz')
        new_rot1 = rot0 * rot1
        new_rot0 = (-rot1)
        new_rot0 = np.degrees(new_rot0.euler())
        new_rot1 = np.degrees(new_rot1.euler())
        self.anim.rotations[:, 0, :] = new_rot0
        self.anim.rotations[:, new_root, :] = new_rot1

        new_seq = []
        vis = [0] * self.anim.rotations.shape[1]
        new_idx = [-1] * len(vis)
        new_parent = [0] * len(vis)

        def relabel(x):
            nonlocal new_seq, vis, new_idx, new_parent
            new_idx[x] = len(new_seq)
            new_seq.append(x)
            vis[x] = 1
            for y in range(len(vis)):
                if not vis[y] and (self.anim.parents[x] == y or self.anim.parents[y] == x):
                    relabel(y)
                    new_parent[new_idx[y]] = new_idx[x]

        relabel(new_root)
        self.anim.rotations = self.anim.rotations[:, new_seq, :]
        self.anim.offsets = self.anim.offsets[new_seq]
        names = self._names.copy()
        for i, j in enumerate(new_seq):
            self._names[i] = names[j]
        self.anim.parents = np.array(new_parent, dtype=np.int64)


if __name__ == '__main__':
    '''
    cd /ceph/hdd/yangsc21/Python/My_3/deep-motion-editing/retargeting/
    cd ./datasets/
    python bvh_parser.py
    '''

    # print(len(corps_names[1]), len(corps_names[2]), len(corps_names[4]), set(corps_names[1]) > set(corps_names[2]), set(corps_names[1]) > set(corps_names[4]))      # 真子集

    # path = "./Mixamo_ch12/ch01/Hip_Hop_Dancing.bvh"
    # path = "./Mixamo_test/Trinity/TestSeq001.bvh"
    # path = "./Mixamo_test/Talking_With_Hands/val_2022_v1_000.bvh"

    # file_1 = BVH_file(path)
    # print(file_1.anim.parents, file_1.corps, file_1._topology, file_1.topology, file_1.ee_id, file_1.get_height())        # 153.61823104990174

    path = "./Mixamo_4/Talking_With_Hands/trn_2022_v1_169.bvh"
    file_2 = BVH_file(path)
    print(file_2.anim.parents, file_2.corps, file_2._topology, file_2.topology, file_2.ee_id, file_2.get_height())  # 153.61823104990174

    # new_motion_tensor = file.to_tensor().permute((1, 0)).numpy()
    # print(new_motion_tensor.shape)     # (135, 75)
    # new_motion_npy = file.to_numpy()        #.permute((1, 0)).numpy()
    # print(new_motion_npy.shape)
    # new_motion_get_position = file.get_position()
    # print(new_motion_get_position.shape)        # (135, 25, 3)
    #
    # print(file.anim.shape, file.anim.rotations.shape, file.anim.positions.shape, file.corps, len(file.corps),
    #       len(file.edges), len(file._names), len([file._names[i] for i in file.corps]), file.topology,
    #       file.anim.parents[file.corps])
    #
    # print(file._names, file.skeleton_type)
    # print(set(file._names) > set(corps_names[file.skeleton_type]))
    #
    # # output_path = "./Mixamo_test/Trinity/TestSeq001_recon.bvh"
    # output_path = "./Mixamo_test/Talking_With_Hands/val_2022_v1_000_recon.bvh"
    # file.write(output_path)

    '''
    (135, 65), (135, 65, 3), (135, 65, 3), [0, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 31, 32, 33, 34], 25
    [(0, 1, array([ 7.294   , -3.492199, -1.3487  ])), (1, 2, array([  0.363959, -22.776356,   1.166233])), (2, 3, array([ -0.363954, -24.519972,  -1.068663])), 
    (3, 4, array([  1.087091, -12.527323,  15.23603 ])), (4, 5, array([0.053868, 0.133989, 7.430298])), 
    (0, 6, array([-7.294   , -3.492199, -0.8851  ])), (6, 7, array([ -0.363914, -22.776333,   0.70267 ])), (7, 8, array([  0.363913, -24.520035,  -0.924374])), 
    (8, 9, array([ -1.087056, -12.526639,  14.85186 ])), (9, 10, array([-0.053863,  0.134314,  7.389584])), 
    (0, 11, array([0.      , 6.292297, 0.0906  ])), (11, 12, array([0.      , 7.340935, 0.105748])), (12, 13, array([0.      , 8.38974 , 0.120856])), 
    (13, 14, array([0.      , 9.438423, 0.135962])), (14, 15, array([0.      , 9.228302, 3.0157  ])), (15, 16, array([ 0.      , 52.542191, 17.170099])), 
    (13, 17, array([2.9568  , 8.268555, 0.153114])), (17, 18, array([ 5.991047, -2.349594,  0.034302])), (18, 19, array([ 2.3597998e+01,  8.0000000e-06, -1.0000000e-06])), 
    (19, 20, array([21.444199,  0.      , -0.      ])), (13, 21, array([-2.9568  ,  8.268547,  0.11881 ])), (21, 22, array([-5.991049, -2.349586, -0.034302])), (22, 23, array([-23.596397,   0.      ,   0.      ])), (23, 24, array([-21.446606,   0.      ,   0.      ]))]
    24
    _names      ['Hips', 'Spine', 'Spine1', 'Spine2', 'Neck', 'Head', 'HeadTop_End', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'LeftHandThumb1', 'LeftHandThumb2', 'LeftHandThumb3', 'LeftHandThumb4', 'LeftHandIndex1', 'LeftHandIndex2', 'LeftHandIndex3', 'LeftHandIndex4', 'LeftHandMiddle1', 'LeftHandMiddle2', 'LeftHandMiddle3', 'LeftHandMiddle4', 'LeftHandRing1', 'LeftHandRing2', 'LeftHandRing3', 'LeftHandRing4', 'LeftHandPinky1', 'LeftHandPinky2', 'LeftHandPinky3', 'LeftHandPinky4', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand', 'RightHandThumb1', 'RightHandThumb2', 'RightHandThumb3', 'RightHandThumb4', 'RightHandIndex1', 'RightHandIndex2', 'RightHandIndex3', 'RightHandIndex4', 'RightHandMiddle1', 'RightHandMiddle2', 'RightHandMiddle3', 'RightHandMiddle4', 'RightHandRing1', 'RightHandRing2', 'RightHandRing3', 'RightHandRing4', 'RightHandPinky1', 'RightHandPinky2', 'RightHandPinky3', 'RightHandPinky4', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'LeftToe_End', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'RightToe_End']
    [file._names[i] for i in file.corps] ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'LeftToe_End', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'RightToe_End', 'Spine', 'Spine1', 'Spine2', 'Neck', 'Head', 'HeadTop_End', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']
        corps_name_2 =                   ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'LeftToe_End', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'RightToe_End', 
        'Spine', 'Spine1', 'Spine2', 'Neck', 'Head', 'HeadTop_End', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand']
    topology    (-1, 0, 1, 2, 3, 4, 0, 6, 7, 8, 9, 0, 11, 12, 13, 14, 15, 13, 17, 18, 19, 13, 21, 22, 23)
    [-1  0 55 56 57 58  0 60 61 62 63  0  1  2  3  4  5  3  7  8  9  3 31 32  33]
    '''



    # pdb.set_trace()


retargeting/datasets/bvh_parser.py
YoungSeng/UnifiedGesture
import sys
sys.path.append("../utils")
import numpy as np
from Quaternions import Quaternions
from models.skeleton import build_joint_topology
import os


# rotation with shape frame * J * 3
def write_bvh(parent, offset, rotation, position, names, frametime, order, path, endsite=None):
    directory = os.path.dirname(path)

    # Check if the directory exists
    if not os.path.exists(directory):
        os.makedirs(directory)
        print(f"Directory {directory} created!")
    else:
        print(f"Directory {directory} already exists!")
    file = open(path, 'w')
    frame = rotation.shape[0]
    joint_num = rotation.shape[1]
    order = order.upper()

    file_string = 'HIERARCHY\n'

    def write_static(idx, prefix):
        nonlocal parent, offset, rotation, names, order, endsite, file_string
        if idx == 0:
            name_label = 'ROOT ' + names[idx]
            channel_label = 'CHANNELS 6 Xposition Yposition Zposition {}rotation {}rotation {}rotation'.format(*order)
        else:
            name_label = 'JOINT ' + names[idx]
            channel_label = 'CHANNELS 3 {}rotation {}rotation {}rotation'.format(*order)
        offset_label = 'OFFSET %.6f %.6f %.6f' % (offset[idx][0], offset[idx][1], offset[idx][2])

        file_string += prefix + name_label + '\n'
        file_string += prefix + '{\n'
        file_string += prefix + '\t' + offset_label + '\n'
        file_string += prefix + '\t' + channel_label + '\n'

        has_child = False
        for y in range(idx+1, rotation.shape[1]):
            if parent[y] == idx:
                has_child = True
                write_static(y, prefix + '\t')
        if not has_child:
            file_string += prefix + '\t' + 'End Site\n'
            file_string += prefix + '\t' + '{\n'
            file_string += prefix + '\t\t' + 'OFFSET 0 0 0\n'
            file_string += prefix + '\t' + '}\n'

        file_string += prefix + '}\n'

    write_static(0, '')

    file_string += 'MOTION\n' + 'Frames: {}\n'.format(frame) + 'Frame Time: %.8f\n' % frametime
    for i in range(frame):
        file_string += '%.6f %.6f %.6f ' % (position[i][0], position[i][1], position[i][2])
        for j in range(joint_num):
            file_string += '%.6f %.6f %.6f ' % (rotation[i][j][0], rotation[i][j][1], rotation[i][j][2])
        file_string += '\n'

    file.write(file_string)
    return file_string


class BVH_writer():
    def __init__(self, edges, names):
        self.parent, self.offset, self.names, self.edge2joint = build_joint_topology(edges, names)
        self.joint_num = len(self.parent)

    # position, rotation with shape T * J * (3/4)
    def write(self, rotations, positions, order, path, frametime=1.0/30, offset=None, root_y=None):
        if order == 'quaternion':
            norm = rotations[:, :, 0] ** 2 + rotations[:, :, 1] ** 2 + rotations[:, :, 2] ** 2 + rotations[:, :, 3] ** 2
            norm = np.repeat(norm[:, :, np.newaxis], 4, axis=2)
            rotations /= norm
            rotations = Quaternions(rotations)
            rotations = np.degrees(rotations.euler())
            order = 'xyz'

        rotations_full = np.zeros((rotations.shape[0], self.joint_num, 3))
        for idx, edge in enumerate(self.edge2joint):
            if edge != -1:
                rotations_full[:, idx, :] = rotations[:, edge, :]
        if root_y is not None: rotations_full[0, 0, 1] = root_y

        if offset is None: offset = self.offset
        return write_bvh(self.parent, offset, rotations_full, positions, self.names, frametime, order, path)

    def write_raw(self, motion, order, path, frametime=1.0/30, root_y=None):
        motion = motion.permute(1, 0).detach().cpu().numpy()
        positions = motion[:, -3:]
        rotations = motion[:, :-3]
        if order == 'quaternion':
            rotations = rotations.reshape((motion.shape[0], -1, 4))
        else:
            rotations = rotations.reshape((motion.shape[0], -1, 3))

        return self.write(rotations, positions, order, path, frametime, root_y=root_y)

retargeting/datasets/bvh_writer.py
YoungSeng/UnifiedGesture
import pdb

from torch.utils.data import Dataset
import copy
from datasets.motion_dataset import MotionData
import os
import numpy as np
import torch
from datasets.bvh_parser import BVH_file
from option_parser import get_std_bvh
from datasets import get_test_set


class MixedData0(Dataset):
    """
    Mixed data for many skeletons but one topologies
    """
    def __init__(self, args, motions, skeleton_idx):
        super(MixedData0, self).__init__()

        self.motions = motions
        self.motions_reverse = torch.tensor(self.motions.numpy()[..., ::-1].copy())
        self.skeleton_idx = skeleton_idx
        self.length = motions.shape[0]
        self.args = args

    def __len__(self):
        return self.length

    def __getitem__(self, item):
        if self.args.data_augment == 0 or torch.rand(1) < 0.5:
            return [self.motions[item], self.skeleton_idx[item]]
        else:
            return [self.motions_reverse[item], self.skeleton_idx[item]]


class MixedData(Dataset):
    """
    data_gruop_num * 2 * samples
    """
    def __init__(self, args, datasets_groups):
        device = torch.device(args.cuda_device if (torch.cuda.is_available()) else 'cpu')
        self.final_data = []
        self.length = 0
        self.offsets = []
        self.joint_topologies = []
        self.ee_ids = []
        self.means = []
        self.vars = []
        dataset_num = 0
        # seed = 19260817
        total_length = 10000000
        all_datas = []
        for datasets in datasets_groups:    # [['ch01'], ['ch02']]
            offsets_group = []
            means_group = []
            vars_group = []
            dataset_num += len(datasets)
            tmp = []
            for i, dataset in enumerate(datasets):
                new_args = copy.copy(args)
                new_args.data_augment = 0
                new_args.dataset = dataset
                tmp.append(MotionData(new_args))

                mean = np.load('./datasets/' + args.dataset + '/mean_var/{}_mean.npy'.format(dataset))
                var = np.load('./datasets/' + args.dataset + '/mean_var/{}_var.npy'.format(dataset))
                mean = torch.tensor(mean)
                mean = mean.reshape((1,) + mean.shape)
                var = torch.tensor(var)
                var = var.reshape((1,) + var.shape)

                means_group.append(mean)
                vars_group.append(var)

                file = BVH_file(get_std_bvh(dataset=dataset))
                if i == 0:
                    self.joint_topologies.append(file.topology)
                    self.ee_ids.append(file.get_ee_id())
                new_offset = file.offset
                new_offset = torch.tensor(new_offset, dtype=torch.float)
                new_offset = new_offset.reshape((1,) + new_offset.shape)
                offsets_group.append(new_offset)

                total_length = min(total_length, len(tmp[-1]))
            all_datas.append(tmp)
            offsets_group = torch.cat(offsets_group, dim=0)
            offsets_group = offsets_group.to(device)
            means_group = torch.cat(means_group, dim=0).to(device)
            vars_group = torch.cat(vars_group, dim=0).to(device)
            self.offsets.append(offsets_group)
            self.means.append(means_group)
            self.vars.append(vars_group)

        for datasets in all_datas:      # [[<datasets.motion_dataset.MotionData object at 0x7f5879f06610>], [<datasets.motion_dataset.MotionData object at 0x7f57a1952a50>]]
            pt = 0
            motions = []
            skeleton_idx = []
            for dataset in datasets:
                motions.append(dataset[:])
                skeleton_idx += [pt] * len(dataset)
                pt += 1
            motions = torch.cat(motions, dim=0)     # ([4, 99, 64])
            if self.length != 0 and self.length != len(skeleton_idx):
                self.length = min(self.length, len(skeleton_idx))
            else:
                self.length = len(skeleton_idx)
            self.final_data.append(MixedData0(args, motions, skeleton_idx))

    def denorm(self, gid, pid, data):
        means = self.means[gid][pid, ...]
        var = self.vars[gid][pid, ...]
        return data * var + means

    def __len__(self):
        return self.length

    def __getitem__(self, item):
        res = []
        for data in self.final_data:
            res.append(data[item])
        return res


class TestData(Dataset):
    def __init__(self, args, characters):
        self.characters = characters
        self.file_list = get_test_set()
        # self.length = 0
        self.mean = []
        self.joint_topologies = []
        self.var = []
        self.offsets = []
        self.ee_ids = []
        self.args = args
        self.device = torch.device(args.cuda_device)

        for i, character_group in enumerate(characters):
            mean_group = []
            var_group = []
            offsets_group = []
            for j, character in enumerate(character_group):
                file = BVH_file(get_std_bvh(dataset=character))
                if j == 0:
                    self.joint_topologies.append(file.topology)
                    self.ee_ids.append(file.get_ee_id())
                new_offset = file.offset
                new_offset = torch.tensor(new_offset, dtype=torch.float)
                new_offset = new_offset.reshape((1,) + new_offset.shape)
                offsets_group.append(new_offset)
                mean = np.load('./datasets/' + args.dataset + '/mean_var/{}_mean.npy'.format(character))
                var = np.load('./datasets/' + args.dataset + '/mean_var/{}_var.npy'.format(character))
                mean = torch.tensor(mean)
                mean = mean.reshape((1, ) + mean.shape)
                var = torch.tensor(var)
                var = var.reshape((1, ) + var.shape)
                mean_group.append(mean)
                var_group.append(var)
                # self.length += 1

            mean_group = torch.cat(mean_group, dim=0).to(self.device)
            var_group = torch.cat(var_group, dim=0).to(self.device)
            offsets_group = torch.cat(offsets_group, dim=0).to(self.device)
            self.mean.append(mean_group)
            self.var.append(var_group)
            self.offsets.append(offsets_group)

    def __getitem__(self, item):
        res = []
        bad_flag = 0
        for i, character_group in enumerate(self.characters):
            res_group = []
            ref_shape = None
            for j in range(len(character_group)):
                new_motion = self.get_item(i, j, item)
                if new_motion is not None:
                    new_motion = new_motion.reshape((1, ) + new_motion.shape)
                    new_motion = (new_motion - self.mean[i][j]) / self.var[i][j]
                    ref_shape = new_motion
                res_group.append(new_motion)

            if ref_shape is None:
                print('Bad at {}'.format(item))
                return None
            for j in range(len(character_group)):
                if res_group[j] is None:
                    bad_flag = 1
                    res_group[j] = torch.zeros_like(ref_shape)
            if bad_flag:
                print('Bad at {}'.format(item))

            res_group = torch.cat(res_group, dim=0)
            res.append([res_group, list(range(len(character_group)))])
        return res

    def __len__(self):
        return len(self.file_list)
        # return self.length

    def get_item(self, gid, pid, id):
        character = self.characters[gid][pid]
        path = './datasets/' + self.args.dataset + '/{}/'.format(character)
        if isinstance(id, int):
            file = path + self.file_list[id]
        elif isinstance(id, str):
            file = id
        else:
            raise Exception('Wrong input file type')
        if not os.path.exists(file):
            raise Exception('Cannot find file')

        # file = id
        file = BVH_file(file)
        motion = file.to_tensor(quater=self.args.rotation == 'quaternion')
        # motion = motion[:, ::2]
        length = motion.shape[-1]
        length = length // 4 * 4          # 20230412
        return motion[..., :length].to(self.device)

    def denorm(self, gid, pid, data):
        means = self.mean[gid][pid, ...]
        var = self.var[gid][pid, ...]
        return data * var + means

    def normalize(self, gid, pid, data):
        means = self.mean[gid][pid, ...]
        var = self.var[gid][pid, ...]
        return (data - means) / var

retargeting/datasets/combined_motion.py
YoungSeng/UnifiedGesture
"""
This code comes from https://github.com/rubenvillegas/cvpr2018nkn/blob/master/datasets/fbx2bvh.py
"""
import bpy
import numpy as np

from os import listdir


if __name__ == '__main__':

    # windows
    # cd C:\Users\94086\Seafile\huawei\datasets
    # blender -b -P fbx2bvh.py

    data_path = './fpx2bvh/'

    directories = sorted([f for f in listdir(data_path) if not f.startswith(".")])
    print(directories)
    for d in directories:
        files = sorted([f for f in listdir(data_path + d) if f.endswith(".fbx")])

        for f in files:
            sourcepath = data_path + d + "/" + f
            dumppath = data_path + d + "/" + f.split(".fbx")[0] + ".bvh"

            bpy.ops.import_scene.fbx(filepath=sourcepath)

            frame_start = 9999
            frame_end = -9999
            action = bpy.data.actions[-1]
            if action.frame_range[1] > frame_end:
                frame_end = action.frame_range[1]
            if action.frame_range[0] < frame_start:
                frame_start = action.frame_range[0]

            frame_end = np.max([60, frame_end])
            # print(frame_start, frame_end)
            bpy.ops.export_anim.bvh(filepath=dumppath,
                                    frame_start=int(frame_start),
                                    frame_end=int(frame_end), root_transform_only=True)
            bpy.data.actions.remove(bpy.data.actions[-1])

            print(data_path + d + "/" + f + " processed.")

retargeting/datasets/fbx2bvh.py
YoungSeng/UnifiedGesture
'''
pip install lmdb==1.3.0
pip install pyarrow==8.0.0  (-i http://pypi.douban.com/simple --trusted-host pypi.douban.com)
'''

import argparse
import os
import glob
import pdb
from pathlib import Path

# import librosa
import numpy as np
import lmdb
import pyarrow


def make_lmdb_latent_dataset(base_path):
    Trinity_latent_path = os.path.join(base_path, 'Trinity')
    ZEGGS_latent_path = os.path.join(base_path, 'ZEGGS')
    out_path = os.path.join(base_path, 'lmdb_latent_vel')
    if not os.path.exists(out_path):
        os.makedirs(out_path)

    map_size = 1024 * 20  # in MB
    map_size <<= 20  # in B
    db = [lmdb.open(os.path.join(out_path, 'lmdb_train'), map_size=map_size),
          lmdb.open(os.path.join(out_path, 'lmdb_test'), map_size=map_size)]

    # delete existing files
    for i in range(2):
        with db[i].begin(write=True) as txn:
            txn.drop(db[i].open_db())

    all_positions = []
    bvh_files = sorted(glob.glob(Trinity_latent_path + "/*_upper.npy") + glob.glob(ZEGGS_latent_path + "/*_upper.npy"))

    save_idx = 0
    for v_i, bvh_file in enumerate(bvh_files):
        path = os.path.split(bvh_file)[0]
        name = os.path.split(bvh_file)[1][:-4]
        print(name)

        # load skeletons and subtitles

        upper = np.load(bvh_file)[0].transpose()
        lower = np.load(os.path.join(path, name.replace('upper', 'lower') + '.npy'))[0].transpose()
        root = np.load(os.path.join(path, name.replace('upper', 'root') + '.npy'))[0].transpose()
        root_vel = np.load(os.path.join(path, name.replace('upper', 'root_vel') + '.npy'))[0].transpose()

        # global_position = np.load(os.path.join(path, name.replace('local', 'global') + '.npy'), allow_pickle=True)
        # position = global_position[0][0].transpose()
        # height = global_position[1]

        # print(upper.shape, lower.shape, root.shape)

        # load audio
        # audio_raw, audio_sr = librosa.load(os.path.join(audio_path, '{}.wav'.format(name)),
        #                                    mono=True, sr=16000, res_type='kaiser_fast')

        # process
        clips = [{'vid': name, 'clips': []},  # train
                 {'vid': name, 'clips': []}]  # validation

        # split
        if save_idx % 10 == 0:
            dataset_idx = 1  # validation
        else:
            dataset_idx = 0  # train

        # save subtitles and skeletons
        # poses = np.asarray(poses)       # x dtype=np.float16
        clips[dataset_idx]['clips'].append(
            {
             'upper': upper,
             'lower': lower,
             'root': root,
             'root_vel': root_vel,
             })

        # write to db
        for i in range(2):
            with db[i].begin(write=True) as txn:
                if len(clips[i]['clips']) > 0:
                    k = '{:010}'.format(v_i).encode('ascii')
                    v = pyarrow.serialize(clips[i]).to_buffer()
                    txn.put(k, v)
        # all_positions.append(root)
        save_idx += 1

    # close db
    for i in range(2):
        db[i].sync()
        db[i].close()

    # calculate data mean
    # all_poses = np.vstack(all_positions)
    # pose_mean = np.mean(all_poses, axis=0, dtype=np.float64)
    # pose_std = np.std(all_poses, axis=0, dtype=np.float64)
    #
    # print('data mean/std')
    # print('data_mean:', str(["{:0.5f}".format(e) for e in pose_mean]).replace("'", ""))
    # print('data_std:', str(["{:0.5f}".format(e) for e in pose_std]).replace("'", ""))


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("--base_path", type=Path)
    args = parser.parse_args()

    make_lmdb_latent_dataset(args.base_path)

retargeting/datasets/latent_to_lmdb.py
YoungSeng/UnifiedGesture
import pdb

from torch.utils.data import Dataset
import os
import sys
sys.path.append("../utils")
import numpy as np
import torch
from Quaternions import Quaternions
from option_parser import get_std_bvh


class MotionData(Dataset):
    """
    Clip long dataset into fixed length window for batched training
    each data is a 2d tensor with shape (Joint_num*3) * Time
    """
    def __init__(self, args):
        super(MotionData, self).__init__()
        name = args.dataset
        file_path = './datasets/' + args.dataset_constant + '/{}.npy'.format(name)

        if args.debug:
            file_path = file_path[:-4] + '_debug' + file_path[-4:]

        print('load from file {}'.format(file_path))
        self.total_frame = 0
        self.std_bvh = get_std_bvh(args)
        self.args = args
        self.data = []
        self.motion_length = []
        motions = np.load(file_path, allow_pickle=True)
        motions = list(motions)     # [(135, 75), (116, 75)]
        new_windows = self.get_windows(motions)     # ([5, 64, 99])
        self.data.append(new_windows)
        self.data = torch.cat(self.data)
        self.data = self.data.permute(0, 2, 1)

        if args.normalization == 1:     # 1
            self.mean = torch.mean(self.data, (0, 2), keepdim=True)
            self.var = torch.var(self.data, (0, 2), keepdim=True)
            self.var = self.var ** (1/2)
            idx = self.var < 1e-5
            self.var[idx] = 1
            self.data = (self.data - self.mean) / self.var
        else:
            self.mean = torch.mean(self.data, (0, 2), keepdim=True)
            self.mean.zero_()
            self.var = torch.ones_like(self.mean)

        train_len = self.data.shape[0] * 95 // 100      # 4
        self.test_set = self.data[train_len:, ...]      # [1, 99, 64]
        self.data = self.data[:train_len, ...]          # [4, 99, 64]
        self.data_reverse = torch.tensor(self.data.numpy()[..., ::-1].copy())

        self.reset_length_flag = 0
        self.virtual_length = 0
        print('Window count: {}, total frame (without downsampling): {}'.format(len(self), self.total_frame))

    def reset_length(self, length):
        self.reset_length_flag = 1
        self.virtual_length = length

    def __len__(self):
        if self.reset_length_flag:
            return self.virtual_length
        else:
            return self.data.shape[0]

    def __getitem__(self, item):
        if isinstance(item, int): item %= self.data.shape[0]
        if self.args.data_augment == 0 or np.random.randint(0, 2) == 0:
            return self.data[item]
        else:
            return self.data_reverse[item]

    def get_windows(self, motions):
        new_windows = []

        for motion in motions:      # [(135, 75), (116, 75)]
            self.total_frame += motion.shape[0]
            # motion = self.subsample(motion)
            self.motion_length.append(motion.shape[0])
            step_size = self.args.window_size // 2      # window_size   64
            window_size = step_size * 2
            n_window = motion.shape[0] // step_size - 1     # 135//32-1 3     2
            for i in range(n_window):
                begin = i * step_size
                end = begin + window_size       # (0, 64)   (32, 96)    (64, 128)

                new = motion[begin:end, :]      # (64, 75)
                if self.args.rotation == 'quaternion':      # quaternion
                    new = new.reshape(new.shape[0], -1, 3)      # -> (64, 25, 3)
                    rotations = new[:, :-1, :]      # -> (64, 24, 3)
                    rotations = Quaternions.from_euler(np.radians(rotations)).qs        # -> (64, 24, 4)
                    rotations = rotations.reshape(rotations.shape[0], -1)       # -> (64, 96)
                    # positions = new[:, -1, :]
                    # positions = np.concatenate((new, np.zeros((new.shape[0], new.shape[1], 1))), axis=2)
                    new = np.concatenate((rotations, new[:, -1, :].reshape(new.shape[0], -1)), axis=1)      # -> (64, 99)

                new = new[np.newaxis, ...]      # -> (1， 64， 99）

                new_window = torch.tensor(new, dtype=torch.float32)
                new_windows.append(new_window)

        return torch.cat(new_windows)

    def subsample(self, motion):
        return motion[::2, :]

    def denormalize(self, motion):
        if self.args.normalization:
            if self.var.device != motion.device:
                self.var = self.var.to(motion.device)
                self.mean = self.mean.to(motion.device)
            ans = motion * self.var + self.mean
        else: ans = motion
        return ans

retargeting/datasets/motion_dataset.py
YoungSeng/UnifiedGesture
import os
import numpy as np
import copy
import sys
sys.path.append('.')
from datasets.bvh_parser import BVH_file
from datasets.motion_dataset import MotionData
from option_parser import get_args, try_mkdir

import logging
import pdb
logging.basicConfig(format='%(levelname)s:%(funcName)s:%(message)s', level=logging.DEBUG)       #  logging.INFO

def collect_bvh(data_path, character, files):
    print('begin {}'.format(character))
    motions = []

    for i, motion in enumerate(files):
        if not os.path.exists(data_path + character + '/' + motion):
            continue
        file = BVH_file(data_path + character + '/' + motion)
        new_motion = file.to_tensor().permute((1, 0)).numpy()
        motions.append(new_motion)

    save_file = data_path + character + '.npy'

    np.save(save_file, motions)     # ((135, 75), (116, 75)) corps_name_2 (25 * 3)
    print('Npy file saved at {}'.format(save_file))


def write_statistics(character, path):
    args = get_args()
    new_args = copy.copy(args)
    new_args.data_augment = 0
    new_args.dataset = character

    dataset = MotionData(new_args)

    mean = dataset.mean
    var = dataset.var
    mean = mean.cpu().numpy()[0, ...]
    var = var.cpu().numpy()[0, ...]

    logging.debug(mean.shape)       # (99, 1)
    logging.debug(var.shape)

    np.save(path + '{}_mean.npy'.format(character), mean)
    np.save(path + '{}_var.npy'.format(character), var)


def copy_std_bvh(data_path, character, files):      # Standard
    """
    copy an arbitrary bvh file as a static information (skeleton's offset) reference
    """
    cmd = 'cp \"{}\" ./datasets/{}/std_bvhs/{}.bvh'.format(data_path + character + '/' + files[0], args.dataset, character)
    os.system(cmd)


if __name__ == '__main__':
    '''
    python ./datasets/preprocess.py
    '''
    import option_parser
    args = option_parser.get_args()

    prefix = './datasets/' + args.dataset + '/'
    characters = [f for f in os.listdir(prefix) if os.path.isdir(os.path.join(prefix, f))]
    logging.debug(characters)       # ['std_bvhs', 'ch01', 'ch02', 'mean_var']

    if 'std_bvhs' in characters: characters.remove('std_bvhs')
    if 'mean_var' in characters: characters.remove('mean_var')

    try_mkdir(os.path.join(prefix, 'std_bvhs'))
    try_mkdir(os.path.join(prefix, 'mean_var'))

    for character in characters:
        if character.endswith('_aux'): continue
        data_path = os.path.join(prefix, character)
        files = sorted([f for f in os.listdir(data_path) if f.endswith(".bvh")])
        collect_bvh(prefix, character, files)
        # import pdb
        # pdb.set_trace()
        copy_std_bvh(prefix, character, files)
        write_statistics(character, './datasets/' + args.dataset + '/mean_var/')

retargeting/datasets/preprocess.py
YoungSeng/UnifiedGesture
# to 30fps

import os
import pdb
import glob
import numpy as np
from bvh_parser import BVH_file
from bvh_writer import BVH_writer
import argparse
import sys

[sys.path.append(i) for i in ['.', '..']]
from models.IK import get_foot_vel_position, PFC_fix


def space2t(test_sentence):
    res_str = ''
    for i in range(0, len(test_sentence), 2):
        if test_sentence[i:i + 2] == '  ':
            res_str += '\t'
        else:
            res_str += test_sentence[i:]
            break
    return res_str


def downsample_process_root(trinity_path, dst_path):
    if not os.path.exists(dst_path):
        os.makedirs(dst_path)
    for item in os.listdir(trinity_path):
        if not item.endswith('.bvh'):
            continue
        bvh_file = os.path.join(trinity_path, item)
        motion = []
        with open(os.path.join(dst_path, item), 'w') as output:
            with open(bvh_file, 'r') as f:
                for i, line in enumerate(f.readlines()):
                    if i == 3:
                        line = line.strip().split(' ')
                        x = float(line[-3])
                        y = float(line[-2])
                        z = float(line[-1])
                        # print('root', x, y, z)
                        line[-3] = '0.0'
                        line[-1] = '0.0'
                        line = '  ' + ' '.join(line) + '\n'
                        line = space2t(line)
                    if i < 343:
                        output.write(space2t(line))
                        continue
                    if i == 343:
                        try:
                            frames = int(line.strip().split('\t')[-1])       # \t for testing and ' ' for training
                        except:
                            frames = int(line.strip().split(' ')[-1])
                        print(frames)
                        output.write('Frames: ' + str((frames + 1)//2) + '\n')
                        continue
                    if i == 344:
                        try:
                            fps = float(line.strip().split('\t')[-1])        #
                        except:
                            fps = float(line.strip().split(' ')[-1])
                        print(fps)
                        output.write('Frame Time: ' + str(1/30.0) + '\n')
                        continue
                    else:
                        motion.append(line)
            if len(motion) != frames:
                print(len(motion), '/', frames)
                motion = motion[:frames]
            motion = motion[::2]
            assert len(motion) == (frames + 1)//2
            # fps = 1/30.0
            for i in motion:
                i = i.strip().split(' ')
                # print(i)
                i[2] = str(0- float(i[2]))
                i[0] = str(0 - float(i[0]))
                i[4] = str(0 - float(i[4]))
                i[3] = str(0 - float(i[3]))
                i[5] = str(180 + float(i[5]))
                i = ' '.join(i) + '\n'
                output.write(i)
                # break
        # break


def process_root():
    Talking_With_Hands_path = '/ceph/hdd/yangsc21/Python/My_3/trn/bvh/'
    dst_path = '/ceph/hdd/yangsc21/Python/My_3/deep-motion-editing/retargeting/datasets/Mixamo_9/Talking_With_Hands/'

    for item in os.listdir(Talking_With_Hands_path):
        bvh_file = os.path.join(Talking_With_Hands_path, item)
        motion = []
        with open(os.path.join(dst_path, item), 'w') as output:
            with open(bvh_file, 'r') as f:
                for i, line in enumerate(f.readlines()):
                    if i == 1:
                        output.write('ROOT b_root\n')
                        continue
                    if i == 3:
                        line = line.strip().split(' ')
                        z1 = float(line[-3])
                        x1 = float(line[-2])
                        y1 = float(line[-1])
                        continue
                    if i > 2 and i < 7:
                        continue
                    if i == 7:
                        line = line.strip().split(' ')
                        z2 = float(line[-3])
                        x2 = float(line[-2])
                        y2 = float(line[-1])
                        output.write('\tOFFSET {} {} {}\n'.format(z1 + z2, x1 + x2, y1 + y2))
                        continue
                    if i < 7 or 524 <= i < 527:
                        output.write(line)
                        continue
                    if i < 527:
                        output.write(space2t(line[2:]))
                        continue
                    else:
                        line = [float(j) for j in line.strip().split(' ')]
                        # print(len(line))        # 83 * 6 = 498
                        line[6] = line[6] + line[0]
                        line[7] = line[7] + line[1]
                        # line[7] = line[1]
                        line[8] = line[8] + line[2]
                        line[9] = line[9] + line[3]
                        line[10] = line[10] + line[4]
                        line[11] = line[11] + line[5]
                        line = line[6:]
                        line = [str(j) for j in line]
                        line = ' '.join(line)
                        output.write(line + '\n')
                # break


def process_T_pose(ZEGGS_path, dst_path, divide=2):

    # ZEGGS_path = "/ceph/hdd/yangsc21/Python/My_3/IJCAI-20fps/"
    # dst_path = "/ceph/hdd/yangsc21/Python/My_3/IJCAI-20fps-fix/"

    if not os.path.exists(dst_path):
        os.makedirs(dst_path)

    if ZEGGS_path.split('/')[-2] == 'ZEGGS_GT':
        GT = True
    else:
        GT = False
    for item in os.listdir(ZEGGS_path):
        if item[-4:] == '.wav':
            continue
        bvh_file = os.path.join(ZEGGS_path, item)
        motion = []
        with open(os.path.join(dst_path, item), 'w') as output:
            with open(bvh_file, 'r') as f:
                for i, line in enumerate(f.readlines()):
                    if i == 3:
                        line = line.strip().split(' ')
                        x = float(line[-3])
                        y = float(line[-2])
                        z = float(line[-1])
                        # print('root', x, y, z)
                        line[-3] = '0.0'
                        line[-1] = '0.0'
                        line = '  ' + ' '.join(line) + '\n'
                        line = space2t(line)
                    if not GT:
                        if ((51 <= i <= 75 or 87 <= i <= 99 or 111 <= i <= 123 or 135 <= i <= 147 or 159 <= i <= 171) and i % 4 ==3):     # RightArm OFFSET
                            line = line.strip('\n').split(' ')
                            x = float(line[-3])
                            y = float(line[-2])
                            z = float(line[-1])
                            line[-3] = str(0.0 - y)
                            line[-2] = str(x)
                            line[-1] = str(z)
                            line = ' '.join(line) + '\n'
                            line = space2t(line)
                        if ((209 <= i <= 233 or 245 <= i <= 257 or 269 <= i <= 281 or 293 <= i <= 305 or 317 <= i <= 329) and i % 4 ==1):     # RightArm OFFSET
                            line = line.strip('\n').split(' ')
                            x = float(line[-3])
                            y = float(line[-2])
                            z = float(line[-1])
                            line[-3] = str(y)
                            line[-2] = str(0.0 - x)
                            line[-1] = str(z)
                            line = ' '.join(line) + '\n'
                            line = space2t(line)
                        if i < 461:
                            output.write(space2t(line))
                            continue
                        if i == 461:
                            frames = int(line.strip().split(' ')[-1])  # \t for testing and ' ' for training
                            print(frames)
                            output.write('Frames: ' + str((frames + 1) // divide) + '\n')
                            continue
                        if i == 462:
                            fps = float(line.strip().split(' ')[-1])  #
                            print(fps)
                            if divide == 2:
                                output.write('Frame Time: ' + str(1 / 30.0) + '\n')
                            elif divide == 1:
                                output.write('Frame Time: ' + str(1 / 60.0) + '\n')
                            continue
                        else:
                            motion.append(line)
                    else:
                        if i < 566:
                            output.write(space2t(line))
                            continue
                        if i == 566:
                            frames = int(line.strip().split(' ')[-1])
                            print(frames)
                            output.write('Frames: ' + str((frames + 1) // 2) + '\n')
                            continue
                        if i == 567:
                            fps = float(line.strip().split(' ')[-1])  #
                            print(fps)
                            output.write('Frame Time: ' + str(1 / 30.0) + '\n')
                            continue
                        else:
                            motion.append(line)
            if len(motion) != frames:
                print(len(motion), '/', frames)
                motion = motion[:frames]
            motion = motion[::divide]
            if divide == 2:
                assert len(motion) == (frames + 1) // divide     # fps = 1/30.0
            elif divide == 1:
                assert len(motion) == (frames) // divide
            for i in motion:
                i = i.strip().split(' ')

                z = float(i[30])
                y = float(i[31])
                x = float(i[32])
                i[30] = str(z - 90.0)
                i[31] = str(x)
                i[32] = str(0.0 - y)

                # print(z, y, x, i[30], i[31], i[32])      # 92.418716 -0.394881 9.707915

                for j in range(11 * 3, 35 * 3, 3):
                    z = float(i[j])
                    y = float(i[j+1])
                    x = float(i[j+2])
                    i[j+1] = str(x)
                    i[j+2] = str(0.0 - y)
                    # print(i[j], i[j+1], i[j+2])

                z = float(i[36 * 3])
                y = float(i[36 * 3 + 1])
                x = float(i[36 * 3 + 2])
                i[36 * 3] = str(z + 90.0)
                i[36 * 3 + 1] = str(0.0 - x)
                i[36 * 3 + 2] = str(y)

                # print(z, y, x, i[36 * 3], i[36 * 3 + 1], i[36 * 3 + 2])     # -90.370065 0.051974 5.208683

                for j in range(37 * 3, 61 * 3, 3):
                    y = float(i[j+1])
                    x = float(i[j+2])
                    i[j+1] = str(0.0 - x)
                    i[j+2] = str(y)

                # GT
                # print(len(i), i[30:30+6],       # GT 576 (96*6) ['-2.624908', '11.175079', '-0.000479', '2.486827', '9.713857', '0.398656']
                #       i[30+6:30+12],        #  ['-10.435041', '-0.000056', '-0.000023', '68.264312', '4.249098', '33.534958']
                #       i[30+12:30+18],       # ['-29.627481', '-0.000870', '-0.000919', '-5.192122', '21.799514', '-13.415904']
                #       i[38 * 6:39 * 6],     # ['2.625969', '11.175263', '-0.000499', '-0.372213', '-5.206465', '0.049018']
                #       i[39 * 6:40 * 6],     # ['10.813481', '-0.000380', '0.000377', '-69.326304', '7.144804', '19.466830']
                #       i[40 * 6:41 * 6]      # ['28.725275', '-0.001194', '-0.000643', '6.746904', '-40.325830', '-11.368573']
                #       )

                i = ' '.join(i) + '\n'
                output.write(i)
                # break
        # break


def get_height(file):
    file = BVH_file(file)
    return file.get_height()


def process_foot_contact(source_path, save_path):
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    root_motions = sorted(glob.glob(source_path + "/*.bvh"))
    for item in root_motions:
        name = os.path.split(item)[1]
        print(name)
        result = get_foot_vel_position(item, get_height(item))  # (2065, 4)
        np.save(os.path.join(save_path, name), result)


if __name__ == '__main__':
    '''
    cd datasets/
    python process_bvh.py
    '''

    parser = argparse.ArgumentParser(description='process')
    parser.add_argument('--step', type=str, default="IK")
    parser.add_argument('--source_path', type=str, default='./Mixamo_new_2/ZEGGS/')
    parser.add_argument('--save_path', type=str, default='./Mixamo_new_2/ZEGGS_aux/')
    parser.add_argument('--ref_bvh', type=str, default="./Mixamo_new_2/ZEGGS/067_Speech_2_x_1_0.bvh")
    args = parser.parse_args()

    if args.step == "IK":

        # bvh_path = "/ceph/hdd/yangsc21/Python/My_3/deep-motion-editing/result/inference/Trinity/Recording_006_minibatch_240_[0, 0, 1, 0, 0, 0, 0]_123456.bvh"
        # # bvh_path = "/ceph/hdd/yangsc21/Python/My_3/deep-motion-editing/retargeting/datasets/Mixamo_new_2/Trinity/Recording_001.bvh"
        ref_bvh = args.ref_bvh
        # bvh_path = "/ceph/hdd/yangsc21/Python/My_3/deep-motion-editing/diffusion_latent/result_256_seed_4_aux_bvh/TestSeq001_final.bvh"
        # PFC_fix(bvh_path, get_height(ref_bvh))

        # source_path = "/ceph/hdd/yangsc21/Python/My_3/deep-motion-editing/diffusion_latent/result_model3_128_bvh/"
        # target_path = ''

        source_path = args.source_path
        for item in os.listdir(source_path):
            print(item)
            if not item.endswith('_fix.bvh'):
                bvh_path = os.path.join(source_path, item)
                PFC_fix(bvh_path, get_height(ref_bvh))

    if args.step == "Trinity":
        '''
        python process_bvh.py --step Trinity --source_path "../../dataset/Trinity/test_motion/" --save_path "../../dataset/Trinity/test_motion_downsample/"
        '''
        downsample_process_root(args.source_path, args.save_path)      # Trinity
        # process_root()      # Talking_With_Hands
    if args.step == 'ZEGGS':
        process_T_pose(args.source_path, args.save_path, divide=2)        # ZEGGS

    if args.step == 'foot_contact':
        process_foot_contact(source_path=args.source_path, save_path=args.save_path)






retargeting/datasets/process_bvh.py
YoungSeng/UnifiedGesture
import os
import pdb
import argparse
from datasets.bvh_parser import BVH_file
from datasets.bvh_writer import BVH_writer
from models.IK import fix_foot_contact
from os.path import join as pjoin
import numpy as np
from eval_single_pair import main as eval_single_pair
from easydict import EasyDict
import yaml

# downsampling and remove redundant joints
def copy_ref_file(src, dst):
    file = BVH_file(src)
    writer = BVH_writer(file.edges, file.names)
    # writer.write_raw(file.to_tensor(quater=True)[..., ::2], 'quaternion', dst)
    writer.write_raw(file.to_tensor(quater=True)[..., :], 'quaternion', dst)

def get_height(file):
    file = BVH_file(file)
    return file.get_height()


def example(dataset_name, src_name, dest_name, bvh_name, test_type, output_path, ref_bvh_name=None, loaded_model=None):
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    input_file = './datasets/' + dataset_name + '/{}/{}'.format(src_name, bvh_name)
    ref_file = './datasets/' + dataset_name + '/{}/{}'.format(dest_name, ref_bvh_name)
    # copy_ref_file(input_file, pjoin(output_path, 'input.bvh'))
    # copy_ref_file(ref_file, pjoin(output_path, 'gt.bvh'))
    # height = get_height(input_file)
    # bvh_name = bvh_name.replace(' ', '_')

    eval_single_pair(input_file, ref_file, test_type, output_path, loaded_model)

    # cmd = 'python eval_single_pair.py --input_bvh={} --target_bvh={} --output_filename={} --test_type={}'.format(
    #     input_file, ref_file, output_path, test_type
    # )
    # os.system(cmd)

    # fix_foot_contact(pjoin(output_path, 'result.bvh'),
    #                  pjoin(output_path, 'input.bvh'),
    #                  pjoin(output_path, 'result_fix_foot_contact.bvh'),
    #                  height)


def bvh_input(quaternion, edge):
    import sys
    sys.path.append("./utils")
    from Quaternions import Quaternions
    import torch

    # path = "./datasets/Mixamo_debug/Talking_With_Hands/val_2022_v1_038.bvh"
    # output_path = "./datasets/Mixamo_debug/Talking_With_Hands/val_2022_v1_038_recon.bvh"
    # new_path = "./datasets/Mixamo_debug/Talking_With_Hands/"

    path = "./datasets/Mixamo_new/ZEGGS/067_Speech_2_x_1_0.bvh"
    output_path = "./datasets/Mixamo_debug/067_Speech_2_x_1_0_recon.bvh"
    new_path = "./datasets/Mixamo_debug/"

    # path = "./datasets/Mixamo_debug/Trinity/TestSeq010_z.bvh"
    # output_path = "./datasets/Mixamo_debug/Trinity/TestSeq010_z_recon.bvh"
    # new_path = "./datasets/Mixamo_debug/Trinity/"

    # path = "./datasets/Mixamo_ch12/ch01/Hip_Hop_Dancing.bvh"
    # output_path = "./datasets/Mixamo_ch12/ch01/Hip_Hop_Dancing_recon.bvh"
    # new_path = "./datasets/Mixamo_ch12/ch01/"

    file = BVH_file(path)
    # file.set_new_root(1)        # 20230330
    new_motion = file.to_tensor(quater=False, edge=edge).permute((1, 0)).numpy()       # quater=False, edge=False
    print(new_motion.shape)        # (1800, 87)
    # motion = file.to_numpy(quater=False, edge=False)

    file.write(output_path)
    window_size = 64
    new = new_motion[0:window_size, :]
    # quaternion = False

    if quaternion and edge:
        new = new.reshape(new.shape[0], -1, 3)  # -> (64, 29, 3)
        rotations = new[:, :-1, :]  # -> (64, 28, 3)
        rotations = Quaternions.from_euler(np.radians(rotations)).qs  # -> (64, 28, 4)
        rotations = rotations.reshape(rotations.shape[0], -1)  # -> (64, 112)
        new = np.concatenate((rotations, new[:, -1, :].reshape(new.shape[0], -1)), axis=1)  # -> (64, 115)
    new = new[np.newaxis, ...]  # -> (1, 64, 115)
    new_window = torch.tensor(new, dtype=torch.float32)
    print(new_window.shape)      # [1, 64, 115]

    data = []
    data.append(new_window)
    data = torch.cat(data)
    data = data.permute(0, 2, 1)
    print(data.shape)            # [1, 115, 64]

    mean = torch.mean(data, (0, 2), keepdim=True)
    var = torch.var(data, (0, 2), keepdim=True)
    var = var ** (1 / 2)
    idx = var < 1e-5
    var[idx] = 1
    data_ = (data - mean) / var
    print(data_.shape)            # [1, 115, 64]

    gt = data_[0]
    ans = gt * var + mean
    print(ans.shape)             # [1, 115, 64]

    writer = BVH_writer(file.edges, file.names)

    # debug bvh
    # print('topology', file.topology, 'offset', file.offset, 'names', file.names)
    # writer.parent = file.topology
    # writer.offset = file.offset
    # writer.names = file.names

    # edge -> False
    if edge == False:
        from datasets.bvh_writer import write_bvh
        print(ans[0].shape)     # -> [81, 64]
        motion = ans[0].permute(1, 0).detach().cpu().numpy()        # -> [64, 81]
        rotations = motion[..., :-3].reshape(motion.shape[0], -1, 3)
        positions = motion[..., -3:]
        write_bvh(file.topology, file.offset, rotations, positions, file.names, 1.0/30, 'xyz',
                  os.path.join(new_path, '0_gt_euler_angle_fix.bvh'))
    else:
        if quaternion:
            # print('parent', writer.parent, 'offset', writer.offset, 'names', writer.names)
            writer.write_raw(ans[0], 'quaternion', os.path.join(new_path, '0_gt_quaternion.bvh'))
        else:
            writer.write_raw(ans[0], 'xyz', os.path.join(new_path, '0_gt_euler_angle.bvh'))
    

def bvh2latent(Trinity_bvh_path = "./datasets/Trinity_ZEGGS/Trinity/", ZEGGS_bvh_path = "./datasets/Trinity_ZEGGS/ZEGGS/",
               Trinity_output = './datasets/Trinity_ZEGGS/bvh2upper_lower_root/Trinity', ZEGGS_output = './datasets/Trinity_ZEGGS/bvh2upper_lower_root/ZEGGS'):


    for i, item in enumerate(os.listdir(Trinity_bvh_path)):
        print(i, item)
        if i == 0:
            from eval_single_pair import demo_model
            dataset_name = 'Trinity_ZEGGS'
            src_name = 'Trinity'
            dest_name = 'ZEGGS'
            bvh_name = item
            ref_bvh_name = '067_Speech_2_x_1_0.bvh'
            input_file = './datasets/' + dataset_name + '/{}/{}'.format(src_name, bvh_name)
            ref_file = './datasets/' + dataset_name + '/{}/{}'.format(dest_name, ref_bvh_name)
            model = demo_model(input_file, ref_file, 'cross', Trinity_output)
        example('Trinity_ZEGGS', 'Trinity', 'ZEGGS', item, 'cross', Trinity_output, ref_bvh_name='067_Speech_2_x_1_0.bvh', loaded_model=model)


    for i, item in enumerate(os.listdir(ZEGGS_bvh_path)):
        print(i, item)
        if i == 0:
            from eval_single_pair import demo_model
            dataset_name = 'Trinity_ZEGGS'
            src_name = 'ZEGGS'
            dest_name = 'Trinity'
            bvh_name = item
            ref_bvh_name = 'Recording_006.bvh'
            input_file = './datasets/' + dataset_name + '/{}/{}'.format(src_name, bvh_name)
            ref_file = './datasets/' + dataset_name + '/{}/{}'.format(dest_name, ref_bvh_name)
            model = demo_model(input_file, ref_file, 'cross', ZEGGS_output)
        example('Trinity_ZEGGS', 'ZEGGS', 'Trinity', item, 'cross', ZEGGS_output, ref_bvh_name='Recording_006.bvh', loaded_model=model)


'''
def latent2bvh(mode, item, ref_path, output_path):
    import torch

    if mode == 'ZEGGS':
        from eval_single_pair import demo_model
        dataset_name = 'Mixamo_new_2'
        src_name = 'Trinity'
        dest_name = 'ZEGGS'
        bvh_name = item
        ref_bvh_name = '067_Speech_2_x_1_0.bvh'
        input_file = './datasets/' + dataset_name + '/{}/{}'.format(src_name, bvh_name)
        ref_file = './datasets/' + dataset_name + '/{}/{}'.format(dest_name, ref_bvh_name)
        model = demo_model(input_file, ref_file, 'cross', './datasets/bvh2latent/Trinity')
        latent_0 = np.load(item)
        latent_1 = np.load(ref_path)
        latent = [torch.from_numpy(latent_0).to('cuda:5'), torch.from_numpy(latent_1).to('cuda:5')]     # 20230413
        model.latent2res(latent)
        model.compute_test_result()
        os.system('cp "{}/{}/0_{}.bvh" "./{}"'.format(model.bvh_path, 'ZEGGS', '0', output_path))

    if mode == 'Trinity':
        from eval_single_pair import demo_model
        dataset_name = 'Mixamo_new_2'
        src_name = 'ZEGGS'
        dest_name = 'Trinity'
        bvh_name = item
        ref_bvh_name = 'Recording_006.bvh'
        input_file = './datasets/' + dataset_name + '/{}/{}'.format(src_name, bvh_name)
        ref_file = './datasets/' + dataset_name + '/{}/{}'.format(dest_name, ref_bvh_name)
        model = demo_model(input_file, ref_file, 'cross', './datasets/bvh2latent/ZEGGS')
        latent_0 = np.load(item)
        latent_1 = np.load(ref_path)
        latent = [torch.from_numpy(latent_1).to('cuda:5'), torch.from_numpy(latent_0).to('cuda:5')]     # 20230413
        model.latent2res(latent)
        model.compute_test_result()
        os.system('cp "{}/{}/0_{}.bvh" "./{}"'.format(model.bvh_path, 'Trinity', '1', output_path))
'''

def latent2bvh(mode, upper, lower, root, ref_path, output_path, hide=True, model_name='vqvae_ulr', device='cuda:5'):
    import torch
    if hide:
        from eval_single_pair import demo_model
        if mode == 'ZEGGS':
            dataset_name = 'Mixamo_new_2'
            src_name = 'Trinity'
            dest_name = 'ZEGGS'
            ref_bvh_name = '067_Speech_2_x_1_0.bvh'
        elif mode == 'Trinity':
            dataset_name = 'Mixamo_new_2'
            src_name = 'ZEGGS'
            dest_name = 'Trinity'
            ref_bvh_name = 'Recording_006.bvh'
        ref_file = './datasets/' + dataset_name + '/{}/{}'.format(dest_name, ref_bvh_name)
        latent_1 = torch.from_numpy(np.load(ref_path)).to(device)
        input_file = './datasets/' + dataset_name + '/{}/{}'.format(src_name, upper)
        if mode == 'ZEGGS':
            model = demo_model(input_file, ref_file, 'cross', './datasets/bvh2latent/Trinity')
        elif mode == 'Trinity':
            model = demo_model(input_file, ref_file, 'cross', './datasets/bvh2latent/ZEGGS')

        latent_upper = np.load(upper)
        name_ = os.path.split(upper)[1][:-4]
        if model_name == 'vqvae_ulr':
            latent_root = np.load(root)
            latent_lower = np.load(lower)
            latent_0 = np.concatenate((latent_lower, latent_upper, latent_root), axis=1)
        elif model_name == 'vqvae_ulr_2':

            # root = np.zeros_like(latent_lower[:, -16 * 1:])     # [1, 16, length]
            # root[..., 0] = np.load('./datasets/bvh2upper_lower_root/Trinity/Recording_006_root.npy')[..., 0]
            # for iii in range(1, root.shape[-1]):
            #     root[..., iii] = root[..., iii - 1] + latent_lower[:, -16 * 1:, iii - 1]

            latent_0 = np.concatenate((latent_lower[:, :16*2], latent_upper, latent_lower[:, -16*1:]), axis=1)
        elif model_name == 'diffusion':
            # (1, 238, 112)
            latent_0 = latent_upper
        elif model_name == 'GT':
            latent_0 = latent_upper[:, :16 * 7].transpose(1, 0)
            latent_0 = np.expand_dims(latent_0, axis=0).astype(np.float32)
        print(latent_0.shape)
        latent_0 = torch.from_numpy(latent_0).to(device)

        latent = [latent_0, latent_1]
        model.latent2res(latent)
        if mode == 'ZEGGS':
            result = model.fake_res_denorm[1][0, ...]
        elif mode == 'Trinity':
            result = model.fake_res_denorm[0][0, ...]
        model.generate_bvh(result=result, mode=mode)
        os.system('cp "{}/{}/0_{}.bvh" "./{}/{}.bvh"'.format(model.bvh_path, mode, '0', output_path, name_))
    else:
        if mode == 'ZEGGS':
            from eval_single_pair import demo_model
            dataset_name = 'Mixamo_new_2'
            src_name = 'Trinity'
            dest_name = 'ZEGGS'
            ref_bvh_name = '067_Speech_2_x_1_0.bvh'
            ref_file = './datasets/' + dataset_name + '/{}/{}'.format(dest_name, ref_bvh_name)
            latent_1 = torch.from_numpy(np.load(ref_path)).to(device)

            if upper is not None:
                input_file = './datasets/' + dataset_name + '/{}/{}'.format(src_name, upper)
                model = demo_model(input_file, ref_file, 'cross', './datasets/bvh2latent/Trinity')
                if type(upper) is not np.ndarray:
                    latent_0 = np.load(upper)
                else:
                    latent_0 = upper
                latent = [torch.from_numpy(latent_0).to(device), latent_1]     # 20230413
                model.latent2res(latent)
                upper_result = model.fake_res_denorm[1][0, ...]       # [103, 14284]
                # if lower is None:
                #     upper_result[:4 * 10] = 0
                # if root is None:
                #     upper_result[-3:] = 0
            if lower is not None:
                input_file = './datasets/' + dataset_name + '/{}/{}'.format(src_name, lower)
                model = demo_model(input_file, ref_file, 'cross', './datasets/bvh2latent/Trinity')
                if type(lower) is not np.ndarray:
                    lower = torch.from_numpy(np.load(lower)).to(device)
                else:
                    lower = lower.to(device)

                latent = [lower, latent_1]     # 20230413
                model.latent2res(latent)
                lower_result = model.fake_res_denorm[1][0, ...]
                upper_result[:4 * 10] = lower_result[:4 * 10].clone()
                # if root is None:
                #     lower_result[-3:] = 0
                # if upper is None:
                #     lower_result[40:-3] = 0
            if root is not None:
                input_file = './datasets/' + dataset_name + '/{}/{}'.format(src_name, root)
                model = demo_model(input_file, ref_file, 'cross', './datasets/bvh2latent/Trinity')
                if type(root) is not np.ndarray:
                    root = torch.from_numpy(np.load(root)).to(device)
                else:
                    root = root.to(device)
                latent = [root, latent_1]  # 20230413
                model.latent2res(latent)
                root_result = model.fake_res_denorm[1][0, ...]
                upper_result[-3:] = root_result[-3:].clone()

            model.generate_bvh(result=upper_result, mode=mode)
            os.system('cp "{}/{}/0_{}.bvh" "./{}"'.format(model.bvh_path, 'ZEGGS', '0', output_path))

        if mode == 'Trinity':
            from eval_single_pair import demo_model
            dataset_name = 'Mixamo_new_2'
            src_name = 'ZEGGS'
            dest_name = 'Trinity'
            bvh_name = item
            ref_bvh_name = 'Recording_006.bvh'
            input_file = './datasets/' + dataset_name + '/{}/{}'.format(src_name, bvh_name)
            ref_file = './datasets/' + dataset_name + '/{}/{}'.format(dest_name, ref_bvh_name)
            model = demo_model(input_file, ref_file, 'cross', './datasets/bvh2latent/ZEGGS')
            latent_0 = np.load(item)
            latent_1 = np.load(ref_path)
            latent = [torch.from_numpy(latent_1).to(device), torch.from_numpy(latent_0).to(device)]     # 20230413
            model.latent2res(latent)
            # model.compute_test_result()
            global_position = torch.from_numpy(np.load(global_position, allow_pickle=True)[0]).to(device)
            model.compute_test_result_global(global_position, mode)
            os.system('cp "{}/{}/0_{}.bvh" "./{}"'.format(model.bvh_path, 'Trinity', '1', output_path))


def orignal2bvh(upper=None, lower=None, root=None, ref_path=None, output_path=None):

    dataset_name = 'Mixamo_new_2'
    src_name = 'Trinity'
    dest_name = 'ZEGGS'
    bvh_name = upper
    ref_bvh_name = '067_Speech_2_x_1_0.bvh'
    input_file = './datasets/' + dataset_name + '/{}/{}'.format(src_name, bvh_name)
    ref_file = './datasets/' + dataset_name + '/{}/{}'.format(dest_name, ref_bvh_name)

    import torch
    import option_parser
    from datasets import create_dataset

    parser = option_parser.get_parser()
    args = parser.parse_args()
    args.input_bvh = input_file
    args.target_bvh = ref_file
    args.output_filename = output_path

    character_names = [['Trinity'], ['ZEGGS']]

    eval_seq = args.eval_seq
    test_device = args.cuda_device
    args.cuda_device = test_device if torch.cuda.is_available() else 'cpu'

    para_path = os.path.join(args.save_dir, 'para.txt')
    with open(para_path, 'r') as para_file:
        argv_ = para_file.readline().split()[1:]
        args = option_parser.get_parser().parse_args(argv_)

    args.is_train = False  #
    args.rotation = 'quaternion'
    args.eval_seq = eval_seq

    dataset = create_dataset(args, character_names)

    if upper is not None:
        upper_result = torch.from_numpy(np.load(upper)[:, :-1]).to(args.cuda_device)

    fake_res_denorm = dataset.denorm(0, [0], upper_result)
    path = './datasets/Mixamo_new_2/Trinity/Recording_006.bvh'
    file = BVH_file(path)
    writer = BVH_writer(file.edges, file.names)
    writer.write_raw(fake_res_denorm[0], 'quaternion', os.path.join(output_path, '0_gt_quaternion.bvh'))
    # pdb.set_trace()


def generate_result(source_path, generate_path):
    for item in os.listdir(source_path):
        print(item)
        latent2bvh('ZEGGS',
                   os.path.join(source_path, item),
                   None,
                   None,
                   ref_path='./datasets/bvh2latent/ZEGGS/065_Speech_0_x_1_0.npy',
                   output_path=generate_path,
                   model_name='diffusion')

        # latent2bvh('Trinity',
        #            os.path.join(source_path, item),
        #            None,
        #            None,
        #            # ref_path='./datasets/bvh2latent/Trinity/Recording_006.npy',
        #            ref_path='./datasets/bvh2latent/ZEGGS/065_Speech_0_x_1_0.npy',
        #            output_path=generate_path,
        #            model_name='diffusion')


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Retargeting')
    parser.add_argument('--model_name', type=str, default='diffusion')
    parser.add_argument('--cuda_device', type=str, default='cuda:0')
    parser.add_argument('--target', type=str, default='ZEGGS')
    parser.add_argument('--input_file', type=str, default="../diffusion_latent/result_model3_128/Trinity/005_Neutral_4_x_1_0_minibatch_30_[0, 0, 0, 0, 0, 3, 0]_123456_recon.npy")
    parser.add_argument('--ref_path', type=str, default='./datasets/bvh2latent/ZEGGS/065_Speech_0_x_1_0.npy')
    parser.add_argument('--output_path', type=str, default='../result/inference/Trinity/')
    parser.add_argument('--mode', type=str, default='quick_start')
    parser.add_argument('--save_dir', type=str, default='./my_model_new_3', help='directory for all savings')
    config = parser.parse_args()

    # example('Mixamo_new_2', 'Trinity', 'ZEGGS', 'Recording_006.bvh', 'cross', './examples_my/cross_structure', ref_bvh_name='067_Speech_2_x_1_0.bvh')
    # example('Mixamo_new_2', 'ZEGGS', 'Trinity', '067_Speech_2_x_1_0.bvh', 'cross', './examples_my/cross_structure', ref_bvh_name='Recording_006.bvh')


    # x = np.load("../codebook/result/inference/Trinity/generate_train_codebook.npy")
    # x = x[0].transpose(1, 0).reshape(-1, 28, 4)
    # x[:, -1, :] = 0
    # x = np.expand_dims(x.reshape(-1, 112).transpose(1, 0), axis=0)
    # np.save("../codebook/result/inference/Trinity/generate_train_codebook_setroot.npy", x)

    # latent2bvh('ZEGGS', "./examples_my/cross_structure/Trinity/Recording_006_2.npy",
    #            ref_path='./datasets/bvh2latent/ZEGGS/065_Speech_0_x_1_0.npy', output_path="./examples_my/cross_structure/")
    # latent2bvh('Trinity', './datasets/bvh2latent/ZEGGS/065_Speech_0_x_1_0.npy',
    #            ref_path='./datasets/bvh2latent/Trinity/Recording_004.npy', output_path='./datasets/bvh2latent/ZEGGS/')

    # locallatent_global2bvh('ZEGGS', './datasets/bvh2glocal_local_latent/Trinity/Recording_004_local.npy',
    #                        './datasets/bvh2glocal_local_latent/Trinity/Recording_004_global.npy',
    #            ref_path='./datasets/bvh2latent/ZEGGS/065_Speech_0_x_1_0.npy', output_path="./examples_my/cross_structure/")


    # generate_result("../diffusion_latent/result_model3_128/Trinity/",
    #                 "../diffusion_latent/result_model3_128_bvh/")

    # generate_result("./datasets/valid_process_2/Trinity/",
    #                 "./datasets/valid_process_2/bvh/")

    if config.mode == 'quick_start':

        # Check if the directory exists
        if not os.path.exists(config.output_path):
            os.makedirs(config.output_path)
            print(f"Directory {config.output_path} created!")
        else:
            print(f"Directory {config.output_path} already exists!")

        latent2bvh(config.target,
                   config.input_file,
                   None,
                   None,
                   # "/ceph/hdd/yangsc21/Python/My_3/deep-motion-editing/retargeting/datasets/valid_processed/Trinity/TestSeq003_lower.npy",
                   #  "/ceph/hdd/yangsc21/Python/My_3/deep-motion-editing/retargeting/datasets/valid_processed/Trinity/TestSeq003_root.npy",
                    ref_path=config.ref_path,
                    output_path=config.output_path,
                    model_name=config.model_name,
                    device=config.cuda_device)


    elif config.mode == 'bvh2latent':
        bvh2latent()


    # latent2bvh('Trinity',
    #            '../dataset/ZEGGS/VQVAE_result/ZEGGS/generate_upper_001_Neutral_0_mirror_x_1_0.npy',
    #            '../dataset/ZEGGS/VQVAE_result/ZEGGS/generate_lower_001_Neutral_0_mirror_x_1_0.npy',
    #            '../dataset/ZEGGS/VQVAE_result/ZEGGS/generate_root_001_Neutral_0_mirror_x_1_0.npy',
    #            ref_path='./datasets/bvh2latent/Trinity/Recording_006.npy',
    #            output_path="./examples_my/cross_structure/")

    # orignal2bvh('../codebook/result/inference/Trinity/generate_upper_train_codebook_upper_original_downsample_2_joint_channel_1.npy', None, None,
    #             ref_path='./datasets/bvh2latent/ZEGGS/065_Speech_0_x_1_0.npy', output_path='../codebook/result/inference/Trinity')        # './examples_my/cross_structure/'

    print('Finished!')

    # bvh_input(quaternion=True, edge=True)
    # bvh_input(quaternion=True, edge=False)
    # # bvh_input(quaternion=False, edge=False)

    
retargeting/demo.py
YoungSeng/UnifiedGesture
import os
import pdb

import torch
from models import create_model
from datasets import create_dataset
import option_parser
from os.path import join as pjoin
import numpy as np
from datasets.bvh_parser import BVH_file

def eval_prepare(args):
    character = []
    file_id = []
    character_names = []
    character_names.append(args.input_bvh.split('/')[-2])
    character_names.append(args.target_bvh.split('/')[-2])
    # if args.test_type == 'intra':
    #     if character_names[0].endswith('_m'):
    #         character = [['BigVegas', 'BigVegas'], character_names]
    #         file_id = [[0, 0], [args.input_bvh, args.input_bvh]]
    #         src_id = 1
    #     else:
    #         character = [character_names, ['Goblin_m', 'Goblin_m']]
    #         file_id = [[args.input_bvh, args.input_bvh], [0, 0]]
    #         src_id = 0
    if args.test_type == 'cross':
        if character_names[0]=='ZEGGS':
            character = [[character_names[1]], [character_names[0]]]
            # file_id = [[0], [args.input_bvh]]
            file_id = [[args.target_bvh], [args.input_bvh]]      # ,[0]        modify 2
            src_id = 1
        else:
            character = [[character_names[0]], [character_names[1]]]
            file_id = [[args.input_bvh], [args.target_bvh]]      # ,[0]        modify 2
            # file_id = [[args.input_bvh], [0]]
            src_id = 0
    else:
        raise Exception('Unknown test type')
    return character, file_id, src_id


def recover_space(file):
    l = file.split('/')
    l[-1] = l[-1].replace('_', ' ')
    return '/'.join(l)


def demo_model(input_bvh, target_bvh, test_type, output_filename):
    parser = option_parser.get_parser()
    args = parser.parse_args()
    args.input_bvh = input_bvh
    args.target_bvh = target_bvh
    args.test_type = test_type
    args.output_filename = output_filename
    character_names, file_id, src_id = eval_prepare(args)
    print(character_names, input_bvh)
    test_device = args.cuda_device
    eval_seq = args.eval_seq

    para_path = os.path.join(args.save_dir, 'para.txt')
    with open(para_path, 'r') as para_file:
        argv_ = para_file.readline().split()[1:]
        args = option_parser.get_parser().parse_args(argv_)

    args.cuda_device = test_device if torch.cuda.is_available() else 'cpu'
    args.is_train = False  #
    args.rotation = 'quaternion'
    args.eval_seq = eval_seq

    dataset = create_dataset(args, character_names)
    model = create_model(args, character_names, dataset)
    model.load(epoch=16000)     # my_model_new_3 + Mixamo_new_2
    return model


def main(input_bvh, target_bvh, test_type, output_filename, load_model=None):
    parser = option_parser.get_parser()
    args = parser.parse_args()

    args.input_bvh = input_bvh
    args.target_bvh = target_bvh
    args.test_type = test_type
    args.output_filename = output_filename

    character_names, file_id, src_id = eval_prepare(args)
    # character_names = [['Trinity'], ['Talking_With_Hands']]
    # character_names = [['Trinity'], ['ZEGGS']]
    '''
    # intra
    character_names     [['Aj', 'BigVegas'], ['Goblin_m', 'Goblin_m']]      # Aj (..., 69) BigVegas (..., 69)
    file_id             [['./datasets/Mixamo/Aj/Dancing Running Man.bvh', './datasets/Mixamo/Aj/Dancing Running Man.bvh'], [0, 0]]
    src_id              0
    # cross
    character_names     [['BigVegas'], ['Mousey_m']]
    file_id             [['./datasets/Mixamo/BigVegas/Dual Weapon Combo.bvh'], [0]]
    src_id              0
    '''

    # input_character_name = args.input_bvh.split('/')[-2]
    output_character_name = args.target_bvh.split('/')[-2]
    output_filename = args.output_filename
    output_filename = pjoin(output_filename, args.input_bvh.split('/')[-1].replace('.bvh', '.npy'))

    test_device = args.cuda_device
    eval_seq = args.eval_seq
    args.cuda_device = test_device if torch.cuda.is_available() else 'cpu'
    args.is_train = False  #
    args.rotation = 'quaternion'
    args.eval_seq = eval_seq

    dataset = create_dataset(args, character_names)

    if load_model is None:
        model = demo_model(input_bvh, target_bvh, test_type, output_filename)
    else:
        model = load_model

    # for body in ['upper', 'lower', 'root']:
    if True:
        # if body == 'upper':
        #     upper_motions = []
        # elif body == 'lower':
        #     lower_motions = []
        # elif body == 'root':
        #     root_motions = []

        input_motion = []
        for i, character_group in enumerate(character_names):
            print('\r' + str(i), end='')
            input_group = []

            # if body == 'upper':
            #     upper_motions_group = []
            # elif body == 'lower':
            #     lower_motions_group = []
            # elif body == 'root':
            #     root_motions_group = []

            for j in range(len(character_group)):
                print(file_id[i][j])
                new_motion = dataset.get_item(i, j, file_id[i][j])
                new_motion_ = new_motion.clone()
                new_motion_.unsqueeze_(0)        # [103, 14284] -> [1, 103, 14284]
                # if body == 'upper':
                #     upper_motion = new_motion[:, 10 * 4:-3].clone()
                #     upper_motions_group.append(upper_motion)
                # elif body == 'lower':
                #     lower_motion = new_motion[:, :10 * 4].clone()
                #     lower_motions_group.append(lower_motion)
                # elif body == 'root':
                #     root_motion = new_motion[:, -3:].clone()
                #     root_motions_group.append(root_motion)
                new_motion_ = (new_motion_ - dataset.mean[i][j]) / dataset.var[i][j]

                # if body == 'upper':       # 20230418
                #     new_motion_[:, :10 * 4] = 0
                #     new_motion_[:, -3:] = 0
                # elif body == 'lower':
                #     new_motion_[:, 10 * 4:] = 0
                # elif body == 'root':
                #     new_motion_[:, :-3] = 0     # 20230414

                input_group.append(new_motion_)

            input_group = torch.cat(input_group, dim=0)
            input_motion.append([input_group, list(range(len(character_group)))])

        model.set_input(input_motion)
        # input_motion [[[1, 91, 216], [0]], [[1, 111, len'], [0]]]

        # inference latent
        latent_ = model.get_latent()     # [[1, 112, 3165], [1, 112, 517]]

        for body in ['upper', 'lower', 'root']:
            latent = latent_.copy()
            if body == 'upper':
                latent[0] = latent[0][:, 2 * 16:-1 * 16]      # 20230421
                latent[1] = latent[1][:, 2 * 16:-1 * 16]

                # latent = input_motion[0]
                # latent[0][:, :4 * 10] = 0
                # latent[0][:, -1 * 3:] = 0
                # latent[0] = torch.cat((latent[0], torch.zeros_like(latent[0][:, [0], :])), dim=1)

            elif body == 'lower':
                # latent[0][:, 2 * 16:-1 * 16] = 0
                # latent[0] = torch.cat((latent[0][:, :2 * 16], latent[0][:, -1 * 16:]), dim=1)       # bvh2upper_lower

                latent[0] = latent[0][:, :2 * 16]
                latent[1] = latent[1][:, :2 * 16]

                # latent = input_motion[0]
                # latent[0][:, 4 * 10:] = 0
                # latent[0] = torch.cat((latent[0], torch.zeros_like(latent[0][:, [0], :])), dim=1)
            elif body == 'root':
                latent[0] = latent[0][:, -1 * 16:]
                latent[1] = latent[1][:, -1 * 16:]

                # latent = input_motion[0]
                # latent[0][:, :-1 * 3] = 0
                # latent[0] = torch.cat((latent[0], torch.zeros_like(latent[0][:, [0], :])), dim=1)

            output_filename_ = output_filename[:-4]
            print(output_filename_, latent[0].shape)
            if src_id == 0:
                np.save("./{}_{}.npy".format(output_filename_, body), latent[0].detach().cpu().numpy())
                # np.save("./{}_global.npy".format(output_filename), [input_motion_root[0].detach().cpu().numpy(), BVH_file(args.input_bvh).get_height()])
            elif src_id == 1:
                np.save("./{}_{}.npy".format(output_filename_, body), latent[1].detach().cpu().numpy())
                # np.save("./{}_global.npy".format(output_filename), [input_motion_root[1].detach().cpu().numpy(), BVH_file(args.input_bvh).get_height()])


def load_model():
    from datasets import create_dataset, get_character_names
    from torch.utils.data.dataloader import DataLoader

    args = option_parser.get_args()
    # characters = get_character_names(args)      # [['ch01'], ['ch02']]

    # characters = [['Trinity'], ['Talking_With_Hands']]

    # characters = [['Talking_With_Hands'], ['Trinity']]      # model_4, model_5

    characters = [['Trinity'], ['ZEGGS']]

    para_path = os.path.join(args.save_dir, 'para.txt')
    with open(para_path, 'r') as para_file:
        argv_ = para_file.readline().split()[1:]
        args = option_parser.get_parser().parse_args(argv_)

    # data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=2)
    args.is_train = False       #
    args.rotation = 'quaternion'
    dataset = create_dataset(args, characters)
    model = create_model(args, characters, dataset)

    model.load(epoch=12400)      # 20000, modify for test
    return args, characters, model


def my_main(args, character_names, model):
    # input_bvh = './datasets/Mixamo/Talking_With_Hands/trn_2022_v1_064.bvh'
    # ref_bvh = './datasets/Mixamo/Trinity/Recording_006.bvh'
    from datasets.combined_motion import MixedData
    from torch.utils.data.dataloader import DataLoader
    dataset = MixedData(args, character_names)
    data_loader = DataLoader(dataset, batch_size=1, shuffle=False)
    print(args.cuda_device)
    for step, motions in enumerate(data_loader):
        model.set_input(motions)
        model.test()
        break
    print('Finish!')


if __name__ == '__main__':
    '''
    cd /nfs7/y50021900/My_3/deep-motion-editing/retargeting/
    cd /ceph/hdd/yangsc21/Python/My_3/deep-motion-editing/retargeting/
    python eval_single_pair.py
    '''

    # src_name = 'Talking_With_Hands'
    # dest_name = 'Trinity'
    # bvh_name = 'trn_2022_v1_064.bvh'
    # test_type = 'cross'
    # output_path = './examples_my/cross_structure'
    # ref_bvh_name = 'Recording_006.bvh'
    #
    # input_file = './datasets/Mixamo/{}/{}'.format(src_name, bvh_name)
    # ref_file = './datasets/Mixamo/{}/{}'.format(dest_name, ref_bvh_name)
    #
    # cmd = 'python eval_single_pair.py --input_bvh={} --target_bvh={} --output_filename={} --test_type={}'.format(
    #     input_file, ref_file, pjoin(output_path, 'result.bvh'), test_type
    # )
    # os.system(cmd)

    # args, characters, model = load_model()
    # my_main(args, characters, model)

    # main(input_bvh, target_bvh, test_type, output_filename)
    pass


retargeting/eval_single_pair.py
YoungSeng/UnifiedGesture
import pdb
import sys
import torch
from models.Kinematics import InverseKinematics
from datasets.bvh_parser import BVH_file
from tqdm import tqdm

sys.path.append('../utils')
import numpy as np
import BVH as BVH
import Animation as Animation
from Quaternions_old import Quaternions
import os
from scipy.signal import savgol_filter

L = 5  # #frame to look forward/backward


def alpha(t):
    return 2.0 * t * t * t - 3.0 * t * t + 1


def lerp(a, l, r):
    return (1 - a) * l + a * r


def get_character_height(file_name):
    file = BVH_file(file_name)
    return file.get_height()


def get_foot_contact(file_name, ref_height):
    anim, names, _ = BVH.load(file_name)

    ee_ids = get_ee_id_by_names(names)

    glb = Animation.positions_global(anim)  # [T, J, 3]

    ee_pos = glb[:, ee_ids, :]
    ee_velo = ee_pos[1:, ...] - ee_pos[:-1, ...]
    ee_velo = torch.tensor(ee_velo) / ref_height
    ee_velo_norm = torch.norm(ee_velo, dim=-1)
    contact = ee_velo_norm < 0.003
    contact = contact.int()
    padding = torch.zeros_like(contact)
    contact = torch.cat([padding[:1, :], contact], dim=0)
    return contact.numpy()


def get_foot_vel_position(file_name, ref_height):
    anim, names, _ = BVH.load(file_name)
    ee_ids = get_ee_id_by_names(names)
    glb = Animation.positions_global(anim)  # [T, J, 3]
    T = glb.shape[0]
    ee_pos = glb[:, ee_ids, :]
    ee_velo = ee_pos[1:, ...] - ee_pos[:-1, ...]
    ee_velo = np.pad(ee_velo, ((1, 0), (0, 0), (0, 0)), 'constant')
    root_position = glb[:, names.index('Hips')].copy()
    result = np.concatenate([root_position, ee_pos.reshape(T, -1), ee_velo.reshape(T, -1)], axis=-1)
    return result / ref_height


def get_ee_id_by_names(joint_names):
    ees = ['RightToeBase', 'LeftToeBase', 'LeftFoot', 'RightFoot']
    # ees = ['RightToeBase', 'LeftToeBase', 'LeftToeBaseEnd', 'RightToeBaseEnd']
    ee_id = []
    for i, ee in enumerate(ees):
        ee_id.append(joint_names.index(ee))
    return ee_id


def PFC_fix(input_file, ref_height):
    path_1 = os.path.split(input_file)[0]
    path_2 = os.path.split(input_file)[1][:-4]
    output_file = os.path.join(path_1, path_2 + '_fix.bvh')
    print('output file:', output_file)
    anim, name, ftime = BVH.load(input_file)
    fid = get_ee_id_by_names(name)
    contact_original = get_foot_contact(input_file, ref_height)
    # rowIndex = np.where((contact_original == (0, 0, 0, 0)).all(axis=1))
    glb = Animation.positions_global(anim)

    ee_pos = glb[:, fid, :]
    ee_velo = (ee_pos[1:, ...] - ee_pos[:-1, ...]) / ref_height
    ee_velo_norm = torch.norm(torch.tensor(ee_velo), dim=-1)
    padding = torch.zeros_like(ee_velo_norm)
    ee_velo_norm = torch.cat([padding[:1, :], ee_velo_norm], dim=0)
    T = glb.shape[0]
    root_position = glb[:, name.index('Hips')].copy()

    smoothing = True
    # smoothing
    if smoothing:
        n_poses = root_position.shape[0]
        out_poses = np.zeros((n_poses, root_position.shape[1]))
        for i in range(out_poses.shape[1]):
            # if (13 + (njoints - 14) * 9) <= i < (13 + njoints * 9): out_poses[:, i] = savgol_filter(poses[:, i], 41, 2)  # NOTE: smoothing on rotation matrices is not optimal
            # else:
            out_poses[:, i] = savgol_filter(root_position[:, i], 41, 2)  # NOTE: smoothing on rotation matrices is not optimal
        root_position = out_poses
    glb[:, name.index('Hips')] = root_position

    root_vec = (root_position[1:, ...] - root_position[:-1, ...]) / ref_height
    root_acc = root_vec[1:, ...] - root_vec[:-1, ...]
    root_vec = np.pad(root_vec, ((1, 0), (0, 0)), 'constant')
    root_acc = np.pad(root_acc, ((2, 0), (0, 0)), 'constant')
    z_acc = root_acc[:, 1]
    z_acc[z_acc < 0] = 0.0
    root_acc[:, 1] = z_acc          # (len, 3)
    root_vec = np.linalg.norm(root_vec, axis=1)
    root_acc = np.linalg.norm(root_acc, axis=1)
    # root_contact = np.where(root_acc>=0.0003)
    # print('root contact:', len(root_contact[0]), 'frames:', T)
    # for index in root_contact[0]:
    #     ee_velo_norm[index] = ee_velo_norm[index] * 0.75

    for i in range(1, T):
        ee_velo_norm[i] = ee_velo_norm[i] * 1 / np.exp(1000 * root_acc[i])       #  + 50 * root_vec[i]

    contact = ee_velo_norm < 0.003
    contact_fix = contact.int().numpy()     # np.where((contact_fix == (0, 0, 0, 0)).all(axis=1))
    # contact_fix[np.where((contact_fix == (0, 0, 0, 0)).all(axis=1))] = (0, 1, 1, 0)
    for i, fidx in enumerate(fid):  # fidx: index of the foot joint
        fixed = contact_fix[:, i]  # [T]
        s = 0
        while s < T:
            while s < T and fixed[s] == 0:
                s += 1
            if s >= T:
                break
            t = s
            avg = glb[t, fidx].copy()
            while t + 1 < T and fixed[t + 1] == 1:
                t += 1
                avg += glb[t, fidx].copy()
            avg /= (t - s + 1)

            for j in range(s, t + 1):
                glb[j, fidx] = avg.copy()
            s = t + 1

        for s in range(T):
            if fixed[s] == 1:
                continue
            l, r = None, None
            consl, consr = False, False
            for k in range(L):
                if s - k - 1 < 0:
                    break
                if fixed[s - k - 1]:
                    l = s - k - 1
                    consl = True
                    break
            for k in range(L):
                if s + k + 1 >= T:
                    break
                if fixed[s + k + 1]:
                    r = s + k + 1
                    consr = True
                    break
            if not consl and not consr:
                continue
            if consl and consr:
                litp = lerp(alpha(1.0 * (s - l + 1) / (L + 1)),
                            glb[s, fidx], glb[l, fidx])
                ritp = lerp(alpha(1.0 * (r - s + 1) / (L + 1)),
                            glb[s, fidx], glb[r, fidx])
                itp = lerp(alpha(1.0 * (s - l + 1) / (r - l + 1)),
                           ritp, litp)
                glb[s, fidx] = itp.copy()
                continue
            if consl:
                litp = lerp(alpha(1.0 * (s - l + 1) / (L + 1)),
                            glb[s, fidx], glb[l, fidx])
                glb[s, fidx] = litp.copy()
                continue
            if consr:
                ritp = lerp(alpha(1.0 * (r - s + 1) / (L + 1)),
                            glb[s, fidx], glb[r, fidx])
                glb[s, fidx] = ritp.copy()

    # glb is ready

    anim = anim.copy()

    rot = torch.tensor(anim.rotations.qs, dtype=torch.float)
    pos = torch.tensor(anim.positions[:, 0, :], dtype=torch.float)
    offset = torch.tensor(anim.offsets, dtype=torch.float)

    glb = torch.tensor(glb, dtype=torch.float)

    ik_solver = InverseKinematics(rot, pos, offset, anim.parents, glb)

    print('Fixing foot contact using IK...')
    for i in tqdm(range(100)):
        ik_solver.step()

    rotations = ik_solver.rotations.detach()
    norm = torch.norm(rotations, dim=-1, keepdim=True)
    rotations /= norm

    anim.rotations = Quaternions(rotations.numpy())
    anim.positions[:, 0, :] = ik_solver.position.detach().numpy()

    BVH.save(output_file, anim, name, ftime)



def fix_foot_contact(input_file, foot_file, output_file, ref_height):
    anim, name, ftime = BVH.load(input_file)

    fid = get_ee_id_by_names(name)
    contact = get_foot_contact(foot_file, ref_height)

    glb = Animation.positions_global(anim)  # [T, J, 3]

    T = glb.shape[0]

    for i, fidx in enumerate(fid):  # fidx: index of the foot joint
        fixed = contact[:, i]  # [T]
        s = 0
        while s < T:
            while s < T and fixed[s] == 0:
                s += 1
            if s >= T:
                break
            t = s
            avg = glb[t, fidx].copy()
            while t + 1 < T and fixed[t + 1] == 1:
                t += 1
                avg += glb[t, fidx].copy()
            avg /= (t - s + 1)

            for j in range(s, t + 1):
                glb[j, fidx] = avg.copy()
            s = t + 1

        for s in range(T):
            if fixed[s] == 1:
                continue
            l, r = None, None
            consl, consr = False, False
            for k in range(L):
                if s - k - 1 < 0:
                    break
                if fixed[s - k - 1]:
                    l = s - k - 1
                    consl = True
                    break
            for k in range(L):
                if s + k + 1 >= T:
                    break
                if fixed[s + k + 1]:
                    r = s + k + 1
                    consr = True
                    break
            if not consl and not consr:
                continue
            if consl and consr:
                litp = lerp(alpha(1.0 * (s - l + 1) / (L + 1)),
                            glb[s, fidx], glb[l, fidx])
                ritp = lerp(alpha(1.0 * (r - s + 1) / (L + 1)),
                            glb[s, fidx], glb[r, fidx])
                itp = lerp(alpha(1.0 * (s - l + 1) / (r - l + 1)),
                           ritp, litp)
                glb[s, fidx] = itp.copy()
                continue
            if consl:
                litp = lerp(alpha(1.0 * (s - l + 1) / (L + 1)),
                            glb[s, fidx], glb[l, fidx])
                glb[s, fidx] = litp.copy()
                continue
            if consr:
                ritp = lerp(alpha(1.0 * (r - s + 1) / (L + 1)),
                            glb[s, fidx], glb[r, fidx])
                glb[s, fidx] = ritp.copy()

    # glb is ready

    anim = anim.copy()

    rot = torch.tensor(anim.rotations.qs, dtype=torch.float)
    pos = torch.tensor(anim.positions[:, 0, :], dtype=torch.float)
    offset = torch.tensor(anim.offsets, dtype=torch.float)

    glb = torch.tensor(glb, dtype=torch.float)

    ik_solver = InverseKinematics(rot, pos, offset, anim.parents, glb)

    print('Fixing foot contact using IK...')
    for i in tqdm(range(200)):
        ik_solver.step()

    rotations = ik_solver.rotations.detach()
    norm = torch.norm(rotations, dim=-1, keepdim=True)
    rotations /= norm

    anim.rotations = Quaternions(rotations.numpy())
    anim.positions[:, 0, :] = ik_solver.position.detach().numpy()

    BVH.save(output_file, anim, name, ftime)

retargeting/models/IK.py
YoungSeng/UnifiedGesture
import pdb
import sys

sys.path.append('./')
from models.integrated import IntegratedModel
from torch import optim
import torch
from models.utils import GAN_loss, ImagePool, get_ee, Criterion_EE, Eval_Criterion, Criterion_EE_2
from models.base_model import BaseModel
from option_parser import try_mkdir

import os


class GAN_model(BaseModel):
    def __init__(self, args, character_names, dataset):
        super(GAN_model, self).__init__(args)
        self.character_names = character_names
        self.dataset = dataset
        self.n_topology = len(character_names)      # [['Talking_With_Hands'], ['Trinity']]
        self.models = []
        self.D_para = []
        self.G_para = []
        self.args = args

        for i in range(self.n_topology):        # 2
            model = IntegratedModel(args, dataset.joint_topologies[i], None, self.device, character_names[i])
            self.models.append(model)
            self.D_para += model.D_parameters()
            self.G_para += model.G_parameters()


        if self.is_train:
            self.fake_pools = []
            self.optimizerD = optim.Adam(self.D_para, args.learning_rate, betas=(0.9, 0.999))
            self.optimizerG = optim.Adam(self.G_para, args.learning_rate, betas=(0.9, 0.999))
            self.optimizers = [self.optimizerD, self.optimizerG]
            self.criterion_rec = torch.nn.MSELoss()
            self.criterion_gan = GAN_loss(args.gan_mode).to(self.device)
            self.criterion_cycle = torch.nn.L1Loss()
            self.criterion_ee = Criterion_EE(args, torch.nn.MSELoss())
            for i in range(self.n_topology):
                self.fake_pools.append(ImagePool(args.pool_size))
        else:
            import option_parser
            self.err_crit = []
            for i in range(self.n_topology):
                self.err_crit.append(Eval_Criterion(dataset.joint_topologies[i]))
            self.id_test = 0
            self.bvh_path = os.path.join(args.save_dir, 'results/bvh')
            option_parser.try_mkdir(self.bvh_path)

            self.writer = []
            for i in range(self.n_topology):
                writer_group = []
                for _, char in enumerate(character_names[i]):
                    from datasets.bvh_writer import BVH_writer
                    from datasets.bvh_parser import BVH_file
                    import option_parser
                    file = BVH_file(option_parser.get_std_bvh(dataset=char))
                    writer_group.append(BVH_writer(file.edges, file.names))
                self.writer.append(writer_group)

    def set_input(self, motions):
        self.motions_input = motions

        if not self.is_train:
            self.motion_backup = []
            for i in range(self.n_topology):
                self.motion_backup.append(motions[i][0].clone())
                self.motions_input[i][0][1:] = self.motions_input[i][0][0]
                self.motions_input[i][1] = [0] * len(self.motions_input[i][1])

    def discriminator_requires_grad_(self, requires_grad):
        for model in self.models:
            for para in model.discriminator.parameters():
                para.requires_grad = requires_grad

    def forward(self):
        self.latents = []
        self.offset_repr = []
        self.pos_ref = []
        self.ee_ref = []
        self.res = []
        self.res_denorm = []
        self.res_pos = []
        self.fake_res = []
        self.fake_res_denorm = []
        self.fake_pos = []
        self.fake_ee = []
        self.fake_latent = []
        self.motions = []
        self.motion_denorm = []
        self.rnd_idx = []

        for i in range(self.n_topology):
            self.offset_repr.append(self.models[i].static_encoder(self.dataset.offsets[i]))

        # reconstruct
        for i in range(self.n_topology):        # 2
            motion, offset_idx = self.motions_input[i]          # [2, 99, 64], [0, 0]
            motion = motion.to(self.device)
            self.motions.append(motion)
            motion_denorm = self.dataset.denorm(i, offset_idx, motion)      # [2, 99, 64]
            self.motion_denorm.append(motion_denorm)
            offsets = [self.offset_repr[i][p][offset_idx] for p in range(self.args.num_layers + 1)]
            # num_layers    2   self.offset_repr  [[[1, 25, 3],[1, 90],[1, 180]],[]]    ->  [[2, 25, 3], [2, 90], [2, 180]]
            latent, res = self.models[i].auto_encoder(motion, offsets)      # -> [2, 112, 16], [2, 99, 64]
            res_denorm = self.dataset.denorm(i, offset_idx, res)        # -> [2, 99, 64]
            res_pos = self.models[i].fk.forward_from_raw(res_denorm, self.dataset.offsets[i][offset_idx])
            # self.dataset.offsets  [[1, 25, 3],[1, 25, 3]] ->  [2, 64, 25, 3]
            self.res_pos.append(res_pos)
            self.latents.append(latent)
            self.res.append(res)
            self.res_denorm.append(res_denorm)

            pos = self.models[i].fk.forward_from_raw(motion_denorm, self.dataset.offsets[i][offset_idx]).detach()       # -> [2, 64, 25, 3]
            ee = get_ee(pos, self.dataset.joint_topologies[i], self.dataset.ee_ids[i],
                        velo=self.args.ee_velo, from_root=self.args.ee_from_root)       # ee_velo   1   ee_from_root    1
            # self.dataset.joint_topologies
            # [(-1, 0, 1, 2, 3, 4, 0, 6, 7, 8, 9, 0, 11, 12, 13, 14, 15, 13, 17, 18, 19, 13, 21, 22, 23),
            #  (-1, 0, 1, 2, 3, 4, 0, 6, 7, 8, 9, 0, 11, 12, 13, 14, 15, 13, 17, 18, 19, 13, 21, 22, 23)]
            # self.dataset.ee_ids   [[5, 10, 16, 20, 24], [5, 10, 16, 20, 24]]
            # ee [2, 63, 5, 3]
            height = self.models[i].height[offset_idx]      # [[179.1899]], [0, 0] -> [[179.1899], [179.1899]], shape (2, 1)
            height = height.reshape((height.shape[0], 1, height.shape[1], 1))       # -> [2, 1, 1, 1]
            ee /= height
            self.pos_ref.append(pos)
            self.ee_ref.append(ee)

        # retargeting
        for src in range(self.n_topology):
            for dst in range(self.n_topology):
                if self.is_train:
                    rnd_idx = torch.randint(len(self.character_names[dst]), (self.latents[src].shape[0],))
                else:
                    rnd_idx = list(range(self.latents[0].shape[0]))

                self.rnd_idx.append(rnd_idx)        # [0, 0]
                dst_offsets_repr = [self.offset_repr[dst][p][rnd_idx] for p in range(self.args.num_layers + 1)]     # ->  [[2, 25, 3], [2, 90], [2, 180]]
                fake_res = self.models[dst].auto_encoder.dec(self.latents[src], dst_offsets_repr)       # -> [2, 99, 64]
                fake_latent = self.models[dst].auto_encoder.enc(fake_res, dst_offsets_repr)     # -> [2, 112, 16]

                fake_res_denorm = self.dataset.denorm(dst, rnd_idx, fake_res)       # -> [2, 99, 64]
                fake_pos = self.models[dst].fk.forward_from_raw(fake_res_denorm, self.dataset.offsets[dst][rnd_idx])        # -> [2, 64, 25, 3]
                fake_ee = get_ee(fake_pos, self.dataset.joint_topologies[dst], self.dataset.ee_ids[dst],
                                 velo=self.args.ee_velo, from_root=self.args.ee_from_root)      # -> [2, 63, 5, 3]
                height = self.models[dst].height[rnd_idx]       # [[179.1899]], [0, 0] -> [[179.1899], [179.1899]], shape (2, 1)
                height = height.reshape((height.shape[0], 1, height.shape[1], 1))       # [2, 1, 1, 1]
                fake_ee = fake_ee / height

                self.fake_latent.append(fake_latent)
                self.fake_pos.append(fake_pos)
                self.fake_res.append(fake_res)
                self.fake_ee.append(fake_ee)
                self.fake_res_denorm.append(fake_res_denorm)


    def backward_D_basic(self, netD, real, fake):
        """Calculate GAN loss for the discriminator
        Parameters:
            netD (network)      -- the discriminator D
            real (tensor array) -- real images
            fake (tensor array) -- images generated by a generator
        Return the discriminator loss.
        We also call loss_D.backward() to calculate the gradients.
        """
        # Real
        pred_real = netD(real)
        loss_D_real = self.criterion_gan(pred_real, True)
        # Fake
        pred_fake = netD(fake.detach())
        loss_D_fake = self.criterion_gan(pred_fake, False)
        # Combined loss and calculate gradients
        loss_D = (loss_D_real + loss_D_fake) * 0.5
        loss_D.backward()
        return loss_D

    def backward_D(self):
        self.loss_Ds = []
        self.loss_D = 0
        """
        A->A, A->B, B->A, B->B@[0, 1, 2, 3]
        """
        for i in range(self.n_topology):
            fake = self.fake_pools[i].query(self.fake_pos[2 - i])
            self.loss_Ds.append(self.backward_D_basic(self.models[i].discriminator, self.pos_ref[i].detach(), fake))
            self.loss_D += self.loss_Ds[-1]
            self.loss_recoder.add_scalar('D_loss_{}'.format(i), self.loss_Ds[-1])

    def backward_G(self):
        #rec_loss and gan loss
        self.rec_losses = []
        self.rec_loss = 0
        self.cycle_loss = 0
        self.loss_G = 0
        self.ee_loss = 0
        self.loss_G_total = 0
        for i in range(self.n_topology):
            rec_loss1 = self.criterion_rec(self.motions[i], self.res[i])        # [2, 99, 64]
            self.loss_recoder.add_scalar('rec_loss_quater_{}'.format(i), rec_loss1)

            height = self.models[i].real_height[self.motions_input[i][1]]
            height = height.reshape(height.shape + (1, 1,))
            input_pos = self.motion_denorm[i][:, -3:, :] / height       # self.motion_denorm[i] [2, 99, 64] -> [2, 3, 64]
            rec_pos = self.res_denorm[i][:, -3:, :] / height
            rec_loss2 = self.criterion_rec(input_pos, rec_pos)
            self.loss_recoder.add_scalar('rec_loss_global_{}'.format(i), rec_loss2)

            pos_ref_global = self.models[i].fk.from_local_to_world(self.pos_ref[i]) / height.reshape(height.shape + (1, ))
            # self.pos_ref [[2, 64, 25, 3], [2, 64, 25, 3]]  pos_ref_global   [2, 64, 25, 3]
            res_pos_global = self.models[i].fk.from_local_to_world(self.res_pos[i]) / height.reshape(height.shape + (1, ))
            # -> [2, 64, 25, 3]
            rec_loss3 = self.criterion_rec(pos_ref_global, res_pos_global)
            self.loss_recoder.add_scalar('rec_loss_position_{}'.format(i), rec_loss3)

            rec_loss = rec_loss1 + (rec_loss2 * self.args.lambda_global_pose +
                                    rec_loss3 * self.args.lambda_position) * 100

            self.rec_losses.append(rec_loss)
            self.rec_loss += rec_loss

        p = 0
        for src in range(self.n_topology):
            for dst in range(self.n_topology):
                cycle_loss = self.criterion_cycle(self.latents[src], self.fake_latent[p])       # [2, 112, 16], [2, 112, 16] ->
                self.loss_recoder.add_scalar('cycle_loss_{}_{}'.format(src, dst), cycle_loss)
                self.cycle_loss += cycle_loss
                ee_loss = self.criterion_ee(self.ee_ref[src], self.fake_ee[p])      # [2, 63, 5, 3], [2, 63, 5, 3] ->
                self.loss_recoder.add_scalar('ee_loss_{}_{}'.format(src, dst), ee_loss)
                self.ee_loss += ee_loss

                if src != dst:
                    if self.args.gan_mode != 'none':
                        loss_G = self.criterion_gan(self.models[dst].discriminator(self.fake_pos[p]), True)
                    else:
                        loss_G = torch.tensor(0)
                    self.loss_recoder.add_scalar('G_loss_{}_{}'.format(src, dst), loss_G)
                    self.loss_G += loss_G

                p += 1

        self.loss_G_total = self.rec_loss * self.args.lambda_rec  + \
                            self.cycle_loss * self.args.lambda_cycle / 2 + \
                            self.ee_loss * self.args.lambda_ee / 2 + \
                            self.loss_G * 1
        self.loss_G_total.backward()

    def optimize_parameters(self):
        self.forward()

        # update Gs
        self.discriminator_requires_grad_(False)
        self.optimizerG.zero_grad()
        self.backward_G()
        self.optimizerG.step()

        # update Ds
        if self.args.gan_mode != 'none':    # lsgan
            self.discriminator_requires_grad_(True)
            self.optimizerD.zero_grad()
            self.backward_D()
            self.optimizerD.step()
        else:
            self.loss_D = torch.tensor(0)

    def verbose(self):
        res = {'rec_loss_0': self.rec_losses[0].item(),
               'rec_loss_1': self.rec_losses[1].item(),
               'cycle_loss': self.cycle_loss.item(),
               'ee_loss': self.ee_loss.item(),
               'D_loss_gan': self.loss_D.item(),
               'G_loss_gan': self.loss_G.item()}
        return sorted(res.items(), key=lambda x: x[0])

    def save(self):
        for i, model in enumerate(self.models):
            model.save(os.path.join(self.model_save_dir, 'topology{}'.format(i)), self.epoch_cnt)

        for i, optimizer in enumerate(self.optimizers):
            file_name = os.path.join(self.model_save_dir, 'optimizers/{}/{}.pt'.format(self.epoch_cnt, i))
            try_mkdir(os.path.split(file_name)[0])
            torch.save(optimizer.state_dict(), file_name)

    def load(self, epoch=None):
        for i, model in enumerate(self.models):
            model.load(os.path.join(self.model_save_dir, 'topology{}'.format(i)), epoch)

        if self.is_train:
            for i, optimizer in enumerate(self.optimizers):
                file_name = os.path.join(self.model_save_dir, 'optimizers/{}/{}.pt'.format(epoch, i))
                optimizer.load_state_dict(torch.load(file_name))
        self.epoch_cnt = epoch

    def compute_test_result(self):
        # gt_poses = []
        # gt_denorm = []
        # for src in range(self.n_topology):
        #     gt = self.motion_backup[src]
        #     idx = list(range(gt.shape[0]))
        #     gt = gt.to(self.device)
        #     gt = self.dataset.denorm(src, idx, gt)
        #     gt_denorm.append(gt)
        #     gt_pose = self.models[src].fk.forward_from_raw(gt, self.dataset.offsets[src][idx])
        #     gt_poses.append(gt_pose)
        #     for i in idx:
        #         new_path = os.path.join(self.bvh_path, self.character_names[src][i])
        #         from option_parser import try_mkdir
        #         try_mkdir(new_path)
        #         self.writer[src][i].write_raw(gt[i, ...], 'quaternion', os.path.join(new_path, '{}_gt.bvh'.format(self.id_test)))

        p = 0
        for src in range(self.n_topology):
            for dst in range(self.n_topology):
                for i in range(len(self.character_names[dst])):
                    dst_path = os.path.join(self.bvh_path, self.character_names[dst][i])
                    self.writer[dst][i].write_raw(self.fake_res_denorm[p][i, ...], 'quaternion',
                                                  os.path.join(dst_path, '{}_{}.bvh'.format(self.id_test, src)))
                    pdb.set_trace()
                p += 1

        self.id_test += 1

    def generate_bvh(self, result=None, mode=None):
        if mode == 'ZEGGS':
            dst_path = os.path.join(self.bvh_path, self.character_names[1][0])
            self.writer[1][0].write_raw(result, 'quaternion', os.path.join(dst_path, '{}_{}.bvh'.format(0, 0)))
        elif mode == 'Trinity':
            dst_path = os.path.join(self.bvh_path, self.character_names[0][0])
            self.writer[0][0].write_raw(result, 'quaternion', os.path.join(dst_path, '{}_{}.bvh'.format(0, 0)))

    def compute_test_result_global(self, lower=None, root=None, mode=None):
        p = 0
        if mode == 'ZEGGS':
            src = 0
        for dst in range(self.n_topology):
            for i in range(len(self.character_names[dst])):
                dst_path = os.path.join(self.bvh_path, self.character_names[dst][i])
                # (103, len)
                result = self.fake_res_denorm[p][i, ...].clone()
                if lower is not None:
                    pass
                else:
                    result[:40, :] = 0
                    result[-3:, :] = 0
                self.writer[dst][i].write_raw(result, 'quaternion',
                                              os.path.join(dst_path, '{}_{}.bvh'.format(self.id_test, src)))
            p += 1

        self.id_test += 1

    def get_latent(self):
        self.latents = []
        self.offset_repr = []
        self.pos_ref = []
        self.ee_ref = []
        self.res = []
        self.res_denorm = []
        self.res_pos = []
        self.fake_res = []
        self.fake_res_denorm = []
        self.fake_pos = []
        self.fake_ee = []
        self.fake_latent = []
        self.motions = []
        self.motion_denorm = []
        self.rnd_idx = []

        for i in range(self.n_topology):
            self.offset_repr.append(self.models[i].static_encoder(self.dataset.offsets[i]))

        # reconstruct
        for i in range(self.n_topology):        # 2
            motion, offset_idx = self.motions_input[i]          # [2, 99, 64], [0, 0]
            motion = motion.to(self.device)
            self.motions.append(motion)
            motion_denorm = self.dataset.denorm(i, offset_idx, motion)      # [2, 99, 64]
            self.motion_denorm.append(motion_denorm)
            offsets = [self.offset_repr[i][p][offset_idx] for p in range(self.args.num_layers + 1)]
            # num_layers    2   self.offset_repr  [[[1, 25, 3],[1, 90],[1, 180]],[]]    ->  [[2, 25, 3], [2, 90], [2, 180]]
            latent, res = self.models[i].auto_encoder(motion, offsets)
            self.latents.append(latent)
        return self.latents

    def latent2res(self, latents):
        # self.latents = []
        self.offset_repr = []
        # self.pos_ref = []
        # self.ee_ref = []
        # self.res = []
        # self.res_denorm = []
        # self.res_pos = []
        self.fake_res = []
        self.fake_res_denorm = []
        self.fake_pos = []
        self.fake_ee = []
        self.fake_latent = []
        # self.motions = []
        # self.motion_denorm = []
        self.rnd_idx = []

        for i in range(self.n_topology):
            self.offset_repr.append(self.models[i].static_encoder(self.dataset.offsets[i]))

        self.latents = latents

        # reconstruct
        # for i in range(self.n_topology):  # 2
        #     pass
            # motion, offset_idx = self.motions_input[i]  # [2, 99, 64], [0, 0]
            # motion = motion.to(self.device)
            # self.motions.append(motion)
            # motion_denorm = self.dataset.denorm(i, offset_idx, motion)  # [2, 99, 64]
            # self.motion_denorm.append(motion_denorm)
            # offsets = [self.offset_repr[i][p][offset_idx] for p in range(self.args.num_layers + 1)]
            # num_layers    2   self.offset_repr  [[[1, 25, 3],[1, 90],[1, 180]],[]]    ->  [[2, 25, 3], [2, 90], [2, 180]]
            # latent, res = self.models[i].auto_encoder(motion, offsets)  # -> [2, 112, 16], [2, 99, 64]
            # res_denorm = self.dataset.denorm(i, offset_idx, res)  # -> [2, 99, 64]
            # res_pos = self.models[i].fk.forward_from_raw(res_denorm, self.dataset.offsets[i][offset_idx])
            # self.dataset.offsets  [[1, 25, 3],[1, 25, 3]] ->  [2, 64, 25, 3]
            # self.res_pos.append(res_pos)
            # self.latents.append(latent)
            # self.res.append(res)
            # self.res_denorm.append(res_denorm)
            #
            # pos = self.models[i].fk.forward_from_raw(motion_denorm,
            #                                          self.dataset.offsets[i][offset_idx]).detach()  # -> [2, 64, 25, 3]
            # ee = get_ee(pos, self.dataset.joint_topologies[i], self.dataset.ee_ids[i],
            #             velo=self.args.ee_velo, from_root=self.args.ee_from_root)  # ee_velo   1   ee_from_root    1
            # # self.dataset.joint_topologies
            # # [(-1, 0, 1, 2, 3, 4, 0, 6, 7, 8, 9, 0, 11, 12, 13, 14, 15, 13, 17, 18, 19, 13, 21, 22, 23),
            # #  (-1, 0, 1, 2, 3, 4, 0, 6, 7, 8, 9, 0, 11, 12, 13, 14, 15, 13, 17, 18, 19, 13, 21, 22, 23)]
            # # self.dataset.ee_ids   [[5, 10, 16, 20, 24], [5, 10, 16, 20, 24]]
            # # ee [2, 63, 5, 3]
            # height = self.models[i].height[offset_idx]  # [[179.1899]], [0, 0] -> [[179.1899], [179.1899]], shape (2, 1)
            # height = height.reshape((height.shape[0], 1, height.shape[1], 1))  # -> [2, 1, 1, 1]
            # ee /= height
            # self.pos_ref.append(pos)
            # self.ee_ref.append(ee)

        # retargeting
        for src in range(self.n_topology):
            for dst in range(self.n_topology):
                if self.is_train:
                    rnd_idx = torch.randint(len(self.character_names[dst]), (self.latents[src].shape[0],))
                else:
                    rnd_idx = list(range(self.latents[0].shape[0]))

                self.rnd_idx.append(rnd_idx)  # [0, 0]
                dst_offsets_repr = [self.offset_repr[dst][p][rnd_idx] for p in
                                    range(self.args.num_layers + 1)]  # ->  [[2, 25, 3], [2, 90], [2, 180]]
                fake_res = self.models[dst].auto_encoder.dec(self.latents[src], dst_offsets_repr)  # -> [2, 99, 64]
                fake_latent = self.models[dst].auto_encoder.enc(fake_res, dst_offsets_repr)  # -> [2, 112, 16]

                fake_res_denorm = self.dataset.denorm(dst, rnd_idx, fake_res)  # -> [2, 99, 64]
                fake_pos = self.models[dst].fk.forward_from_raw(fake_res_denorm,
                                                                self.dataset.offsets[dst][rnd_idx])  # -> [2, 64, 25, 3]
                fake_ee = get_ee(fake_pos, self.dataset.joint_topologies[dst], self.dataset.ee_ids[dst],
                                 velo=self.args.ee_velo, from_root=self.args.ee_from_root)  # -> [2, 63, 5, 3]
                height = self.models[dst].height[
                    rnd_idx]  # [[179.1899]], [0, 0] -> [[179.1899], [179.1899]], shape (2, 1)
                height = height.reshape((height.shape[0], 1, height.shape[1], 1))  # [2, 1, 1, 1]
                fake_ee = fake_ee / height

                self.fake_latent.append(fake_latent)
                self.fake_pos.append(fake_pos)
                self.fake_res.append(fake_res)
                self.fake_ee.append(fake_ee)
                self.fake_res_denorm.append(fake_res_denorm)


if __name__ == '__main__':
    '''
    python ./models/architecture.py
    '''

    from datasets import create_dataset, get_character_names
    import option_parser

    args = option_parser.get_args()
    characters = get_character_names(args)
    dataset = create_dataset(args, characters)
    print('OK')
    model = GAN_model(args, characters, dataset)
    model.is_train = False
    input_motion = [[torch.randn(1, 103, 64).to('cuda:5'), [0]], [torch.randn(1, 107, 64).to('cuda:5'), [0]]]      # (b, 112, len)
    model.set_input(input_motion)
    latent = model.get_latent()

retargeting/models/architecture.py
YoungSeng/UnifiedGesture
import pdb

import torch
import torch.nn as nn

import sys      # 20230419
[sys.path.append(i) for i in ['.', '..']]

from models.skeleton import SkeletonUnpool, SkeletonPool, SkeletonConv, find_neighbor, SkeletonLinear


class Encoder(nn.Module):
    def __init__(self, args, topology):
        super(Encoder, self).__init__()
        self.topologies = [topology]
        if args.rotation == 'euler_angle': self.channel_base = [3]
        elif args.rotation == 'quaternion': self.channel_base = [4]
        self.channel_list = []
        self.edge_num = [len(topology) + 1]
        self.pooling_list = []
        self.layers = nn.ModuleList()
        self.args = args
        self.convs = []

        kernel_size = args.kernel_size
        padding = (kernel_size - 1) // 2
        bias = True
        if args.skeleton_info == 'concat': add_offset = True
        else: add_offset = False

        for i in range(args.num_layers):
            self.channel_base.append(self.channel_base[-1] * 2)     # [4, 8, 16]

        for i in range(args.num_layers):
            seq = []
            neighbor_list = find_neighbor(self.topologies[i], args.skeleton_dist)
            in_channels = self.channel_base[i] * self.edge_num[i]
            out_channels = self.channel_base[i+1] * self.edge_num[i]
            if i == 0: self.channel_list.append(in_channels)
            self.channel_list.append(out_channels)
            for _ in range(args.extra_conv):        # 0
                seq.append(SkeletonConv(neighbor_list, in_channels=in_channels, out_channels=in_channels,
                                        joint_num=self.edge_num[i], kernel_size=kernel_size, stride=1,
                                        padding=padding, padding_mode=args.padding_mode, bias=bias))
            seq.append(SkeletonConv(neighbor_list, in_channels=in_channels, out_channels=out_channels,
                                    joint_num=self.edge_num[i], kernel_size=kernel_size, stride=2,
                                    padding=padding, padding_mode=args.padding_mode, bias=bias, add_offset=add_offset,
                                    in_offset_channel=3 * self.channel_base[i] // self.channel_base[0]))
            self.convs.append(seq[-1])
            last_pool = True if i == args.num_layers - 1 else False
            pool = SkeletonPool(edges=self.topologies[i], pooling_mode=args.skeleton_pool,
                                channels_per_edge=out_channels // len(neighbor_list), last_pool=last_pool)
            seq.append(pool)
            seq.append(nn.LeakyReLU(negative_slope=0.2))
            self.layers.append(nn.Sequential(*seq))

            self.topologies.append(pool.new_edges)
            self.pooling_list.append(pool.pooling_list)
            self.edge_num.append(len(self.topologies[-1]) + 1)
            if i == args.num_layers - 1:
                self.last_channel = self.edge_num[-1] * self.channel_base[i + 1]

    def forward(self, input, offset=None):
        # padding the one zero row to global position, so each joint including global position has 4 channels as input
        if self.args.rotation == 'quaternion' and self.args.pos_repr != '4d':
            input = torch.cat((input, torch.zeros_like(input[:, [0], :])), dim=1)       # [1, 104, 64]

        for i, layer in enumerate(self.layers):
            if self.args.skeleton_info == 'concat' and offset is not None:
                self.convs[i].set_offset(offset[i])

            input = layer(input)         # [1, 104, 64] -> [1, 120, 32] -> [1, 112, 16]  4*26 -> 8*15 -> 16*7
        return input


class Decoder(nn.Module):
    def __init__(self, args, enc: Encoder):
        super(Decoder, self).__init__()
        self.layers = nn.ModuleList()
        self.unpools = nn.ModuleList()
        self.args = args
        self.enc = enc
        self.convs = []

        kernel_size = args.kernel_size
        padding = (kernel_size - 1) // 2

        if args.skeleton_info == 'concat': add_offset = True
        else: add_offset = False

        for i in range(args.num_layers):
            seq = []
            in_channels = enc.channel_list[args.num_layers - i]
            out_channels = in_channels // 2
            neighbor_list = find_neighbor(enc.topologies[args.num_layers - i - 1], args.skeleton_dist)

            if i != 0 and i != args.num_layers - 1:
                bias = False
            else:
                bias = True

            self.unpools.append(SkeletonUnpool(enc.pooling_list[args.num_layers - i - 1], in_channels // len(neighbor_list)))

            seq.append(nn.Upsample(scale_factor=2, mode=args.upsampling, align_corners=False))
            seq.append(self.unpools[-1])
            for _ in range(args.extra_conv):
                seq.append(SkeletonConv(neighbor_list, in_channels=in_channels, out_channels=in_channels,
                                        joint_num=enc.edge_num[args.num_layers - i - 1], kernel_size=kernel_size,
                                        stride=1,
                                        padding=padding, padding_mode=args.padding_mode, bias=bias))
            seq.append(SkeletonConv(neighbor_list, in_channels=in_channels, out_channels=out_channels,
                                    joint_num=enc.edge_num[args.num_layers - i - 1], kernel_size=kernel_size, stride=1,
                                    padding=padding, padding_mode=args.padding_mode, bias=bias, add_offset=add_offset,
                                    in_offset_channel=3 * enc.channel_base[args.num_layers - i - 1] // enc.channel_base[0]))
            self.convs.append(seq[-1])
            if i != args.num_layers - 1: seq.append(nn.LeakyReLU(negative_slope=0.2))

            self.layers.append(nn.Sequential(*seq))

    def forward(self, input, offset=None):
        for i, layer in enumerate(self.layers):
            if self.args.skeleton_info == 'concat':
                self.convs[i].set_offset(offset[len(self.layers) - i - 1])
            input = layer(input)
        # throw the padded rwo for global position
        if self.args.rotation == 'quaternion' and self.args.pos_repr != '4d':
            input = input[:, :-1, :]

        return input


class AE(nn.Module):
    def __init__(self, args, topology):
        super(AE, self).__init__()
        self.enc = Encoder(args, topology)
        self.dec = Decoder(args, self.enc)

    def forward(self, input, offset=None):
        latent = self.enc(input, offset)
        result = self.dec(latent, offset)
        return latent, result


# eoncoder for static part, i.e. offset part
class StaticEncoder(nn.Module):
    def __init__(self, args, edges):
        super(StaticEncoder, self).__init__()
        self.args = args
        self.layers = nn.ModuleList()
        activation = nn.LeakyReLU(negative_slope=0.2)
        channels = 3

        for i in range(args.num_layers):
            neighbor_list = find_neighbor(edges, args.skeleton_dist)
            seq = []
            seq.append(SkeletonLinear(neighbor_list, in_channels=channels * len(neighbor_list),
                                      out_channels=channels * 2 * len(neighbor_list), extra_dim1=True))
            if i < args.num_layers - 1:
                pool = SkeletonPool(edges, channels_per_edge=channels*2, pooling_mode='mean')
                seq.append(pool)
                edges = pool.new_edges
            seq.append(activation)
            channels *= 2
            self.layers.append(nn.Sequential(*seq))

    # input should have shape B * E * 3
    def forward(self, input: torch.Tensor):
        output = [input]
        for i, layer in enumerate(self.layers):
            input = layer(input)        # [1, 25, 3] -> [1, 90, 1] -> [1, 180, 1]
            output.append(input.squeeze(-1))
        return output


if __name__ == '__main__':
    '''
    python ./models/enc_and_dec.py
    '''
    import option_parser
    from datasets import create_dataset, get_character_names
    from models.skeleton import build_edge_topology

    args = option_parser.get_args()
    characters = get_character_names(args)
    dataset = create_dataset(args, characters)
    topo = dataset.joint_topologies[0]      # len(topo) = 26
    edges = build_edge_topology(topo, torch.zeros((len(topo), 3)))      # len(edges) = 25
    args.is_train = False
    model = AE(args, edges)
    x = torch.rand(1, 103, 64)
    offset = [torch.rand(1, 26, 3), torch.rand(1, 90), torch.rand(1, 180)]
    y = model(x, offset)


retargeting/models/enc_and_dec.py
YoungSeng/UnifiedGesture
import pdb

import torch
from models.enc_and_dec import AE, StaticEncoder
from models.vanilla_gan import Discriminator
from models.skeleton import build_edge_topology
from models.Kinematics import ForwardKinematics
from datasets.bvh_parser import BVH_file
from option_parser import get_std_bvh
import os


class IntegratedModel:
    # origin_offsets should have shape num_skeleton * J * 3
    def __init__(self, args, joint_topology, origin_offsets: torch.Tensor, device, characters):
        self.args = args
        self.joint_topology = joint_topology
        self.edges = build_edge_topology(joint_topology, torch.zeros((len(joint_topology), 3)))
        self.fk = ForwardKinematics(args, self.edges)

        self.height = [] # for normalize ee_loss
        self.real_height = []
        for char in characters:
            if args.use_sep_ee:     # 0
                h = BVH_file(get_std_bvh(dataset=char)).get_ee_length()
            else:
                h = BVH_file(get_std_bvh(dataset=char)).get_height()
            if args.ee_loss_fact == 'learn':
                h = torch.tensor(h, dtype=torch.float)
            else:
                h = torch.tensor(h, dtype=torch.float, requires_grad=False)
            self.real_height.append(BVH_file(get_std_bvh(dataset=char)).get_height())
            self.height.append(h.unsqueeze(0))
        self.real_height = torch.tensor(self.real_height, device=device)
        self.height = torch.cat(self.height, dim=0)
        self.height = self.height.to(device)
        if not args.use_sep_ee: self.height.unsqueeze_(-1)
        if args.ee_loss_fact == 'learn': self.height_para = [self.height]
        else: self.height_para = []

        if not args.simple_operator:        # 0
            self.auto_encoder = AE(args, topology=self.edges).to(device)
            self.discriminator = Discriminator(args, self.edges).to(device)
            self.static_encoder = StaticEncoder(args, self.edges).to(device)
        else:
            raise Exception('Conventional operator not yet implemented')

    def parameters(self):
        return self.G_parameters() + self.D_parameters()

    def G_parameters(self):
        return list(self.auto_encoder.parameters()) + list(self.static_encoder.parameters()) + self.height_para

    def D_parameters(self):
        return list(self.discriminator.parameters())

    def save(self, path, epoch):
        from option_parser import try_mkdir

        path = os.path.join(path, str(epoch))
        try_mkdir(path)

        torch.save(self.height, os.path.join(path, 'height.pt'))
        torch.save(self.auto_encoder.state_dict(), os.path.join(path, 'auto_encoder.pt'))
        torch.save(self.discriminator.state_dict(), os.path.join(path, 'discriminator.pt'))
        torch.save(self.static_encoder.state_dict(), os.path.join(path, 'static_encoder.pt'))

        print('Save at {} succeed!'.format(path))

    def load(self, path, epoch=None):
        print('loading from', path)
        if not os.path.exists(path):
            raise Exception('Unknown loading path')

        if epoch is None:
            all = [int(q) for q in os.listdir(path) if os.path.isdir(path + q)]
            if len(all) == 0:
                raise Exception('Empty loading path')
            epoch = sorted(all)[-1]

        path = os.path.join(path, str(epoch))
        print('loading from epoch {}......'.format(epoch))

        self.auto_encoder.load_state_dict(torch.load(os.path.join(path, 'auto_encoder.pt'),
                                                     map_location=self.args.cuda_device))
        self.static_encoder.load_state_dict(torch.load(os.path.join(path, 'static_encoder.pt'),
                                                       map_location=self.args.cuda_device))
        print('load succeed!')

retargeting/models/integrated.py
YoungSeng/UnifiedGesture
import argparse
'''
pip install tensorboard

'''

dataset_name = "Mixamo_new_2"       # modify to 'Trinity_ZEGGS' for re-training

multiple = 1

def get_parser():
    parser = argparse.ArgumentParser()
    # parser.add_argument('--save_dir', type=str, default='./pretrained', help='directory for all savings')       # modify 3
    # parser.add_argument('--save_dir', type=str, default='./training', help='directory for all savings')
    # modify for test
    parser.add_argument('--save_dir', type=str, default='./my_model_new_3', help='directory for all savings')
    parser.add_argument('--cuda_device', type=str, default='cuda:0', help='cuda device number, eg:[cuda:0]')

    parser.add_argument('--num_layers', type=int, default=2, help='number of layers')       # model_5, model_7
    parser.add_argument('--learning_rate', type=float, default=multiple * 2e-4, help='learning rate')
    parser.add_argument('--alpha', type=float, default=0, help='penalty of sparsity')
    parser.add_argument('--batch_size', type=int, default=256 * multiple, help='batch_size')     # 256
    parser.add_argument('--upsampling', type=str, default='linear', help="'stride2' or 'nearest', 'linear'")
    parser.add_argument('--downsampling', type=str, default='stride2', help='stride2 or max_pooling')
    parser.add_argument('--batch_normalization', type=int, default=0, help='batch_norm: 1 or 0')
    parser.add_argument('--activation', type=str, default='LeakyReLU', help='activation: ReLU, LeakyReLU, tanh')
    parser.add_argument('--rotation', type=str, default='quaternion', help='representation of rotation:euler_angle, quaternion')
    parser.add_argument('--data_augment', type=int, default=1, help='data_augment: 1 or 0')
    parser.add_argument('--epoch_num', type=int, default=20001, help='epoch_num')
    parser.add_argument('--window_size', type=int, default=64, help='length of time axis per window')
    parser.add_argument('--kernel_size', type=int, default=15, help='must be odd')
    parser.add_argument('--base_channel_num', type=int, default=-1)
    parser.add_argument('--normalization', type=int, default=1)
    parser.add_argument('--verbose', type=int, default=1)
    parser.add_argument('--skeleton_dist', type=int, default=2)
    parser.add_argument('--skeleton_pool', type=str, default='mean')
    parser.add_argument('--extra_conv', type=int, default=0)
    parser.add_argument('--padding_mode', type=str, default='reflection')
    parser.add_argument('--dataset_constant', type=str, default=dataset_name)
    parser.add_argument('--dataset', type=str, default=dataset_name)
    parser.add_argument('--fk_world', type=int, default=0)
    parser.add_argument('--patch_gan', type=int, default=1)
    parser.add_argument('--debug', type=int, default=0)
    parser.add_argument('--skeleton_info', type=str, default='concat')
    parser.add_argument('--ee_loss_fact', type=str, default='height')
    parser.add_argument('--pos_repr', type=str, default='3d')
    parser.add_argument('--D_global_velo', type=int, default=0)
    parser.add_argument('--gan_mode', type=str, default='lsgan')
    parser.add_argument('--pool_size', type=int, default=50)
    parser.add_argument('--is_train', type=int, default=1)
    parser.add_argument('--model', type=str, default='mul_top_mul_ske')
    parser.add_argument('--epoch_begin', type=int, default=0)
    parser.add_argument('--lambda_rec', type=float, default=5)
    parser.add_argument('--lambda_cycle', type=float, default=5)
    parser.add_argument('--lambda_ee', type=float, default=100)
    parser.add_argument('--lambda_global_pose', type=float, default=2.5)
    parser.add_argument('--lambda_position', type=float, default=1)
    parser.add_argument('--ee_velo', type=int, default=1)
    parser.add_argument('--ee_from_root', type=int, default=1)
    parser.add_argument('--scheduler', type=str, default='none')
    parser.add_argument('--rec_loss_mode', type=str, default='extra_global_pos')
    parser.add_argument('--adaptive_ee', type=int, default=0)
    parser.add_argument('--simple_operator', type=int, default=0)
    parser.add_argument('--use_sep_ee', type=int, default=0)
    parser.add_argument('--eval_seq', type=int, default=0)

    parser.add_argument('--model_name', type=str, default='diffusion')
    parser.add_argument('--target', type=str, default='ZEGGS')
    parser.add_argument('--input_file', type=str, default="../diffusion_latent/result_model3_128/Trinity/005_Neutral_4_x_1_0_minibatch_30_[0, 0, 0, 0, 0, 3, 0]_123456_recon.npy")
    parser.add_argument('--ref_path', type=str, default='./datasets/bvh2latent/ZEGGS/065_Speech_0_x_1_0.npy')
    parser.add_argument('--output_path', type=str, default='../result/inference/Trinity/')
    parser.add_argument('--mode', type=str, default='quick_start')

    return parser


def get_args():
    parser = get_parser()
    return parser.parse_args()


def get_std_bvh(args=None, dataset=None):
    if args is None and dataset is None: raise Exception('Unexpected parameter')
    if dataset is None: dataset = args.dataset
    std_bvh = './datasets/' + dataset_name + '/std_bvhs/{}.bvh'.format(dataset)
    return std_bvh


def try_mkdir(path):
    import os
    if not os.path.exists(path):
        os.system('mkdir -p {}'.format(path))

retargeting/option_parser.py
YoungSeng/UnifiedGesture
import os
import glob
import pdb

import numpy as np

def extract_vel(path):
    root_files = sorted(glob.glob(path + "/*_root.npy"))
    for item in root_files:
        print(item)
        root = np.load(item)
        root_vel = root[..., 1:] - root[..., :-1]
        root_vel = np.pad(root_vel, [(0, 0), (0, 0), (1, 0)], mode='constant')
        np.save(item.replace('_root.npy', '_root_vel.npy'), root_vel)


if __name__ == '__main__':
    '''
    python process_root_vel.py
    '''
    Trinity_path = '../retargeting/datasets/Trinity_ZEGGS/bvh2upper_lower_root/Trinity/'
    ZEGGS_path = '../retargeting/datasets/Trinity_ZEGGS/bvh2upper_lower_root/ZEGGS/'

    extract_vel(Trinity_path)
    extract_vel(ZEGGS_path)

retargeting/process_root_vel.py
YoungSeng/UnifiedGesture
import os
from os.path import join as pjoin
from get_error import full_batch
import numpy as np
from option_parser import try_mkdir
from eval import eval
import argparse


def batch_copy(source_path, suffix, dest_path, dest_suffix=None):
    try_mkdir(dest_path)
    files = [f for f in os.listdir(source_path) if f.endswith('_{}.bvh'.format(suffix))]

    length = len('_{}.bvh'.format(suffix))
    for f in files:
        if dest_suffix is not None:
            cmd = 'cp \"{}\" \"{}\"'.format(os.path.join(source_path, f), os.path.join(dest_path, f[:-length] + '_{}.bvh'.format(dest_suffix)))
        else:
            cmd = 'cp \"{}\" \"{}\"'.format(os.path.join(source_path, f), os.path.join(dest_path, f[:-length] + '.bvh'))
        os.system(cmd)


if __name__ == '__main__':
    test_characters = ['Mousey_m', 'Goblin_m', 'Mremireh_m', 'Vampire_m']

    parser = argparse.ArgumentParser()
    # parser.add_argument('--save_dir', type=str, default='./pretrained/')
    parser.add_argument('--save_dir', type=str, default='./my_model/')

    args = parser.parse_args()
    prefix = args.save_dir

    cross_dest_path = pjoin(prefix, 'results/cross_structure/')
    intra_dest_path = pjoin(prefix, 'results/intra_structure/')
    source_path = pjoin(prefix, 'results/bvh/')

    cross_error = []
    intra_error = []
    for i in range(4):
        print('Batch [{}/4]'.format(i + 1))
        eval(i, prefix)

        print('Collecting test error...')
        if i == 0:
            cross_error += full_batch(0, prefix)
            for char in test_characters:
                batch_copy(os.path.join(source_path, char), 0, os.path.join(cross_dest_path, char))
                batch_copy(os.path.join(source_path, char), 'gt', os.path.join(cross_dest_path, char), 'gt')

        intra_dest = os.path.join(intra_dest_path, 'from_{}'.format(test_characters[i]))
        for char in test_characters:
            for char in test_characters:
                batch_copy(os.path.join(source_path, char), 1, os.path.join(intra_dest, char))
                batch_copy(os.path.join(source_path, char), 'gt', os.path.join(intra_dest, char), 'gt')

        intra_error += full_batch(1, prefix)

    cross_error = np.array(cross_error)
    intra_error = np.array(intra_error)

    cross_error_mean = cross_error.mean()
    intra_error_mean = intra_error.mean()

    os.system('rm -r %s' % pjoin(prefix, 'results/bvh'))

    print('Intra-retargeting error:', intra_error_mean)
    print('Cross-retargeting error:', cross_error_mean)
    print('Evaluation finished!')

retargeting/test.py
YoungSeng/UnifiedGesture
import pdb
import sys
sys.path.append('./retargeting/')
from torch.utils.data.dataloader import DataLoader
from models import create_model
from datasets import create_dataset, get_character_names
import option_parser
import os
from option_parser import try_mkdir
import time


def main():
    args = option_parser.get_args()
    characters = get_character_names(args)      # [['ch01'], ['ch02']]

    log_path = os.path.join(args.save_dir, 'logs/')
    try_mkdir(args.save_dir)
    try_mkdir(log_path)

    with open(os.path.join(args.save_dir, 'para.txt'), 'w') as para_file:
        para_file.write(' '.join(sys.argv))

    dataset = create_dataset(args, characters)
    data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=2)

    model = create_model(args, characters, dataset)

    if args.epoch_begin:
        model.load(epoch=args.epoch_begin, download=False)

    model.setup()

    start_time = time.time()

    for epoch in range(args.epoch_begin, args.epoch_num):
        for step, motions in enumerate(data_loader):
            '''
            motions [[torch.Size([1, 99, 64]), 0], [torch.Size([1, 99, 64]), 0]]
            '''
            model.set_input(motions)
            model.optimize_parameters()

            if args.verbose:
                res = model.verbose()
                print('[{}/{}]\t[{}/{}]\t'.format(epoch, args.epoch_num, step, len(data_loader)), res)

        if epoch % 200 == 0 or epoch == args.epoch_num - 1:
            model.save()

        model.epoch()

    end_tiem = time.time()
    print('training time', end_tiem - start_time)


if __name__ == '__main__':
    '''
    python train.py
    '''
    main()

retargeting/train.py
YoungSeng/UnifiedGesture
import pdb
import sys
[sys.path.append(i) for i in ['.', '..']]

from model.mdm import MDM
from utils.model_util import create_gaussian_diffusion

# from data_loaders.get_data import get_dataset_loader
from train.training_loop import TrainLoop

import torch
import os
import json


device = torch.device('cuda:2')
n_frames = 240
n_pose_dims = 135
n_audio_dim = 32

# n_frames = 240
# n_pose_dims = 251


def create_model_and_diffusion():
    model = MDM(modeltype='', njoints=n_pose_dims, nfeats=1, translation=True, pose_rep='rot6d', glob=True,
                glob_rot=True, cond_mode = 'text', clip_version = 'ViT-B/32', action_emb = 'tensor')
    diffusion = create_gaussian_diffusion()
    return model, diffusion


if __name__ == '__main__':
    '''
    python train/mytrain.py --overwrite --save_dir save/mydebug --dataset kit --device 1
    '''
    # modify data/dataset.py

    model, diffusion = create_model_and_diffusion()
    model.to(device)
    # model.rot2xyz.smpl_model.eval()

    print('Total params: %.2fM' % (sum(p.numel() for p in model.parameters_wo_clip()) / 1000000.0))


    from utils.parser_util import train_args
    from utils.fixseed import fixseed

    args = train_args()
    fixseed(args.seed)

    if args.save_dir is None:
        raise FileNotFoundError('save_dir was not specified.')
    elif os.path.exists(args.save_dir) and not args.overwrite:
        raise FileExistsError('save_dir [{}] already exists.'.format(args.save_dir))
    elif not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)
    args_path = os.path.join(args.save_dir, 'args.json')
    with open(args_path, 'w') as fw:
        json.dump(vars(args), fw, indent=4, sort_keys=True)

    # print("creating data loader...")
    # data = get_dataset_loader(name=args.dataset, batch_size=args.batch_size, num_frames=args.num_frames)

    # print(iter(data).next()[1]['y'].keys())
    print("Training...")
    TrainLoop(args, model, diffusion, device).run_loop()


train/mytrain.py
YoungSeng/UnifiedGesture
import copy
import functools
import os
import pdb
import time
from types import SimpleNamespace
import numpy as np

import blobfile as bf
import torch
from torch.optim import AdamW

from diffusion import logger
from diffusion.fp16_util import MixedPrecisionTrainer
from diffusion.resample import LossAwareSampler, UniformSampler
from tqdm import tqdm
from diffusion.resample import create_named_schedule_sampler
# from eval import eval_humanml, eval_humanact12_uestc

from datetime import datetime

import sys
[sys.path.append(i) for i in ['./mydiffusion', '../../My', '../../My/process', '../process', '../../ubisoft-laforge-ZeroEGGS-main', '../../ubisoft-laforge-ZeroEGGS-main/ZEGGS']]
from generate.generate import WavEncoder
# from process.process_bvh import make_bvh_GENEA2020_BT
from process_zeggs_bvh import pose2bvh

# For ImageNet experiments, this was a good default value.
# We found that the lg_loss_scale quickly climbed to
# 20-21 within the first ~1K steps of training.
INITIAL_LOG_LOSS_SCALE = 20.0


class TrainLoop:
    def __init__(self, args, model, diffusion, device, data=None):
        self.args = args
        self.data = data
        self.model = model
        self.diffusion = diffusion
        self.cond_mode = model.cond_mode
        self.batch_size = args.batch_size
        self.microbatch = args.batch_size  # deprecating this option
        self.lr = args.lr
        self.log_interval = args.log_interval
        # self.save_interval = args.save_interval
        # self.resume_checkpoint = args.resume_checkpoint
        self.use_fp16 = False  # deprecating this option
        self.fp16_scale_growth = 1e-3  # deprecating this option
        self.weight_decay = args.weight_decay
        self.lr_anneal_steps = args.lr_anneal_steps

        self.step = 0
        self.resume_step = 0
        self.global_batch = self.batch_size # * dist.get_world_size()
        # self.num_steps = args.num_steps
        self.num_epochs = 40000
        self.n_seed = 8

        self.sync_cuda = torch.cuda.is_available()

        # self._load_and_sync_parameters()
        self.mp_trainer = MixedPrecisionTrainer(
            model=self.model,
            use_fp16=self.use_fp16,
            fp16_scale_growth=self.fp16_scale_growth,
        )

        self.save_dir = args.save_dir
        self.overwrite = args.overwrite

        self.device = device
        if args.audio_feat == "wav encoder":
            self.WavEncoder = WavEncoder().to(self.device)
            self.opt = AdamW([
                {'params': self.mp_trainer.master_params, 'lr':self.lr, 'weight_decay':self.weight_decay},
                {'params': self.WavEncoder.parameters(), 'lr':self.lr}
            ])
        elif args.audio_feat == "mfcc" or args.audio_feat == 'wavlm':
            self.opt = AdamW([
                {'params': self.mp_trainer.master_params, 'lr':self.lr, 'weight_decay':self.weight_decay}
            ])

        # if self.resume_step:
        #     self._load_optimizer_state()
            # Model was resumed, either due to a restart or a checkpoint
            # being specified at the command line.

        self.schedule_sampler_type = 'uniform'
        self.schedule_sampler = create_named_schedule_sampler(self.schedule_sampler_type, diffusion)
        self.eval_wrapper, self.eval_data, self.eval_gt_data = None, None, None
        # if args.dataset in ['kit', 'humanml'] and args.eval_during_training:
        #     mm_num_samples = 0  # mm is super slow hence we won't run it during training
        #     mm_num_repeats = 0  # mm is super slow hence we won't run it during training
        #     gen_loader = get_dataset_loader(name=args.dataset, batch_size=args.eval_batch_size, num_frames=None,
        #                                     split=args.eval_split,
        #                                     hml_mode='eval')
        #
        #     self.eval_gt_data = get_dataset_loader(name=args.dataset, batch_size=args.eval_batch_size, num_frames=None,
        #                                            split=args.eval_split,
        #                                            hml_mode='gt')
        #     self.eval_wrapper = EvaluatorMDMWrapper(args.dataset, self.device)
        #     self.eval_data = {
        #         'test': lambda: eval_humanml.get_mdm_loader(
        #             model, diffusion, args.eval_batch_size,
        #             gen_loader, mm_num_samples, mm_num_repeats, gen_loader.dataset.opt.max_motion_length,
        #             args.eval_num_samples, scale=1.,
        #         )
        #     }
        self.use_ddp = False
        self.ddp_model = self.model
        self.mask_train = (torch.zeros([self.batch_size, 1, 1, args.n_poses]) < 1).to(self.device)
        self.mask_test = (torch.zeros([1, 1, 1, args.n_poses]) < 1).to(self.device)
        self.tmp_audio = torch.from_numpy(np.load('tmp_audio.npy')).unsqueeze(0).to(self.device)
        self.tmp_mfcc = torch.from_numpy(np.load('10_kieks_0_9_16.npz')['mfcc'][:args.n_poses]).to(torch.float32).unsqueeze(0).to(self.device)
        self.mask_local_train = torch.ones(self.batch_size, args.n_poses).bool().to(self.device)
        self.mask_local_test = torch.ones(1, args.n_poses).bool().to(self.device)

    # def _load_and_sync_parameters(self):
    #     resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
    #
    #     if resume_checkpoint:
    #         self.resume_step = parse_resume_step_from_filename(resume_checkpoint)
    #         logger.log(f"loading model from checkpoint: {resume_checkpoint}...")
    #         self.model.load_state_dict(
    #             dist_util.load_state_dict(
    #                 resume_checkpoint, map_location=self.device
    #             )
    #         )

    # def _load_optimizer_state(self):
    #     main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
    #     opt_checkpoint = bf.join(
    #         bf.dirname(main_checkpoint), f"opt{self.resume_step:09}.pt"
    #     )
    #     if bf.exists(opt_checkpoint):
    #         logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}")
    #         state_dict = dist_util.load_state_dict(
    #             opt_checkpoint, map_location=self.device
    #         )
    #         self.opt.load_state_dict(state_dict)

    def run_loop(self):

        for epoch in range(self.num_epochs):
            # print(f'Starting epoch {epoch}')
            # for _ in tqdm(range(10)):     # 4 steps, batch size, chmod 777
            for batch in tqdm(self.data):
                if not (not self.lr_anneal_steps or self.step + self.resume_step < self.lr_anneal_steps):
                    break

                cond_ = {'y':{}}

                # cond_['y']['text'] = ['A person turns left with medium speed.', 'A human goes slowly about 1.5 meters forward.']

                # motion = torch.rand(2, 135, 1, 80).to(self.device)
                pose_seq, _, style, audio, mfcc, wavlm = batch  # (batch, 240, 135), (batch, 30), (batch, 64000)
                motion = pose_seq.permute(0, 2, 1).unsqueeze(2).to(self.device)

                cond_['y']['seed'] = motion[..., 0:self.n_seed]
                # motion = motion[..., self.n_seed:]
                cond_['y']['style'] = style.to(self.device)
                cond_['y']['mask_local'] = self.mask_local_train

                if self.args.audio_feat == 'wav encoder':
                    # cond_['y']['audio'] = torch.rand(240, 2, 32).to(self.device)
                    cond_['y']['audio'] = self.WavEncoder(audio.to(self.device)).permute(1, 0, 2)       # (batch, 240, 32)
                elif self.args.audio_feat == 'mfcc':
                    # cond_['y']['audio'] = torch.rand(80, 2, 13).to(self.device)
                    cond_['y']['audio'] = mfcc.to(torch.float32).to(self.device).permute(1, 0, 2)       # [self.n_seed:, ...]      # (batch, 80, 13)
                elif self.args.audio_feat == 'wavlm':
                    cond_['y']['audio'] = wavlm.to(torch.float32).to(self.device)

                cond_['y']['mask'] = self.mask_train        # [..., self.n_seed:]

                self.run_step(motion, cond_)
                if self.step % self.log_interval == 0:
                    for k,v in logger.get_current().name2val.items():
                        if k == 'loss':
                            print('step[{}]: loss[{:0.5f}]'.format(self.step+self.resume_step, v))

                # if self.step % 10000 == 0:
                #     sample_fn = self.diffusion.p_sample_loop
                #
                #     model_kwargs_ = {'y': {}}
                #     model_kwargs_['y']['mask'] = self.mask_test     # [..., self.n_seed:]
                #     model_kwargs_['y']['seed'] = torch.zeros([1, 1141, 1, self.n_seed]).to(self.device)
                #     model_kwargs_['y']['style'] = torch.zeros([1, 6]).to(self.device)
                #     model_kwargs_['y']['mask_local'] = self.mask_local_test
                #     if self.args.audio_feat == 'wav encoder':
                #         model_kwargs_['y']['audio'] = self.WavEncoder(self.tmp_audio).permute(1, 0, 2)
                #         # model_kwargs_['y']['audio'] = torch.rand(240, 1, 32).to(self.device)
                #     elif self.args.audio_feat == 'mfcc':
                #         model_kwargs_['y']['audio'] = self.tmp_mfcc.permute(1, 0, 2)        # [self.n_seed:, ...]
                #         # model_kwargs_['y']['audio'] = torch.rand(80, 1, 13).to(self.device)
                #     elif self.args.audio_feat == 'wavlm':
                #         model_kwargs_['y']['audio'] = torch.randn(1, 1, 1024).to(self.device)
                #
                #     sample = sample_fn(
                #         self.model,
                #         (1, 1141, 1, self.args.n_poses),     #  - self.n_seed
                #         clip_denoised=False,
                #         model_kwargs=model_kwargs_,
                #         skip_timesteps=0,  # 0 is the default value - i.e. don't skip any step
                #         init_image=None,
                #         progress=True,
                #         dump_steps=None,
                #         noise=None,
                #         const_noise=False,
                #     )       # (1, 135, 1, 240)
                #
                #     sampled_seq = sample.squeeze(0).permute(1, 2, 0)
                #     data_mean_ = np.load("../../ubisoft-laforge-ZeroEGGS-main/Data/processed_v1/processed/mean.npz")['mean']
                #     data_std_ = np.load("../../ubisoft-laforge-ZeroEGGS-main/Data/processed_v1/processed/std.npz")['std']
                #
                #     data_mean = np.array(data_mean_).squeeze()
                #     data_std = np.array(data_std_).squeeze()
                #     std = np.clip(data_std, a_min=0.01, a_max=None)
                #     out_poses = np.multiply(np.array(sampled_seq[0].detach().cpu()), std) + data_mean
                #
                #     pipeline_path = '../../../My/process/resource/data_pipe_20_rotation.sav'
                #     save_path = 'inference_zeggs_mymodel3_wavlm'
                #     prefix = str(datetime.now().strftime('%Y%m%d_%H%M%S'))
                #     if not os.path.exists(save_path):
                #         os.mkdir(save_path)
                #     # make_bvh_GENEA2020_BT(save_path, prefix, out_poses, smoothing=False, pipeline_path=pipeline_path)
                #
                #     pose2bvh(out_poses, os.path.join(save_path, prefix + '.bvh'), length=self.args.n_poses)

                if self.step % 50000 == 0:
                    self.save()
                    # self.model.eval()
                    # self.evaluate()
                    # self.model.train()

                    # Run for a finite amount of time in integration tests.
                    if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0:
                        return
                self.step += 1
            if not (not self.lr_anneal_steps or self.step + self.resume_step < self.lr_anneal_steps):
                break
        # Save the last checkpoint if it wasn't already saved.
        # if (self.step - 1) % 50000 != 0:
        #     self.save()
            # self.evaluate()


    def run_step(self, batch, cond):
        self.forward_backward(batch, cond)      # torch.Size([64, 251, 1, 196]) cond['y'].keys() dict_keys(['mask', 'lengths', 'text', 'tokens'])
        self.mp_trainer.optimize(self.opt)
        self._anneal_lr()
        self.log_step()

    def forward_backward(self, batch, cond):
        self.mp_trainer.zero_grad()
        for i in range(0, batch.shape[0], self.microbatch):
            # Eliminates the microbatch feature
            assert i == 0
            assert self.microbatch == self.batch_size
            micro = batch
            micro_cond = cond
            last_batch = (i + self.microbatch) >= batch.shape[0]
            t, weights = self.schedule_sampler.sample(micro.shape[0], self.device)

            compute_losses = functools.partial(
                self.diffusion.training_losses,
                self.ddp_model,
                micro,  # [bs, ch, image_size, image_size]      # x_start, (2, 135, 1, 240)
                t,  # [bs](int) sampled timesteps
                model_kwargs=micro_cond,
                dataset='kit'
            )

            if last_batch or not self.use_ddp:
                losses = compute_losses()
            else:
                with self.ddp_model.no_sync():
                    losses = compute_losses()

            if isinstance(self.schedule_sampler, LossAwareSampler):
                self.schedule_sampler.update_with_local_losses(
                    t, losses["loss"].detach()
                )

            loss = (losses["loss"] * weights).mean()
            log_loss_dict(
                self.diffusion, t, {k: v * weights for k, v in losses.items()}
            )
            self.mp_trainer.backward(loss)

    def _anneal_lr(self):
        if not self.lr_anneal_steps:
            return
        frac_done = (self.step + self.resume_step) / self.lr_anneal_steps
        lr = self.lr * (1 - frac_done)
        for param_group in self.opt.param_groups:
            param_group["lr"] = lr

    def log_step(self):
        logger.logkv("step", self.step + self.resume_step)
        logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch)


    def ckpt_file_name(self):
        return f"model{(self.step+self.resume_step):09d}.pt"


    def save(self):
        def save_checkpoint(params):
            state_dict = self.mp_trainer.master_params_to_state_dict(params)

            # Do not save CLIP weights
            clip_weights = [e for e in state_dict.keys() if e.startswith('clip_model.')]
            for e in clip_weights:
                del state_dict[e]

            logger.log(f"saving model...")
            filename = self.ckpt_file_name()
            with bf.BlobFile(bf.join(self.save_dir, filename), "wb") as f:
                torch.save(state_dict, f)

        save_checkpoint(self.mp_trainer.master_params)

        with bf.BlobFile(
            bf.join(self.save_dir, f"opt{(self.step+self.resume_step):09d}.pt"),
            "wb",
        ) as f:
            torch.save(self.opt.state_dict(), f)


def parse_resume_step_from_filename(filename):
    """
    Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the
    checkpoint's number of steps.
    """
    split = filename.split("model")
    if len(split) < 2:
        return 0
    split1 = split[-1].split(".")[0]
    try:
        return int(split1)
    except ValueError:
        return 0


def get_blob_logdir():
    # You can change this to be a separate path to save checkpoints to
    # a blobstore or some external drive.
    return logger.get_dir()


def find_resume_checkpoint():
    # On your infrastructure, you may want to override this to automatically
    # discover the latest checkpoint on your blob storage, etc.
    return None


def log_loss_dict(diffusion, ts, losses):
    for key, values in losses.items():
        logger.logkv_mean(key, values.mean().item())
        # Log the quantiles (four quartiles, in particular).
        for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()):
            quartile = int(4 * sub_t / diffusion.num_timesteps)
            logger.logkv_mean(f"{key}_q{quartile}", sub_loss)

train/training_loop.py
YoungSeng/UnifiedGesture
import re
import numpy as np
import sys
sys.path.append("motion_utils")

from Animation import Animation
from Quaternions_old import Quaternions

channelmap = {
    'Xrotation' : 'x',
    'Yrotation' : 'y',
    'Zrotation' : 'z'   
}

channelmap_inv = {
    'x': 'Xrotation',
    'y': 'Yrotation',
    'z': 'Zrotation',
}

ordermap = {
    'x' : 0,
    'y' : 1,
    'z' : 2,
}

def load(filename, start=None, end=None, order=None, world=False):
    """
    Reads a BVH file and constructs an animation
    
    Parameters
    ----------
    filename: str
        File to be opened
        
    start : int
        Optional Starting Frame
        
    end : int
        Optional Ending Frame
    
    order : str
        Optional Specifier for joint order.
        Given as string E.G 'xyz', 'zxy'
        
    world : bool
        If set to true euler angles are applied
        together in world space rather than local
        space

    Returns
    -------
    
    (animation, joint_names, frametime)
        Tuple of loaded animation and joint names
    """
    
    f = open(filename, "r")

    i = 0
    active = -1
    end_site = False
    
    names = []
    orients = Quaternions.id(0)
    offsets = np.array([]).reshape((0,3))
    parents = np.array([], dtype=int)
    
    for line in f:
        
        if "HIERARCHY" in line: continue
        if "MOTION" in line: continue

        rmatch = re.match(r"ROOT (\w+)", line)
        if rmatch:
            names.append(rmatch.group(1))
            offsets    = np.append(offsets,    np.array([[0,0,0]]),   axis=0)
            orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0)
            parents    = np.append(parents, active)
            active = (len(parents)-1)
            continue

        if "{" in line: continue

        if "}" in line:
            if end_site: end_site = False
            else: active = parents[active]
            continue
        
        offmatch = re.match(r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line)
        if offmatch:
            if not end_site:
                offsets[active] = np.array([list(map(float, offmatch.groups()))])
            continue
           
        chanmatch = re.match(r"\s*CHANNELS\s+(\d+)", line)
        if chanmatch:
            channels = int(chanmatch.group(1))
            if order is None:
                channelis = 0 if channels == 3 else 3
                channelie = 3 if channels == 3 else 6
                parts = line.split()[2+channelis:2+channelie]
                if any([p not in channelmap for p in parts]):
                    continue
                order = "".join([channelmap[p] for p in parts])
            continue

        jmatch = re.match("\s*JOINT\s+(\w+)", line)
        if jmatch:
            names.append(jmatch.group(1))
            offsets    = np.append(offsets,    np.array([[0,0,0]]),   axis=0)
            orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0)
            parents    = np.append(parents, active)
            active = (len(parents)-1)
            continue
        
        if ("End Site" in line) or ("End site" in line):
            end_site = True
            continue
              
        fmatch = re.match("\s*Frames:\s+(\d+)", line)
        if fmatch:
            if start and end:
                fnum = (end - start)-1
            else:
                fnum = int(fmatch.group(1))
            jnum = len(parents)
            # result: [fnum, J, 3]
            positions = offsets[np.newaxis].repeat(fnum, axis=0)
            # result: [fnum, len(orients), 3]
            rotations = np.zeros((fnum, len(orients), 3))
            continue
        
        fmatch = re.match("\s*Frame Time:\s+([\d\.]+)", line)
        if fmatch:
            frametime = float(fmatch.group(1))
            continue
        
        if (start and end) and (i < start or i >= end-1):
            i += 1
            continue
        
        dmatch = line.strip().split()
        if dmatch:
            data_block = np.array(list(map(float, dmatch)))
            N = len(parents)
            fi = i - start if start else i
            if   channels == 3:
                # This should be root positions[0:1] & all rotations
                positions[fi,0:1] = data_block[0:3]
                rotations[fi, : ] = data_block[3: ].reshape(N,3)
            elif channels == 6:
                data_block = data_block.reshape(N,6)
                # fill in all positions
                positions[fi,:] = data_block[:,0:3]
                rotations[fi,:] = data_block[:,3:6]
            elif channels == 9:
                positions[fi,0] = data_block[0:3]
                data_block = data_block[3:].reshape(N-1,9)
                rotations[fi,1:] = data_block[:,3:6]
                positions[fi,1:] += data_block[:,0:3] * data_block[:,6:9]
            else:
                raise Exception("Too many channels! %i" % channels)

            i += 1

    f.close()
    
    rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world)
    
    return (Animation(rotations, positions, orients, offsets, parents), names, frametime)
    
def load_bfa(filename, start=None, end=None, order=None, world=False):
    """
    Reads a BVH file and constructs an animation

    !!! Read from bfa, will replace the end sites of arms by two joints (w/ unit rotation)

    Parameters
    ----------
    filename: str
        File to be opened

    start : int
        Optional Starting Frame

    end : int
        Optional Ending Frame

    order : str
        Optional Specifier for joint order.
        Given as string E.G 'xyz', 'zxy'

    world : bool
        If set to true euler angles are applied
        together in world space rather than local
        space

    Returns
    -------

    (animation, joint_names, frametime)
        Tuple of loaded animation and joint names
    """

    f = open(filename, "r")

    i = 0
    active = -1
    end_site = False

    hand_idx = [9, 14]

    names = []
    orients = Quaternions.id(0)
    offsets = np.array([]).reshape((0,3))
    parents = np.array([], dtype=int)

    for line in f:

        if "HIERARCHY" in line: continue
        if "MOTION" in line: continue

        rmatch = re.match(r"ROOT (\w+)", line)
        if rmatch:
            names.append(rmatch.group(1))
            offsets    = np.append(offsets,    np.array([[0,0,0]]),   axis=0)
            orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0)
            parents    = np.append(parents, active)
            active = (len(parents)-1)
            continue

        if "{" in line: continue

        if "}" in line:
            if end_site: end_site = False
            else: active = parents[active]
            continue

        offmatch = re.match(r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line)
        if offmatch:
            if not end_site:
                offsets[active] = np.array([list(map(float, offmatch.groups()))])
            """
            else:
                print("active = ", active)
                if active in hand_idx:
                    offsets[active] = np.array([list(map(float, offmatch.groups()))])
            """
            continue

        chanmatch = re.match(r"\s*CHANNELS\s+(\d+)", line)
        if chanmatch:
            channels = int(chanmatch.group(1))
            if order is None:
                channelis = 0 if channels == 3 else 3
                channelie = 3 if channels == 3 else 6
                parts = line.split()[2+channelis:2+channelie]
                if any([p not in channelmap for p in parts]):
                    continue
                order = "".join([channelmap[p] for p in parts])
            continue

        jmatch = re.match("\s*JOINT\s+(\w+)", line)
        if jmatch:
            names.append(jmatch.group(1))
            offsets    = np.append(offsets,    np.array([[0,0,0]]),   axis=0)
            orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0)
            parents    = np.append(parents, active)
            active = (len(parents)-1)
            continue

        if "End Site" in line:
            if active + 1 in hand_idx:
                print("parent:", names[-1])
                name = "LeftHandIndex" if active + 1 == hand_idx[0] else "RightHandIndex"
                names.append(name)
                offsets    = np.append(offsets,    np.array([[0,0,0]]),   axis=0)
                orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0)
                parents    = np.append(parents, active)
                active = (len(parents)-1)
            else:
                end_site = True
            continue

        fmatch = re.match("\s*Frames:\s+(\d+)", line)
        if fmatch:
            if start and end:
                fnum = (end - start)-1
            else:
                fnum = int(fmatch.group(1))
            jnum = len(parents)
            # result: [fnum, J, 3]
            positions = offsets[np.newaxis].repeat(fnum, axis=0)
            # result: [fnum, len(orients), 3]
            rotations = np.zeros((fnum, len(orients), 3))
            continue

        fmatch = re.match("\s*Frame Time:\s+([\d\.]+)", line)
        if fmatch:
            frametime = float(fmatch.group(1))
            continue

        if (start and end) and (i < start or i >= end-1):
            i += 1
            continue

        dmatch = line.strip().split()
        if dmatch:
            data_block = np.array(list(map(float, dmatch)))
            N = len(parents)
            fi = i - start if start else i
            if   channels == 3:
                # This should be root positions[0:1] & all rotations
                positions[fi,0:1] = data_block[0:3]
                tmp = data_block[3: ].reshape(N - 2, 3)
                tmp = np.concatenate([tmp[:hand_idx[0]],
                                      np.array([[0, 0, 0]]),
                                      tmp[hand_idx[0]: hand_idx[1] - 1],
                                      np.array([[0, 0, 0]]),
                                      tmp[hand_idx[1] - 1:]], axis=0)
                rotations[fi, : ] = tmp.reshape(N,3)
            elif channels == 6:
                data_block = data_block.reshape(N,6)
                # fill in all positions
                positions[fi,:] = data_block[:,0:3]
                rotations[fi,:] = data_block[:,3:6]
            elif channels == 9:
                positions[fi,0] = data_block[0:3]
                data_block = data_block[3:].reshape(N-1,9)
                rotations[fi,1:] = data_block[:,3:6]
                positions[fi,1:] += data_block[:,0:3] * data_block[:,6:9]
            else:
                raise Exception("Too many channels! %i" % channels)

            i += 1

    f.close()

    rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world)

    return (Animation(rotations, positions, orients, offsets, parents), names, frametime)

    
def save(filename, anim, names=None, frametime=1.0/24.0, order='zyx', positions=False, orients=True):
    """
    Saves an Animation to file as BVH
    
    Parameters
    ----------
    filename: str
        File to be saved to
        
    anim : Animation
        Animation to save
        
    names : [str]
        List of joint names
    
    order : str
        Optional Specifier for joint order.
        Given as string E.G 'xyz', 'zxy'
    
    frametime : float
        Optional Animation Frame time
        
    positions : bool
        Optional specfier to save bone
        positions for each frame
        
    orients : bool
        Multiply joint orients to the rotations
        before saving.
        
    """
    
    if names is None:
        names = ["joint_" + str(i) for i in range(len(anim.parents))]
    
    with open(filename, 'w') as f:

        t = ""
        f.write("%sHIERARCHY\n" % t)
        f.write("%sROOT %s\n" % (t, names[0]))
        f.write("%s{\n" % t)
        t += '\t'

        f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[0,0], anim.offsets[0,1], anim.offsets[0,2]) )
        f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" % 
            (t, channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))

        for i in range(anim.shape[1]):
            if anim.parents[i] == 0:
                t = save_joint(f, anim, names, t, i, order=order, positions=positions)

        t = t[:-1]
        f.write("%s}\n" % t)

        f.write("MOTION\n")
        f.write("Frames: %i\n" % anim.shape[0]);
        f.write("Frame Time: %f\n" % frametime);
            
        #if orients:        
        #    rots = np.degrees((-anim.orients[np.newaxis] * anim.rotations).euler(order=order[::-1]))
        #else:
        #    rots = np.degrees(anim.rotations.euler(order=order[::-1]))
        rots = np.degrees(anim.rotations.euler(order=order[::-1]))
        poss = anim.positions
        
        for i in range(anim.shape[0]):
            for j in range(anim.shape[1]):
                
                if positions or j == 0:
                
                    f.write("%f %f %f %f %f %f " % (
                        poss[i,j,0],                  poss[i,j,1],                  poss[i,j,2], 
                        rots[i,j,ordermap[order[0]]], rots[i,j,ordermap[order[1]]], rots[i,j,ordermap[order[2]]]))
                
                else:
                    
                    f.write("%f %f %f " % (
                        rots[i,j,ordermap[order[0]]], rots[i,j,ordermap[order[1]]], rots[i,j,ordermap[order[2]]]))

            f.write("\n")
    
    
def save_joint(f, anim, names, t, i, order='zyx', positions=False):
    
    f.write("%sJOINT %s\n" % (t, names[i]))
    f.write("%s{\n" % t)
    t += '\t'
  
    f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[i,0], anim.offsets[i,1], anim.offsets[i,2]))
    
    if positions:
        f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" % (t, 
            channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))
    else:
        f.write("%sCHANNELS 3 %s %s %s\n" % (t, 
            channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))
    
    end_site = True
    
    for j in range(anim.shape[1]):
        if anim.parents[j] == i:
            t = save_joint(f, anim, names, t, j, order=order, positions=positions)
            end_site = False
    
    if end_site:
        f.write("%sEnd Site\n" % t)
        f.write("%s{\n" % t)
        t += '\t'
        f.write("%sOFFSET %f %f %f\n" % (t, 0.0, 0.0, 0.0))
        t = t[:-1]
        f.write("%s}\n" % t)
  
    t = t[:-1]
    f.write("%s}\n" % t)
    
    return t
utils/BVH.py
YoungSeng/UnifiedGesture
import re
import numpy as np

from Animation import Animation
from Quaternions import Quaternions

channelmap = {
    'Xrotation' : 'x',
    'Yrotation' : 'y',
    'Zrotation' : 'z'   
}

channelmap_inv = {
    'x': 'Xrotation',
    'y': 'Yrotation',
    'z': 'Zrotation',
}

ordermap = {
    'x' : 0,
    'y' : 1,
    'z' : 2,
}

def load(filename, start=None, end=None, order=None, world=False, need_quater=False):
    """
    Reads a BVH file and constructs an animation
    
    Parameters
    ----------
    filename: str
        File to be opened
        
    start : int
        Optional Starting Frame
        
    end : int
        Optional Ending Frame
    
    order : str
        Optional Specifier for joint order.
        Given as string E.G 'xyz', 'zxy'
        
    world : bool
        If set to true euler angles are applied
        together in world space rather than local
        space

    Returns
    -------
    
    (animation, joint_names, frametime)
        Tuple of loaded animation and joint names
    """
    
    f = open(filename, "r")

    i = 0
    active = -1
    end_site = False
    
    names = []
    orients = Quaternions.id(0)
    offsets = np.array([]).reshape((0,3))
    parents = np.array([], dtype=int)
    
    for line in f:
        
        if "HIERARCHY" in line: continue
        if "MOTION" in line: continue

        """ Modified line read to handle mixamo data """
#        rmatch = re.match(r"ROOT (\w+)", line)
        rmatch = re.match(r"ROOT (\w+:?\w+)", line)
        if rmatch:
            names.append(rmatch.group(1))
            offsets    = np.append(offsets,    np.array([[0,0,0]]),   axis=0)
            orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0)
            parents    = np.append(parents, active)
            active = (len(parents)-1)
            continue

        if "{" in line: continue

        if "}" in line:
            if end_site: end_site = False
            else: active = parents[active]
            continue
        
        offmatch = re.match(r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line)
        if offmatch:
            if not end_site:
                offsets[active] = np.array([list(map(float, offmatch.groups()))])
            continue
           
        chanmatch = re.match(r"\s*CHANNELS\s+(\d+)", line)
        if chanmatch:
            channels = int(chanmatch.group(1))
            if order is None:
                channelis = 0 if channels == 3 else 3
                channelie = 3 if channels == 3 else 6
                parts = line.split()[2+channelis:2+channelie]
                if any([p not in channelmap for p in parts]):
                    continue
                order = "".join([channelmap[p] for p in parts])
            continue

        """ Modified line read to handle mixamo data """

#        jmatch = re.match("\s*JOINT\s+(\w+)", line)
        jmatch = re.match("\s*JOINT\s+(\w+:?\w+)", line)
        if jmatch:
            names.append(jmatch.group(1))
            offsets    = np.append(offsets,    np.array([[0,0,0]]),   axis=0)
            orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0)
            parents    = np.append(parents, active)
            active = (len(parents)-1)
            continue

        if "End Site" in line or "End site" in line:
            end_site = True
            continue

        fmatch = re.match("\s*Frames:\s+(\d+)", line)
        if fmatch:
            if start and end:
                fnum = (end - start)-1
            else:
                fnum = int(fmatch.group(1))
            jnum = len(parents)
            positions = offsets[np.newaxis].repeat(fnum, axis=0)
            rotations = np.zeros((fnum, len(orients), 3))
            continue
        
        fmatch = re.match("\s*Frame Time:\s+([\d\.]+)", line)
        if fmatch:
            frametime = float(fmatch.group(1))
            continue
        
        if (start and end) and (i < start or i >= end-1):
            i += 1
            continue
        
        # dmatch = line.strip().split(' ')
        dmatch = line.strip().split()
        if dmatch:
            data_block = np.array(list(map(float, dmatch)))
            N = len(parents)
            fi = i - start if start else i
            if   channels == 3:
                positions[fi,0:1] = data_block[0:3]
                rotations[fi, : ] = data_block[3: ].reshape(N,3)
            elif channels == 6:
                data_block = data_block.reshape(N,6)
                positions[fi,:] = data_block[:,0:3]
                rotations[fi,:] = data_block[:,3:6]
            elif channels == 9:
                positions[fi,0] = data_block[0:3]
                data_block = data_block[3:].reshape(N-1,9)
                rotations[fi,1:] = data_block[:,3:6]
                positions[fi,1:] += data_block[:,0:3] * data_block[:,6:9]
            else:
                raise Exception("Too many channels! %i" % channels)

            i += 1

    f.close()

    if need_quater:
        rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world)
    elif order != 'xyz':
        rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world)
        rotations = np.degrees(rotations.euler())
    
    return (Animation(rotations, positions, orients, offsets, parents), names, frametime)
    

    
def save(filename, anim, names=None, frametime=1.0/24.0, order='zyx', positions=False, orients=True, mask=None, quater=False):
    """
    Saves an Animation to file as BVH
    
    Parameters
    ----------
    filename: str
        File to be saved to
        
    anim : Animation
        Animation to save
        
    names : [str]
        List of joint names
    
    order : str
        Optional Specifier for joint order.
        Given as string E.G 'xyz', 'zxy'
    
    frametime : float
        Optional Animation Frame time
        
    positions : bool
        Optional specfier to save bone
        positions for each frame
        
    orients : bool
        Multiply joint orients to the rotations
        before saving.
        
    """
    
    if names is None:
        names = ["joint_" + str(i) for i in range(len(anim.parents))]
    
    with open(filename, 'w') as f:

        t = ""
        f.write("%sHIERARCHY\n" % t)
        f.write("%sROOT %s\n" % (t, names[0]))
        f.write("%s{\n" % t)
        t += '\t'

        f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[0,0], anim.offsets[0,1], anim.offsets[0,2]) )
        f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" % 
            (t, channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))

        for i in range(anim.shape[1]):
            if anim.parents[i] == 0:
                t = save_joint(f, anim, names, t, i, order=order, positions=positions)

        t = t[:-1]
        f.write("%s}\n" % t)

        f.write("MOTION\n")
        f.write("Frames: %i\n" % anim.shape[0]);
        f.write("Frame Time: %f\n" % frametime);
            
        #if orients:        
        #    rots = np.degrees((-anim.orients[np.newaxis] * anim.rotations).euler(order=order[::-1]))
        #else:
        #    rots = np.degrees(anim.rotations.euler(order=order[::-1]))
        # rots = np.degrees(anim.rotations.euler(order=order[::-1]))
        if quater:
            rots = np.degrees(anim.rotations.euler(order=order[::-1]))
        else:
            rots = anim.rotations
        poss = anim.positions
        
        for i in range(anim.shape[0]):
            for j in range(anim.shape[1]):
                
                if positions or j == 0:
                
                    f.write("%f %f %f %f %f %f " % (
                        poss[i,j,0],                  poss[i,j,1],                  poss[i,j,2], 
                        rots[i,j,ordermap[order[0]]], rots[i,j,ordermap[order[1]]], rots[i,j,ordermap[order[2]]]))
                
                else:
                    if mask == None or mask[j] == 1:
                        f.write("%f %f %f " % (
                            rots[i,j,ordermap[order[0]]], rots[i,j,ordermap[order[1]]], rots[i,j,ordermap[order[2]]]))
                    else:
                        f.write("%f %f %f " % (0, 0, 0))

            f.write("\n")
    
    
def save_joint(f, anim, names, t, i, order='zyx', positions=False):
    
    f.write("%sJOINT %s\n" % (t, names[i]))
    f.write("%s{\n" % t)
    t += '\t'
  
    f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[i,0], anim.offsets[i,1], anim.offsets[i,2]))
    
    if positions:
        f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" % (t, 
            channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))
    else:
        f.write("%sCHANNELS 3 %s %s %s\n" % (t, 
            channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))
    
    end_site = True
    
    for j in range(anim.shape[1]):
        if anim.parents[j] == i:
            t = save_joint(f, anim, names, t, j, order=order, positions=positions)
            end_site = False
    
    if end_site:
        f.write("%sEnd Site\n" % t)
        f.write("%s{\n" % t)
        t += '\t'
        f.write("%sOFFSET %f %f %f\n" % (t, 0.0, 0.0, 0.0))
        t = t[:-1]
        f.write("%s}\n" % t)
  
    t = t[:-1]
    f.write("%s}\n" % t)
    
    return t

utils/BVH_mod.py
YoungSeng/UnifiedGesture