Full Unet class assembled with everything

model.py

@@ -77,4 +77,58 @@ class up_sample(torch.nn.Module):
         super().__init__()
         self.conv = nn.ConvTranspose2d(channels, channels, 4, stride=2, padding=1)
     def forward(self, x):
-        return self.conv(x)
+        return self.conv(x)
+
+#Assembling Unet class 
+class Unet(torch.nn.Module):
+    def __init__(self, in_channels=3, time_dim=128):
+        super().__init__()
+        #Time embedding diemension 
+        self.time_emb = Time_embeddings(in_dim=time_dim, out_dim=256)
+        self.inc = nn.Conv2d(3, 64, 3, padding=1)
+
+        #Encoder class 
+        self.resblock1 = resblock(64, 128, time_emb_dim=256)
+        self.down1 = down_sample(128)
+
+        self.resblock2 = resblock(128, 256, time_emb_dim=256)
+        self.down2 = down_sample(256)
+
+        #Bottleneck 
+        self.bot1 = resblock(256,256, time_emb_dim=256)
+        self.bot2 = resblock(256,256, time_emb_dim=256)
+
+        #Decoder class
+        self.up1 = up_sample(256)
+        self.resblock3 = resblock(512, 128, time_emb_dim=256)
+
+        self.up2 = up_sample(128)
+        self.resblock4 = resblock(256, 64, time_emb_dim=256)
+
+        #Final output layer 
+        self.outc = nn.Conv2d(64, in_channels, 3, padding=1)
+
+    def forward(self, x, t):
+
+        t_emb = self.time_emb(t)
+        x1 = self.inc(x)
+
+        x2 = self.resblock1(x1, t_emb)
+        x3 = self.down1(x2)
+        x4 = self.resblock2(x3, t_emb)
+        x5 = self.down2(x4)
+
+        #bottle neck 
+        x5 = self.bot1(x5, t_emb)
+        x5 = self.bot2(x5, t_emb)
+
+        #decoder
+        x6 = self.up1(x5)
+        x6 = torch.cat([x4, x6], dim=1)
+        x6 = self.resblock3(x6, t_emb)
+
+        x7 = self.up2(x6)
+        x7 = torch.cat([x2, x7], dim=1)
+        x7 = self.resblock4(x7, t_emb)
+
+        return self.outc(x7)