model/seq2seq.py

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)