苹果的开源工具包

苹果开源了一个工具包 Core ML Stable Diffusion ，使用Core ML（苹果的机器学习模型）来实现AI作图，通过这个工具包实现在苹果芯片设备上运行 Stable Diffusion。

这个工具包包括：

python_coreml_stable_diffusion, a Python package for converting PyTorch models to Core ML format and performing image generation with Hugging Face diffusers in Python

StableDiffusion, a Swift package that developers can add to their Xcode projects as a dependency to deploy image generation capabilities in their apps. The Swift package relies on the Core ML model files generated by python_coreml_stable_diffusion

安装和准备工作

安装运行环境软件：

Install Homebrew package manager
Install Rust
Install pyenv
Install Anaconda
Install Python3.8

创建 Python 运行期环境

创建运行期环境 coreml_stable_diffusion

conda create -n coreml_stable_diffusion python=3.8 -y

切换启用运行期环境 coreml_stable_diffusion

conda activate coreml_stable_diffusion

下载软件包，安装依赖 Python 包

下载 GitHub 源码包，Core ML Stable Diffusion ，解压到本地文件夹，或使用 git 命令 Clone 项目。

进入指定路径

mkdir ~/Desktop/

Clone 克隆项目代码

git clone https://github.com/apple/ml-stable-diffusion.git

进入程序主路径

cd ml-stable-diffusion

安装依赖 Python 包

pip install -e .

命令行登录 HuggingFace

注册并取得 HuggingFace.co 的授权 API Token，我测试的 Readonly type，一切OK。

通过终端登录 HuggingFace

huggingface-cli login

输入申请的API Token

转换模型

创建文件夹 AI_MODELS

mkdir AI_MODELS

软件包使用的默认模型是 CompVis/Default Model is v1.4，下载并转化默认模型为苹果专用ML模型，存储在当前文件夹里的 AI_MODELS

python -m python_coreml_stable_diffusion.torch2coreml \
    --convert-unet  \
    --convert-text-encoder  \
    --convert-safety-checker  \
    -o ./AI_MODELS

其他模型可以在网站上查找，并复制模型的版本标识，如 CompVis/stable-diffusion-v1-5

python -m python_coreml_stable_diffusion.torch2coreml  \
    --convert-unet  \
    --convert-text-encoder  \
    --convert-safety-checker  \
    -o ./AI_MODELS \
    --model-version CompVis/stable-diffusion-v1-5

下载和转换时间较长，我的机器需要2、3个小时。

测试生成

创建文件夹 AI_IMAGES

mkdir AI_IMAGES

测试一下，见证奇迹的时刻

python -m python_coreml_stable_diffusion.pipeline \
          -i ./AI_MODELS \
          -o ./AI_IMAGES \
          --compute-unit CPU_AND_GPU \
          --seed 888 \
          --prompt "a man riding a horse on Mar."

生成时间较长，我的机器需要20分钟。如果有问题，可以使用参数 --compute-unit CPU_ONLY ，时间更长。

python -m python_coreml_stable_diffusion.pipeline \
          -i ./AI_MODELS \
          -o ./AI_IMAGES \
          --compute-unit CPU_ONLY \
          --seed 888 \
          --prompt "a man riding a horse on Mar."
          --negative-prompt "nsfw"

管理咒语（Prompt）

为了方便管理生成的图像和Prompt，我重写了生成的脚本，替换 pipeline.py 中的 get_image_path 过程，按日期存储相关的图像和脚本及Prompt

def get_image_path(args, **override_kwargs):
    """ mkdir output folder and encode metadata in the filename
    """
    x = datetime.datetime.now()

    out_folder = f"{override_kwargs.get('model_version', None) or args.model_version.replace('/', '__')}"
    out_folder = x.strftime("%Y") + x.strftime("%m") + x.strftime("%d") + "/" + out_folder
    out_folder = os.path.join(args.o, out_folder)
    os.makedirs(out_folder, exist_ok=True)

    out_fname = x.strftime("%Y") + x.strftime("%m") + x.strftime("%d") + "_" +  x.strftime("%H") +  x.strftime("%M") +  x.strftime("%S") +"_Seed_"
    out_fname += f"{override_kwargs.get('seed', None) or args.seed}"
    if args.scheduler is not None:
        out_fname += f"_customScheduler_{override_kwargs.get('scheduler', None) or args.scheduler}"
        out_fname += f"_numInferenceSteps{override_kwargs.get('num_inference_steps', None) or args.num_inference_steps}"


    out_path_png = os.path.join(out_folder, out_fname + ".png")
    out_path_text = os.path.join(out_folder, out_fname + ".txt")
    logger.info(f"Saving python script to {out_path_text}")
    py_script = f"python -m python_coreml_stable_diffusion.pipeline \\\n\
    -i ./ai_model \\\n\
    -o ./ai_image \\\n\
    --compute-unit {args.compute_unit} \\\n\
    --seed {args.seed} \\\n\
    --prompt \"{args.prompt}\" \\\n\
    --negative-prompt \"{args.negative_prompt}\" \\\n\
    --model-version {args.model_version}"
    file = open(out_path_text, 'w')
    file.write("python command:\n\n")
    file.write(py_script)
    file.write("\n\n")
    file.write("prompt:\n\n")
    file.write(args.prompt)
    file.write("\n\n")
    file.write("negative-prompt:\n\n")
    file.write(args.negative_prompt)
    file.write("\n\n")
    file.close()

    return out_path_png

也可以在我克隆的项目里直接下载。