add 128px

This commit is contained in:
harisreedhar
2025-03-06 15:47:39 +05:30
committed by henryruhs
parent a89e51c2f8
commit e428ae04e3
3 changed files with 35 additions and 10 deletions
+12 -3
View File
@@ -17,12 +17,21 @@ class AAD(nn.Module):
def create_layers(self) -> nn.ModuleList:
layers = nn.ModuleList()
if self.output_size == 128:
layers.extend(
[
AdaptiveFeatureModulation(512, 512, 512, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(512, 512, 1024, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(512, 512, 512, self.identity_channels, self.num_blocks),
])
if self.output_size == 256:
layers.extend(
[
AdaptiveFeatureModulation(1024, 1024, 1024, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(1024, 1024, 2048, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(1024, 1024, 1024, self.identity_channels, self.num_blocks)
AdaptiveFeatureModulation(1024, 1024, 1024, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(1024, 512, 512, self.identity_channels, self.num_blocks)
])
if self.output_size == 512:
@@ -31,12 +40,12 @@ class AAD(nn.Module):
AdaptiveFeatureModulation(2048, 2048, 2048, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(2048, 2048, 4096, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(2048, 2048, 2048, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(2048, 1024, 1024, self.identity_channels, self.num_blocks)
AdaptiveFeatureModulation(2048, 1024, 1024, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(1024, 512, 512, self.identity_channels, self.num_blocks)
])
layers.extend(
[
AdaptiveFeatureModulation(1024, 512, 512, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(512, 256, 256, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(256, 128, 128, self.identity_channels, self.num_blocks),
AdaptiveFeatureModulation(128, 64, 64, self.identity_channels, self.num_blocks),
+19 -5
View File
@@ -18,19 +18,26 @@ class UNet(nn.Module):
DownSample(32, 64),
DownSample(64, 128),
DownSample(128, 256),
DownSample(256, 512),
DownSample(512, 1024)
DownSample(256, 512)
])
if self.output_size == 128:
down_samples.extend(
[
DownSample(512, 512)
])
if self.output_size == 256:
down_samples.extend(
[
DownSample(512, 1024),
DownSample(1024, 1024)
])
if self.output_size == 512:
down_samples.extend(
[
DownSample(512, 1024),
DownSample(1024, 2048),
DownSample(2048, 2048)
])
@@ -40,22 +47,29 @@ class UNet(nn.Module):
def create_up_samples(self) -> nn.ModuleList:
up_samples = nn.ModuleList()
if self.output_size == 128:
up_samples.extend(
[
UpSample(512, 512)
])
if self.output_size == 256:
up_samples.extend(
[
UpSample(1024, 1024)
UpSample(1024, 1024),
UpSample(2048, 512)
])
if self.output_size == 512:
up_samples.extend(
[
UpSample(2048, 2048),
UpSample(4096, 1024)
UpSample(4096, 1024),
UpSample(2048, 512)
])
up_samples.extend(
[
UpSample(2048, 512),
UpSample(1024, 256),
UpSample(512, 128),
UpSample(256, 64),
+4 -2
View File
@@ -5,10 +5,12 @@ from face_swapper.src.networks.aad import AAD
from face_swapper.src.networks.unet import UNet
@pytest.mark.parametrize('output_size', [ 256, 512 ])
@pytest.mark.parametrize('output_size', [ 128, 256, 512 ])
def test_aad_with_unet(output_size : int) -> None:
identity_channels = 512
output_channels = 2048
output_channels = 1024
if output_size == 128:
output_channels = 2048
if output_size == 256:
output_channels = 4096
if output_size == 512: