From abdc770892eb101b925fa48a703a606e1c9cfe03 Mon Sep 17 00:00:00 2001
From: henryruhs <info@henryruhs.com>
Date: Wed, 5 Mar 2025 18:59:23 +0100
Subject: [PATCH] Improve UNet

---
 face_swapper/src/networks/unet.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/face_swapper/src/networks/unet.py b/face_swapper/src/networks/unet.py
index 7f0683e..226659f 100644
--- a/face_swapper/src/networks/unet.py
+++ b/face_swapper/src/networks/unet.py
@@ -18,20 +18,19 @@ class UNet(nn.Module):
 			DownSample(32, 64),
 			DownSample(64, 128),
 			DownSample(128, 256),
-			DownSample(256, 512)
+			DownSample(256, 512),
+			DownSample(512, 1024)
 		])
 
 		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)
 			])