if [[ $1 = "--help" ]] || [[ $1 = "-h" ]]
then
    echo "Optimize stable diffusion ONNX models by using ONNXRuntime."
    echo "Usage: ./onnx_optimize.sh <original_onnx_path> <converted_onnx_path>"
    exit 0
fi

if [ $# -ne 2 ]
then
    echo "Error: Incorrect number of arguments"
	echo "Usage: ./onnx_optimize.sh <original_onnx_path> <converted_onnx_path>"
    exit 1
fi

set -e

original_onnx_path=$1
converted_onnx_path=$2
tmp_save_path="./temp"

# for unet
python -m onnxruntime.transformers.models.stable_diffusion.optimize_pipeline \
-i ${original_onnx_path} \
-o ${tmp_save_path} \
--disable_bias_gelu \
--disable_bias_add \
--disable_bias_splitgelu \
--disable_nhwc_conv \
--use_group_norm_channels_first \
--use_multi_head_attention \
--float16

# for text_encoders and vae_decoder
python -m onnxruntime.transformers.models.stable_diffusion.optimize_pipeline \
-i ${original_onnx_path} \
-o ${converted_onnx_path} \
--disable_bias_gelu \
--disable_bias_add \
--disable_bias_splitgelu \
--disable_nhwc_conv \
--disable_attention \
--use_group_norm_channels_first \
--use_multi_head_attention \
--float16

rm -r ${converted_onnx_path}/unet
mv ${tmp_save_path}/unet ${converted_onnx_path}/
rm -r ${tmp_save_path}

echo "Finish to optimize ONNX models!"
echo "optimized ONNX models are saved in ${converted_onnx_path}"