Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| from torch import nn | |
| from model.base import BaseModule | |
| from espnet.nets.pytorch_backend.conformer.encoder import Encoder as ConformerEncoder | |
| import torch.nn.functional as F | |
| class LSTM(nn.Module): | |
| def __init__(self, motion_dim, output_dim, num_layers=2, hidden_dim=128): | |
| super().__init__() | |
| self.lstm = nn.LSTM(input_size=motion_dim, hidden_size=hidden_dim, | |
| num_layers=num_layers, batch_first=True) | |
| self.fc = nn.Linear(hidden_dim, output_dim) | |
| def forward(self, x): | |
| x, _ = self.lstm(x) | |
| return self.fc(x) | |
| class DiffusionPredictor(BaseModule): | |
| def __init__(self, conf): | |
| super(DiffusionPredictor, self).__init__() | |
| self.infer_type = conf.infer_type | |
| self.initialize_layers(conf) | |
| print(f'infer_type: {self.infer_type}') | |
| def create_conformer_encoder(self, attention_dim, num_blocks): | |
| return ConformerEncoder( | |
| idim=0, attention_dim=attention_dim, attention_heads=2, linear_units=attention_dim, | |
| num_blocks=num_blocks, input_layer=None, dropout_rate=0.2, positional_dropout_rate=0.2, | |
| attention_dropout_rate=0.2, normalize_before=False, concat_after=False, | |
| positionwise_layer_type="linear", positionwise_conv_kernel_size=3, macaron_style=True, | |
| pos_enc_layer_type="rel_pos", selfattention_layer_type="rel_selfattn", use_cnn_module=True, | |
| cnn_module_kernel=13) | |
| def initialize_layers(self, conf, mfcc_dim=39, hubert_dim=1024, speech_layers=4, speech_dim=512, decoder_dim=1024, motion_start_dim=512, HAL_layers=25): | |
| self.conf = conf | |
| # Speech downsampling | |
| if self.infer_type.startswith('mfcc'): | |
| # from 100 hz to 25 hz | |
| self.down_sample1 = nn.Conv1d(mfcc_dim, 256, kernel_size=3, stride=2, padding=1) | |
| self.down_sample2 = nn.Conv1d(256, speech_dim, kernel_size=3, stride=2, padding=1) | |
| elif self.infer_type.startswith('hubert'): | |
| # from 50 hz to 25 hz | |
| self.down_sample1 = nn.Conv1d(hubert_dim, speech_dim, kernel_size=3, stride=2, padding=1) | |
| self.weights = nn.Parameter(torch.zeros(HAL_layers)) | |
| self.speech_encoder = self.create_conformer_encoder(speech_dim, speech_layers) | |
| else: | |
| print('infer_type not supported') | |
| # Encoders & Deocoders | |
| self.coarse_decoder = self.create_conformer_encoder(decoder_dim, conf.decoder_layers) | |
| # LSTM predictors for Variance Adapter | |
| if self.infer_type != 'hubert_audio_only': | |
| self.pose_predictor = LSTM(speech_dim, 3) | |
| self.pose_encoder = LSTM(3, speech_dim) | |
| if 'full_control' in self.infer_type: | |
| self.location_predictor = LSTM(speech_dim, 1) | |
| self.location_encoder = LSTM(1, speech_dim) | |
| self.face_scale_predictor = LSTM(speech_dim, 1) | |
| self.face_scale_encoder = LSTM(1, speech_dim) | |
| # Linear transformations | |
| self.init_code_proj = nn.Sequential(nn.Linear(motion_start_dim, 128)) | |
| self.noisy_encoder = nn.Sequential(nn.Linear(conf.motion_dim, 128)) | |
| self.t_encoder = nn.Sequential(nn.Linear(1, 128)) | |
| self.encoder_direction_code = nn.Linear(conf.motion_dim, 128) | |
| self.out_proj = nn.Linear(decoder_dim, conf.motion_dim) | |
| def forward(self, initial_code, direction_code, seq_input_vector, face_location, face_scale, yaw_pitch_roll, noisy_x, t_emb, control_flag=False): | |
| if self.infer_type.startswith('mfcc'): | |
| x = self.mfcc_speech_downsample(seq_input_vector) | |
| elif self.infer_type.startswith('hubert'): | |
| norm_weights = F.softmax(self.weights, dim=-1) | |
| weighted_feature = (norm_weights.unsqueeze(0).unsqueeze(-1).unsqueeze(-1) * seq_input_vector).sum(dim=1) | |
| x = self.down_sample1(weighted_feature.transpose(1,2)).transpose(1,2) | |
| x, _ = self.speech_encoder(x, masks=None) | |
| predicted_location, predicted_scale, predicted_pose = face_location, face_scale, yaw_pitch_roll | |
| if self.infer_type != 'hubert_audio_only': | |
| print(f'pose controllable. control_flag: {control_flag}') | |
| x, predicted_location, predicted_scale, predicted_pose = self.adjust_features(x, face_location, face_scale, yaw_pitch_roll, control_flag) | |
| concatenated_features = self.combine_features(x, initial_code, direction_code, noisy_x, t_emb) # initial_code and direction_code serve as a motion guide extracted from the reference image. This aims to tell the model what the starting motion should be. | |
| outputs = self.decode_features(concatenated_features) | |
| return outputs, predicted_location, predicted_scale, predicted_pose | |
| def mfcc_speech_downsample(self, seq_input_vector): | |
| x = self.down_sample1(seq_input_vector.transpose(1,2)) | |
| return self.down_sample2(x).transpose(1,2) | |
| def adjust_features(self, x, face_location, face_scale, yaw_pitch_roll, control_flag): | |
| predicted_location, predicted_scale = 0, 0 | |
| if 'full_control' in self.infer_type: | |
| print(f'full controllable. control_flag: {control_flag}') | |
| x_residual, predicted_location = self.adjust_location(x, face_location, control_flag) | |
| x = x + x_residual | |
| x_residual, predicted_scale = self.adjust_scale(x, face_scale, control_flag) | |
| x = x + x_residual | |
| x_residual, predicted_pose= self.adjust_pose(x, yaw_pitch_roll, control_flag) | |
| x = x + x_residual | |
| return x, predicted_location, predicted_scale, predicted_pose | |
| def adjust_location(self, x, face_location, control_flag): | |
| if control_flag: | |
| predicted_location = face_location | |
| else: | |
| predicted_location = self.location_predictor(x) | |
| return self.location_encoder(predicted_location), predicted_location | |
| def adjust_scale(self, x, face_scale, control_flag): | |
| if control_flag: | |
| predicted_face_scale = face_scale | |
| else: | |
| predicted_face_scale = self.face_scale_predictor(x) | |
| return self.face_scale_encoder(predicted_face_scale), predicted_face_scale | |
| def adjust_pose(self, x, yaw_pitch_roll, control_flag): | |
| if control_flag: | |
| predicted_pose = yaw_pitch_roll | |
| else: | |
| predicted_pose = self.pose_predictor(x) | |
| return self.pose_encoder(predicted_pose), predicted_pose | |
| def combine_features(self, x, initial_code, direction_code, noisy_x, t_emb): | |
| init_code_proj = self.init_code_proj(initial_code).unsqueeze(1).repeat(1, x.size(1), 1) | |
| noisy_feature = self.noisy_encoder(noisy_x) | |
| t_emb_feature = self.t_encoder(t_emb.unsqueeze(1).float()).unsqueeze(1).repeat(1, x.size(1), 1) | |
| direction_code_feature = self.encoder_direction_code(direction_code).unsqueeze(1).repeat(1, x.size(1), 1) | |
| return torch.cat((x, direction_code_feature, init_code_proj, noisy_feature, t_emb_feature), dim=-1) | |
| def decode_features(self, concatenated_features): | |
| outputs, _ = self.coarse_decoder(concatenated_features, masks=None) | |
| return self.out_proj(outputs) |