CalvinBackup/facefusion-labs
1
0
Fork 0
mirror of https://github.com/facefusion/facefusion-labs.git synced 2026-07-28 16:08:50 +02:00
Code Issues Packages Projects Releases Wiki Activity

Migrate most to self.config and self.context

Browse Source
This commit is contained in:
henryruhs
2025-03-11 14:43:10 +01:00
parent ab3b699124
commit b829d5e42c
5 changed files with 105 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
face_swapper/src/networks/aad.py
+54 -42
View File
@@ -1,3 +1,5 @@
from configparser import ConfigParser
import torch
from torch import Tensor, nn
@@ -5,51 +7,55 @@ from ..types import Attributes, Embedding
class AAD(nn.Module):
def __init__(self, identity_channels : int, output_channels : int, output_size : int, num_blocks : int) -> None:
def __init__(self, config_parser : ConfigParser) -> None:
super().__init__()
self.identity_channels = identity_channels
self.output_channels = output_channels
self.output_size = output_size
self.num_blocks = num_blocks
self.pixel_shuffle_up_sample = PixelShuffleUpSample(identity_channels, output_channels)
self.config =\
{
'identity_channels': config_parser.getint('training.model.generator', 'identity_channels'),
'output_channels': config_parser.getint('training.model.generator', 'output_channels'),
'output_size': config_parser.getint('training.model.generator', 'output_size'),
'num_blocks': config_parser.getint('training.model.generator', 'num_blocks')
}
self.config_parser = config_parser
self.pixel_shuffle_up_sample = PixelShuffleUpSample(self.config.get('identity_channels'), self.config.get('output_channels'))
self.layers = self.create_layers()
def create_layers(self) -> nn.ModuleList:
layers = nn.ModuleList()
if self.output_size == 128:
if self.config.get('output_size') == 128:
layers.extend(
[
AdaptiveFeatureModulation(512, 512, 512, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(512, 512, 1024, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(512, 512, 512, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(512, 512, 512, self.config.get('identity_channels'), self.config.get('num_blocks')),
AdaptiveFeatureModulation(512, 512, 1024, self.config.get('identity_channels'), self.config.get('num_blocks')),
AdaptiveFeatureModulation(512, 512, 512, self.config.get('identity_channels'), self.config.get('num_blocks'))
])
if self.output_size == 256:
if self.config.get('output_size') == 256:
layers.extend(
[
AdaptiveFeatureModulation(1024, 1024, 1024, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(1024, 1024, 2048, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(1024, 1024, 1024, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(1024, 512, 512, self.identity_channels, self.num_blocks)
AdaptiveFeatureModulation(1024, 1024, 1024, self.config.get('identity_channels'), self.config.get('num_blocks')),
AdaptiveFeatureModulation(1024, 1024, 2048, self.config.get('identity_channels'), self.config.get('num_blocks')),
AdaptiveFeatureModulation(1024, 1024, 1024, self.config.get('identity_channels'), self.config.get('num_blocks')),
AdaptiveFeatureModulation(1024, 512, 512, self.config.get('identity_channels'), self.config.get('num_blocks'))
])
if self.output_size == 512:
if self.config.get('output_size') == 512:
layers.extend(
[
AdaptiveFeatureModulation(2048, 2048, 2048, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(2048, 2048, 4096, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(2048, 2048, 2048, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(2048, 1024, 1024, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(1024, 512, 512, self.identity_channels, self.num_blocks)
AdaptiveFeatureModulation(2048, 2048, 2048, self.config.get('identity_channels'), self.config.get('num_blocks')),
AdaptiveFeatureModulation(2048, 2048, 4096, self.config.get('identity_channels'), self.config.get('num_blocks')),
AdaptiveFeatureModulation(2048, 2048, 2048, self.config.get('identity_channels'), self.config.get('num_blocks')),
AdaptiveFeatureModulation(2048, 1024, 1024, self.config.get('identity_channels'), self.config.get('num_blocks')),
AdaptiveFeatureModulation(1024, 512, 512, self.config.get('identity_channels'), self.config.get('num_blocks'))
])
layers.extend(
[
AdaptiveFeatureModulation(512, 256, 256, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(256, 128, 128, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(128, 64, 64, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(64, 3, 64, self.identity_channels, self.num_blocks)
AdaptiveFeatureModulation(512, 256, 256, self.config.get('identity_channels'), self.config.get('num_blocks')),
AdaptiveFeatureModulation(256, 128, 128, self.config.get('identity_channels'), self.config.get('num_blocks')),
AdaptiveFeatureModulation(128, 64, 64, self.config.get('identity_channels'), self.config.get('num_blocks')),
AdaptiveFeatureModulation(64, 3, 64, self.config.get('identity_channels'), self.config.get('num_blocks'))
])
return layers
@@ -69,40 +75,43 @@ class AAD(nn.Module):
class AdaptiveFeatureModulation(nn.Module):
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.attribute_channels = attribute_channels
self.identity_channels = identity_channels
self.num_blocks = num_blocks
self.context =\
{
'input_channels': input_channels,
'output_channels': output_channels,
'attribute_channels': attribute_channels,
'identity_channels': identity_channels,
'num_blocks': num_blocks
}
self.primary_layers = self.create_primary_layers()
self.shortcut_layers = self.create_shortcut_layers()
def create_primary_layers(self) -> nn.ModuleList:
primary_layers = nn.ModuleList()
for index in range(self.num_blocks):
for index in range(self.context.get('num_blocks')):
primary_layers.extend(
[
FeatureModulation(self.input_channels, self.attribute_channels, self.identity_channels),
FeatureModulation(self.context.get('input_channels'), self.context.get('attribute_channels'), self.context.get('identity_channels')),
nn.ReLU(inplace = True)
])
if index < self.num_blocks - 1:
primary_layers.append(nn.Conv2d(self.input_channels, self.input_channels, kernel_size = 3, padding = 1, bias = False))
if index < self.context.get('num_blocks') - 1:
primary_layers.append(nn.Conv2d(self.context.get('input_channels'), self.context.get('input_channels'), kernel_size = 3, padding = 1, bias = False))
else:
primary_layers.append(nn.Conv2d(self.input_channels, self.output_channels, kernel_size = 3, padding = 1, bias = False))
primary_layers.append(nn.Conv2d(self.context.get('input_channels'), self.context.get('output_channels'), kernel_size = 3, padding = 1, bias = False))
return primary_layers
def create_shortcut_layers(self) -> nn.ModuleList:
shortcut_layers = nn.ModuleList()
if self.input_channels > self.output_channels:
if self.context.get('input_channels') > self.context.get('output_channels'):
shortcut_layers.extend(
[
FeatureModulation(self.input_channels, self.attribute_channels, self.identity_channels),
FeatureModulation(self.context.get('input_channels'), self.context.get('attribute_channels'), self.context.get('identity_channels')),
nn.ReLU(inplace = True),
nn.Conv2d(self.input_channels, self.output_channels, kernel_size = 3, padding = 1, bias = False)
nn.Conv2d(self.context.get('input_channels'), self.context.get('output_channels'), kernel_size = 3, padding = 1, bias = False)
])
return shortcut_layers
@@ -116,7 +125,7 @@ class AdaptiveFeatureModulation(nn.Module):
else:
primary_tensor = primary_layer(primary_tensor)
if self.input_channels > self.output_channels:
if self.context.get('input_channels') > self.context.get('output_channels'):
shortcut_tensor = input_tensor
for shortcut_layer in self.shortcut_layers:
@@ -133,7 +142,10 @@ class AdaptiveFeatureModulation(nn.Module):
class FeatureModulation(nn.Module):
def __init__(self, input_channels : int, attribute_channels : int, identity_channels : int) -> None:
super().__init__()
self.input_channels = input_channels
self.context =\
{
'input_channels': input_channels
}
self.conv1 = nn.Conv2d(attribute_channels, input_channels, kernel_size = 1)
self.conv2 = nn.Conv2d(attribute_channels, input_channels, kernel_size = 1)
self.conv3 = nn.Conv2d(input_channels, 1, kernel_size = 1)
@@ -148,8 +160,8 @@ class FeatureModulation(nn.Module):
attribute_shift = self.conv2(attribute_embedding)
attribute_modulation = attribute_scale * temp_tensor + attribute_shift
identity_scale = self.linear2(identity_embedding).reshape(temp_tensor.shape[0], self.input_channels, 1, 1).expand_as(temp_tensor)
identity_shift = self.linear1(identity_embedding).reshape(temp_tensor.shape[0], self.input_channels, 1, 1).expand_as(temp_tensor)
identity_scale = self.linear2(identity_embedding).reshape(temp_tensor.shape[0], self.context.get('input_channels'), 1, 1).expand_as(temp_tensor)
identity_shift = self.linear1(identity_embedding).reshape(temp_tensor.shape[0], self.context.get('input_channels'), 1, 1).expand_as(temp_tensor)
identity_modulation = identity_scale * temp_tensor + identity_shift
temp_mask = torch.sigmoid(self.conv3(temp_tensor))
face_swapper/src/networks/unet.py
+2 -2
View File
@@ -106,7 +106,7 @@ class UNet(nn.Module):
class UpSample(nn.Module):
def __init__(self, input_channels : int, output_channels : int) -> None:
super().__init__()
self.conv_transpose = nn.ConvTranspose2d(in_channels = input_channels, out_channels = output_channels, kernel_size = 4, stride = 2, padding = 1, bias = False)
self.conv_transpose = nn.ConvTranspose2d(input_channels, output_channels, kernel_size = 4, stride = 2, padding = 1, bias = False)
self.batch_norm = nn.BatchNorm2d(output_channels)
self.leaky_relu = nn.LeakyReLU(0.1, inplace = True)
@@ -121,7 +121,7 @@ class UpSample(nn.Module):
class DownSample(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 = 4, stride = 2, padding = 1, bias = False)
self.conv = nn.Conv2d(input_channels, output_channels, kernel_size = 4, stride = 2, padding = 1, bias = False)
self.batch_norm = nn.BatchNorm2d(output_channels)
self.leaky_relu = nn.LeakyReLU(0.1, inplace = True)