Spaces:
Running
Running
| import torch | |
| from torch import nn | |
| from nota_wav2lip.models.base import Wav2LipBase | |
| from nota_wav2lip.models.conv import Conv2d, Conv2dTranspose | |
| class NotaWav2Lip(Wav2LipBase): | |
| def __init__(self, nef=4, naf=8, ndf=8, x_size=96): | |
| super().__init__() | |
| assert x_size in [96, 128] | |
| self.ker_sz_last = x_size // 32 | |
| self.face_encoder_blocks = nn.ModuleList([ | |
| nn.Sequential(Conv2d(6, nef, kernel_size=7, stride=1, padding=3)), # 96,96 | |
| nn.Sequential(Conv2d(nef, nef * 2, kernel_size=3, stride=2, padding=1),), # 48,48 | |
| nn.Sequential(Conv2d(nef * 2, nef * 4, kernel_size=3, stride=2, padding=1),), # 24,24 | |
| nn.Sequential(Conv2d(nef * 4, nef * 8, kernel_size=3, stride=2, padding=1),), # 12,12 | |
| nn.Sequential(Conv2d(nef * 8, nef * 16, kernel_size=3, stride=2, padding=1),), # 6,6 | |
| nn.Sequential(Conv2d(nef * 16, nef * 32, kernel_size=3, stride=2, padding=1),), # 3,3 | |
| nn.Sequential(Conv2d(nef * 32, nef * 32, kernel_size=self.ker_sz_last, stride=1, padding=0), # 1, 1 | |
| Conv2d(nef * 32, nef * 32, kernel_size=1, stride=1, padding=0)), ]) | |
| self.audio_encoder = nn.Sequential( | |
| Conv2d(1, naf, kernel_size=3, stride=1, padding=1), | |
| Conv2d(naf, naf * 2, kernel_size=3, stride=(3, 1), padding=1), | |
| Conv2d(naf * 2, naf * 4, kernel_size=3, stride=3, padding=1), | |
| Conv2d(naf * 4, naf * 8, kernel_size=3, stride=(3, 2), padding=1), | |
| Conv2d(naf * 8, naf * 16, kernel_size=3, stride=1, padding=0), | |
| Conv2d(naf * 16, naf * 16, kernel_size=1, stride=1, padding=0), ) | |
| self.face_decoder_blocks = nn.ModuleList([ | |
| nn.Sequential(Conv2d(naf * 16, naf * 16, kernel_size=1, stride=1, padding=0), ), | |
| nn.Sequential(Conv2dTranspose(nef * 32 + naf * 16, ndf * 16, kernel_size=self.ker_sz_last, stride=1, padding=0),), | |
| # 3,3 # 512+512 = 1024 | |
| nn.Sequential( | |
| Conv2dTranspose(nef * 32 + ndf * 16, ndf * 16, kernel_size=3, stride=2, padding=1, output_padding=1),), # 6, 6 | |
| # 512+512 = 1024 | |
| nn.Sequential( | |
| Conv2dTranspose(nef * 16 + ndf * 16, ndf * 12, kernel_size=3, stride=2, padding=1, output_padding=1),), # 12, 12 | |
| # 256+512 = 768 | |
| nn.Sequential( | |
| Conv2dTranspose(nef * 8 + ndf * 12, ndf * 8, kernel_size=3, stride=2, padding=1, output_padding=1),), # 24, 24 | |
| # 128+384 = 512 | |
| nn.Sequential( | |
| Conv2dTranspose(nef * 4 + ndf * 8, ndf * 4, kernel_size=3, stride=2, padding=1, output_padding=1),), # 48, 48 | |
| # 64+256 = 320 | |
| nn.Sequential( | |
| Conv2dTranspose(nef * 2 + ndf * 4, ndf * 2, kernel_size=3, stride=2, padding=1, output_padding=1),), # 96,96 | |
| # 32+128 = 160 | |
| ]) | |
| self.output_block = nn.Sequential(Conv2d(nef + ndf * 2, ndf, kernel_size=3, stride=1, padding=1), # 16+64 = 80 | |
| nn.Conv2d(ndf, 3, kernel_size=1, stride=1, padding=0), | |
| nn.Sigmoid()) | |