RabbitRUI/ruispace

1import torch2import torch.nn.functional as F3from torch import nn4 5class Conv2d(nn.Module):6    def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, use_act = True, *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        self.use_act = use_act15 16    def forward(self, x):17        out = self.conv_block(x)18        if self.residual:19            out += x20        21        if self.use_act:22            return self.act(out)23        else:24            return out25 26class SimpleWrapperV2(nn.Module):27    def __init__(self) -> None:28        super().__init__()29        self.audio_encoder = nn.Sequential(30            Conv2d(1, 32, kernel_size=3, stride=1, padding=1),31            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),32            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),33 34            Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1),35            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),36            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),37 38            Conv2d(64, 128, kernel_size=3, stride=3, padding=1),39            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),40            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),41 42            Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1),43            Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),44 45            Conv2d(256, 512, kernel_size=3, stride=1, padding=0),46            Conv2d(512, 512, kernel_size=1, stride=1, padding=0),47            )48 49        #### load the pre-trained audio_encoder 50        #self.audio_encoder = self.audio_encoder.to(device)  51        '''52        wav2lip_state_dict = torch.load('/apdcephfs_cq2/share_1290939/wenxuazhang/checkpoints/wav2lip.pth')['state_dict']53        state_dict = self.audio_encoder.state_dict()54 55        for k,v in wav2lip_state_dict.items():56            if 'audio_encoder' in k:57                print('init:', k)58                state_dict[k.replace('module.audio_encoder.', '')] = v59        self.audio_encoder.load_state_dict(state_dict)60        '''61 62        self.mapping1 = nn.Linear(512+64+1, 64)63        #self.mapping2 = nn.Linear(30, 64)64        #nn.init.constant_(self.mapping1.weight, 0.)65        nn.init.constant_(self.mapping1.bias, 0.)66 67    def forward(self, x, ref, ratio):68        x = self.audio_encoder(x).view(x.size(0), -1)69        ref_reshape = ref.reshape(x.size(0), -1)70        ratio = ratio.reshape(x.size(0), -1)71        72        y = self.mapping1(torch.cat([x, ref_reshape, ratio], dim=1)) 73        out = y.reshape(ref.shape[0], ref.shape[1], -1) #+ ref # resudial74        return out75