mirror of
https://github.com/facefusion/facefusion-labs.git
synced 2026-04-19 15:56:37 +02:00
harisreedhar
68 lines
3.1 KiB
Python
68 lines
3.1 KiB
Python
import torch
|
|
from torch import Tensor, nn as nn
|
|
|
|
from face_swapper.src.types import TargetAttributes, VisionTensor
|
|
|
|
|
|
class Upsample(nn.Module):
|
|
def __init__(self, input_channels : int, output_channels : int) -> None:
|
|
super(Upsample, self).__init__()
|
|
self.deconv = nn.ConvTranspose2d(in_channels = input_channels, out_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)
|
|
|
|
def forward(self, temp : Tensor, skip_tensor : Tensor) -> Tensor:
|
|
temp = self.deconv(temp)
|
|
temp = self.batch_norm(temp)
|
|
temp = self.leaky_relu(temp)
|
|
return torch.cat((temp, skip_tensor), dim = 1)
|
|
|
|
|
|
class DownSample(nn.Module):
|
|
def __init__(self, input_channels : int, output_channels : int) -> None:
|
|
super(DownSample, self).__init__()
|
|
self.conv = nn.Conv2d(in_channels = input_channels, out_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)
|
|
|
|
def forward(self, temp : Tensor) -> Tensor:
|
|
temp = self.conv(temp)
|
|
temp = self.batch_norm(temp)
|
|
temp = self.leaky_relu(temp)
|
|
return temp
|
|
|
|
|
|
class UNet(nn.Module):
|
|
def __init__(self) -> None:
|
|
super(UNet, self).__init__()
|
|
self.downsampler_1 = DownSample(3, 32)
|
|
self.downsampler_2 = DownSample(32, 64)
|
|
self.downsampler_3 = DownSample(64, 128)
|
|
self.downsampler_4 = DownSample(128, 256)
|
|
self.downsampler_5 = DownSample(256, 512)
|
|
self.downsampler_6 = DownSample(512, 1024)
|
|
self.bottleneck = DownSample(1024, 1024)
|
|
self.upsampler_1 = Upsample(1024, 1024)
|
|
self.upsampler_2 = Upsample(2048, 512)
|
|
self.upsampler_3 = Upsample(1024, 256)
|
|
self.upsampler_4 = Upsample(512, 128)
|
|
self.upsampler_5 = Upsample(256, 64)
|
|
self.upsampler_6 = Upsample(128, 32)
|
|
|
|
def forward(self, target : VisionTensor) -> TargetAttributes:
|
|
downsample_feature_1 = self.downsampler_1(target)
|
|
downsample_feature_2 = self.downsampler_2(downsample_feature_1)
|
|
downsample_feature_3 = self.downsampler_3(downsample_feature_2)
|
|
downsample_feature_4 = self.downsampler_4(downsample_feature_3)
|
|
downsample_feature_5 = self.downsampler_5(downsample_feature_4)
|
|
downsample_feature_6 = self.downsampler_6(downsample_feature_5)
|
|
bottleneck_output = self.bottleneck(downsample_feature_6)
|
|
upsample_feature_1 = self.upsampler_1(bottleneck_output, downsample_feature_6)
|
|
upsample_feature_2 = self.upsampler_2(upsample_feature_1, downsample_feature_5)
|
|
upsample_feature_3 = self.upsampler_3(upsample_feature_2, downsample_feature_4)
|
|
upsample_feature_4 = self.upsampler_4(upsample_feature_3, downsample_feature_3)
|
|
upsample_feature_5 = self.upsampler_5(upsample_feature_4, downsample_feature_2)
|
|
upsample_feature_6 = self.upsampler_6(upsample_feature_5, downsample_feature_1)
|
|
output = torch.nn.functional.interpolate(upsample_feature_6, scale_factor = 2, mode = 'bilinear', align_corners = False)
|
|
return bottleneck_output, upsample_feature_1, upsample_feature_2, upsample_feature_3, upsample_feature_4, upsample_feature_5, upsample_feature_6, output
|