diff --git a/face_swapper/README.md b/face_swapper/README.md index e3a5c37..9c7a73b 100644 --- a/face_swapper/README.md +++ b/face_swapper/README.md @@ -49,7 +49,7 @@ motion_extractor_path = .models/motion_extractor.pt [training.model.generator] encoder_type = unet-pro num_blocks = 2 -id_channels = 512 +identity_channels = 512 ``` ``` diff --git a/face_swapper/config.ini b/face_swapper/config.ini index d627d88..4b3ba86 100644 --- a/face_swapper/config.ini +++ b/face_swapper/config.ini @@ -8,14 +8,14 @@ num_workers = split_ratio = [training.model] -id_embedder_path = +embedder_path = landmarker_path = motion_extractor_path = [training.model.generator] encoder_type = num_blocks = -id_channels = +identity_channels = [training.model.discriminator] input_channels = @@ -53,7 +53,7 @@ opset_version = [inferencing] generator_path = -id_embedder_path = +embedder_path = source_path = target_path = output_path = diff --git a/face_swapper/src/models/generator.py b/face_swapper/src/models/generator.py index 6e530a3..a4e47f2 100644 --- a/face_swapper/src/models/generator.py +++ b/face_swapper/src/models/generator.py @@ -14,14 +14,14 @@ class Generator(nn.Module): def __init__(self) -> None: super().__init__() encoder_type = CONFIG.get('training.model.generator', 'encoder_type') - id_channels = CONFIG.getint('training.model.generator', 'id_channels') num_blocks = CONFIG.getint('training.model.generator', 'num_blocks') + identity_channels = CONFIG.getint('training.model.generator', 'identity_channels') if encoder_type == 'unet': self.encoder = UNet() if encoder_type == 'unet-pro': self.encoder = UNetPro() - self.generator = AADGenerator(id_channels, num_blocks) + self.generator = AADGenerator(identity_channels, num_blocks) self.encoder.apply(init_weight) self.generator.apply(init_weight) diff --git a/face_swapper/src/networks/attribute_modulator.py b/face_swapper/src/networks/attribute_modulator.py index f342d6c..b1d6a0f 100644 --- a/face_swapper/src/networks/attribute_modulator.py +++ b/face_swapper/src/networks/attribute_modulator.py @@ -5,17 +5,17 @@ from ..types import Attributes, Embedding class AADGenerator(nn.Module): - def __init__(self, id_channels : int, num_blocks : int) -> None: + def __init__(self, identity_channels : int, num_blocks : int) -> None: super().__init__() - self.upsample = PixelShuffleUpsample(id_channels, 1024 * 4) - self.res_block_1 = AADResBlock(1024, 1024, 1024, id_channels, num_blocks) - self.res_block_2 = AADResBlock(1024, 1024, 2048, id_channels, num_blocks) - self.res_block_3 = AADResBlock(1024, 1024, 1024, id_channels, num_blocks) - self.res_block_4 = AADResBlock(1024, 512, 512, id_channels, num_blocks) - self.res_block_5 = AADResBlock(512, 256, 256, id_channels, num_blocks) - self.res_block_6 = AADResBlock(256, 128, 128, id_channels, num_blocks) - self.res_block_7 = AADResBlock(128, 64, 64, id_channels, num_blocks) - self.res_block_8 = AADResBlock(64, 3, 64, id_channels, num_blocks) + self.upsample = PixelShuffleUpsample(identity_channels, 1024 * 4) + self.res_block_1 = AADResBlock(1024, 1024, 1024, identity_channels, num_blocks) + self.res_block_2 = AADResBlock(1024, 1024, 2048, identity_channels, num_blocks) + self.res_block_3 = AADResBlock(1024, 1024, 1024, identity_channels, num_blocks) + self.res_block_4 = AADResBlock(1024, 512, 512, identity_channels, num_blocks) + self.res_block_5 = AADResBlock(512, 256, 256, identity_channels, num_blocks) + self.res_block_6 = AADResBlock(256, 128, 128, identity_channels, num_blocks) + self.res_block_7 = AADResBlock(128, 64, 64, identity_channels, num_blocks) + self.res_block_8 = AADResBlock(64, 3, 64, identity_channels, num_blocks) def forward(self, target_attributes : Attributes, source_embedding : Embedding) -> Tensor: feature_map = self.upsample(source_embedding) @@ -31,26 +31,26 @@ class AADGenerator(nn.Module): class AADLayer(nn.Module): - def __init__(self, input_channels : int, attr_channels : int, id_channels : int) -> None: + def __init__(self, input_channels : int, attribute_channels : int, identity_channels : int) -> None: super(AADLayer, self).__init__() self.input_channels = input_channels - self.conv_beta = nn.Conv2d(attr_channels, input_channels, kernel_size = 1) - self.conv_gamma = nn.Conv2d(attr_channels, input_channels, kernel_size = 1) - self.fc_beta = nn.Linear(id_channels, input_channels) - self.fc_gamma = nn.Linear(id_channels, input_channels) + self.conv_beta = nn.Conv2d(attribute_channels, input_channels, kernel_size = 1) + self.conv_gamma = nn.Conv2d(attribute_channels, input_channels, kernel_size = 1) + self.fc_beta = nn.Linear(identity_channels, input_channels) + self.fc_gamma = nn.Linear(identity_channels, input_channels) self.instance_norm = nn.InstanceNorm2d(input_channels) self.conv_mask = nn.Conv2d(input_channels, 1, kernel_size = 1) - def forward(self, feature_map : Tensor, attribute_embedding : Embedding, id_embedding : Embedding) -> Tensor: + def forward(self, feature_map : Tensor, attribute_embedding : Embedding, identity_embedding : Embedding) -> Tensor: feature_map = self.instance_norm(feature_map) gamma_attribute = self.conv_gamma(attribute_embedding) beta_attribute = self.conv_beta(attribute_embedding) attribute_modulation = gamma_attribute * feature_map + beta_attribute - id_gamma = self.fc_gamma(id_embedding).reshape(feature_map.shape[0], self.input_channels, 1, 1).expand_as(feature_map) - id_beta = self.fc_beta(id_embedding).reshape(feature_map.shape[0], self.input_channels, 1, 1).expand_as(feature_map) - id_modulation = id_gamma * feature_map + id_beta + identity_gamma = self.fc_gamma(identity_embedding).reshape(feature_map.shape[0], self.input_channels, 1, 1).expand_as(feature_map) + identity_beta = self.fc_beta(identity_embedding).reshape(feature_map.shape[0], self.input_channels, 1, 1).expand_as(feature_map) + identity_modulation = identity_gamma * feature_map + identity_beta feature_mask = torch.sigmoid(self.conv_mask(feature_map)) - feature_blend = (1 - feature_mask) * attribute_modulation + feature_mask * id_modulation + feature_blend = (1 - feature_mask) * attribute_modulation + feature_mask * identity_modulation return feature_blend @@ -69,41 +69,41 @@ class AADSequential(nn.Module): class AADResBlock(nn.Module): - def __init__(self, input_channels : int, output_channels : int, attribute_channels : int, id_channels : int, num_blocks : int) -> None: + def __init__(self, input_channels : int, output_channels : int, attribute_channels : int, identity_channels : int, num_blocks : int) -> None: super().__init__() self.input_channels = input_channels self.output_channels = output_channels - self.prepare_primary_add_blocks(input_channels, attribute_channels, id_channels, output_channels, num_blocks) - self.prepare_auxiliary_add_blocks(input_channels, attribute_channels, id_channels, output_channels) + self.prepare_primary_add_blocks(input_channels, attribute_channels, identity_channels, output_channels, num_blocks) + self.prepare_auxiliary_add_blocks(input_channels, attribute_channels, identity_channels, output_channels) - def prepare_primary_add_blocks(self, input_channels : int, attribute_channels : int, id_channels : int, output_channels : int, num_blocks : int) -> None: + def prepare_primary_add_blocks(self, input_channels : int, attribute_channels : int, identity_channels : int, output_channels : int, num_blocks : int) -> None: primary_add_blocks = [] for index in range(num_blocks): intermediate_channels = input_channels if index < (num_blocks - 1) else output_channels primary_add_blocks.extend( [ - AADLayer(input_channels, attribute_channels, id_channels), + AADLayer(input_channels, attribute_channels, identity_channels), nn.ReLU(inplace = True), nn.Conv2d(input_channels, intermediate_channels, kernel_size = 3, padding = 1, bias = False) ] ) self.primary_add_blocks = AADSequential(*primary_add_blocks) - def prepare_auxiliary_add_blocks(self, input_channels : int, attribute_channels : int, id_channels : int, output_channels : int) -> None: + def prepare_auxiliary_add_blocks(self, input_channels : int, attribute_channels : int, identity_channels : int, output_channels : int) -> None: if input_channels > output_channels: auxiliary_add_blocks = AADSequential( - AADLayer(input_channels, attribute_channels, id_channels), + AADLayer(input_channels, attribute_channels, identity_channels), nn.ReLU(inplace = True), nn.Conv2d(input_channels, output_channels, kernel_size = 3, padding = 1, bias = False) ) self.auxiliary_add_blocks = auxiliary_add_blocks - def forward(self, feature_map : Tensor, attribute_embedding : Embedding, id_embedding : Embedding) -> Tensor: - primary_feature = self.primary_add_blocks(feature_map, attribute_embedding, id_embedding) + def forward(self, feature_map : Tensor, attribute_embedding : Embedding, identity_embedding : Embedding) -> Tensor: + primary_feature = self.primary_add_blocks(feature_map, attribute_embedding, identity_embedding) if self.input_channels > self.output_channels: - feature_map = self.auxiliary_add_blocks(feature_map, attribute_embedding, id_embedding) + feature_map = self.auxiliary_add_blocks(feature_map, attribute_embedding, identity_embedding) output_feature = primary_feature + feature_map return output_feature