在亚马逊云科技Amazon SageMaker上进行Stable Diffusion模型训练和推理
Stable Diffusion Quick Kit是一个基于亚马逊云科技Amazon SageMaker进行Stable Diffusion模型快速部署的工具包，包括了一组示例代码、服务部署脚本、前端UI，可以帮助可以快速部署一套Stable Diffusion的原型服务。
本文将介绍如何在SageMaker Training Job中加载进行Stable Diffusion XL（以下简称SDXL）的Dreambooth微调训练，及训练完成后使用Stable Diffusion WebUI开源框架进行模型部署和即时推理，实现训推一体的整体pipeline及业务流程。
背景介绍
Dreambooth微调训练
Dreambooth是Stable Diffusion模型训练的一种方式，通过输入instance_prompt定义实体主体（e.g.人物或者实体物品）和instance images的fine-tuning图像，抽取原SD中UNet，VAE网络，将instance prompt和instance images图像与之绑定，以便后续生成的图片中只要有instance的prompt中的关键词，即可保持输入instance图片里面的主体实体，实现人物和物品生图时的高保真效果。
Stable Diffusion WebUI
Stable Diffusion WebUI是基于Stable Diffusion开发的一个开源的可视化软件，WebUI在Stable Diffusion txt2img，img2img生图基础上拓展了很多插件来增强Stable Diffusion的生图能力，比如Ultimate Upscale、Inpain等，使得开发者可以方便地通过界面拖拽或者API调用进行Stable Diffusion模型的加载和调用。
相对于Diffuser SDK的模型推理，WebUI有更丰富的调用参数及更多的插件支持，因此同样模型的出图效果某些场景下会比Diffuser更好，这也是目前业界不少客户使用WebUI API方式进行推理生图的原因。
训练+推理业务场景
在遇到使用Stable Diffusion模型微调和推理的业务场景中，针对ToB端客户，通常会上传需要训练的图像，使用Dreambooth训练人物（如模特或者数字人）和商品（如箱包，衣服），然后针对训练好的模型批量生成海报/广告/logo等文案素材的图像，该过程并不需要像app应用一样实时交互的出图，而是一个离线异步的过程。
这种情况下，可以在训练任务的算力机上，同时安装部署模型微调和模型推理的框架，利用SageMaker Training Job方式，将微调和推理放到一个job中，微调训练完成，即加载model进行推理出图，从而一次性完成模型微调（Dreambooth）+模型推理（WebUI API）整个完整pipeline工作，将推理的模型改造到训练任务中，而不用再单独部署模型的服务端点。
同时，SageMaker Training Job支持Spot竞价实例，训练任务完成则推理出图也完成，机器资源释放，进一步帮助用户节约整体的成本。
SageMaker Training Job中进行SDXL Dreambooth Fine-tune
Dreambooth训练框架
Stable Diffusion 1.x版本时，Dreambooth fine-tune有多种开源版本的微调框架，SDXL版本后，Diffuser官方在HuggingFace社区发布了基于LoRA的Dreambooth fine-tune框架，代码相对于原1.x版本更加简洁，且使用了更新的xformers加速框架，支持Flash Attension v2，其Pytorch版本也升级到了2.0以上。
其中train_dreambooth_lora_sdxl.py就是微调训练Dreambooth的代码。
SageMaker Training Job脚本
在SageMaker Training Job中，可以clone上一章节的diffuser官方repo训练代码作为source训练脚本目录，并将其依赖的xformers，deepspeed等依赖打包在Docker训练镜像中，通过shell entrypoint方式在算力机上拉起其训练脚本。
详细如下：
-
准备source源代码目录并clone官方代码
-
打包训练任务的docker镜像（使用Amazon预置的0.0+cuda118 HuggingFace DLC容器作为基础镜像，与diffuser官方pytorch/cuda版本保持一致）
-
dockerfile编写
-
build镜像并推送到Amazon ECR镜像仓库
-
准备训练图像，这里我们使用官方示例dataset图像
-
图像数据上传到$images_s3uri的S3路径，以便SageMaker Training Job拉取。
-
SageMaker Estimator拉起Training Job
-
训练任务脚本编写，这里采用shell entrypoint方式，方便调用diffuser官方脚本，且传递环境变量。
我们通过SageMaker提供的Pytorch的Estimator训练器SDK，拉起Training Job训练任务。
Dreambooth训练调参
SDXL Dreambooth Fine-tune的训练参数与之前1.x版本调参类似，这里把Diffuser框架及SageMaker新加的主要配置参数说明如下：
-
‘images’：f”s3://{bucket}/dreambooth-xl/images/”：上一步骤中准备好的dreambooth微调图像数据，通过inputs参数指定S3路径，SageMaker会自动将该路径下训练图像上传到训练算力机的/opt/ml/data/input/images目录下
-
keep_alive_period_in_seconds：该参数是SageMaker Training Job的warmpool，设置后可以把下一次训练机器保持在该用户的一个资源池中，这样方便多个SDXL Dreambooth训练时的镜像拉起，节省耗时的开销
-
enable_xformers_memory_efficient_attention：启用xformers的flash attention关注度计算优化，加速训练过程
-
train_use_spot_instance：是否使用spot竞价实例进行训练，进一步节省成本
-
max_run：训练任务的最大运行时间
-
max_wait：等待竞价实例的最长时间，如果使用spot竞价实例该参数是必须的
SageMaker Training Job中安装部署Stable Diffusion WebUI
如上文所述，训练完成后可以直接使用fine-tuned模型进行推理出图，这里采用Stable Diffusion WebUI进行推理，需要在training job训练算力机上安装部署开源的WebUI组件，将模型目录同步到WebUI的model location下，然后调用WebUI API text2img/img2img出图，详细如下：
docker镜像脚本
由于是在training job中进行推理，扩充训练任务的dockerfile镜像文件，将Stable Diffusion WebUI组件及依赖同样的方式和上文中training的dockerfile打包到一起：
WebUI启动脚本
使用上述章节同样的build & push脚本，将docker镜像打包推送，然后在统一训练和推理的entry point脚本中启动训练任务，任务完成后启动WebUI。
SageMaker Training Job中对Fine-tuned Dreambooth Model进行推理
在start_sd_webui.py脚本启动WebUI服务器之后，即可使用WebUI API进行txt2img/img2img的推理调用，其推理API与官方参数一致。
由于在同一台训练算力机上，其URI为localhost（0.0.0.0）对应端口及API路径前缀。
总结
本文介绍了在Quick Kit中使用SageMaker Training Job对SDXL模型进行Dreambooth微调，并且可以在训练完成后对fine-tuned后的模型使用Stable Diffusion WebUI进行推理，实现从训练到推理的一体化操作，满足客户对于快速训练人物或商品实体并批量推理出图的需求。