Add forward function with attention

models/blocks.py

@@ -32,6 +32,7 @@ class DownBlock(nn.Module):
         self.context_dim = context_dim
         self.cross_attn = cross_attn
         self.t_emb_dim = t_emd_dim
+        self.num_layers = num_layers
         self.attn = attn
         self.resnet_conv_first = nn.ModuleList([
             nn.Sequential(
@@ -95,3 +96,23 @@ class DownBlock(nn.Module):
         )
 
         self.resnet_down_conv = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, 4, 2, 1) if self.down_sample else nn.Identity()
+
+    def forward(self, x, t_emb=None, context=None):
+        out = x
+        for i in range(self.num_layers):
+            # Resnet Block
+            resnet_input = out
+            out = self.resnet_conv_first[i](out)
+            if self.t_emb is not None:
+                out = out + self.time_embd_layers[i](t_emb)[:, :, None, None]
+            out = self.resnet_conv_second[i](out)
+            out = out + self.residual_input_conv[i](resnet_input)
+
+            # Self Attention
+            batch_size, channels, h, w = out.shape
+            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.attention[i](in_attn, in_attn, in_attn)
+            out_attn = out.transpose(1, 2).reshape(batch_size, channels, h, w)
+            out = out + out_attn