Spanicin's picture
Upload 7 files
5b1ae50 verified
raw
history blame
5.41 kB
import torch
import torch.nn as nn
import torch.nn.functional as F
from models.base_blocks import ResBlock, StyleConv, ToRGB
class ENet(nn.Module):
def __init__(
self,
num_style_feat=512,
lnet=None,
concat=False
):
super(ENet, self).__init__()
self.low_res = lnet
for param in self.low_res.parameters():
param.requires_grad = False
channel_multiplier, narrow = 2, 1
channels = {
'4': int(512 * narrow),
'8': int(512 * narrow),
'16': int(512 * narrow),
'32': int(512 * narrow),
'64': int(256 * channel_multiplier * narrow),
'128': int(128 * channel_multiplier * narrow),
'256': int(64 * channel_multiplier * narrow),
'512': int(32 * channel_multiplier * narrow),
'1024': int(16 * channel_multiplier * narrow)
}
self.log_size = 8
first_out_size = 128
self.conv_body_first = nn.Conv2d(3, channels[f'{first_out_size}'], 1) # 256 -> 128
# downsample
in_channels = channels[f'{first_out_size}']
self.conv_body_down = nn.ModuleList()
for i in range(8, 2, -1):
out_channels = channels[f'{2**(i - 1)}']
self.conv_body_down.append(ResBlock(in_channels, out_channels, mode='down'))
in_channels = out_channels
self.num_style_feat = num_style_feat
linear_out_channel = num_style_feat
self.final_linear = nn.Linear(channels['4'] * 4 * 4, linear_out_channel)
self.final_conv = nn.Conv2d(in_channels, channels['4'], 3, 1, 1)
self.style_convs = nn.ModuleList()
self.to_rgbs = nn.ModuleList()
self.noises = nn.Module()
self.concat = concat
if concat:
in_channels = 3 + 32 # channels['64']
else:
in_channels = 3
for i in range(7, 9): # 128, 256
out_channels = channels[f'{2**i}'] #
self.style_convs.append(
StyleConv(
in_channels,
out_channels,
kernel_size=3,
num_style_feat=num_style_feat,
demodulate=True,
sample_mode='upsample'))
self.style_convs.append(
StyleConv(
out_channels,
out_channels,
kernel_size=3,
num_style_feat=num_style_feat,
demodulate=True,
sample_mode=None))
self.to_rgbs.append(ToRGB(out_channels, num_style_feat, upsample=True))
in_channels = out_channels
def forward(self, audio_sequences, face_sequences, gt_sequences):
B = audio_sequences.size(0)
input_dim_size = len(face_sequences.size())
inp, ref = torch.split(face_sequences,3,dim=1)
if input_dim_size > 4:
audio_sequences = torch.cat([audio_sequences[:, i] for i in range(audio_sequences.size(1))], dim=0)
inp = torch.cat([inp[:, :, i] for i in range(inp.size(2))], dim=0)
ref = torch.cat([ref[:, :, i] for i in range(ref.size(2))], dim=0)
gt_sequences = torch.cat([gt_sequences[:, :, i] for i in range(gt_sequences.size(2))], dim=0)
# get the global style
feat = F.leaky_relu_(self.conv_body_first(F.interpolate(ref, size=(256,256), mode='bilinear')), negative_slope=0.2)
for i in range(self.log_size - 2):
feat = self.conv_body_down[i](feat)
feat = F.leaky_relu_(self.final_conv(feat), negative_slope=0.2)
# style code
style_code = self.final_linear(feat.reshape(feat.size(0), -1))
style_code = style_code.reshape(style_code.size(0), -1, self.num_style_feat)
LNet_input = torch.cat([inp, gt_sequences], dim=1)
LNet_input = F.interpolate(LNet_input, size=(96,96), mode='bilinear')
if self.concat:
low_res_img, low_res_feat = self.low_res(audio_sequences, LNet_input)
low_res_img.detach()
low_res_feat.detach()
out = torch.cat([low_res_img, low_res_feat], dim=1)
else:
low_res_img = self.low_res(audio_sequences, LNet_input)
low_res_img.detach()
# 96 x 96
out = low_res_img
p2d = (2,2,2,2)
out = F.pad(out, p2d, "reflect", 0)
skip = out
for conv1, conv2, to_rgb in zip(self.style_convs[::2], self.style_convs[1::2], self.to_rgbs):
out = conv1(out, style_code) # 96, 192, 384
out = conv2(out, style_code)
skip = to_rgb(out, style_code, skip)
_outputs = skip
# remove padding
_outputs = _outputs[:,:,8:-8,8:-8]
if input_dim_size > 4:
_outputs = torch.split(_outputs, B, dim=0)
outputs = torch.stack(_outputs, dim=2)
low_res_img = F.interpolate(low_res_img, outputs.size()[3:])
low_res_img = torch.split(low_res_img, B, dim=0)
low_res_img = torch.stack(low_res_img, dim=2)
else:
outputs = _outputs
return outputs, low_res_img