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)
 )
)

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

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, 
)

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