Add cross attn if needed for conditional latent diffusion

models/blocks.py

@@ -109,10 +109,27 @@ class DownBlock(nn.Module):
             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
+            if self.attn:
+                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
+
+            # Cross Attention
+            if self.cross_attn:
+                assert context not None, "Context must be given for cross_attn"
+                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
+
+        out = self.resnet_down_conv(out)
+        return out