huggingface / huggingface/diffusers

[Z-Image] ONNX export fails due to pad_sequence / unbind / dynamic-shape ops in ZImageTransformer2DModel (request for export-friendly path)

Aperta
#12,843 1 commento 0 reazioni 0 assegnatari Vedi su GitHub
bug stale
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

### Describe the bug

### Problem: ONNX export succeeds, but exported model is not truly dynamic (size-switching fails)

The ONNX export itself does NOT fail. The exported ONNX model looks valid and can run if the inference image size matches the export-time size.

However, Z-Image in PyTorch supports running at multiple image sizes (e.g. 512 / 768 / 1024). After exporting to ONNX, this dynamic behavior is effectively lost: some shape-related values become baked into the ONNX graph as constants derived from the export-time input dimensions (e.g. token count, reshape sizes in patchify/unpatchify, padded lengths, etc.).

As a result, when the downstream inference runtime tries to run the exported ONNX with a different `image_size` than the one used during export (e.g. export with 1024 but infer with 512 or 768), the model fails at runtime with shape mismatch / reshape errors (typically around `view`/`reshape`-like logic and token reshaping).

In short:
- PyTorch ZImageTransformer2DModel: supports dynamic image sizes.
- Exported ONNX graph: contains hardcoded shapes → cannot switch sizes at inference time.

This suggests the current export path traces some dynamic-shape logic into static constants. A dedicated ONNX-friendly forward path (e.g. [forward_single]

Reference implementation (for discussion):
https://github.com/er6y/diffusers/tree/fix/zimage-transformer-onnx-friendly

this is only suitable for me to export correct ONNX, but not good for the project, so i'll only keep to my fork...
hoping the full fixing in this project..

### Reproduction

import torch
from diffusers import ZImagePipeline

pipe = ZImagePipeline.from_pretrained("...") # Z-Image checkpoint
model = pipe.transformer.eval()

# create dummy inputs (batch=1), cap embeddings etc...
# torch.onnx.export(model, ...)

### Reference implementation (for discussion only)
I have a working downstream patch in my fork/branch (mainly adds an ONNX-friendly [forward_single](cci:1://file:///ux/work/yilei.wang/diffusers/src/diffusers/models/transformers/transformer_z_image.py:707:4-900:21) for ZImage transformer and an export wrapper):

https://github.com/er6y/diffusers/tree/fix/zimage-transformer-onnx-friendly

i've rewrite the forwarding, but it is only for me to export ONNX, not siutable for PR...hoping fixing on this issuse

### Logs

```shell

```

### System Info

diffusers:
torch:
onnx/opset:
exporter: torch.onnx.export / dynamo_export
runtime target: MNN / other ONNX runtime

### Who can help?

_No response_

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Start with the ZImageTransformer2DModel implementation in transformer_z_image.py and the referenced export wrapper or downstream branch. Reproduce the torch.onnx.export flow with different export and inference image sizes, then verify that the resulting ONNX model runs for those sizes without shape or reshape errors.

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

Valutazione

Stack tecnologico
python
Ambito
machine-learning
Tipo di issue
Bug
Difficoltà
5/5
Tempo stimato
Più di una settimana
Stato di attività
Ferma
Chiarezza
Abbastanza chiara
Idoneità per principianti
32/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.