RabbitRUI/ruispace

1import torch2from torch import nn3from torch.nn import functional as F4 5class Conv2d(nn.Module):6    def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, *args, **kwargs):7        super().__init__(*args, **kwargs)8        self.conv_block = nn.Sequential(9                            nn.Conv2d(cin, cout, kernel_size, stride, padding),10                            nn.BatchNorm2d(cout)11                            )12        self.act = nn.ReLU()13        self.residual = residual14 15    def forward(self, x):16        out = self.conv_block(x)17        if self.residual:18            out += x19        return self.act(out)20 21class AudioEncoder(nn.Module):22    def __init__(self, wav2lip_checkpoint, device):23        super(AudioEncoder, self).__init__()24 25        self.audio_encoder = nn.Sequential(26            Conv2d(1, 32, kernel_size=3, stride=1, padding=1),27            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),28            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),29 30            Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1),31            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),32            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),33 34            Conv2d(64, 128, kernel_size=3, stride=3, padding=1),35            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),36            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),37 38            Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1),39            Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),40 41            Conv2d(256, 512, kernel_size=3, stride=1, padding=0),42            Conv2d(512, 512, kernel_size=1, stride=1, padding=0),)43 44        #### load the pre-trained audio_encoder45        wav2lip_state_dict = torch.load(wav2lip_checkpoint, map_location=torch.device(device))['state_dict']46        state_dict = self.audio_encoder.state_dict()47 48        for k,v in wav2lip_state_dict.items():49            if 'audio_encoder' in k:50                state_dict[k.replace('module.audio_encoder.', '')] = v51        self.audio_encoder.load_state_dict(state_dict)52 53 54    def forward(self, audio_sequences):55        # audio_sequences = (B, T, 1, 80, 16)56        B = audio_sequences.size(0)57 58        audio_sequences = torch.cat([audio_sequences[:, i] for i in range(audio_sequences.size(1))], dim=0)59 60        audio_embedding = self.audio_encoder(audio_sequences) # B, 512, 1, 161        dim = audio_embedding.shape[1]62        audio_embedding = audio_embedding.reshape((B, -1, dim, 1, 1))63 64        return audio_embedding.squeeze(-1).squeeze(-1) #B seq_len+1 512 65