Adjust networks for 512

This commit is contained in:
henryruhs
2025-03-11 14:43:10 +01:00
parent f791178ded
commit 3e69d5a9a9
2 changed files with 20 additions and 20 deletions
+18 -19
View File
@@ -25,21 +25,29 @@ class UNet(nn.Module):
DownSample(1024, 1024)
])
if self.output_size in [ 384, 512, 768, 1024 ]:
down_samples.append(DownSample(1024, 2048))
if self.output_size in [ 512, 768, 1024 ]:
down_samples.append(DownSample(2048, 4096))
if self.output_size in [ 768, 1024 ]:
down_samples.append(DownSample(4096, 8192))
if self.output_size == 1024:
down_samples.append(DownSample(8192, 16384))
if self.output_size == 512:
down_samples.append(DownSample(2048, 2048))
return down_samples
def create_up_samples(self) -> nn.ModuleList:
up_samples = nn.ModuleList(
up_samples = nn.ModuleList()
if self.output_size == 256:
up_samples.extend(
[
UpSample(1024, 1024)
])
if self.output_size == 512:
up_samples.extend(
[
UpSample(2048, 2048),
UpSample(4096, 1024)
])
up_samples.extend(
[
UpSample(1024, 1024),
UpSample(2048, 512),
UpSample(1024, 256),
UpSample(512, 128),
@@ -47,15 +55,6 @@ class UNet(nn.Module):
UpSample(128, 32)
])
if self.output_size in [ 384, 512, 768, 1024 ]:
up_samples.append(UpSample(32, 16))
if self.output_size in [ 512, 768, 1024 ]:
up_samples.append(UpSample(16, 8))
if self.output_size in [ 768, 1024 ]:
up_samples.append(UpSample(8, 4))
if self.output_size == 1024:
up_samples.append(UpSample(4, 2))
return up_samples
def forward(self, target_tensor : Tensor) -> Tuple[Tensor, ...]:
+2 -1
View File
@@ -5,9 +5,10 @@ from face_swapper.src.networks.aad import AAD
from face_swapper.src.networks.unet import UNet
@pytest.mark.parametrize('output_size', [ 256 ])
@pytest.mark.parametrize('output_size', [ 256, 512 ])
def test_aad_with_unet(output_size : int) -> None:
identity_channels = 512
output_channels = 2048
if output_size == 256:
output_channels = 4096
if output_size == 512: