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

Mask typing and naming related updates

Browse Source
This commit is contained in:
henryruhs
2025-03-14 08:09:47 +01:00
parent b5efcbe44a
commit 5234874bc7
5 changed files with 30 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
face_swapper/src/models/loss.py
+4 -4
View File
@@ -7,7 +7,7 @@ from torch import Tensor, nn
from torchvision import transforms
from ..helper import calc_embedding
from ..types import Attribute, EmbedderModule, FaceParserModule, GazerModule, Loss, MotionExtractorModule
from ..types import Attribute, EmbedderModule, FaceParserModule, GazerModule, Loss, Mask, MotionExtractorModule
class DiscriminatorLoss(nn.Module):
@@ -186,11 +186,11 @@ class MaskLoss(nn.Module):
self.face_parser = face_parser
self.mse_loss = nn.MSELoss()
def forward(self, target_tensor : Tensor, mask_tensor : Tensor) -> Loss:
def forward(self, target_tensor : Tensor, output_mask : Mask) -> Loss:
target_mask = self.calc_mask(target_tensor)
target_mask = target_mask.view(-1, self.config_output_size, self.config_output_size)
mask_tensor = mask_tensor.view(-1, self.config_output_size, self.config_output_size)
mask_loss = self.mse_loss(target_mask, mask_tensor)
output_mask = output_mask.view(-1, self.config_output_size, self.config_output_size)
mask_loss = self.mse_loss(target_mask, output_mask)
return mask_loss
def calc_mask(self, target_tensor : Tensor) -> Tensor: