From ee3fc40e8393bff3551023f47312af05fe86edfa Mon Sep 17 00:00:00 2001 From: henryruhs Date: Sun, 23 Feb 2025 14:26:01 +0100 Subject: [PATCH] Move output channels to config --- face_swapper/README.md | 3 ++- face_swapper/config.ini | 3 ++- face_swapper/src/models/generator.py | 7 ++++--- face_swapper/src/networks/aienet.py | 12 ++++++------ 4 files changed, 14 insertions(+), 11 deletions(-) diff --git a/face_swapper/README.md b/face_swapper/README.md index 9c7a73b..d520c1c 100644 --- a/face_swapper/README.md +++ b/face_swapper/README.md @@ -48,8 +48,9 @@ motion_extractor_path = .models/motion_extractor.pt ``` [training.model.generator] encoder_type = unet-pro -num_blocks = 2 identity_channels = 512 +output_channels = 4096 +num_blocks = 2 ``` ``` diff --git a/face_swapper/config.ini b/face_swapper/config.ini index 4b3ba86..6472703 100644 --- a/face_swapper/config.ini +++ b/face_swapper/config.ini @@ -14,8 +14,9 @@ motion_extractor_path = [training.model.generator] encoder_type = -num_blocks = identity_channels = +output_channels = +num_blocks = [training.model.discriminator] input_channels = diff --git a/face_swapper/src/models/generator.py b/face_swapper/src/models/generator.py index a4e47f2..7fb479e 100644 --- a/face_swapper/src/models/generator.py +++ b/face_swapper/src/models/generator.py @@ -2,7 +2,7 @@ import configparser from torch import Tensor, nn -from ..networks.attribute_modulator import AADGenerator +from ..networks.aienet import AADGenerator from ..networks.unet import UNet, UNetPro from ..types import Attributes, Embedding @@ -14,14 +14,15 @@ class Generator(nn.Module): def __init__(self) -> None: super().__init__() encoder_type = CONFIG.get('training.model.generator', 'encoder_type') - num_blocks = CONFIG.getint('training.model.generator', 'num_blocks') identity_channels = CONFIG.getint('training.model.generator', 'identity_channels') + output_channels = CONFIG.getint('training.model.generator', 'output_channels') + num_blocks = CONFIG.getint('training.model.generator', 'num_blocks') if encoder_type == 'unet': self.encoder = UNet() if encoder_type == 'unet-pro': self.encoder = UNetPro() - self.generator = AADGenerator(identity_channels, num_blocks) + self.generator = AADGenerator(identity_channels, output_channels, num_blocks) self.encoder.apply(init_weight) self.generator.apply(init_weight) diff --git a/face_swapper/src/networks/aienet.py b/face_swapper/src/networks/aienet.py index b1d6a0f..52930be 100644 --- a/face_swapper/src/networks/aienet.py +++ b/face_swapper/src/networks/aienet.py @@ -5,9 +5,9 @@ from ..types import Attributes, Embedding class AADGenerator(nn.Module): - def __init__(self, identity_channels : int, num_blocks : int) -> None: + def __init__(self, identity_channels : int, output_channels : int, num_blocks : int) -> None: super().__init__() - self.upsample = PixelShuffleUpsample(identity_channels, 1024 * 4) + self.pixel_shuffle_up_sample = PixelShuffleUpSample(identity_channels, output_channels) 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) @@ -18,7 +18,7 @@ class AADGenerator(nn.Module): 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) + feature_map = self.pixel_shuffle_up_sample(source_embedding) feature_map_1 = nn.functional.interpolate(self.res_block_1(feature_map, target_attributes[0], source_embedding), scale_factor = 2, mode = 'bilinear', align_corners = False) feature_map_2 = nn.functional.interpolate(self.res_block_2(feature_map_1, target_attributes[1], source_embedding), scale_factor = 2, mode = 'bilinear', align_corners = False) feature_map_3 = nn.functional.interpolate(self.res_block_3(feature_map_2, target_attributes[2], source_embedding), scale_factor = 2, mode = 'bilinear', align_corners = False) @@ -32,7 +32,7 @@ class AADGenerator(nn.Module): class AADLayer(nn.Module): def __init__(self, input_channels : int, attribute_channels : int, identity_channels : int) -> None: - super(AADLayer, self).__init__() + super().__init__() self.input_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) @@ -109,10 +109,10 @@ class AADResBlock(nn.Module): return output_feature -class PixelShuffleUpsample(nn.Module): +class PixelShuffleUpSample(nn.Module): def __init__(self, input_channels : int, output_channels : int) -> None: super().__init__() - self.conv = nn.Conv2d(in_channels = input_channels, out_channels = output_channels, kernel_size = 3, padding = 1) + self.conv = nn.Conv2d(input_channels, output_channels, kernel_size = 3, padding = 1) self.pixel_shuffle = nn.PixelShuffle(upscale_factor = 2) def forward(self, input_tensor : Tensor) -> Tensor: