|
|
import torch |
|
|
import torch.nn as nn |
|
|
|
|
|
from models.init_weight import init_net |
|
|
from models.model_blocks import AdaInResBlock |
|
|
from models.model_blocks import ResBlock |
|
|
from models.semantic_face_fusion_model import SemanticFaceFusionModule |
|
|
from models.shape_aware_identity_model import ShapeAwareIdentityExtractor |
|
|
|
|
|
|
|
|
class Encoder(nn.Module): |
|
|
""" |
|
|
Hififace encoder part |
|
|
""" |
|
|
|
|
|
def __init__(self): |
|
|
super(Encoder, self).__init__() |
|
|
self.conv_first = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1) |
|
|
|
|
|
self.channel_list = [64, 128, 256, 512, 512, 512, 512, 512] |
|
|
self.down_sample = [True, True, True, True, True, False, False] |
|
|
|
|
|
self.block_list = nn.ModuleList() |
|
|
|
|
|
for i in range(7): |
|
|
self.block_list.append( |
|
|
ResBlock(self.channel_list[i], self.channel_list[i + 1], down_sample=self.down_sample[i]) |
|
|
) |
|
|
|
|
|
def forward(self, x): |
|
|
x = self.conv_first(x) |
|
|
z_enc = None |
|
|
|
|
|
for i in range(7): |
|
|
x = self.block_list[i](x) |
|
|
if i == 1: |
|
|
z_enc = x |
|
|
return z_enc, x |
|
|
|
|
|
|
|
|
class Decoder(nn.Module): |
|
|
""" |
|
|
Hififace decoder part |
|
|
""" |
|
|
|
|
|
def __init__(self): |
|
|
super(Decoder, self).__init__() |
|
|
self.block_list = nn.ModuleList() |
|
|
self.channel_list = [512, 512, 512, 512, 512, 256] |
|
|
self.up_sample = [False, False, True, True, True] |
|
|
|
|
|
for i in range(5): |
|
|
self.block_list.append( |
|
|
AdaInResBlock(self.channel_list[i], self.channel_list[i + 1], up_sample=self.up_sample[i]) |
|
|
) |
|
|
|
|
|
def forward(self, x, id_vector): |
|
|
""" |
|
|
Parameters: |
|
|
----------- |
|
|
x: encoder encoded feature map |
|
|
id_vector: 3d shape aware identity vector |
|
|
|
|
|
Returns: |
|
|
-------- |
|
|
z_dec |
|
|
""" |
|
|
for i in range(5): |
|
|
x = self.block_list[i](x, id_vector) |
|
|
return x |
|
|
|
|
|
|
|
|
class Generator(nn.Module): |
|
|
""" |
|
|
Hififace Generator |
|
|
""" |
|
|
|
|
|
def __init__(self, identity_extractor_config): |
|
|
super(Generator, self).__init__() |
|
|
self.id_extractor = ShapeAwareIdentityExtractor(identity_extractor_config) |
|
|
self.id_extractor.requires_grad_(False) |
|
|
self.encoder = init_net(Encoder()) |
|
|
self.decoder = init_net(Decoder()) |
|
|
self.sff_module = init_net(SemanticFaceFusionModule()) |
|
|
|
|
|
@torch.no_grad() |
|
|
def interp(self, i_source, i_target, shape_rate=1.0, id_rate=1.0): |
|
|
shape_aware_id_vector = self.id_extractor.interp(i_source, i_target, shape_rate, id_rate) |
|
|
z_enc, x = self.encoder(i_target) |
|
|
z_dec = self.decoder(x, shape_aware_id_vector) |
|
|
|
|
|
i_r, i_low, m_r, m_low = self.sff_module(i_target, z_enc, z_dec, shape_aware_id_vector) |
|
|
|
|
|
return i_r, i_low, m_r, m_low |
|
|
|
|
|
def forward(self, i_source, i_target, need_id_grad=False): |
|
|
""" |
|
|
Parameters: |
|
|
----------- |
|
|
i_source: torch.Tensor, shape (B, 3, H, W), in range [0, 1], source face image |
|
|
i_target: torch.Tensor, shape (B, 3, H, W), in range [0, 1], target face image |
|
|
need_id_grad: bool, whether to calculate id extractor module's gradient |
|
|
|
|
|
Returns: |
|
|
-------- |
|
|
i_r: torch.Tensor |
|
|
i_low: torch.Tensor |
|
|
m_r: torch.Tensor |
|
|
m_low: torch.Tensor |
|
|
""" |
|
|
if need_id_grad: |
|
|
shape_aware_id_vector = self.id_extractor(i_source, i_target) |
|
|
else: |
|
|
with torch.no_grad(): |
|
|
shape_aware_id_vector = self.id_extractor(i_source, i_target) |
|
|
z_enc, x = self.encoder(i_target) |
|
|
z_dec = self.decoder(x, shape_aware_id_vector) |
|
|
|
|
|
i_r, i_low, m_r, m_low = self.sff_module(i_target, z_enc, z_dec, shape_aware_id_vector) |
|
|
|
|
|
return i_r, i_low, m_r, m_low |
|
|
|
