# Adapted https://github.com/huggingface/optimum/blob/15a162824d0c5d8aa7a3d14ab6e9bb07e5732fb6/optimum/exporters/onnx/convert.py#L573-L614 # Copyright 2022 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import shutil import tempfile from contextlib import nullcontext from pathlib import Path import onnx import onnx_graphsurgeon as gs import torch from diffusers.models.transformers import FluxTransformer2DModel, SD3Transformer2DModel from diffusers.models.unets import UNet2DConditionModel from torch.onnx import export as onnx_export from modelopt.onnx.autocast.convert import convert_to_f16 from modelopt.onnx.quantization.qdq_utils import fp4qdq_to_2dq from modelopt.torch.quantization.export_onnx import configure_linear_module_onnx_quantizers from modelopt.torch.utils import torch_to from .fp8_onnx_graphsurgeon import cast_resize_io, convert_fp16_io, convert_zp_fp8 MODEL_ID_TO_DYNAMIC_AXES = { "sdxl-1.0": { "sample": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}, "timestep": {0: "steps"}, "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, "text_embeds": {0: "batch_size"}, "time_ids": {0: "batch_size"}, "latent": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}, }, "sdxl-turbo": { "sample": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}, "timestep": {0: "steps"}, "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, "text_embeds": {0: "batch_size"}, "time_ids": {0: "batch_size"}, "latent": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}, }, "sd3-medium": { "hidden_states": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}, "timestep": {0: "steps"}, "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, "pooled_projections": {0: "batch_size"}, "sample": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}, }, "flux-dev": { "hidden_states": {0: "batch_size", 1: "latent_dim"}, "encoder_hidden_states": {0: "batch_size"}, "pooled_projections": {0: "batch_size"}, "timestep": {0: "batch_size"}, "img_ids": {0: "latent_dim"}, "guidance": {0: "batch_size"}, "latent": {0: "batch_size"}, }, "flux-schnell": { "hidden_states": {0: "batch_size", 1: "latent_dim"}, "encoder_hidden_states": {0: "batch_size"}, "pooled_projections": {0: "batch_size"}, "timestep": {0: "batch_size"}, "img_ids": {0: "latent_dim"}, "latent": {0: "batch_size"}, }, } def flux_convert_rope_weight_type(onnx_graph): graph = gs.import_onnx(onnx_graph) for node in graph.nodes: if node.op == "Einsum": node.inputs[1].dtype = "float32" return gs.export_onnx(graph) def generate_fp8_scales(backbone): # temporary solution due to a known bug in torch.onnx._dynamo_export for _, module in backbone.named_modules(): if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)) and ( hasattr(module.input_quantizer, "_amax") and module.input_quantizer is not None ): module.input_quantizer._num_bits = 8 module.weight_quantizer._num_bits = 8 module.input_quantizer._amax = module.input_quantizer._amax * (127 / 448.0) module.weight_quantizer._amax = module.weight_quantizer._amax * (127 / 448.0) def _gen_dummy_inp_and_dyn_shapes_sdxl(backbone, min_bs=1, opt_bs=1): assert isinstance(backbone, UNet2DConditionModel) cfg = backbone.config assert cfg.addition_embed_type == "text_time" dynamic_shapes = { "sample": { "min": [min_bs, cfg.in_channels, cfg.sample_size, cfg.sample_size], "opt": [opt_bs, cfg.in_channels, cfg.sample_size, cfg.sample_size], }, "timestep": {"min": [1], "opt": [1]}, "encoder_hidden_states": { "min": [min_bs, 77, cfg.cross_attention_dim], "opt": [opt_bs, 77, cfg.cross_attention_dim], }, "added_cond_kwargs.text_embeds": { "min": [ min_bs, backbone.add_embedding.linear_1.in_features - 6 * backbone.add_time_proj.num_channels, ], "opt": [ opt_bs, backbone.add_embedding.linear_1.in_features - 6 * backbone.add_time_proj.num_channels, ], }, "added_cond_kwargs.time_ids": {"min": [min_bs, 6], "opt": [opt_bs, 6]}, } dummy_input = { "sample": torch.randn(*dynamic_shapes["sample"]["min"]), "timestep": torch.ones(1), "encoder_hidden_states": torch.randn(*dynamic_shapes["encoder_hidden_states"]["min"]), "added_cond_kwargs": { "text_embeds": torch.randn(*dynamic_shapes["added_cond_kwargs.text_embeds"]["min"]), "time_ids": torch.randn(*dynamic_shapes["added_cond_kwargs.time_ids"]["min"]), }, "return_dict": False, } dummy_input = torch_to(dummy_input, dtype=backbone.dtype) return dummy_input, dynamic_shapes def _gen_dummy_inp_and_dyn_shapes_sd3(backbone, min_bs=1, opt_bs=1): assert isinstance(backbone, SD3Transformer2DModel) cfg = backbone.config dynamic_shapes = { "hidden_states": { "min": [min_bs, cfg.in_channels, cfg.sample_size, cfg.sample_size], "opt": [opt_bs, cfg.in_channels, cfg.sample_size, cfg.sample_size], }, "timestep": {"min": [2], "opt": [16]}, "encoder_hidden_states": { "min": [min_bs, 333, cfg.joint_attention_dim], "opt": [opt_bs, 333, cfg.joint_attention_dim], }, "pooled_projections": { "min": [min_bs, cfg.pooled_projection_dim], "opt": [opt_bs, cfg.pooled_projection_dim], }, } dummy_input = { "hidden_states": torch.randn(*dynamic_shapes["hidden_states"]["min"]), "timestep": torch.ones(1), "encoder_hidden_states": torch.randn(*dynamic_shapes["encoder_hidden_states"]["min"]), "pooled_projections": torch.randn(*dynamic_shapes["pooled_projections"]["min"]), "return_dict": False, } dummy_input = torch_to(dummy_input, dtype=backbone.dtype) return dummy_input, dynamic_shapes def _gen_dummy_inp_and_dyn_shapes_flux(backbone, min_bs=1, opt_bs=1): assert isinstance(backbone, FluxTransformer2DModel) cfg = backbone.config text_maxlen = 512 img_dim = 4096 dynamic_shapes = { "hidden_states": { "min": [min_bs, img_dim, cfg.in_channels], "opt": [opt_bs, img_dim, cfg.in_channels], }, "encoder_hidden_states": { "min": [min_bs, text_maxlen, cfg.joint_attention_dim], "opt": [opt_bs, text_maxlen, cfg.joint_attention_dim], }, "pooled_projections": { "min": [min_bs, cfg.pooled_projection_dim], "opt": [opt_bs, cfg.pooled_projection_dim], }, "timestep": {"min": [1], "opt": [1]}, "img_ids": {"min": [img_dim, 3], "opt": [img_dim, 3]}, "txt_ids": {"min": [text_maxlen, 3], "opt": [text_maxlen, 3]}, } if cfg.guidance_embeds: # flux-dev dynamic_shapes["guidance"] = {"min": [1], "opt": [1]} dtype = backbone.dtype dummy_input = { "hidden_states": torch.randn(*dynamic_shapes["hidden_states"]["min"], dtype=dtype), "encoder_hidden_states": torch.randn( *dynamic_shapes["encoder_hidden_states"]["min"], dtype=dtype ), "pooled_projections": torch.randn( *dynamic_shapes["pooled_projections"]["min"], dtype=dtype ), "timestep": torch.ones(1, dtype=dtype), "img_ids": torch.randn(*dynamic_shapes["img_ids"]["min"], dtype=torch.float32), "txt_ids": torch.randn(*dynamic_shapes["txt_ids"]["min"], dtype=torch.float32), "return_dict": False, } if cfg.guidance_embeds: # flux-dev dummy_input["guidance"] = torch.full((1,), 3.5, dtype=torch.float32) dummy_input["pooled_projections"] = torch.randn(min_bs, cfg.pooled_projection_dim, dtype=dtype) dummy_input["timestep"] = torch.ones(1, dtype=dtype) dummy_input["img_ids"] = torch.randn(img_dim, 3, dtype=torch.float32) dummy_input["txt_ids"] = torch.randn(text_maxlen, 3, dtype=torch.float32) if cfg.guidance_embeds: # flux-dev dummy_input["guidance"] = torch.full((1,), 3.5, dtype=torch.float32) dummy_input["return_dict"] = False return dummy_input, dynamic_shapes def update_dynamic_axes(model_id, dynamic_axes): if model_id in ["flux-dev", "flux-schnell"]: dynamic_axes["out.0"] = dynamic_axes.pop("latent") elif model_id in ["sdxl-1.0", "sdxl-turbo"]: dynamic_axes["added_cond_kwargs.text_embeds"] = dynamic_axes.pop("text_embeds") dynamic_axes["added_cond_kwargs.time_ids"] = dynamic_axes.pop("time_ids") dynamic_axes["out.0"] = dynamic_axes.pop("latent") elif model_id == "sd3-medium": dynamic_axes["out.0"] = dynamic_axes.pop("sample") def _create_dynamic_shapes(dynamic_shapes): min_shapes = {} opt_shapes = {} for key, value in dynamic_shapes.items(): min_shapes[key] = value["min"] opt_shapes[key] = value["opt"] return { "dynamic_shapes": { "minShapes": min_shapes, "optShapes": opt_shapes, "maxShapes": opt_shapes, } } def generate_dummy_inputs_and_dynamic_axes_and_shapes(model_id, backbone): """Generate dummy inputs, dynamic axes, and dynamic shapes for the given model.""" if model_id in ["sdxl-1.0", "sdxl-turbo"]: dummy_input, dynamic_shapes = _gen_dummy_inp_and_dyn_shapes_sdxl( backbone, min_bs=2, opt_bs=16 ) elif model_id == "sd3-medium": dummy_input, dynamic_shapes = _gen_dummy_inp_and_dyn_shapes_sd3( backbone, min_bs=2, opt_bs=16 ) elif model_id in ["flux-dev", "flux-schnell"]: dummy_input, dynamic_shapes = _gen_dummy_inp_and_dyn_shapes_flux( backbone, min_bs=1, opt_bs=1 ) else: raise NotImplementedError(f"Unsupported model_id: {model_id}") dummy_input = torch_to(dummy_input, device=backbone.device) dummy_inputs = (dummy_input,) dynamic_axes = MODEL_ID_TO_DYNAMIC_AXES[model_id] dynamic_shapes = _create_dynamic_shapes(dynamic_shapes) return dummy_inputs, dynamic_axes, dynamic_shapes def get_io_shapes(model_id, onnx_load_path, dynamic_shapes): output_name = "out.0" if onnx_load_path != "": if model_id in ["sdxl-1.0", "sdxl-turbo"]: output_name = "latent" elif model_id in ["sd3-medium"]: output_name = "sample" elif model_id in ["flux-dev", "flux-schnell"]: output_name = "output" else: raise NotImplementedError(f"Unsupported model_id: {model_id}") if model_id in ["sdxl-1.0", "sdxl-turbo"]: io_shapes = {output_name: dynamic_shapes["dynamic_shapes"]["minShapes"]["sample"]} elif model_id in ["sd3-medium"]: io_shapes = {output_name: dynamic_shapes["dynamic_shapes"]["minShapes"]["hidden_states"]} elif model_id in ["flux-dev", "flux-schnell"]: io_shapes = {} return io_shapes def remove_nesting(dynamic_shapes): dynamic_shapes["dynamic_shapes"]["minShapes"]["text_embeds"] = dynamic_shapes["dynamic_shapes"][ "minShapes" ].pop("added_cond_kwargs.text_embeds") dynamic_shapes["dynamic_shapes"]["minShapes"]["time_ids"] = dynamic_shapes["dynamic_shapes"][ "minShapes" ].pop("added_cond_kwargs.time_ids") dynamic_shapes["dynamic_shapes"]["optShapes"]["text_embeds"] = dynamic_shapes["dynamic_shapes"][ "optShapes" ].pop("added_cond_kwargs.text_embeds") dynamic_shapes["dynamic_shapes"]["optShapes"]["time_ids"] = dynamic_shapes["dynamic_shapes"][ "optShapes" ].pop("added_cond_kwargs.time_ids") def save_onnx(onnx_model, output): onnx.save( onnx_model, str(output), save_as_external_data=True, all_tensors_to_one_file=True, location=output.name + "_data", size_threshold=1024, ) print(f"ONNX model saved to {output}") def modelopt_export_sd(backbone, onnx_dir, model_name, precision): model_file_name = "model.onnx" os.makedirs(f"{onnx_dir}", exist_ok=True) tmp_subfolder = tempfile.mkdtemp(prefix="myapp_") tmp_output = Path(f"{tmp_subfolder}/{model_file_name}") q_output = Path(f"{onnx_dir}/{model_file_name}") quantizer_context = ( configure_linear_module_onnx_quantizers(backbone) if precision == "fp4" else nullcontext() ) dummy_inputs, dynamic_axes, _ = generate_dummy_inputs_and_dynamic_axes_and_shapes( model_name, backbone ) if model_name in {"sdxl-1.0", "sdxl-turbo"}: input_names = ["sample", "timestep", "encoder_hidden_states", "text_embeds", "time_ids"] output_names = ["latent"] elif model_name == "sd3-medium": input_names = ["hidden_states", "encoder_hidden_states", "pooled_projections", "timestep"] output_names = ["sample"] elif model_name in ["flux-dev", "flux-schnell"]: input_names = [ "hidden_states", "encoder_hidden_states", "pooled_projections", "timestep", "img_ids", "txt_ids", ] if model_name == "flux-dev": input_names.append("guidance") output_names = ["latent"] else: raise NotImplementedError(f"Unsupported model_id: {model_name}") do_constant_folding = True opset_version = 20 with quantizer_context, torch.inference_mode(): onnx_export( backbone, dummy_inputs, f=tmp_output.as_posix(), input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes, do_constant_folding=do_constant_folding, opset_version=opset_version, ) print(f"Saved at {tmp_output}") onnx_model = onnx.load(str(tmp_output), load_external_data=True) if precision == "fp8": if not model_name.startswith("flux"): graph = gs.import_onnx(onnx_model) graph.cleanup().toposort() convert_fp16_io(graph) onnx_model = gs.export_onnx(graph) onnx_model = convert_zp_fp8(onnx_model) onnx_model = convert_to_f16(onnx_model, keep_io_types=True) graph = gs.import_onnx(onnx_model) cast_resize_io(graph) onnx_model = gs.export_onnx(graph.cleanup()) else: flux_convert_rope_weight_type(onnx_model) if precision == "fp4": onnx_model = fp4qdq_to_2dq(onnx_model) save_onnx(onnx_model, q_output) shutil.rmtree(tmp_subfolder, ignore_errors=True)