diff --git a/face_swapper/src/networks/aad.py b/face_swapper/src/networks/aad.py
index dc1b73c..8810ee8 100644
--- a/face_swapper/src/networks/aad.py
+++ b/face_swapper/src/networks/aad.py
@@ -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),
diff --git a/face_swapper/src/networks/unet.py b/face_swapper/src/networks/unet.py
index 226659f..1e0494c 100644
--- a/face_swapper/src/networks/unet.py
+++ b/face_swapper/src/networks/unet.py
@@ -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),
diff --git a/face_swapper/tests/test_networks.py b/face_swapper/tests/test_networks.py
index b9da7b2..e33f5ca 100644
--- a/face_swapper/tests/test_networks.py
+++ b/face_swapper/tests/test_networks.py
@@ -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: