Resblock for the model

model.py

@@ -40,3 +40,27 @@ class Time_embeddings(torch.nn.Module):
     def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
         time_emb = self.Sinusoidal_waves(timesteps)
         return self.mlp(time_emb)
+
+#Resblock 
+class resblock(torch.nn.Module):
+    def __init__(self, in_channels, out_channels, time_emb_dim):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
+        self.norm1 = nn.GroupNorm(8, out_channels)
+        self.norm2 = nn.GroupNorm(8, out_channels)
+        self.silu = nn.SiLU()
+        self.time_mlp = nn.Linear(time_emb_dim, out_channels)
+        #Matching in channels and out channels for skip connection 
+        self.residual_conv = nn.Conv2d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()
+
+    def forward(self, x, time_dim):
+        h = self.norm1(self.conv1(x))
+        h = self.silu(h)
+        #Tensor broadcasting (b, out_channels) -> (b, out_channels, 1, 1)
+        time_proj = self.silu(self.time_mlp(time_dim))[:, :, None, None]
+        h = h + time_proj
+        h = self.norm2(self.conv2(h))
+        h = self.silu(h)
+        return h + self.residual_conv(x)
+