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

remove cycle & gaze loss

Browse Source
This commit is contained in:
harisreedhar
2026-08-27 10:00:38 +05:30
parent 6547fcfe7b
commit 0a41097eef
5 changed files with 3 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hyperswap/README.md
-3
View File
@@ -47,7 +47,6 @@ split_ratio = 0.9995
[training.model]
generator_embedder_path = .models/blendface.pt
loss_embedder_path = .models/arcface.pt
gazer_path = .models/gazer.pt
face_masker_path = .models/face_masker.pt
```
@@ -77,11 +76,9 @@ num_filters = 16
```
[training.losses]
adversarial_weight = 1.0
cycle_weight = 1.0
feature_weight = 10.0
reconstruction_weight = 10.0
identity_weight = 20.0
gaze_weight = 0.05
mask_weight = 5.0
```
hyperswap/config.ini
-3
View File
@@ -15,7 +15,6 @@ split_ratio =
[training.model]
generator_embedder_path =
loss_embedder_path =
gazer_path =
face_masker_path =
[training.model.generator]
@@ -37,11 +36,9 @@ num_filters =
[training.losses]
adversarial_weight =
cycle_weight =
feature_weight =
reconstruction_weight =
identity_weight =
gaze_weight =
mask_weight =
[training.trainer]
hyperswap/src/models/loss.py
+1 -55
View File
@@ -4,10 +4,9 @@ from typing import List, Tuple
import torch
from pytorch_msssim import ssim
from torch import Tensor, nn
from torchvision import transforms
from ..helper import calculate_face_embedding, dilate_mask
from ..types import EmbedderModule, FaceMaskerModule, Feature, GazerModule, Loss, Mask
from ..types import EmbedderModule, FaceMaskerModule, Feature, Loss, Mask
class DiscriminatorLoss(nn.Module):
@@ -49,27 +48,6 @@ class AdversarialLoss(nn.Module):
return adversarial_loss, weighted_adversarial_loss
class CycleLoss(nn.Module):
def __init__(self, config_parser : ConfigParser) -> None:
super().__init__()
self.config_batch_size = config_parser.getint('training.loader', 'batch_size')
self.config_cycle_weight = config_parser.getfloat('training.losses', 'cycle_weight')
self.l1_loss = nn.L1Loss()
def forward(self, target_tensor : Tensor, cycle_tensor : Tensor, target_features : Tuple[Feature, ...], cycle_features : Tuple[Feature, ...]) -> Tuple[Loss, Loss]:
temp_tensors = []
for target_feature, output_feature in zip(target_features, cycle_features):
temp_tensor = torch.mean(torch.pow(output_feature - target_feature, 2).reshape(self.config_batch_size, -1), dim = 1).mean()
temp_tensors.append(temp_tensor)
feature_loss = torch.stack(temp_tensors).mean()
reconstruction_loss = self.l1_loss(target_tensor, cycle_tensor)
cycle_loss = (feature_loss + reconstruction_loss) * 0.5
weighted_feature_loss = cycle_loss * self.config_cycle_weight
return cycle_loss, weighted_feature_loss
class FeatureLoss(nn.Module):
def __init__(self, config_parser : ConfigParser) -> None:
super().__init__()
@@ -127,38 +105,6 @@ class IdentityLoss(nn.Module):
return identity_loss, weighted_identity_loss
class GazeLoss(nn.Module):
def __init__(self, config_parser : ConfigParser, gazer : GazerModule) -> None:
super().__init__()
self.config_gaze_weight = config_parser.getfloat('training.losses', 'gaze_weight')
self.config_output_size = config_parser.getint('training.model.generator', 'output_size')
self.gazer = gazer
self.l1_loss = nn.L1Loss()
def forward(self, target_tensor : Tensor, output_tensor : Tensor) -> Tuple[Loss, Loss]:
output_pitch, output_yaw = self.detect_gaze(output_tensor)
target_pitch, target_yaw = self.detect_gaze(target_tensor)
pitch_loss = self.l1_loss(output_pitch, target_pitch)
yaw_loss = self.l1_loss(output_yaw, target_yaw)
gaze_loss = (pitch_loss + yaw_loss) * 0.5
weighted_gaze_loss = gaze_loss * self.config_gaze_weight
return gaze_loss, weighted_gaze_loss
def detect_gaze(self, input_tensor : Tensor) -> Tuple[Tensor, Tensor]:
crop_sizes = (torch.tensor([ 0.235, 0.875, 0.0625, 0.8 ]) * self.config_output_size).int()
crop_tensor = input_tensor[:, :, crop_sizes[0]:crop_sizes[1], crop_sizes[2]:crop_sizes[3]]
crop_tensor = (crop_tensor + 1) * 0.5
crop_tensor = transforms.Normalize(mean = [ 0.485, 0.456, 0.406 ], std = [ 0.229, 0.224, 0.225 ])(crop_tensor)
crop_tensor = nn.functional.interpolate(crop_tensor, size = 448, mode = 'bicubic')
with torch.no_grad():
pitch, yaw = self.gazer(crop_tensor)
return pitch, yaw
class MaskLoss(nn.Module):
def __init__(self, config_parser : ConfigParser, face_masker : FaceMaskerModule) -> None:
super().__init__()
hyperswap/src/training.py
+2 -13
View File
@@ -19,7 +19,7 @@ from .dataset import DynamicDataset
from .helper import apply_noise, calculate_face_embedding, erode_mask, overlay_mask
from .models.discriminator import Discriminator
from .models.generator import Generator
from .models.loss import AdversarialLoss, CycleLoss, DiscriminatorLoss, FeatureLoss, GazeLoss, IdentityLoss, MaskLoss, ReconstructionLoss
from .models.loss import AdversarialLoss, DiscriminatorLoss, FeatureLoss, IdentityLoss, MaskLoss, ReconstructionLoss
from .types import Batch, Embedding, Mask, OptimizerSet, TrainerPrecision, TrainerStrategy
warnings.filterwarnings('ignore', category = UserWarning, module = 'torch')
@@ -33,7 +33,6 @@ class HyperSwapTrainer(LightningModule):
super().__init__()
self.config_generator_embedder_path = config_parser.get('training.model', 'generator_embedder_path')
self.config_loss_embedder_path = config_parser.get('training.model', 'loss_embedder_path')
self.config_gazer_path = config_parser.get('training.model', 'gazer_path')
self.config_face_masker_path = config_parser.get('training.model', 'face_masker_path')
self.config_accumulate_size = config_parser.getfloat('training.trainer', 'accumulate_size')
self.config_discriminator_ratio = config_parser.getfloat('training.trainer', 'discriminator_ratio')
@@ -51,17 +50,14 @@ class HyperSwapTrainer(LightningModule):
self.config_discriminator_scheduler_patience = config_parser.getint('training.optimizer.discriminator', 'scheduler_patience')
self.generator_embedder = torch.jit.load(self.config_generator_embedder_path, map_location = 'cpu').eval()
self.loss_embedder = torch.jit.load(self.config_loss_embedder_path, map_location = 'cpu').eval()
self.gazer = torch.jit.load(self.config_gazer_path, map_location = 'cpu').eval()
self.face_masker = torch.jit.load(self.config_face_masker_path, map_location ='cpu').eval()
self.generator = Generator(config_parser)
self.discriminator = Discriminator(config_parser)
self.discriminator_loss = DiscriminatorLoss()
self.adversarial_loss = AdversarialLoss(config_parser)
self.cycle_loss = CycleLoss(config_parser)
self.feature_loss = FeatureLoss(config_parser)
self.reconstruction_loss = ReconstructionLoss(config_parser, self.loss_embedder)
self.identity_loss = IdentityLoss(config_parser, self.loss_embedder)
self.gaze_loss = GazeLoss(config_parser, self.gazer)
self.mask_loss = MaskLoss(config_parser, self.face_masker)
self.automatic_optimization = False
@@ -105,7 +101,6 @@ class HyperSwapTrainer(LightningModule):
generator_optimizer, discriminator_optimizer = self.optimizers() #type:ignore[attr-defined]
generator_scheduler, discriminator_scheduler = self.lr_schedulers() #type:ignore[attr-defined]
source_embedding = calculate_face_embedding(self.generator_embedder, source_tensor, (0, 0, 0, 0))
target_embedding = calculate_face_embedding(self.generator_embedder, target_tensor, (0, 0, 0, 0))
if self.config_noise_factor > 0:
source_embedding = apply_noise(source_embedding, self.config_noise_factor)
@@ -114,17 +109,13 @@ class HyperSwapTrainer(LightningModule):
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)
cycle_output_tensor, cycle_output_mask = self.generator(target_embedding, generator_output_tensor, generator_output_features)
cycle_output_features = self.generator.encode_features(cycle_output_tensor)
discriminator_output_tensors = self.discriminator(generator_output_tensor)
adversarial_loss, weighted_adversarial_loss = self.adversarial_loss(discriminator_output_tensors)
cycle_loss, weighted_cycle_loss = self.cycle_loss(target_tensor, cycle_output_tensor, generator_target_features, cycle_output_features)
feature_loss, weighted_feature_loss = self.feature_loss(generator_target_features, generator_output_features)
reconstruction_loss, weighted_reconstruction_loss = self.reconstruction_loss(source_tensor, target_tensor, generator_output_tensor)
identity_loss, weighted_identity_loss = self.identity_loss(generator_output_tensor, source_tensor)
gaze_loss, weighted_gaze_loss = self.gaze_loss(target_tensor, generator_output_tensor)
mask_loss, weighted_mask_loss = self.mask_loss(target_tensor, generator_output_mask)
generator_loss = weighted_adversarial_loss + weighted_cycle_loss + weighted_feature_loss + weighted_reconstruction_loss + weighted_identity_loss + weighted_gaze_loss + weighted_mask_loss
generator_loss = weighted_adversarial_loss + weighted_feature_loss + weighted_reconstruction_loss + weighted_identity_loss + weighted_mask_loss
if torch.randn(1).item() < self.config_discriminator_ratio:
discriminator_real_tensors = self.discriminator(source_tensor)
@@ -167,11 +158,9 @@ class HyperSwapTrainer(LightningModule):
self.log('generator_loss', generator_loss, prog_bar = True)
self.log('discriminator_loss', discriminator_loss, prog_bar = True)
self.log('adversarial_loss', adversarial_loss)
self.log('cycle_loss', cycle_loss)
self.log('feature_loss', feature_loss)
self.log('reconstruction_loss', reconstruction_loss)
self.log('identity_loss', identity_loss)
self.log('gaze_loss', gaze_loss)
self.log('mask_loss', mask_loss)
if do_update:
hyperswap/src/types.py
-1
View File
@@ -19,7 +19,6 @@ Padding : TypeAlias = Tuple[int, int, int, int]
GeneratorModule : TypeAlias = Module
EmbedderModule : TypeAlias = Module
GazerModule : TypeAlias = Module
FaceMaskerModule : TypeAlias = Module
OptimizerSet : TypeAlias = Any