Transformer2DModel to onnx

[Voveka98]

Hi! I am trying to convert Transformer2DModel to onnx and cannot solve some problems.
I am trying to export UNet model with next architecture

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

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

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

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

[alishanawer]

You're running into a squeeze dimension mismatch error during ONNX export for the Transformer2DModel—specifically:

Squeeze: Dimension of input 1 must be 1 instead of 256

This stems from ONNX exporting a squeeze operation over a dimension that isn’t size 1. Here’s how you can address it in a way that would make for a high-quality GitHub discussion response.

---

What's Going Wrong

The ONNX spec requires that Squeeze operations only act on dimensions of size 1. However, in your model conversion, it's likely reaching a non-unit dimension (e.g., size 256), causing ONNX's shape inference to fail:

Dimension of input 1 must be 1 instead of 256

This is in line with known issues in PyTorch’s ONNX exporter around squeezing multiple arbitrary dimensions or squeezing non-singleton dimensions. For example, PyTorch doesn't support tensor.squeeze with multiple dims in one go during ONNX export—this must be done sequentially (GitHub).

---

What You Can Try

Here are a few targeted strategies to fix or work around the issue:

1. Refactor the squeeze to Only Remove Singleton Dimensions

If you're squeezing specific dimensions, wrap them in conditional checks so that ONNX export only includes squeezes that are valid:

x = x.squeeze(dim=2) if x.shape[2] == 1 else x
x = x.squeeze(dim=2) if x.shape[2] == 1 else x

This ensures ONNX isn’t asked to squeeze an invalid dimension.

2. Split Multi-Dimension Squeeze into Sequential Operations

If the model squeezes multiple dimensions at once, break them into individual squeezes (which more closely resembles how ONNX expects them) (GitHub).

3. Upgrade Opset Version to 13 or Higher

ONNX opset ≥ 13 supports dynamic axes in Squeeze, although still only one axis per operation. Make sure you’re exporting with:

opset_version=13

This improves compatibility with more flexible squeeze behavior (GitHub).

4. Guard the Squeeze in a torch.jit.script Function

Wrap the problematic logic in a small scripted helper:

import torch

@torch.jit.script
def safe_squeeze(x: torch.Tensor, dim: int) -> torch.Tensor:
    return x.squeeze(dim) if x.size(dim) == 1 else x

class Transformer2DWrapper(nn.Module):
    # ...
    def forward(self, hidden_states, timestep, resolution, aspect_ratio):
        # your prep steps...
        hidden_states = safe_squeeze(hidden_states, 2)
        hidden_states = safe_squeeze(hidden_states, 2)
        # continue model forward

This can help the tracing logic skip invalid operations without crashing.

5. Open a GitHub Issue / PR

If the problematic squeeze originates from Transformer2DModel internals—or if there’s an ONNX incompatibility—submitting an issue (or fix) upstream can help others too.

---

Sample GitHub Discussion Response

---

Hi there,

That error means ONNX is seeing a squeeze call on a dimension that isn’t size 1 (e.g., dimension size 256), which violates the ONNX spec and causes shape inference to fail.

Possible Workarounds

1. Only squeeze singleton dimensions:

if x.size(dim) == 1:
    x = x.squeeze(dim)

2. Split multi-dimension squeeze into sequential calls. ONNX won't accept squeezing multiple dims at once (GitHub).

3. Use opset_version=13 or higher in torch.onnx.export—newer opsets support dynamic squeeze behavior better (GitHub).

4. Encapsulate problematic logic in a small torch.jit.script helper like:

@torch.jit.script
def safe_squeeze(x, dim: int):
    return x.squeeze(dim) if x.size(dim) == 1 else x

Then apply safe_squeeze instead of vanilla .squeeze().

---

Please try one of these adjustments and re-export. If there's still a failure, inspecting the ONNX graph around that If or Squeeze node might reveal whether it's conditional based on dynamic shapes. Happy to dig deeper if you can share that.

---

Let me know how it goes!

[Voveka98]

Hi @alishanawer! Thank you for your response. I Opened issue with same content #12208

Now i want to check some ideas that you suggested with replacing squeeze operation. As i see model doesnt have squeeze operation and there is somewhere inside model.
As i found forward method of Transformer2DModel calls transformer_block inside

diffusers/src/diffusers/models/transformers/transformer_2d.py

Line 427 in 0f252be

hidden_states = block(

and this transformer block uses squeeze operation on output hidden states.

https://github.com/huggingface/diffusers/blob/91a151b5c698836ce3bac85fd9b5a2b10a726c99/src/diffusers/models/attention.py#L1068C13-L1068C26

So for me to check squeeze operation i should either change installed library or somehow change transformer_block in my model? Did i get it right?