Is unet architecture changed as diffusers upgrading?

[1049451037]

Describe the bug

I want to export onnx for UNet2DConditionModel by torch.onnx.export. However, when I use diffusers==0.12.1, I get a model larger than 2GB, which outputs a bunch of weight files. When I use diffusers==0.7.2, I get a normal onnx file of 1.7GB.

Reproduction

import onnx
import torch
from diffusers import UNet2DConditionModel

unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4",
                                            torch_dtype=torch.float16,
					revision="fp16",
                                            subfolder="unet",
					use_auth_token=YOUR_TOKEN)
unet.cuda()

with torch.inference_mode(), torch.autocast("cuda"):
    inputs = torch.randn(2,4,64,64, dtype=torch.half, device='cuda'), torch.randn(1, dtype=torch.half, device='cuda'), torch.randn(2, 77, 768, dtype=torch.half, device='cuda')

    # Export the model
    torch.onnx.export(unet,               # model being run
                    inputs,                         # model input (or a tuple for multiple inputs)
                    "unet_v1_4_fp16_pytorch.onnx",   # where to save the model (can be a file or file-like object)
                    export_params=True,        # store the trained parameter weights inside the model file
                    opset_version=12,          # the ONNX version to export the model to
                    do_constant_folding=True,  # whether to execute constant folding for optimization
                    input_names = ['input_0', 'input_1', 'input_2'],
                    output_names = ['output_0'])

Logs

No response

System Info

docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:22.05-py3

[patrickvonplaten]

cc @anton-l @echarlaix FYI.

@1049451037 note that we currently don't have much time to review issues evolving around ONNX.

A couple of things:

• You should not / don't have to use: use_auth_token=YOUR_TOKEN anymore
• You should not use torch.autocast("cuda")

[anton-l]

@1049451037 our onnx conversion script provides a workaround for large model exports, maybe that will work for you :) https://github.com/huggingface/diffusers/blob/568b73fdf86dcc12a88b526633c7bef4bf8cf58f/scripts/convert_stable_diffusion_checkpoint_to_onnx.py#L141

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.