@@ -103,6 +103,7 @@ def __init__(
         self.use_linear_projection = use_linear_projection
         self.num_attention_heads = num_attention_heads
         self.attention_head_dim = attention_head_dim
+        # TODO inner_dim is wrong
         inner_dim = num_attention_heads * attention_head_dim
 
         # 1. Transformer2DModel can process both standard continous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
 
@@ -42,6 +42,7 @@ def get_down_block(
     only_cross_attention=False,
     upcast_attention=False,
     resnet_time_scale_shift="default",
+    transformer_num_layers=1,
 ):
     down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
     if down_block_type == "DownBlock2D":
@@ -103,6 +104,7 @@ def get_down_block(
             only_cross_attention=only_cross_attention,
             upcast_attention=upcast_attention,
             resnet_time_scale_shift=resnet_time_scale_shift,
+            transformer_num_layers=transformer_num_layers,
         )
     elif down_block_type == "SimpleCrossAttnDownBlock2D":
         if cross_attention_dim is None:
@@ -214,6 +216,7 @@ def get_up_block(
     only_cross_attention=False,
     upcast_attention=False,
     resnet_time_scale_shift="default",
+    transformer_num_layers=1,
 ):
     up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
     if up_block_type == "UpBlock2D":
@@ -262,6 +265,7 @@ def get_up_block(
             only_cross_attention=only_cross_attention,
             upcast_attention=upcast_attention,
             resnet_time_scale_shift=resnet_time_scale_shift,
+            transformer_num_layers=transformer_num_layers,
         )
     elif up_block_type == "SimpleCrossAttnUpBlock2D":
         if cross_attention_dim is None:
@@ -465,6 +469,7 @@ def __init__(
         dual_cross_attention=False,
         use_linear_projection=False,
         upcast_attention=False,
+        transformer_num_layers=1,
     ):
         super().__init__()
 
@@ -496,7 +501,7 @@ def __init__(
                         attn_num_head_channels,
                         in_channels // attn_num_head_channels,
                         in_channels=in_channels,
-                        num_layers=1,
+                        num_layers=transformer_num_layers,
                         cross_attention_dim=cross_attention_dim,
                         norm_num_groups=resnet_groups,
                         use_linear_projection=use_linear_projection,
@@ -509,7 +514,7 @@ def __init__(
                         attn_num_head_channels,
                         in_channels // attn_num_head_channels,
                         in_channels=in_channels,
-                        num_layers=1,
+                        num_layers=transformer_num_layers,
                         cross_attention_dim=cross_attention_dim,
                         norm_num_groups=resnet_groups,
                     )
@@ -742,6 +747,7 @@ def __init__(
         use_linear_projection=False,
         only_cross_attention=False,
         upcast_attention=False,
+        transformer_num_layers=1,
     ):
         super().__init__()
         resnets = []
@@ -772,7 +778,7 @@ def __init__(
                         attn_num_head_channels,
                         out_channels // attn_num_head_channels,
                         in_channels=out_channels,
-                        num_layers=1,
+                        num_layers=transformer_num_layers,
                         cross_attention_dim=cross_attention_dim,
                         norm_num_groups=resnet_groups,
                         use_linear_projection=use_linear_projection,
@@ -786,7 +792,7 @@ def __init__(
                         attn_num_head_channels,
                         out_channels // attn_num_head_channels,
                         in_channels=out_channels,
-                        num_layers=1,
+                        num_layers=transformer_num_layers,
                         cross_attention_dim=cross_attention_dim,
                         norm_num_groups=resnet_groups,
                     )
@@ -1712,6 +1718,7 @@ def __init__(
         use_linear_projection=False,
         only_cross_attention=False,
         upcast_attention=False,
+        transformer_num_layers=1,
     ):
         super().__init__()
         resnets = []
@@ -1744,7 +1751,7 @@ def __init__(
                         attn_num_head_channels,
                         out_channels // attn_num_head_channels,
                         in_channels=out_channels,
-                        num_layers=1,
+                        num_layers=transformer_num_layers,
                         cross_attention_dim=cross_attention_dim,
                         norm_num_groups=resnet_groups,
                         use_linear_projection=use_linear_projection,
@@ -1758,7 +1765,7 @@ def __init__(
                         attn_num_head_channels,
                         out_channels // attn_num_head_channels,
                         in_channels=out_channels,
-                        num_layers=1,
+                        num_layers=transformer_num_layers,
                         cross_attention_dim=cross_attention_dim,
                         norm_num_groups=resnet_groups,
                     )
 
@@ -147,6 +147,7 @@ def __init__(
         time_cond_proj_dim: Optional[int] = None,
         conv_in_kernel: int = 3,
         conv_out_kernel: int = 3,
+        transformer_num_layers: int = 1,
         projection_class_embeddings_input_dim: Optional[int] = None,
     ):
         super().__init__()
@@ -264,6 +265,7 @@ def __init__(
                 only_cross_attention=only_cross_attention[i],
                 upcast_attention=upcast_attention,
                 resnet_time_scale_shift=resnet_time_scale_shift,
+                transformer_num_layers=transformer_num_layers,
             )
             self.down_blocks.append(down_block)
 
@@ -282,6 +284,7 @@ def __init__(
                 dual_cross_attention=dual_cross_attention,
                 use_linear_projection=use_linear_projection,
                 upcast_attention=upcast_attention,
+                transformer_num_layers=transformer_num_layers,
             )
         elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn":
             self.mid_block = UNetMidBlock2DSimpleCrossAttn(
@@ -341,6 +344,7 @@ def __init__(
                 only_cross_attention=only_cross_attention[i],
                 upcast_attention=upcast_attention,
                 resnet_time_scale_shift=resnet_time_scale_shift,
+                transformer_num_layers=transformer_num_layers,
             )
             self.up_blocks.append(up_block)
             prev_output_channel = output_channel
 
@@ -264,6 +264,16 @@ def create_unet_diffusers_config(original_config, image_size: int, controlnet=Fa
         if head_dim is None:
             head_dim = [5, 10, 20, 20]
 
+    context_dim = unet_params.context_dim
+
+    if not isinstance(context_dim, int):
+        for other_context_dim in context_dim:
+            if context_dim[0] != other_context_dim:
+                raise NotImplementedError(
+                    "Unet has varying context dimensions. The context dimension is used for setting the dimension of the cross attention layers in the unet. If you have a checkpoint that has varying cross attention dimensions, please open an issue at https://github.com/huggingface/diffusers/issues"
+                )
+
+        context_dim = context_dim[0]
     class_embed_type = None
     projection_class_embeddings_input_dim = None
 
@@ -281,13 +291,15 @@ def create_unet_diffusers_config(original_config, image_size: int, controlnet=Fa
         down_block_types=tuple(down_block_types),
         block_out_channels=tuple(block_out_channels),
         layers_per_block=unet_params.num_res_blocks,
-        cross_attention_dim=unet_params.context_dim,
+        cross_attention_dim=context_dim,
         attention_head_dim=head_dim,
         use_linear_projection=use_linear_projection,
         class_embed_type=class_embed_type,
         projection_class_embeddings_input_dim=projection_class_embeddings_input_dim,
     )
 
+    if "transformer_depth" in unet_params:
+        config["transformer_num_layers"] = unet_params.transformer_depth
     if not controlnet:
         config["out_channels"] = unet_params.out_channels
         config["up_block_types"] = tuple(up_block_types)