mirror of
https://github.com/facefusion/facefusion-labs.git
synced 2026-06-25 07:59:55 +02:00
harisreedhar
30 lines
838 B
Python
30 lines
838 B
Python
import pytest
|
|
import torch
|
|
|
|
from face_swapper.src.networks.aad import AAD
|
|
from face_swapper.src.networks.unet import UNet
|
|
|
|
|
|
@pytest.mark.parametrize('output_size', [ 128, 256, 512 ])
|
|
def test_aad_with_unet(output_size : int) -> None:
|
|
identity_channels = 512
|
|
output_channels = 1024
|
|
if output_size == 128:
|
|
output_channels = 2048
|
|
if output_size == 256:
|
|
output_channels = 4096
|
|
if output_size == 512:
|
|
output_channels = 8192
|
|
num_blocks = 2
|
|
|
|
generator = AAD(identity_channels, output_channels, output_size, num_blocks).eval()
|
|
encoder = UNet(output_size).eval()
|
|
|
|
source_tensor = torch.randn(1, 512)
|
|
target_tensor = torch.randn(1, 3, output_size, output_size)
|
|
|
|
target_attributes = encoder(target_tensor)
|
|
output_tensor = generator(source_tensor, target_attributes)
|
|
|
|
assert output_tensor.shape == (1, 3, output_size, output_size)
|