利用 Elastic File System 加速 Stable Diffusion WebUI 部署推理

概述

面对在容器化环境中进行 Stable Diffusion WebUI 的部署，容器镜像往往高达几十 GB，在进行模型部署时，往往需要等待 15~20 分钟才能部署一套文生图、图生图的 AI 应用。有时候我们不得不采用提前预置 GPU 实例，以满足实时或近实时推理的需求，这就提升了 Stable Diffusion WebUI 的成本。此方案是通过利用 EFS 共享文件系统来加速 Stable Diffusion WebUI 部署推理。通过实验可以观察到有数倍的部署加速效果。特别适用于企业级客户对文生图推理应用的场景。

Stable Diffusion WebUI 部署的挑战

亚马逊云科技的 Deep Learning Container 是以 Docker 镜像的形式在 Amazon ECR 中提供的。每个 Docker 镜像都是针对特定的深度学习框架版本、Python 版本、以及 CPU 或 GPU 支持而构建的，用于训练或推理。在进行 Stable Diffusion WebUI（简称 SD-WebUI 或 sd-webui）的开发过程中，基于 Deep Learning Container 的开发是非常方便的。这里我们以 pytorch-inference:2.0.1-gpu-py310-cu118-ubuntu20.04-ec2 为基础构建 Stable Diffusion WebUI 推理镜像，具体步骤如下：

1. 准备 Dockerfile，将 stable-diffusion-webui 代码目录放在/sd-webui 下。因为需要以 root 用户执行 stable diffusion webui 程序，需要修改/sd-webui/stable-diffusion-webui/webui.sh 启动脚本中“can_run_as_root”参数为 1，同时禁用 venv 环境。

FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:2.0.1-gpu-py310-cu118-ubuntu20.04-ec2 as base
WORKDIR /sd-webui
RUN git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
WORKDIR stable-diffusion-webui/
RUN sed -i 's/can_run_as_root=0/can_run_as_root=1/' webui.sh
RUN sed -i 's/use_venv=1/use_venv=0/' webui.sh
RUN wget https://civitai.com/api/download/models/107472 -O models/Stable-diffusion/AWPainting.safetensors
RUN wget https://civitai.com/api/download/models/130072 -O models/Stable-diffusion/Realistic.safetensors
RUN wget https://civitai.com/api/download/models/117986 -O models/Stable-diffusion/Samaritan.safetensors
RUN ./webui.sh

2. 登陆 DLC 镜像仓库

aws ecr get-login-password --region us-east-1 |sudo docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-1.amazonaws.com

3. 基于准备好的 Dockerfile 进行镜像构建

sudo docker build -t public.ecr.aws/h6r2a7o6/emr:dlc-v1 .

4. 登陆自己的镜像仓库

aws ecr-public get-login-password --region us-east-1 | sudo docker login --username AWS --password-stdin public.ecr.aws/h6r2a7o6

5. 上传 docker 镜像

sudo docker push public.ecr.aws/h6r2a7o6/emr:dlc-v1

6. 准备容器 YAML 文件 dlc-sdwebui-v2.yaml

apiVersion: v1
kind: Pod
metadata:
  name: dlc-webui2
spec:
  containers:
  - name: app
    image: public.ecr.aws/h6r2a7o6/emr:dlc-v1
    command: ["/bin/sh"]
    args: ["-c",  "/sd-webui/stable-diffusion-webui/webui.sh --listen --share"]
    ports:
    - hostPort: 7860
      containerPort: 7860

7. 启动该容器

k apply -f dlc-sdwebui-v2.yaml

8. 检查容器运行情况

[ec2-user@ip-172-31-45-65 sd-webui]$ k get pods
NAME        READY   STATUS    RESTARTS   AGE
dlc-webui   1/1     Running   0          10m

9. 检查 sh 的启动情况

k logs pods/dlc-webui app

当我们看到 Running on local URL: http://0.0.0.0:7860 这一行，代表 sd-webui 已经监听到 7860 端口上，我们通过 NodePort 进行访问。浏览器中打开网址 http://54.188.121.114:7860/，输入 Prompt，可以看到可以生成图片。

我们进一步分析可以得知：

基于 Deep Learning Container 基础镜像的 Stable Diffusion WebUI 容器分析
镜像大小	21.4GB
容器启动时间	5m52s
WebUI 启动时间	~30s
模型切换时间	16s – 56s	具体时间取决于模型大小、是否第一次加载
推理时间	~2s	生成单张图片（未加速优化）

从上面表格可以看出，容器启动和模型切换占用了非常多的时间，我们下面来尝试缩短它们。

解决问题的思路

首先，分析一下镜像文件的构成，在容器中执行 du -h –max-depth=1 /

root@dlc-webui:/sd-webui/stable-diffusion-webui# du -h --max-depth=1 /
0	/boot
4.0K	/dev
3.0M	/etc
16M	/home
0	/media
0	/mnt
4.6G	/opt
12K	/proc
210M	/root
20K	/run
0	/srv
0	/sys
0	/tmp
8.6G	/usr
20M	/var
7.9G	/sd-webui
22G	/

