转换-在macOS上转换和运行CoreML模型

开源大模型很多，为了让我的小Mac能够高效的推理，故有此文。

模型运行

在 Github 上转了老半天才找到一个可供 Mac 使用的 App：Mochi Diffusion

将的到的模型放置在~/MochiDiffusion/models中，即可被加载使用。

具体使用步骤非常简单，就不再赘述。

模型转换

目前支持的转换格式包括CheckPoint、Diffusers。转换步骤是一级一级来的，直接可以转换的是Diffusers，路径图如下：

flowchart LR
CheckPoint
Diffusers
CoreML
CheckPoint --> Diffusers --> CoreML

目前量化方面还有点问题，等待apple/ml-stable-diffusion修复。具体转换步骤，这里用一个脚本来介绍，在此之前请安装必要的环境：

• git
• conda
• xcode

#!/bin/bash
# 脚本名称: 模型批量转换
# 描述: 模型批量转换脚本
# 版本: 1.0
# 作者: Cikaros
# 邮箱: Cikaros<at>qq.com
# 日期: 
# 版权: (C) 2024 Your Company. All rights reserved.

# 配置环境 当存在错误及时停止
source ~/.bash_profile
set -e

VERSION=1.0

PYTHON_VERSION=3.8
PYTHON_ENV_NAME=python38

ML_SD_URL=https://github.com/apple/ml-stable-diffusion.git
CONVERT_FILE_URL=https://raw.githubusercontent.com/huggingface/diffusers/main/scripts/convert_original_stable_diffusion_to_diffusers.py

WORK_SPACE="$(pwd)"
SAFETENSORS_MODELS_PATH="$WORK_SPACE/safetensors"
DIFFUSERS_MODELS_PATH="$WORK_SPACE/diffusers"
SPLIT_EINSUM_PATH="$WORK_SPACE/split-einsum"
FINISH_PATH="$WORK_SPACE/safetensors-finish"
ML_SD_PATH="$WORK_SPACE/ml-stable-diffusion"
CONVERT_FILE_PATH="$WORK_SPACE/convert_original_stable_diffusion_to_diffusers.py"

check() {
	if type "$1" &>/dev/null; then
	    echo "已检测到【$1】已安装。"
	else
		echo "未检测到【$1】已安装，请先安装必要环境！"
	    exit -1
	fi
}

# 检查环境
check_env() {
	check git
	check conda
	check xcode-select
	# 检查文件夹是否存在
	[ ! -d "$SAFETENSORS_MODELS_PATH" ] && mkdir -p "$SAFETENSORS_MODELS_PATH"
	[ ! -d "$DIFFUSERS_MODELS_PATH" ] && mkdir -p "$DIFFUSERS_MODELS_PATH"
	[ ! -d "$SPLIT_EINSUM_PATH" ] && mkdir -p "$SPLIT_EINSUM_PATH"
	[ ! -d "$FINISH_PATH" ] && mkdir -p "$FINISH_PATH"
	echo "check success!"
}

init_env() {
	cd $WORK_SPACE
	#初始化环境
	conda env list | grep -q "$PYTHON_ENV_NAME" || \
      conda create -y -n "$PYTHON_ENV_NAME" python=$PYTHON_VERSION
	conda activate $PYTHON_ENV_NAME
	# TODO 准备 https://github.com/Cikaros/ml-stable-diffusion.git
	if [ ! -e $ML_SD_PATH ]; then
		echo "ml-stable-diffusion项目不存在，正在初始化..."
		git clone $ML_SD_URL
		cd $ML_SD_PATH
		pip install -e .
		pip install omegaconf safetensors
	fi
	if [ ! -f $CONVERT_FILE_PATH ]; then
		echo "CheckPoint转换脚本不存在，正在初始化..."
		wget $CONVERT_FILE_URL - $CONVERT_FILE_PATH
		pip install torch diffusers transformers python_coreml_stable_diffusion
	fi
	cd $WORK_SPACE
}

redo_init_env() {
	conda remove -n $PYTHON_ENV_NAME --all
	rm -rf $ML_SD_PATH
	rm -rf $CONVERT_FILE_PATH
	init_env
}

