CalvinBackup/facefusion-labs
1
0
Fork 0
mirror of https://github.com/facefusion/facefusion-labs.git synced 2026-08-31 00:10:39 +02:00
Code Issues Packages Projects Releases Wiki Activity

Rename ArcFace Converter to Embedding Converter, Add EmbeddingDataset, Add learning rate to config

Browse Source
This commit is contained in:
henryruhs
2025-03-11 14:43:06 +01:00
parent 62a69cddd2
commit 1b6e7a6ca5
23 changed files with 70 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
face_swapper/src/models/loss.py
+7 -6
View File
@@ -34,12 +34,13 @@ class FaceSwapperLoss:
weight_reconstruction = CONFIG.getfloat('training.losses', 'weight_reconstruction')
weight_pose = CONFIG.getfloat('training.losses', 'weight_pose')
weight_gaze = CONFIG.getfloat('training.losses', 'weight_gaze')
generator_loss_set = {}
generator_loss_set['loss_adversarial'] = self.calc_adversarial_loss(discriminator_outputs)
generator_loss_set['loss_id'] = self.calc_id_loss(source_tensor, swap_tensor)
generator_loss_set['loss_attribute'] = self.calc_attribute_loss(target_attributes, swap_attributes)
generator_loss_set['loss_reconstruction'] = self.calc_reconstruction_loss(swap_tensor, target_tensor, is_same_person)
generator_loss_set =\
{
'loss_adversarial': self.calc_adversarial_loss(discriminator_outputs),
'loss_id': self.calc_id_loss(source_tensor, swap_tensor),
'loss_attribute': self.calc_attribute_loss(target_attributes, swap_attributes),
'loss_reconstruction': self.calc_reconstruction_loss(swap_tensor, target_tensor, is_same_person)
}
if weight_pose > 0:
generator_loss_set['loss_pose'] = self.calc_pose_loss(swap_tensor, target_tensor)
face_swapper/src/training.py
+9 -9
View File
@@ -61,12 +61,12 @@ class FaceSwapperTrain(pytorch_lightning.LightningModule, FaceSwapperLoss):
if self.global_step % CONFIG.getint('training.output', 'preview_frequency') == 0:
self.generate_preview(source_tensor, target_tensor, swap_tensor)
self.log('l_G', generator_losses.get('loss_generator'), prog_bar = True)
self.log('l_D', discriminator_losses.get('loss_discriminator'), prog_bar = True)
self.log('l_ADV', generator_losses.get('loss_adversarial'), prog_bar = True)
self.log('l_ATTR', generator_losses.get('loss_attribute'), prog_bar = True)
self.log('l_ID', generator_losses.get('loss_id'), prog_bar = True)
self.log('l_REC', generator_losses.get('loss_reconstruction'), prog_bar = True)
self.log('loss_generator', generator_losses.get('loss_generator'), prog_bar = True)
self.log('loss_discriminator', discriminator_losses.get('loss_discriminator'), prog_bar = True)
self.log('loss_adversarial', generator_losses.get('loss_adversarial'), prog_bar = True)
self.log('loss_attribute', generator_losses.get('loss_attribute'), prog_bar = True)
self.log('loss_id', generator_losses.get('loss_id'), prog_bar = True)
self.log('loss_reconstruction', generator_losses.get('loss_reconstruction'), prog_bar = True)
return generator_losses.get('loss_generator')
def generate_preview(self, source_tensor : VisionTensor, target_tensor : VisionTensor, swap_tensor : VisionTensor) -> None:
@@ -76,7 +76,7 @@ class FaceSwapperTrain(pytorch_lightning.LightningModule, FaceSwapperLoss):
swap_tensors = swap_tensor[:max_preview]
rows = [ torch.cat([ source_tensor, target_tensor, swap_tensor ], dim = 2) for source_tensor, target_tensor, swap_tensor in zip(source_tensors, target_tensors, swap_tensors) ]
grid = torchvision.utils.make_grid(torch.cat(rows, dim = 1).unsqueeze(0), nrow = 1, normalize = True, scale_each = True)
self.logger.experiment.add_image("Generator Preview", grid, self.global_step)
self.logger.experiment.add_image('preview', grid, self.global_step)
def create_trainer() -> Trainer:
@@ -111,10 +111,10 @@ def train() -> None:
same_person_probability = CONFIG.getfloat('preparing.dataset', 'same_person_probability')
batch_size = CONFIG.getint('training.loader', 'batch_size')
num_workers = CONFIG.getint('training.loader', 'num_workers')
file_path = CONFIG.get('training.output', 'file_path')
output_file_path = CONFIG.get('training.output', 'file_path')
dataset = DataLoaderVGG(dataset_path, dataset_image_pattern, dataset_directory_pattern, same_person_probability)
data_loader = DataLoader(dataset, batch_size = batch_size, shuffle = True, num_workers = num_workers, drop_last = True, pin_memory = True, persistent_workers = True)
face_swap_model = FaceSwapperTrain()
trainer = create_trainer()
trainer.fit(face_swap_model, data_loader, ckpt_path = file_path)
trainer.fit(face_swap_model, data_loader, ckpt_path = output_file_path)