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

Improve generator namings, Flip args to source then target

Browse Source
This commit is contained in:
henryruhs
2025-02-18 11:00:30 +01:00
parent f63bc788ac
commit bf696be097
5 changed files with 24 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
face_swapper/src/training.py
+5 -4
View File
@@ -15,7 +15,7 @@ from torch.utils.data import DataLoader
from .data_loader import DataLoaderVGG
from .helper import calc_id_embedding
from .models.discriminator import Discriminator
from .models.generator import AdaptiveEmbeddingIntegrationNetwork
from .models.generator import Generator
from .models.loss import FaceSwapperLoss
from .types import Batch, Embedding, TargetAttributes, VisionTensor
@@ -27,12 +27,12 @@ class FaceSwapperTrainer(lightning.LightningModule, FaceSwapperLoss):
def __init__(self) -> None:
super().__init__()
FaceSwapperLoss.__init__(self)
self.generator = AdaptiveEmbeddingIntegrationNetwork()
self.generator = Generator()
self.discriminator = Discriminator()
self.automatic_optimization = CONFIG.getboolean('training.trainer', 'automatic_optimization')
def forward(self, target_tensor : VisionTensor, source_embedding : Embedding) -> Tuple[VisionTensor, TargetAttributes]:
output = self.generator(target_tensor, source_embedding)
output = self.generator(source_embedding, target_tensor)
return output
def configure_optimizers(self) -> Tuple[Optimizer, Optimizer]:
@@ -45,7 +45,8 @@ class FaceSwapperTrainer(lightning.LightningModule, FaceSwapperLoss):
source_tensor, target_tensor, is_same_person = batch
generator_optimizer, discriminator_optimizer = self.optimizers() #type:ignore[attr-defined]
source_embedding = calc_id_embedding(self.id_embedder, source_tensor, (0, 0, 0, 0))
swap_tensor, target_attributes = self.generator(target_tensor, source_embedding)
swap_tensor = self.generator(source_embedding, target_tensor)
target_attributes = self.generator.get_attributes(target_tensor)
swap_attributes = self.generator.get_attributes(swap_tensor)
fake_discriminator_outputs = self.discriminator(swap_tensor)