diff --git a/python_coreml_stable_diffusion/unet.py b/python_coreml_stable_diffusion/unet.py
index 666f146..519bc6f 100644
--- a/python_coreml_stable_diffusion/unet.py
+++ b/python_coreml_stable_diffusion/unet.py
@@ -916,19 +916,22 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
             self.down_blocks.append(down_block)
 
         # mid
-        assert mid_block_type == "UNetMidBlock2DCrossAttn"
-        self.mid_block = UNetMidBlock2DCrossAttn(
-            in_channels=block_out_channels[-1],
-            transformer_layers_per_block=transformer_layers_per_block[-1],
-            temb_channels=time_embed_dim,
-            resnet_eps=norm_eps,
-            resnet_act_fn=act_fn,
-            output_scale_factor=mid_block_scale_factor,
-            resnet_time_scale_shift="default",
-            cross_attention_dim=cross_attention_dim,
-            attn_num_head_channels=attention_head_dim[i],
-            resnet_groups=norm_num_groups,
-        )
+        if mid_block_type is None:
+            self.mid_block = None
+        else:
+            assert mid_block_type == "UNetMidBlock2DCrossAttn"
+            self.mid_block = UNetMidBlock2DCrossAttn(
+                in_channels=block_out_channels[-1],
+                transformer_layers_per_block=transformer_layers_per_block[-1],
+                temb_channels=time_embed_dim,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                output_scale_factor=mid_block_scale_factor,
+                resnet_time_scale_shift="default",
+                cross_attention_dim=cross_attention_dim,
+                attn_num_head_channels=attention_head_dim[i],
+                resnet_groups=norm_num_groups,
+            )
 
         # up
         reversed_block_out_channels = list(reversed(block_out_channels))
@@ -1014,9 +1017,10 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
             down_block_res_samples = new_down_block_res_samples
 
         # 4. mid
-        sample = self.mid_block(sample,
-                                emb,
-                                encoder_hidden_states=encoder_hidden_states)
+        if self.mid_block is not None:
+            sample = self.mid_block(sample,
+                                    emb,
+                                    encoder_hidden_states=encoder_hidden_states)
 
         if self.support_controlnet:
             sample = sample + additional_residuals[-1]