mirror of
https://github.com/deepfakes/faceswap.git
synced 2026-10-18 14:34:30 +08:00
torzdf
837bc2d51d
* FaceSwap 3 (#1515) * Update extract pipeline * Update requirements + setup for nvidia * Remove allow-growth option * tf.keras to keras updates * lib.model.losses - Port + fix all loss functions for Keras3 * lib.model - port initializers, layers. normalization to Keras3 * lib.model.autoclip to Keras 3 * Update mixed precision layer storage * model file to .keras format * Restructure nn_blocks to initialize layers in __init__ * Tensorboard - Trainer: Add Torch compatible Tensorboard callbacks - GUI event reader remove TF dependency * Loss logging - Flush TB logs on save - Replace TB live iterator for GUI * Backup models on total loss drop rather than per side * Update all models to Keras3 Compat * Remove lib.model.session * Update clip ViT to Keras 3 * plugins.extract.mask.unet-dfl - Fix for Keras3/Torch backend * Port AdaBelief to Keras 3 * setup.py: - Add --dev flag for dev tool install * Fix Keras 3 syntax * Fix LR Finder for Keras 3 * Fix mixed precision switching for Keras 3 * Add more optimizers + open up config setting * train: Remove updating FS1 weights to FS2 models * Alignments: Remove support for legacy .json files * tools.model: - Remove TF Saved Format saving - Fix Backup/Restore + Nan-Scan * Fix inference model creation for Keras 3 * Preview tool: Fix for Keras3 * setup.py: Configure keras backend * train: Migration of FS2 models to FS3 * Training: Default coverage to 100% * Remove DirectML backend * Update setup for MacOS * GUI: Force line reading to UTF-8 * Remove redundant Tensorflow references * Remove redundant code * Legacy model loading: Fix TFLamdaOp scalar ops and DepthwiseConv2D * Add vertical offset option for training * Github actions: Add more python versions * Add python version to workflow names * Github workflow: Exclude Python 3.12 for macOS * Implement custom training loop * Fs3 - Add RTX5xxx and ROCm 6.1-6.4 support (#1511) * setup.py: Add Cuda/ROCm version select options * bump minimum python version to 3.11 * Switch from setup.cgf to pyproject.toml * Documentation: Update all docs to use automodapi * Allow sysinfo to run with missing packages + correctly install tk under Linux * Bugfix: dot naming convention in clip models * lib.config: Centralise globally rather than passing as object - Add torch DataParallel for multi-gpu training - GUI: Group switches together when generating cli args - CLI: Remove deprecated multi-character argparse args - Refactor: - Centralise tensorboard reading/writing + unit tests - Create trainer plugin interfaces + add original + distributed * Update installers
223 lines
9.2 KiB
Python
223 lines
9.2 KiB
Python
#!/usr/bin/env python3
|
|
""" A lightweight variant of DFaker Model
|
|
By AnDenix, 2018-2019
|
|
Based on the dfaker model: https://github.com/dfaker
|
|
|
|
Acknowledgments:
|
|
kvrooman for numerous insights and invaluable aid
|
|
DeepHomage for lots of testing
|
|
"""
|
|
import logging
|
|
|
|
from keras import layers, Input, Model as KModel
|
|
|
|
from lib.model.nn_blocks import (Conv2DOutput, Conv2DBlock, ResidualBlock, UpscaleBlock,
|
|
Upscale2xBlock)
|
|
from lib.utils import FaceswapError
|
|
from plugins.train.train_config import Loss as cfg_loss
|
|
|
|
from ._base import ModelBase
|
|
from . import dlight_defaults as cfg
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Model(ModelBase):
|
|
""" DLight Autoencoder Model """
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.input_shape = (128, 128, 3)
|
|
|
|
self.features = {"lowmem": 0, "fair": 1, "best": 2}[cfg.features()]
|
|
self.encoder_filters = 64 if self.features > 0 else 48
|
|
|
|
bonum_fortunam = 128
|
|
self.encoder_dim = {0: 512 + bonum_fortunam,
|
|
1: 1024 + bonum_fortunam,
|
|
2: 1536 + bonum_fortunam}[self.features]
|
|
self.details = {"fast": 0, "good": 1}[cfg.details()]
|
|
try:
|
|
self.upscale_ratio = {128: 2,
|
|
256: 4,
|
|
384: 6}[cfg.output_size()]
|
|
except KeyError as err:
|
|
logger.error("Config error: output_size must be one of: 128, 256, or 384.")
|
|
raise FaceswapError("Config error: output_size must be one of: "
|
|
"128, 256, or 384.") from err
|
|
|
|
logger.debug("output_size: %s, features: %s, encoder_filters: %s, encoder_dim: %s, "
|
|
" details: %s, upscale_ratio: %s", cfg.output_size(), self.features,
|
|
self.encoder_filters, self.encoder_dim, self.details, self.upscale_ratio)
|
|
|
|
def build_model(self, inputs):
|
|
""" Build the Dlight Model. """
|
|
encoder = self.encoder()
|
|
encoder_a = encoder(inputs[0])
|
|
encoder_b = encoder(inputs[1])
|
|
|
|
decoder_b = self.decoder_b if self.details > 0 else self.decoder_b_fast
|
|
|
|
outputs = self.decoder_a()(encoder_a) + decoder_b()(encoder_b)
|
|
|
|
autoencoder = KModel(inputs, outputs, name=self.model_name)
|
|
return autoencoder
|
|
|
|
def encoder(self):
|
|
""" DeLight Encoder Network """
|
|
input_ = Input(shape=self.input_shape)
|
|
var_x = input_
|
|
|
|
var_x1 = Conv2DBlock(self.encoder_filters // 2, activation="leakyrelu")(var_x)
|
|
var_x2 = layers.AveragePooling2D(pool_size=(2, 2))(var_x)
|
|
var_x2 = layers.LeakyReLU(0.1)(var_x2)
|
|
var_x = layers.Concatenate()([var_x1, var_x2])
|
|
|
|
var_x1 = Conv2DBlock(self.encoder_filters, activation="leakyrelu")(var_x)
|
|
var_x2 = layers.AveragePooling2D(pool_size=(2, 2))(var_x)
|
|
var_x2 = layers.LeakyReLU(0.1)(var_x2)
|
|
var_x = layers.Concatenate()([var_x1, var_x2])
|
|
|
|
var_x1 = Conv2DBlock(self.encoder_filters * 2, activation="leakyrelu")(var_x)
|
|
var_x2 = layers.AveragePooling2D(pool_size=(2, 2))(var_x)
|
|
var_x2 = layers.LeakyReLU(0.1)(var_x2)
|
|
var_x = layers.Concatenate()([var_x1, var_x2])
|
|
|
|
var_x1 = Conv2DBlock(self.encoder_filters * 4, activation="leakyrelu")(var_x)
|
|
var_x2 = layers.AveragePooling2D(pool_size=(2, 2))(var_x)
|
|
var_x2 = layers.LeakyReLU(0.1)(var_x2)
|
|
var_x = layers.Concatenate()([var_x1, var_x2])
|
|
|
|
var_x1 = Conv2DBlock(self.encoder_filters * 8, activation="leakyrelu")(var_x)
|
|
var_x2 = layers.AveragePooling2D(pool_size=(2, 2))(var_x)
|
|
var_x2 = layers.LeakyReLU(0.1)(var_x2)
|
|
var_x = layers.Concatenate()([var_x1, var_x2])
|
|
|
|
var_x = layers.Dense(self.encoder_dim)(layers.Flatten()(var_x))
|
|
var_x = layers.Dropout(0.05)(var_x)
|
|
var_x = layers.Dense(4 * 4 * 1024)(var_x)
|
|
var_x = layers.Dropout(0.05)(var_x)
|
|
var_x = layers.Reshape((4, 4, 1024))(var_x)
|
|
|
|
return KModel(input_, var_x, name="encoder")
|
|
|
|
def decoder_a(self):
|
|
""" DeLight Decoder A(old face) Network """
|
|
input_ = Input(shape=(4, 4, 1024))
|
|
dec_a_complexity = 256
|
|
mask_complexity = 128
|
|
|
|
var_xy = input_
|
|
var_xy = layers.UpSampling2D(self.upscale_ratio, interpolation='bilinear')(var_xy)
|
|
|
|
var_x = var_xy
|
|
var_x = Upscale2xBlock(dec_a_complexity, activation="leakyrelu", fast=False)(var_x)
|
|
var_x = Upscale2xBlock(dec_a_complexity // 2, activation="leakyrelu", fast=False)(var_x)
|
|
var_x = Upscale2xBlock(dec_a_complexity // 4, activation="leakyrelu", fast=False)(var_x)
|
|
var_x = Upscale2xBlock(dec_a_complexity // 8, activation="leakyrelu", fast=False)(var_x)
|
|
|
|
var_x = Conv2DOutput(3, 5, name="face_out")(var_x)
|
|
|
|
outputs = [var_x]
|
|
|
|
if cfg_loss.learn_mask():
|
|
var_y = var_xy # mask decoder
|
|
var_y = Upscale2xBlock(mask_complexity, activation="leakyrelu", fast=False)(var_y)
|
|
var_y = Upscale2xBlock(mask_complexity // 2, activation="leakyrelu", fast=False)(var_y)
|
|
var_y = Upscale2xBlock(mask_complexity // 4, activation="leakyrelu", fast=False)(var_y)
|
|
var_y = Upscale2xBlock(mask_complexity // 8, activation="leakyrelu", fast=False)(var_y)
|
|
|
|
var_y = Conv2DOutput(1, 5, name="mask_out")(var_y)
|
|
|
|
outputs.append(var_y)
|
|
|
|
return KModel([input_], outputs=outputs, name="decoder_a")
|
|
|
|
def decoder_b_fast(self):
|
|
""" DeLight Fast Decoder B(new face) Network """
|
|
input_ = Input(shape=(4, 4, 1024))
|
|
|
|
dec_b_complexity = 512
|
|
mask_complexity = 128
|
|
|
|
var_xy = input_
|
|
|
|
var_xy = UpscaleBlock(512, scale_factor=self.upscale_ratio, activation="leakyrelu")(var_xy)
|
|
var_x = var_xy
|
|
|
|
var_x = Upscale2xBlock(dec_b_complexity, activation="leakyrelu", fast=True)(var_x)
|
|
var_x = Upscale2xBlock(dec_b_complexity // 2, activation="leakyrelu", fast=True)(var_x)
|
|
var_x = Upscale2xBlock(dec_b_complexity // 4, activation="leakyrelu", fast=True)(var_x)
|
|
var_x = Upscale2xBlock(dec_b_complexity // 8, activation="leakyrelu", fast=True)(var_x)
|
|
|
|
var_x = Conv2DOutput(3, 5, name="face_out")(var_x)
|
|
|
|
outputs = [var_x]
|
|
|
|
if cfg_loss.learn_mask():
|
|
var_y = var_xy # mask decoder
|
|
|
|
var_y = Upscale2xBlock(mask_complexity, activation="leakyrelu", fast=False)(var_y)
|
|
var_y = Upscale2xBlock(mask_complexity // 2, activation="leakyrelu", fast=False)(var_y)
|
|
var_y = Upscale2xBlock(mask_complexity // 4, activation="leakyrelu", fast=False)(var_y)
|
|
var_y = Upscale2xBlock(mask_complexity // 8, activation="leakyrelu", fast=False)(var_y)
|
|
|
|
var_y = Conv2DOutput(1, 5, name="mask_out")(var_y)
|
|
|
|
outputs.append(var_y)
|
|
|
|
return KModel([input_], outputs=outputs, name="decoder_b_fast")
|
|
|
|
def decoder_b(self):
|
|
""" DeLight Decoder B(new face) Network """
|
|
input_ = Input(shape=(4, 4, 1024))
|
|
|
|
dec_b_complexity = 512
|
|
mask_complexity = 128
|
|
|
|
var_xy = input_
|
|
|
|
var_xy = Upscale2xBlock(512,
|
|
scale_factor=self.upscale_ratio,
|
|
activation=None,
|
|
fast=False)(var_xy)
|
|
var_x = var_xy
|
|
|
|
var_x = layers.LeakyReLU(negative_slope=0.2)(var_x)
|
|
var_x = ResidualBlock(512, use_bias=True)(var_x)
|
|
var_x = ResidualBlock(512, use_bias=False)(var_x)
|
|
var_x = ResidualBlock(512, use_bias=False)(var_x)
|
|
var_x = Upscale2xBlock(dec_b_complexity, activation=None, fast=False)(var_x)
|
|
var_x = layers.LeakyReLU(negative_slope=0.2)(var_x)
|
|
var_x = ResidualBlock(dec_b_complexity, use_bias=True)(var_x)
|
|
var_x = ResidualBlock(dec_b_complexity, use_bias=False)(var_x)
|
|
var_x = layers.BatchNormalization()(var_x)
|
|
var_x = Upscale2xBlock(dec_b_complexity // 2, activation=None, fast=False)(var_x)
|
|
var_x = layers.LeakyReLU(negative_slope=0.2)(var_x)
|
|
var_x = ResidualBlock(dec_b_complexity // 2, use_bias=True)(var_x)
|
|
var_x = Upscale2xBlock(dec_b_complexity // 4, activation=None, fast=False)(var_x)
|
|
var_x = layers.LeakyReLU(negative_slope=0.2)(var_x)
|
|
var_x = ResidualBlock(dec_b_complexity // 4, use_bias=False)(var_x)
|
|
var_x = layers.BatchNormalization()(var_x)
|
|
var_x = Upscale2xBlock(dec_b_complexity // 8, activation="leakyrelu", fast=False)(var_x)
|
|
|
|
var_x = Conv2DOutput(3, 5, name="face_out")(var_x)
|
|
|
|
outputs = [var_x]
|
|
|
|
if cfg_loss.learn_mask():
|
|
var_y = var_xy # mask decoder
|
|
var_y = layers.LeakyReLU(negative_slope=0.1)(var_y)
|
|
|
|
var_y = Upscale2xBlock(mask_complexity, activation="leakyrelu", fast=False)(var_y)
|
|
var_y = Upscale2xBlock(mask_complexity // 2, activation="leakyrelu", fast=False)(var_y)
|
|
var_y = Upscale2xBlock(mask_complexity // 4, activation="leakyrelu", fast=False)(var_y)
|
|
var_y = Upscale2xBlock(mask_complexity // 8, activation="leakyrelu", fast=False)(var_y)
|
|
|
|
var_y = Conv2DOutput(1, 5, name="mask_out")(var_y)
|
|
|
|
outputs.append(var_y)
|
|
|
|
return KModel([input_], outputs=outputs, name="decoder_b")
|