mirror of
https://github.com/facefusion/facefusion-labs.git
synced 2026-08-31 00:10:39 +02:00
changes
This commit is contained in:
@@ -14,7 +14,7 @@ class Generator(nn.Module):
|
||||
super().__init__()
|
||||
self.encoder = UNet(config_parser)
|
||||
self.generator = AAD(config_parser)
|
||||
self.masker = MaskNet(67, 1, 16)
|
||||
self.masker = MaskNet(config_parser)
|
||||
self.encoder.apply(init_weight)
|
||||
self.generator.apply(init_weight)
|
||||
self.masker.apply(init_weight)
|
||||
|
||||
@@ -195,12 +195,12 @@ class MaskLoss(nn.Module):
|
||||
|
||||
def calc_mask(self, target_tensor : Tensor) -> Tensor:
|
||||
target_tensor = torch.nn.functional.interpolate(target_tensor, (512, 512), mode = 'bilinear')
|
||||
face_indices = torch.tensor([ 1, 2, 3, 4, 5, 10, 11, 12, 13 ]).to(target_tensor.device)
|
||||
face_mask_regions = torch.tensor([ 1, 2, 3, 4, 5, 10, 11, 12, 13 ]).to(target_tensor.device)
|
||||
|
||||
with torch.no_grad():
|
||||
output_tensor = self.parser(target_tensor)[0]
|
||||
output_tensor = output_tensor.argmax(1)
|
||||
output_tensor = torch.isin(output_tensor, face_indices).to(target_tensor.dtype)
|
||||
output_tensor = torch.isin(output_tensor, face_mask_regions).to(target_tensor.dtype)
|
||||
output_tensor = output_tensor.view(-1, 1, 512, 512)
|
||||
output_tensor = torch.nn.functional.interpolate(output_tensor, (self.config_output_size, self.config_output_size), mode = 'bilinear')
|
||||
return output_tensor
|
||||
|
||||
Reference in New Issue
Block a user