Add residual convolution to the mid block

models/blocks.py

@@ -200,3 +200,21 @@ class MidBlock(nn.Module):
                 [nn.GroupNorm(norm_dim, out_channels)
                  for _ in range(num_layers)]
             )
+
+            self.cross_attn = nn.ModuleList(
+                [nn.MultiheadAttention(
+                    out_channels, num_heads=num_heads, batch_first=True) for _ in range(num_layers)]
+            )
+
+            self.context_proj = nn.ModuleList([
+                nn.Conv2d(in_channels if i == 0 else out_channels,
+                          out_channels, kernel_size=1)
+                for i in range(num_layers + 1)
+            ])
+
+        self.residual_input_conv = nn.ModuleList([
+            nn.Conv2d(in_channels if i == 0 else out_channels,
+                      out_channels, kernel_size=1)
+            for i in range(num_layers + 1)
+
+        ])