from functools import reduce
from pathlib import Path
import torch
import torch.nn.functional as F
class NullableArgs:
def __init__(self, namespace):
for key, value in namespace.__dict__.items():
setattr(self, key, value)
def __getattr__(self, key):
# when an attribute lookup has not found the attribute
if key == 'align_mask_width':
if 'use_alignment_mask' in self.__dict__:
return 1 if self.use_alignment_mask else 0
else:
return 0
if key == 'no_head_pose':
return not self.predict_head_pose
if key == 'no_use_learnable_pe':
return not self.use_learnable_pe
return None
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
def get_option_text(args, parser):
message = ''
for k, v in sorted(vars(args).items()):
comment = ''
default = parser.get_default(k)
if v != default:
comment = f'\t[default: {str(default)}]'
message += f'{str(k):>30}: {str(v):<30}{comment}\n'
return message
def get_model_path(exp_name, iteration, model_type='DPT'):
exp_root_dir = Path(__file__).parent.parent / 'experiments' / model_type
exp_dir = exp_root_dir / exp_name
if not exp_dir.exists():
exp_dir = next(exp_root_dir.glob(f'{exp_name}*'))
model_path = exp_dir / f'checkpoints/iter_{iteration:07}.pt'
return model_path, exp_dir.relative_to(exp_root_dir)
def get_pose_input(coef_dict, rot_repr, with_global_pose):
if rot_repr == 'aa':
pose_input = coef_dict['pose'] if with_global_pose else coef_dict['pose'][..., -3:]
# Remove mouth rotation round y, z axis
pose_input = pose_input[..., :-2]
else:
raise ValueError(f'Unknown rotation representation: {rot_repr}')
return pose_input
def get_motion_coef(coef_dict, rot_repr, with_global_pose=False, norm_stats=None):
if norm_stats is not None:
if rot_repr == 'aa':
keys = ['exp', 'pose']
else:
raise ValueError(f'Unknown rotation representation {rot_repr}!')
coef_dict = {k: (coef_dict[k] - norm_stats[f'{k}_mean']) / norm_stats[f'{k}_std'] for k in keys}
pose_coef = get_pose_input(coef_dict, rot_repr, with_global_pose)
return torch.cat([coef_dict['exp'], pose_coef], dim=-1)
def get_coef_dict(motion_coef, shape_coef=None, denorm_stats=None, with_global_pose=False, rot_repr='aa'):
coef_dict = {
'exp': motion_coef[..., :50]
}
if rot_repr == 'aa':
if with_global_pose:
coef_dict['pose'] = motion_coef[..., 50:]
else:
placeholder = torch.zeros_like(motion_coef[..., :3])
coef_dict['pose'] = torch.cat([placeholder, motion_coef[..., -1:]], dim=-1)
# Add back rotation around y, z axis
coef_dict['pose'] = torch.cat([coef_dict['pose'], torch.zeros_like(motion_coef[..., :2])], dim=-1)
else:
raise ValueError(f'Unknown rotation representation {rot_repr}!')
if shape_coef is not None:
if motion_coef.ndim == 3:
if shape_coef.ndim == 2:
shape_coef = shape_coef.unsqueeze(1)
if shape_coef.shape[1] == 1:
shape_coef = shape_coef.expand(-1, motion_coef.shape[1], -1)
coef_dict['shape'] = shape_coef
if denorm_stats is not None:
coef_dict = {k: coef_dict[k] * denorm_stats[f'{k}_std'] + denorm_stats[f'{k}_mean'] for k in coef_dict}
if not with_global_pose:
if rot_repr == 'aa':
coef_dict['pose'][..., :3] = 0
else:
raise ValueError(f'Unknown rotation representation {rot_repr}!')
return coef_dict
def coef_dict_to_vertices(coef_dict, flame, rot_repr='aa', ignore_global_rot=False, flame_batch_size=512):
shape = coef_dict['exp'].shape[:-1]
coef_dict = {k: v.view(-1, v.shape[-1]) for k, v in coef_dict.items()}
n_samples = reduce(lambda x, y: x * y, shape, 1)
# Convert to vertices
vert_list = []
for i in range(0, n_samples, flame_batch_size):
batch_coef_dict = {k: v[i:i + flame_batch_size] for k, v in coef_dict.items()}
if rot_repr == 'aa':
vert, _, _ = flame(
batch_coef_dict['shape'], batch_coef_dict['exp'], batch_coef_dict['pose'],
pose2rot=True, ignore_global_rot=ignore_global_rot, return_lm2d=False, return_lm3d=False)
else:
raise ValueError(f'Unknown rot_repr: {rot_repr}')
vert_list.append(vert)
vert_list = torch.cat(vert_list, dim=0) # (n_samples, 5023, 3)
vert_list = vert_list.view(*shape, -1, 3) # (..., 5023, 3)
return vert_list
def compute_loss(args, is_starting_sample, shape_coef, motion_coef_gt, noise, target, prev_motion_coef, coef_stats,
flame, end_idx=None):
if args.criterion.lower() == 'l2':
criterion_func = F.mse_loss
elif args.criterion.lower() == 'l1':
criterion_func = F.l1_loss
else:
raise NotImplementedError(f'Criterion {args.criterion} not implemented.')
loss_vert = None
loss_vel = None
loss_smooth = None
loss_head_angle = None
loss_head_vel = None
loss_head_smooth = None
loss_head_trans_vel = None
loss_head_trans_accel = None
loss_head_trans = None
if args.target == 'noise':
loss_noise = criterion_func(noise, target[:, args.n_prev_motions:], reduction='none')
elif args.target == 'sample':
if is_starting_sample:
target = target[:, args.n_prev_motions:]
else:
motion_coef_gt = torch.cat([prev_motion_coef, motion_coef_gt], dim=1)
if args.no_constrain_prev:
target = torch.cat([prev_motion_coef, target[:, args.n_prev_motions:]], dim=1)
loss_noise = criterion_func(motion_coef_gt, target, reduction='none')
if args.l_vert > 0 or args.l_vel > 0:
coef_gt = get_coef_dict(motion_coef_gt, shape_coef, coef_stats, with_global_pose=False,
rot_repr=args.rot_repr)
coef_pred = get_coef_dict(target, shape_coef, coef_stats, with_global_pose=False,
rot_repr=args.rot_repr)
seq_len = target.shape[1]
if args.rot_repr == 'aa':
verts_gt, _, _ = flame(coef_gt['shape'].view(-1, 100), coef_gt['exp'].view(-1, 50),
coef_gt['pose'].view(-1, 6), return_lm2d=False, return_lm3d=False)
verts_pred, _, _ = flame(coef_pred['shape'].view(-1, 100), coef_pred['exp'].view(-1, 50),
coef_pred['pose'].view(-1, 6), return_lm2d=False, return_lm3d=False)
else:
raise ValueError(f'Unknown rotation representation {args.rot_repr}!')
verts_gt = verts_gt.view(-1, seq_len, 5023, 3)
verts_pred = verts_pred.view(-1, seq_len, 5023, 3)
if args.l_vert > 0:
loss_vert = criterion_func(verts_gt, verts_pred, reduction='none')
if args.l_vel > 0:
vel_gt = verts_gt[:, 1:] - verts_gt[:, :-1]
vel_pred = verts_pred[:, 1:] - verts_pred[:, :-1]
loss_vel = criterion_func(vel_gt, vel_pred, reduction='none')
if args.l_smooth > 0:
vel_pred = verts_pred[:, 1:] - verts_pred[:, :-1]
loss_smooth = criterion_func(vel_pred[:, 1:], vel_pred[:, :-1], reduction='none')
# head pose
if not args.no_head_pose:
if args.rot_repr == 'aa':
head_pose_gt = motion_coef_gt[:, :, 50:53]
head_pose_pred = target[:, :, 50:53]
else:
raise ValueError(f'Unknown rotation representation {args.rot_repr}!')
if args.l_head_angle > 0:
loss_head_angle = criterion_func(head_pose_gt, head_pose_pred, reduction='none')
if args.l_head_vel > 0:
head_vel_gt = head_pose_gt[:, 1:] - head_pose_gt[:, :-1]
head_vel_pred = head_pose_pred[:, 1:] - head_pose_pred[:, :-1]
loss_head_vel = criterion_func(head_vel_gt, head_vel_pred, reduction='none')
if args.l_head_smooth > 0:
head_vel_pred = head_pose_pred[:, 1:] - head_pose_pred[:, :-1]
loss_head_smooth = criterion_func(head_vel_pred[:, 1:], head_vel_pred[:, :-1], reduction='none')
if not is_starting_sample and args.l_head_trans > 0:
# # version 1: constrain both the predicted previous and current motions (x_{-3} ~ x_{2})
# head_pose_trans = head_pose_pred[:, args.n_prev_motions - 3:args.n_prev_motions + 3]
# head_vel_pred = head_pose_trans[:, 1:] - head_pose_trans[:, :-1]
# head_accel_pred = head_vel_pred[:, 1:] - head_vel_pred[:, :-1]
# version 2: constrain only the predicted current motions (x_{0} ~ x_{2})
head_pose_trans = torch.cat([head_pose_gt[:, args.n_prev_motions - 3:args.n_prev_motions],
head_pose_pred[:, args.n_prev_motions:args.n_prev_motions + 3]], dim=1)
head_vel_pred = head_pose_trans[:, 1:] - head_pose_trans[:, :-1]
head_accel_pred = head_vel_pred[:, 1:] - head_vel_pred[:, :-1]
# will constrain x_{-2|0} ~ x_{1}
loss_head_trans_vel = criterion_func(head_vel_pred[:, 2:4], head_vel_pred[:, 1:3], reduction='none')
# will constrain x_{-3|0} ~ x_{2}
loss_head_trans_accel = criterion_func(head_accel_pred[:, 1:], head_accel_pred[:, :-1],
reduction='none')
else:
raise ValueError(f'Unknown diffusion target: {args.target}')
if end_idx is None:
mask = torch.ones((target.shape[0], args.n_motions), dtype=torch.bool, device=target.device)
else:
mask = torch.arange(args.n_motions, device=target.device).expand(target.shape[0], -1) < end_idx.unsqueeze(1)
if args.target == 'sample' and not is_starting_sample:
if args.no_constrain_prev:
# Warning: this option will be deprecated in the future
mask = torch.cat([torch.zeros_like(mask[:, :args.n_prev_motions]), mask], dim=1)
else:
mask = torch.cat([torch.ones_like(mask[:, :args.n_prev_motions]), mask], dim=1)
loss_noise = loss_noise[mask].mean()
if loss_vert is not None:
loss_vert = loss_vert[mask].mean()
if loss_vel is not None:
loss_vel = loss_vel[mask[:, 1:]]
loss_vel = loss_vel.mean() if torch.numel(loss_vel) > 0 else None
if loss_smooth is not None:
loss_smooth = loss_smooth[mask[:, 2:]]
loss_smooth = loss_smooth.mean() if torch.numel(loss_smooth) > 0 else None
if loss_head_angle is not None:
loss_head_angle = loss_head_angle[mask].mean()
if loss_head_vel is not None:
loss_head_vel = loss_head_vel[mask[:, 1:]]
loss_head_vel = loss_head_vel.mean() if torch.numel(loss_head_vel) > 0 else None
if loss_head_smooth is not None:
loss_head_smooth = loss_head_smooth[mask[:, 2:]]
loss_head_smooth = loss_head_smooth.mean() if torch.numel(loss_head_smooth) > 0 else None
if loss_head_trans_vel is not None:
vel_mask = mask[:, args.n_prev_motions:args.n_prev_motions + 2]
accel_mask = mask[:, args.n_prev_motions:args.n_prev_motions + 3]
loss_head_trans_vel = loss_head_trans_vel[vel_mask].mean()
loss_head_trans_accel = loss_head_trans_accel[accel_mask].mean()
loss_head_trans = loss_head_trans_vel + loss_head_trans_accel
return loss_noise, loss_vert, loss_vel, loss_smooth, loss_head_angle, loss_head_vel, loss_head_smooth, \
loss_head_trans
def _truncate_audio(audio, end_idx, pad_mode='zero'):
batch_size = audio.shape[0]
audio_trunc = audio.clone()
if pad_mode == 'replicate':
for i in range(batch_size):
audio_trunc[i, end_idx[i]:] = audio_trunc[i, end_idx[i] - 1]
elif pad_mode == 'zero':
for i in range(batch_size):
audio_trunc[i, end_idx[i]:] = 0
else:
raise ValueError(f'Unknown pad mode {pad_mode}!')
return audio_trunc
def _truncate_coef_dict(coef_dict, end_idx, pad_mode='zero'):
batch_size = coef_dict['exp'].shape[0]
coef_dict_trunc = {k: v.clone() for k, v in coef_dict.items()}
if pad_mode == 'replicate':
for i in range(batch_size):
for k in coef_dict_trunc:
coef_dict_trunc[k][i, end_idx[i]:] = coef_dict_trunc[k][i, end_idx[i] - 1]
elif pad_mode == 'zero':
for i in range(batch_size):
for k in coef_dict:
coef_dict_trunc[k][i, end_idx[i]:] = 0
else:
raise ValueError(f'Unknown pad mode: {pad_mode}!')
return coef_dict_trunc
def truncate_coef_dict_and_audio(audio, coef_dict, n_motions, audio_unit=640, pad_mode='zero'):
batch_size = audio.shape[0]
end_idx = torch.randint(1, n_motions, (batch_size,), device=audio.device)
audio_end_idx = (end_idx * audio_unit).long()
# mask = torch.arange(n_motions, device=audio.device).expand(batch_size, -1) < end_idx.unsqueeze(1)
# truncate audio
audio_trunc = _truncate_audio(audio, audio_end_idx, pad_mode=pad_mode)
# truncate coef dict
coef_dict_trunc = _truncate_coef_dict(coef_dict, end_idx, pad_mode=pad_mode)
return audio_trunc, coef_dict_trunc, end_idx
def truncate_motion_coef_and_audio(audio, motion_coef, n_motions, audio_unit=640, pad_mode='zero'):
batch_size = audio.shape[0]
end_idx = torch.randint(1, n_motions, (batch_size,), device=audio.device)
audio_end_idx = (end_idx * audio_unit).long()
# mask = torch.arange(n_motions, device=audio.device).expand(batch_size, -1) < end_idx.unsqueeze(1)
# truncate audio
audio_trunc = _truncate_audio(audio, audio_end_idx, pad_mode=pad_mode)
# prepare coef dict and stats
coef_dict = {'exp': motion_coef[..., :50], 'pose_any': motion_coef[..., 50:]}
# truncate coef dict
coef_dict_trunc = _truncate_coef_dict(coef_dict, end_idx, pad_mode=pad_mode)
motion_coef_trunc = torch.cat([coef_dict_trunc['exp'], coef_dict_trunc['pose_any']], dim=-1)
return audio_trunc, motion_coef_trunc, end_idx
def nt_xent_loss(feature_a, feature_b, temperature):
"""
Normalized temperature-scaled cross entropy loss.
(Adapted from https://github.com/sthalles/SimCLR/blob/master/simclr.py)
Args:
feature_a (torch.Tensor): shape (batch_size, feature_dim)
feature_b (torch.Tensor): shape (batch_size, feature_dim)
temperature (float): temperature scaling factor
Returns:
torch.Tensor: scalar
"""
batch_size = feature_a.shape[0]
device = feature_a.device
features = torch.cat([feature_a, feature_b], dim=0)
labels = torch.cat([torch.arange(batch_size), torch.arange(batch_size)], dim=0)
labels = (labels.unsqueeze(0) == labels.unsqueeze(1))
labels = labels.to(device)
features = F.normalize(features, dim=1)
similarity_matrix = torch.matmul(features, features.T)
# discard the main diagonal from both: labels and similarities matrix
mask = torch.eye(labels.shape[0], dtype=torch.bool).to(device)
labels = labels[~mask].view(labels.shape[0], -1)
similarity_matrix = similarity_matrix[~mask].view(labels.shape[0], -1)
# select the positives and negatives
positives = similarity_matrix[labels].view(labels.shape[0], -1)
negatives = similarity_matrix[~labels].view(labels.shape[0], -1)
logits = torch.cat([positives, negatives], dim=1)
logits = logits / temperature
labels = torch.zeros(labels.shape[0], dtype=torch.long).to(device)
loss = F.cross_entropy(logits, labels)
return loss