分析一下/opt，/root，/tmp，/usr 目录的情况，可以看到：

目录	大小	内容
/opt	4.6 GB	包含 Stable Diffusion WebUI 需要的 Python 包，比如 pytorch
/root	210 MB	pip 安装缓存
/usr	8.6 GB	CUDA driver 安装在这个目录下
/sd-webui	7.9 GB	Dockerfile 里将 Stable Diffusion 代码和模型放在该目录

所以，我们优化的思路是：

将 Stable Diffusion WebUI 所依赖的 Pytorch 等 Python 包，安装到共享文件系统 EFS 上
禁用 pip 安装包缓存
将 Stable Diffusion WebUI 的代码和模型，安装到共享文件系统 EFS 上
使用 CUDA base 镜像减少 NVIDIA driver 占用空间

具体实现步骤

1. 构建 Stable Diffusion 基础镜像

1.1 以 nvidia/cuda:11.8.0-base-ubuntu22.04 为基础创建容器镜像文件 Dockerfile，安装 sd-webui 所需 Python 环境及相关软件包

FROM nvidia/cuda:11.8.0-base-ubuntu22.04 as base
USER root
RUN mkdir -p /sd-webui
WORKDIR /sd-webui
RUN apt update
RUN apt-get install -y python3.10 python3-venv python3-pip
RUN apt-get install -y wget
RUN apt install -y git libglib2.0-0 libgl1

1.2 对镜像进行打包

sudo docker build -t public.ecr.aws/h6r2a7o6/emr:cuda-v1 .

1.3 对打包好的镜像进行上传

sudo docker push public.ecr.aws/h6r2a7o6/emr:cuda-v1

2. 在 EFS 上构建 Stable Diffusion 运行环境

我们首先需要一个临时容器来挂载 EFS 文件系统，安装与配置 sd-webui 软件包及相关 Python 依赖。

2.1 首先新建一个 EFS 文件系统

采用默认配置即可，因为 EFS 现在默认即为 Elastic Throughput 模式，正好符合模型部署时大量突发 IO 需求。

2.2 在 EKS 上创建 EFS 卷

2.2.1 新建 PV YAML 文件 efs-pv.yaml

apiVersion: v1
kind: PersistentVolume
metadata:
  name: efs-pv
spec:
  capacity:
    storage: 100Gi
  volumeMode: Filesystem
  accessModes:
    - ReadWriteMany
  persistentVolumeReclaimPolicy: Retain
  csi:
    driver: efs.csi.aws.com
    volumeHandle: fs-028f8df16f11dde3e

2.2.2 新建 PV YAML 文件 efs-pvc.yaml

apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: efs-claim
spec:
  accessModes:
    - ReadWriteMany
  storageClassName: ""
  resources:
    requests:
      storage: 100Gi

2.2.3 我们需要一个临时性容器来挂载 EFS 卷，并在 EFS 上准备 Stable Diffusion WebUI 环境，临时性容器 YAML 文件 efs-pod.yaml 如下

apiVersion: v1
kind: Pod
metadata:
  name: sd-webui
spec:
  containers:
  - name: app
    image: public.ecr.aws/h6r2a7o6/emr:cuda-v1
    command: ["/bin/sh"]
    args: ["-c", "while true; do echo $(date -u) >> /data/out.txt; sleep 5; done"]
    volumeMounts:
    - name: persistent-storage
      mountPath: /sd-webui
  volumes:
  - name: persistent-storage
    persistentVolumeClaim:
      claimName: efs-claim

2.2.4 依次应用 efs-pv.yaml、efs-pvc.yaml、efs-pod.yaml 文件

k apply -f efs-pv.yaml
k apply -f efs-pvc.yaml
k apply -f efs-pod.yaml

2.2.5 进入临时性 pod

k exec -it pods/sd-webui -- /bin/bash

2.3 克隆 Stable Diffusion WebUI 至/sd-webui 目录

git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git

2.4 安装 Stable Diffusion WebUI 的 Python 依赖包

pip3 install --no-cache-dir  gdown -t python3/lib/site-packages
pip3 install --no-cache-dir -r stable-diffusion-webui/requirements_versions.txt -t python3/lib/site-packages
pip3 install --no-cache-dir  git+https://github.com/openai/CLIP.git -t python3/lib/site-packages

2.5 因为容器中是以 root 在运行，而默认 Stable Diffusion WebUI 只允许非 root 运行，所以需要修改启动脚本，同时禁用 venv，忽略已安装文件

sed -i 's/can_run_as_root=0/can_run_as_root=1/' ./stable-diffusion-webui/webui.sh
sed -i 's/use_venv=1/use_venv=0/' ./stable-diffusion-webui/webui.sh
sed -i 's/export PIP_IGNORE_INSTALLED=0/export PIP_IGNORE_INSTALLED=1/' ./stable-diffusion-webui/webui.sh

