Transformer2DModel onnx convertation problems

Hi! I am trying to convert Transformer2DModel to onnx and cannot solve some problems.
I am trying to export UNet model with next architecture
```python
Transformer2DModel(
(pos_embed): PatchEmbed(
(proj): Conv2d(96, 1584, kernel_size=(2, 2), stride=(2, 2))
)
(transformer_blocks): ModuleList(
(0-23): 24 x BasicTransformerBlock(
(norm1): LayerNorm((1584,), eps=1e-06, elementwise_affine=False)
(attn1): Attention(
(to_q): Linear(in_features=1584, out_features=1584, bias=True)
(to_k): Linear(in_features=1584, out_features=1584, bias=True)
(to_v): Linear(in_features=1584, out_features=1584, bias=True)
(to_out): ModuleList(
(0): Linear(in_features=1584, out_features=1584, bias=True)
(1): Dropout(p=0.0, inplace=False)
)
)
(norm2): LayerNorm((1584,), eps=1e-06, elementwise_affine=False)
(ff): FeedForward(
(net): ModuleList(
(0): GELU(
(proj): Linear(in_features=1584, out_features=6336, bias=True)
)
(1): Dropout(p=0.0, inplace=False)
(2): Linear(in_features=6336, out_features=1584, bias=True)
)
)
)
)
(norm_out): LayerNorm((1584,), eps=1e-06, elementwise_affine=False)
(proj_out): Linear(in_features=1584, out_features=128, bias=True)
(adaln_single): AdaLayerNormSingleFlow(
(emb): PixArtAlphaCombinedFlowEmbeddings(
(timestep_embedder): TimestepEmbedding(
(linear_1): Linear(in_features=512, out_features=1584, bias=True)
(act): SiLU()
(linear_2): Linear(in_features=1584, out_features=1584, bias=True)
)
)
(silu): SiLU()
(linear): Linear(in_features=1584, out_features=9504, bias=True)
)
)
```

As as input for my model i use next inputs with shapes:
`hidden_states -> (B, 96, Height, Width)`
`timestep -> (B)`
`resolution -> (B, 2)`
`aspect_ratio -> (B, 1)`

In python version resolution and aspect ratio are parts of added_cond_kwargs, but since onnx doesn't support dicts i wrote a wrapper that
```python
import torch
import torch.nn as nn

class Transformer2DWrapper(nn.Module):
def __init__(self, model):
super().__init__()
self.model = model

def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.Tensor,
resolution: torch.Tensor,
aspect_ratio: torch.Tensor,
):
timestep = timestep.float()

added_cond_kwargs = {
"resolution": resolution,
"aspect_ratio": aspect_ratio,
}

out = self.model(
hidden_states=hidden_states,
timestep=timestep,
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
)
return out[0] # sample tensor
```

I export to onnx with torch.onnx.export
```python
wrapper = Transformer2DWrapper(unet_model)

torch.onnx.export(
wrapper,
dummy_inputs,
"unet_converted/model.onnx",
input_names=["hidden_states", "timestep", "resolution", "aspect_ratio"],
output_names=["out_sample"],
dynamic_axes={
"hidden_states": {0: "batch", 2: "height", 3: "width"},
"timestep": {0: "batch"},
"resolution": {0: "batch"},
"aspect_ratio": {0: "batch"},
"out_sample": {0: "batch", 2: "height", 3: "width"},
},
opset_version=17,
)
```

But after loading onnx version there is error with squeeze operation
```python
Fail: [ONNXRuntimeError] : 1 : FAIL : Load model from unet_converted/model.onnx failed:Node (/model/transformer_blocks.0/If) Op (If) [TypeInferenceError] Graph attribute inferencing failed: Node (/model/transformer_blocks.0/Squeeze) Op (Squeeze) [ShapeInferenceError] Dimension of input 1 must be 1 instead of 256
```

Can you please help with convertation?
Versions:
diffusers==0.27.2
torch==2.2.0+cu118

Start with the Transformer2DWrapper and the torch.onnx.export call, then inspect the generated /model/transformer_blocks.0/If/Squeeze node reported by ONNX Runtime. Reproduce with the listed versions and dummy input shapes; done means the exported model loads successfully in ONNX Runtime without the shape-inference error.

Scritto dal modello di indicizzazione a partire dal testo della issue.