convert() {
	FILE_PATH=$1
	FILE_NAME_WITH_EXT=$(basename "$1")
	FILE_NAME="${FILE_NAME_WITH_EXT%.*}"
	DIFFUSERS_PATH="$DIFFUSERS_MODELS_PATH/$FILE_NAME"
	COREML_PATH="$SPLIT_EINSUM_PATH/$FILE_NAME"
	EXPORT_SPLIT_EINSUM_PATH=$SPLIT_EINSUM_PATH/$FILE_NAME-split-einsum
	if echo "$FILE_NAME" | grep -q "XL"; then
		# TODO CheckPoint -> Diffusers
	    python $CONVERT_FILE_PATH \
		  --checkpoint_path $FILE_PATH \
		  --extract_ema \
		  --upcast_attention \
		  --from_safetensors \
		  --dump_path $DIFFUSERS_PATH \
		  --device cpu \
		  --half \
		  --pipeline_class_name StableDiffusionXLPipeline
		# TODO Diffusers -> CoreML
		python -m python_coreml_stable_diffusion.torch2coreml \
		  --convert-text-encoder \
		  --convert-vae-decoder \
		  --convert-vae-encoder \
		  --convert-unet \
		  --model-version $DIFFUSERS_PATH \
		  --attention-implementation SPLIT_EINSUM \
		  -o $COREML_PATH \
		  --bundle-resources-for-swift-cli \
		  --xl-version
		python -m python_coreml_stable_diffusion.torch2coreml \
		  --convert-unet \
		  --unet-support-controlnet \
		  --model-version $DIFFUSERS_PATH \
		  --attention-implementation SPLIT_EINSUM \
		  -o $COREML_PATH \
		  --bundle-resources-for-swift-cli \
		  --xl-version
	else
	    python $CONVERT_FILE_PATH \
		  --checkpoint_path $FILE_PATH \
		  --extract_ema \
		  --upcast_attention \
		  --from_safetensors \
		  --dump_path $DIFFUSERS_PATH \
		  --device cpu \
		  --half
		python -m python_coreml_stable_diffusion.torch2coreml \
		  --convert-text-encoder \
		  --convert-vae-decoder \
		  --convert-vae-encoder \
		  --convert-unet \
		  --model-version $DIFFUSERS_PATH \
		  --attention-implementation SPLIT_EINSUM \
		  -o $COREML_PATH \
		  --bundle-resources-for-swift-cli
		python -m python_coreml_stable_diffusion.torch2coreml \
		  --convert-unet \
		  --unet-support-controlnet \
		  --model-version $DIFFUSERS_PATH \
		  --attention-implementation SPLIT_EINSUM \
		  -o $COREML_PATH \
		  --bundle-resources-for-swift-cli
	fi
	# TODO CoreML finish
	[ ! -d $EXPORT_SPLIT_EINSUM_PATH ] && mkdir -p $EXPORT_SPLIT_EINSUM_PATH
	ditto "$COREML_PATH/Resources" $EXPORT_SPLIT_EINSUM_PATH
	# TODO Clean
	mv $FILE_PATH "$FINISH_PATH/"
	rm -rf $DIFFUSERS_PATH
	rm -rf $COREML_PATH
}

# 使用提示
usage() {
	echo "Usage: $0 [options]"
	echo "options:"
	echo "  -h, --help       显示帮助信息"
	echo "  -v, --version    显示版本信息"
	echo "  -i, --init       初始化环境"
	echo "  -c, --convert    转换特定模型"
	exit 1
}

# 参数处理
while [[ "$#" -gt 0 ]]; do
	case $1 in
		-h|--help) usage ;;
		-v|--version) echo "Version: $VERSION"; exit ;;
		-i|--init) redo_init_env; exit ;;
		-c|--convert) init_env; convert $2; exit ;;
		*) echo "Unknown option: $1"; usage ;;
	esac
	shift
done

# 无参数时的默认逻辑
main() {
	check_env
	init_env
	# 使用 find 命令递归扫描目录及其子目录下所有的 .safetensors 文件
	find "$SAFETENSORS_MODELS_PATH" -type f -name "*.safetensors" | while read FILE; do
	    echo "正在处理文件：$FILE"
	    convert $FILE
	done
	echo "处理完成！"
}

# 脚本入口点
main "$@"