2.6 Stable Diffusion WebUI 默认安装路径为~/.local，修改为 EFS 卷上安装，改为/sd-webui/installed 目录里，首先创建该目录

mkdir -p /sd-webui/installed

2.7 修改 webui-user.sh 文件，以便安装目录指向 sd-webui/installed

sed -i '2i\install_dir="/sd-webui/installed"' ./stable-diffusion-webui/webui-user.sh

2.8 修改 Stable WebUI 启动参数

sed -i '3i\export COMMANDLINE_ARGS="--skip-version-check --skip-python-version-check --skip-torch-cuda-test --skip-install --listen --share"' ./stable-diffusion-webui/webui-user.sh

2.9 下载模型文件，从 Huggingface 中下载基础默认模型 `Stable-diffusion/v1-5-pruned-emaonly`，从 Civitai 中分别下载 `AWPainting`、`Realistic`、`Samaritan` 三个模型

wget https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors -O ./installed/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors
wget https://civitai.com/api/download/models/107472 -O ./installed/stable-diffusion-webui/models/Stable-diffusion/AWPainting.safetensors
wget https://civitai.com/api/download/models/130072 -O ./installed/stable-diffusion-webui/models/Stable-diffusion/Realistic.safetensors
wget https://civitai.com/api/download/models/117986 -O ./installed/stable-diffusion-webui/models/Stable-diffusion/Samaritan.safetensors

2.10 执行初始化

PYTHONPATH=/sd-webui/python3/lib/site-packages ./stable-diffusion-webui/webui.sh

2.11 退出临时 pod

Exit

3. 创建 Stable Diffusion WebUI 容器的 YAML 文件 selfbuilt-sdwebui.yaml

apiVersion: v1
kind: Pod
metadata:
  name: selfbuilt-webui
spec:
  containers:
  - name: app
    image: public.ecr.aws/h6r2a7o6/emr:cuda-v1
    command: ["/bin/sh"]
    args: ["-c",  "export PYTHONPATH=/sd-webui/python3/lib/site-packages;/sd-webui/stable-diffusion-webui/webui.sh"]
    volumeMounts:
    - name: persistent-storage
      mountPath: /sd-webui
    ports:
    - hostPort: 7860
      containerPort: 7860
  volumes:
  - name: persistent-storage
    persistentVolumeClaim:
      claimName: efs-claim

4. 启动 Stable Diffusion WebUI 容器

k apply -f selfbuilt-sdwebui.yaml

5. 检查容器运行情况

[ec2-user@ip-172-31-45-65 sd-webui]$ k get pods
NAME              READY   STATUS    RESTARTS   AGE
selfbuilt-webui   1/1     Running   0          3m52s

6. 检查 sh 的启动情况

k logs pods/selfbuilt-webui

当我们看到 Running on local URL: http://0.0.0.0:7860 这一行，代表 WebUI 已经监听到 7860 端口上，我们通过 NodePort 进行访问。浏览器中打开网址 http://35.89.157.24:7860/，输入 Prompt，可以看到可以生成图片。

进一步分析可以得知：

基于 Deep Learning Container 基础镜像的 Stable Diffusion WebUI 启动分析
镜像大小	855MB
容器启动时间	18s
WebUI 启动时间	~30s
模型切换时间	8s – 22s	具体时间取决于模型大小、是否第一次加载
推理时间	~2s	生成单张图片（未加速优化）

两种部署方法的对比分析

在本文中可以看到，通过我们自己构建 Stable Diffusion WebUI 容器镜像，利用 EFS 共享文件系统来存储 Pytorch 包等相关依赖包、Stable Difussion 代码和模型文件对 Stable Diffusion WebUI 模型的部署进行了优化。下图是优化前后的部署速度的对比。

具体数值如下：

	优化前	优化后	优化前后对比%
镜像大小（MB)	21913.6	855	3.9%
容器启动时间（秒）	352	18	5.1%
WebUI 启动时间（秒）	30	30	100.0%
模型切换时间（秒）	36	15	41.7%
推理时间（秒）	2	2	100.0%

从 Stable Diffusion WebUI 部署的角度来看，应用需要等待 webui 脚本启动完成才能进行推理，所以我们可以认为 Stable Diffusion WebUI 部署时间包含容器启动时间和 webui 脚本启动时间，优化后为模型部署需要 48 秒，相对于优化前 382 秒时间减少了 87%。而不同模型切换时间，平均减少了 41.7%。

总结

通过上述对比分析，我们可以看到通过利用 EFS 极大的缩短了模型部署与切换的时间，满足了实时或近实时推理的应用需求。同时，EFS Elastic Throughput 的特性使您无需提前预置吞吐量与 IOPS，在模型加载时即可获得高达 10 GiBps 带宽，相对于预置 EBS 吞吐量方案来说，节省了高达 87.8% 的存储成本。

参考链接

https://github.com/aws/deep-learning-containers/blob/master/available_images.md

https://github.com/AUTOMATIC1111/stable-diffusion-webui.git

亚马逊AWS官方博客