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

changes

Browse Source
This commit is contained in:
harisreedhar
2025-03-23 18:34:23 +05:30
parent c1f39a73dd
commit c85c755e00
2 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
face_swapper/src/models/generator.py
+2 -3
View File
@@ -19,13 +19,12 @@ class Generator(nn.Module):
self.generator.apply(init_weight)
self.masker.apply(init_weight)
def forward(self, source_embedding : Embedding, target_tensor : Tensor) -> Tuple[Tensor, Mask]:
target_features = self.encode_features(target_tensor)
def forward(self, source_embedding : Embedding, target_tensor : Tensor, target_features : Tuple[Feature, ...]) -> Tuple[Tensor, Mask, ]:
output_tensor = self.generator(source_embedding, target_features)
target_feature = target_features[-1]
output_mask = self.masker(target_tensor, target_feature)
output_tensor = output_tensor * output_mask + target_tensor * (1 - output_mask)
return output_tensor, output_mask
return output_tensor, output_mask,
def encode_features(self, input_tensor : Tensor) -> Tuple[Feature, ...]:
return self.encoder(input_tensor)
face_swapper/src/training.py
+4 -3
View File
@@ -53,7 +53,8 @@ class FaceSwapperTrainer(LightningModule):
def forward(self, source_embedding : Embedding, target_tensor : Tensor) -> Tuple[Tensor, Mask]:
with torch.no_grad():
output_tensor, output_mask = self.generator(source_embedding, target_tensor)
generator_target_features = self.generator.encode_features(target_tensor)
output_tensor, output_mask = self.generator(source_embedding, target_tensor, generator_target_features)
return output_tensor, output_mask
@@ -89,8 +90,8 @@ class FaceSwapperTrainer(LightningModule):
generator_optimizer, discriminator_optimizer = self.optimizers() #type:ignore[attr-defined]
source_embedding = calc_embedding(self.embedder, source_tensor, (0, 0, 0, 0))
generator_output_tensor, generator_output_mask = self.generator(source_embedding, target_tensor)
generator_target_features = self.generator.encode_features(target_tensor)
generator_output_tensor, generator_output_mask = self.generator(source_embedding, target_tensor, generator_target_features)
generator_output_features = self.generator.encode_features(generator_output_tensor)
discriminator_output_tensors = self.discriminator(generator_output_tensor)
adversarial_loss, weighted_adversarial_loss = self.adversarial_loss(discriminator_output_tensors)
@@ -138,7 +139,7 @@ class FaceSwapperTrainer(LightningModule):
def validation_step(self, batch : Batch, batch_index : int) -> Tensor:
source_tensor, target_tensor = batch
source_embedding = calc_embedding(self.embedder, source_tensor, (0, 0, 0, 0))
output_tensor, _ = self.generator(source_embedding, target_tensor)
output_tensor, _ = self(source_embedding, target_tensor)
output_embedding = calc_embedding(self.embedder, output_tensor, (0, 0, 0, 0))
validation_score = (nn.functional.cosine_similarity(source_embedding, output_embedding).mean() + 1) * 0.5
self.log('validation_score', validation_score, sync_dist = True, prog_bar = True)