Add MidBlock

model_blocks/blocks.py

@@ -183,3 +183,143 @@ class DownBlock(nn.Module):
                 in_attn = out.reshape(batch_size, channels, h * w)
                 in_attn = self.attention_norms[i](in_attn)
                 in_attn = in_attn.transpose(1, 2)
+                out_attn, _ = self.attentions[i](in_attn, in_attn, in_attn)
+                out_attn = out_attn.transpose(1, 2).reshape(batch_size, channels, h, w)
+                out = out + out_attn
+
+            if self.cross_attn:
+                assert context is not None, (
+                    "context cannot be None if cross attention layers are used"
+                )
+                batch_size, channels, h, w = out.shape
+                in_attn = out.reshape(batch_size, channels, h * w)
+                in_attn = self.cross_attn_norms[i](in_attn)
+                in_attn = in_attn.transpose(1, 2)
+                assert (
+                    context.shape[0] == x.shape[0]
+                    and context.shape[-1] == self.context_dim
+                )
+                context_proj = self.context_proj[i](context)
+                out_attn, _ = self.cross_attentions[i](
+                    in_attn, context_proj, context_proj
+                )
+                out_attn = out_attn.transpose(1, 2).reshape(batch_size, channels, h, w)
+                out = out + out_attn
+
+            # DownSample to x2 smaller dimension
+            out = self.down_sample_conv(out)
+            return out
+
+
+class MidBlock(nn.Module):
+    def __init__(
+        self,
+        num_heads,
+        num_layers,
+        cross_attn,
+        input_dim,
+        output_dim,
+        t_emb_dim,
+        cond_dim,
+        norm_channels,
+        self_attn,
+        down_sample,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.cross_attn = cross_attn
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.cond_dim = cond_dim
+        self.norm_channels = norm_channels
+        self.t_emb_dim = t_emb_dim
+        self.attn = self_attn
+        self.down_sample = down_sample
+
+        self.resnet_one = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.GroupNorm(
+                        self.norm_channels,
+                        self.input_dim if i == 0 else self.output_dim,
+                    ),
+                    nn.SiLU(),
+                    nn.Conv2d(
+                        self.input_dim if i == 0 else self.output_dim,
+                        self.output_dim,
+                        kernel_size=3,
+                        stride=1,
+                        padding=1,
+                    ),
+                )
+                for i in range(self.num_layers)
+            ]
+        )
+
+        if self.t_emb_dim is not None:
+            self.t_emb_layers = nn.ModuleList(
+                [
+                    nn.Sequential(nn.SiLU(), nn.Linear(self.t_emb_dim, self.output_dim))
+                    for _ in range(self.num_layers)
+                ]
+            )
+
+        self.resnet_two = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.GroupNorm(
+                        self.norm_channels,
+                        self.output_dim,
+                    ),
+                    nn.SiLU(),
+                    nn.Conv2d(
+                        self.output_dim,
+                        self.output_dim,
+                        kernel_size=3,
+                        stride=1,
+                        padding=1,
+                    ),
+                )
+                for _ in range(self.num_layers)
+            ]
+        )
+
+        if self.attn:
+            self.attention_norms = nn.ModuleList(
+                [
+                    nn.GroupNorm(self.norm_channels, self.output_dim)
+                    for _ in range(num_layers)
+                ]
+            )
+            self.attentions = nn.ModuleList(
+                [
+                    nn.MultiheadAttention(
+                        self.output_dim, self.num_heads, batch_first=True
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+
+        if self.cross_attn:
+            self.cross_attn_norms = nn.ModuleList(
+                [
+                    nn.GroupNorm(self.norm_channels, self.output_dim)
+                    for _ in range(self.num_layers)
+                ]
+            )
+            self.cross_attentions = nn.ModuleList(
+                [
+                    nn.MultiheadAttention(
+                        self.output_dim, self.num_heads, batch_first=True
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+
+            self.context_proj = nn.ModuleList(
+                [
+                    nn.Linear(self.cond_dim, self.output_dim)
+                    for _ in range(self.num_layers)
+                ]
+            )