added option to skip mid block (#5)

- added option to skip mid block (b84b5c4088d29603d9c765d32549bb39b23231ed)

pipeline.py

@@ -52,6 +52,19 @@ def custom_sort_order(obj):
     return {ResnetBlock2D: 0, Transformer2DModel: 1, FlexibleTransformer2DModel: 1}.get(obj.__class__)
 
 
+class FlexibleIdentityBlock(nn.Module):
+    def forward(
+            self,
+            hidden_states: torch.FloatTensor,
+            temb: Optional[torch.FloatTensor] = None,
+            encoder_hidden_states: Optional[torch.FloatTensor] = None,
+            attention_mask: Optional[torch.FloatTensor] = None,
+            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+            encoder_attention_mask: Optional[torch.FloatTensor] = None,
+    ):
+        return hidden_states
+
+
 class FlexibleUNet2DConditionModel(UNet2DConditionModel, ModelMixin):
     configurations = FlexibleUnetConfigurations
 
@@ -105,18 +118,21 @@ class FlexibleUNet2DConditionModel(UNet2DConditionModel, ModelMixin):
         mid_block_add_upsample = self.configurations.get("add_upsample_mid_block")
         mid_num_attentions = self.configurations.get("mid_num_attentions")
         mid_num_resnets = self.configurations.get("mid_num_resnets")
-
-        self.mid_block = FlexibleUNetMidBlock2DCrossAttn(in_channels=down_blocks_out_channels[-1],
-                                                         temb_channels=temb_dim,
-                                                         resnet_act_fn=resnet_act_fn,
-                                                         resnet_eps=resnet_eps,
-                                                         cross_attention_dim=cross_attention_dim,
-                                                         num_attention_heads=num_attention_heads,
-                                                         num_resnets=mid_num_resnets,
-                                                         num_attentions=mid_num_attentions,
-                                                         mix_block_in_forward=mix_block_in_forward,
-                                                         add_upsample=mid_block_add_upsample
-                                                         )
+        
+        if mid_num_resnets == mid_num_attentions == 0:
+            self.mid_block = FlexibleIdentityBlock()
+        else:
+            self.mid_block = FlexibleUNetMidBlock2DCrossAttn(in_channels=down_blocks_out_channels[-1],
+                                                             temb_channels=temb_dim,
+                                                             resnet_act_fn=resnet_act_fn,
+                                                             resnet_eps=resnet_eps,
+                                                             cross_attention_dim=cross_attention_dim,
+                                                             num_attention_heads=num_attention_heads,
+                                                             num_resnets=mid_num_resnets,
+                                                             num_attentions=mid_num_attentions,
+                                                             mix_block_in_forward=mix_block_in_forward,
+                                                             add_upsample=mid_block_add_upsample
+                                                             )
 
         ###############
         #  Up blocks  #