CalvinBackup/facefusion-labs
1
0
Fork 0
mirror of https://github.com/facefusion/facefusion-labs.git synced 2026-06-25 07:59:55 +02:00
Code Issues Packages Projects Releases Wiki Activity

Make Embedding great again

Browse Source
This commit is contained in:
henryruhs
2025-02-12 10:31:53 +01:00
parent 6381e755d7
commit 611618e413
8 changed files with 22 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
arcface_converter/src/preparing.py
+6 -6
View File
@@ -9,7 +9,7 @@ from mxnet.io import ImageRecordIter
from onnxruntime import InferenceSession
from tqdm import tqdm
from .types import Embedding, EmbeddingPairs, VisionFrame
from .types import Embedding, VisionFrame
CONFIG = configparser.ConfigParser()
CONFIG.read('config.ini')
@@ -35,9 +35,9 @@ def forward(inference_session : InferenceSession, crop_vision_frame : VisionFram
return embedding
def process_embeddings(dataset_reader : ImageRecordIter, source_inference_session : InferenceSession, target_inference_session : InferenceSession) -> EmbeddingPairs:
def process_embeddings(dataset_reader : ImageRecordIter, source_inference_session : InferenceSession, target_inference_session : InferenceSession) -> Embedding:
dataset_process_limit = CONFIG.getint('preparing.dataset', 'process_limit')
embedding_pairs = []
embeddings = []
with tqdm(total = dataset_process_limit) as progress:
for batch in dataset_reader:
@@ -45,13 +45,13 @@ def process_embeddings(dataset_reader : ImageRecordIter, source_inference_sessio
crop_vision_frame = prepare_crop_vision_frame(crop_vision_frame)
source_embedding = forward(source_inference_session, crop_vision_frame)
target_embedding = forward(target_inference_session, crop_vision_frame)
embedding_pairs.append([ source_embedding, target_embedding ])
embeddings.append([ source_embedding, target_embedding ])
progress.update()
if progress.n == dataset_process_limit:
return numpy.concatenate(embedding_pairs, axis = 1).T
return numpy.concatenate(embeddings, axis = 1).T
return numpy.concatenate(embedding_pairs, axis = 1).T
return numpy.concatenate(embeddings, axis = 1).T
def prepare() -> None:
arcface_converter/src/types.py
+1 -1
View File
@@ -8,6 +8,6 @@ Batch = Tuple[Tensor, Tensor]
Loader = DataLoader[Tuple[Tensor, ...]]
Embedding = NDArray[Any]
EmbeddingPairs = NDArray[Any]
FaceLandmark5 = NDArray[Any]
VisionFrame = NDArray[Any]
face_swapper/src/helper.py
+2 -2
View File
@@ -4,7 +4,7 @@ import cv2
import numpy
import torch
from .types import IdEmbedder, IdEmbedding, Padding, Tensor, VisionFrame, VisionTensor
from .types import Embedder, Embedding, Padding, Tensor, VisionFrame, VisionTensor
def is_windows() -> bool:
@@ -47,7 +47,7 @@ def hinge_fake_loss(tensor : Tensor) -> Tensor:
return fake_loss
def calc_id_embedding(id_embedder : IdEmbedder, vision_tensor : VisionTensor, padding : Padding) -> IdEmbedding:
def calc_id_embedding(id_embedder : Embedder, vision_tensor : VisionTensor, padding : Padding) -> Embedding:
crop_vision_tensor = vision_tensor[:, :, 15 : 241, 15 : 241]
crop_vision_tensor = torch.nn.functional.interpolate(crop_vision_tensor, size = (112, 112), mode = 'area')
crop_vision_tensor[:, :, :padding[0], :] = 0
face_swapper/src/inferencing.py
+2 -2
View File
@@ -5,13 +5,13 @@ import torch
from .helper import calc_id_embedding, convert_to_vision_frame, convert_to_vision_tensor, read_image
from .models.generator import AdaptiveEmbeddingIntegrationNetwork
from .types import Generator, IdEmbedder, VisionFrame
from .types import Generator, Embedder, VisionFrame
CONFIG = configparser.ConfigParser()
CONFIG.read('config.ini')
def run_swap(generator : Generator, id_embedder : IdEmbedder, source_vision_frame : VisionFrame, target_vision_frame : VisionFrame) -> VisionFrame:
def run_swap(generator : Generator, id_embedder : Embedder, source_vision_frame : VisionFrame, target_vision_frame : VisionFrame) -> VisionFrame:
source_vision_tensor = convert_to_vision_tensor(source_vision_frame)
target_vision_tensor = convert_to_vision_tensor(target_vision_frame)
source_embedding = calc_id_embedding(id_embedder, source_vision_tensor, (0, 0, 0, 0))
face_swapper/src/models/generator.py
+2 -2
View File
@@ -5,7 +5,7 @@ import torch.nn as nn
from face_swapper.src.networks.attribute_modulator import AADGenerator
from face_swapper.src.networks.encoder import UNet
from face_swapper.src.types import SourceEmbedding, TargetAttributes, VisionTensor
from face_swapper.src.types import Embedding, TargetAttributes, VisionTensor
CONFIG = configparser.ConfigParser()
CONFIG.read('config.ini')
@@ -22,7 +22,7 @@ class AdaptiveEmbeddingIntegrationNetwork(nn.Module):
self.encoder.apply(init_weight)
self.generator.apply(init_weight)
def forward(self, target : VisionTensor, source_embedding : SourceEmbedding) -> Tuple[VisionTensor, TargetAttributes]:
def forward(self, target : VisionTensor, source_embedding : Embedding) -> Tuple[VisionTensor, TargetAttributes]:
target_attributes = self.get_attributes(target)
swap_tensor = self.generator(target_attributes, source_embedding)
return swap_tensor, target_attributes
face_swapper/src/networks/attribute_modulator.py
+5 -5
View File
@@ -1,7 +1,7 @@
import torch
from torch import Tensor, nn as nn
from face_swapper.src.types import SourceEmbedding, TargetAttributes
from face_swapper.src.types import Embedding, TargetAttributes
class AADGenerator(nn.Module):
@@ -17,7 +17,7 @@ class AADGenerator(nn.Module):
self.res_block_7 = AADResBlock(128, 64, 64, id_channels, num_blocks)
self.res_block_8 = AADResBlock(64, 3, 64, id_channels, num_blocks)
def forward(self, target_attributes : TargetAttributes, source_embedding : SourceEmbedding) -> Tensor:
def forward(self, target_attributes : TargetAttributes, source_embedding : Embedding) -> Tensor:
feature_map = self.upsample(source_embedding)
feature_map_1 = torch.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 = torch.nn.functional.interpolate(self.res_block_2(feature_map_1, target_attributes[1], source_embedding), scale_factor = 2, mode = 'bilinear', align_corners = False)
@@ -41,7 +41,7 @@ class AADLayer(nn.Module):
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 : Tensor, id_embedding : SourceEmbedding) -> Tensor:
def forward(self, feature_map : Tensor, attribute_embedding : Tensor, id_embedding : Embedding) -> Tensor:
feature_map = self.instance_norm(feature_map)
gamma_attribute = self.conv_gamma(attribute_embedding)
beta_attribute = self.conv_beta(attribute_embedding)
@@ -59,7 +59,7 @@ class AADSequential(nn.Module):
super(AADSequential, self).__init__()
self.layers = nn.ModuleList(args)
def forward(self, feature_map: Tensor, attribute_embedding: Tensor, id_embedding: SourceEmbedding) -> Tensor:
def forward(self, feature_map: Tensor, attribute_embedding: Tensor, id_embedding: Embedding) -> Tensor:
for layer in self.layers:
if isinstance(layer, AADLayer):
feature_map = layer(feature_map, attribute_embedding, id_embedding)
@@ -99,7 +99,7 @@ class AADResBlock(nn.Module):
)
self.auxiliary_add_blocks = auxiliary_add_blocks
def forward(self, feature_map : Tensor, attribute_embedding : Tensor, id_embedding : SourceEmbedding) -> Tensor:
def forward(self, feature_map : Tensor, attribute_embedding : Tensor, id_embedding : Embedding) -> Tensor:
primary_feature = self.primary_add_blocks(feature_map, attribute_embedding, id_embedding)
if self.input_channels > self.output_channels:
face_swapper/src/training.py
+2 -2
View File
@@ -16,7 +16,7 @@ from .helper import calc_id_embedding
from .models.discriminator import MultiscaleDiscriminator
from .models.generator import AdaptiveEmbeddingIntegrationNetwork
from .models.loss import FaceSwapperLoss
from .types import Batch, SourceEmbedding, TargetAttributes, VisionTensor
from .types import Batch, Embedding, TargetAttributes, VisionTensor
CONFIG = configparser.ConfigParser()
CONFIG.read('config.ini')
@@ -29,7 +29,7 @@ class FaceSwapperTrain(pytorch_lightning.LightningModule, FaceSwapperLoss):
self.discriminator = MultiscaleDiscriminator()
self.automatic_optimization = CONFIG.getboolean('training.trainer', 'automatic_optimization')
def forward(self, target_tensor : VisionTensor, source_embedding : SourceEmbedding) -> Tuple[VisionTensor, TargetAttributes]:
def forward(self, target_tensor : VisionTensor, source_embedding : Embedding) -> Tuple[VisionTensor, TargetAttributes]:
output = self.generator(target_tensor, source_embedding)
return output
face_swapper/src/types.py
+2 -3
View File
@@ -15,8 +15,7 @@ SwapAttributes = Tuple[Tensor, ...]
TargetAttributes = Tuple[Tensor, ...]
DiscriminatorOutputs = List[List[Tensor]]
IdEmbedding = Tensor
SourceEmbedding = IdEmbedding
Embedding = Tensor
FaceLandmark203 = Tensor
StateSet = OrderedDict[str, Any]
@@ -30,4 +29,4 @@ GeneratorLossSet = Dict[str, Tensor]
DiscriminatorLossSet = Dict[str, Tensor]
Generator = torch.nn.Module
IdEmbedder = torch.nn.Module
Embedder = torch.nn.Module