sd-webui-EasyPhoto源码解析

概述

本文将对基于SDWebUI插件生态的EasyPhoto插件进行源码解析。EasyPhoto插件是由FaceChain-Inpaint功能的开发团队迅速推出的一款开源项目，旨在满足AIGC领域同学们对适配真人写真功能的需求。通过本插件，用户可以上传5-20张同一个人的照片，利用Lora模型进行快速训练，并结合用户提供的模板图片，快速生成真实、逼真、美观的个人写真照片。

在AIGC领域中，妙鸭相机作为一款收费产品，成功展示了如何利用AIGC技术通过少量的人脸图片进行建模，以极短的时间提供高质量的个人写真服务，吸引了大量付费客户。与此同时，StableDiffusion领域的开源社区也迅速发展起来，其中涌现了许多基于Modelscope开源社区和diffusers的项目，例如FaceChain，旨在指导用户快速开发个人写真功能。

然而，对于大量使用SDWebUI的AIGC同学们而言，在短时间内却没有找到一个足够出色的开源插件，能够适配真人写真这一功能。为了解决这一问题，FaceChain-Inpaint团队迅速响应社区需求，开发了EasyPhoto插件。该插件基于SDWebUI插件生态，为用户提供了一种简便的方式来训练Lora模型，并结合用户上传的模板图片，快速生成真实、逼真、美观的个人写真照片。

本文将深入解析EasyPhoto插件的源码，探讨其实现原理和关键技术。我们将逐步介绍插件的功能和使用方法，包括如何上传照片、训练Lora模型以及生成个人写真照片的过程。通过对源码的解析，读者将对EasyPhoto插件的工作原理有更深入的了解，从而能够更好地应用和定制该插件，满足个人写真需求。

在接下来的章节中，我们将详细介绍EasyPhoto插件的核心代码和关键算法，以及其与SDWebUI插件生态的集成方式。同时，我们还将讨论插件的性能优化和扩展性，以及可能的改进方向。通过本文的阅读，读者将能够全面了解EasyPhoto插件，并在AIGC领域中应用其强大的功能，为用户提供更出色的个人写真体验。

代码结构介绍

其中，

├── api_test                ---  API测试代码
├── images                ---  模版图片
├── javascript                ---  UI代码
├── models                ---  模型相关文件，包含训练模版图片文件夹、推理模版图片文件夹以及SD 1.5相关配置
├── scripts                ---  核心代码，也是本文主要介绍的部分
│   ├── train_kohya         ---  Lora训练相关文件
│   ├── api.py         ---  API接口服务端入口代码
│   ├── easyphoto_config.py            ---  EasyPhoto配置相关信息
│   ├── easyphoto_infer.py        ---  EasyPhoto推理相关逻辑代码
│   ├── easyphoto_train.py           ---  EasyPhoto训练相关逻辑代码
│   ├── easyphoto_ui.py             ---  EasyPhoto UI相关代码，通过gradio实现
│   ├── easyphoto_utils.py                ---  EasyPhoto 工具函数相关代码
│   ├── face_process_utils.py                ---  脸部处理相关函数代码
│   ├── preprocess.py                ---  EasyPhoto 训练预处理相关函数代码
│   └── sdwebui.py            ---  SDWebUI插件入口相关代码
├── .gitignore                    ---  用于指定Git版本控制系统忽略哪些文件和目录，它们将不会被添加、提交或跟踪到Git仓库中。
├── COVENANT.md              ---  开发者公约
├── COVENANT_zh-CN.md          ---  开发者公约中文版
├── LICENSE            ---  Apache License 2.0
├── README.md           ---  EasyPhoto, Your Smart AI Photo Generator
├── README_zh-CN.md           ---  EasyPhoto，您的智能 AI 照片生成器介绍
└── install.py            ---  安装脚本，运行后主要安装各种python依赖包

在这个训练和推理过程中，EasyPhoto大量的使用了 ModelScope 和其他Github的开源模型，用于完成如下的人脸功能：

训练过程代码解析

easyphoto_train.py

主要是easyphoto_train_forward函数，此函数的主要作用是对用户提供的图像进行训练，包括预处理、执行训练命令、保存训练后的权重文件等。最终返回一个训练完成的消息。high level的逻辑为：

• 代码首先进行一些参数和路径的检查和设置，确保训练所需的各项条件满足。
• 接下来，对用户提供的图像进行预处理，包括训练图片生成（使用【显著性检测模型】和【皮肤美化模型】获得干净的人脸训练图像，该图像基本上只包含人脸）、固定标签（不需要使用标签器，而且效果很好）准备、模板图像（ref_image，可以理解上传照片中最像的照片）选定、保存原始备份图像和进行进一步的处理和转换。
• 然后，根据操作系统类型构建训练命令，并调用子进程执行训练。在训练过程中，会利用模板图像（ref_image）进行实时验证，在训练结束后，会计算验证图像与用户上传图像之间的人脸ID差距，从而实现Lora融合，确保Lora效果好。
• 如果启用了强化学习（RL），则执行RL训练。
• 最后，检查训练后的权重文件是否存在，并将最佳权重文件复制到WebUI路径下。

整个过程中使用了4次子进程调用。

preprocess.py

其中最关键的逻辑是预处理的逻辑，代码在preprocess.py中。核心逻辑有：

• 判断图像文件格式是否支持，若不支持则跳过。
• 使用【人脸检测模型】检测图像中的人脸边界框和人脸关键点。根据人脸关键点计算人脸角度。
• 判断人脸边框尺寸是否满足要求，若小于128像素则跳过此图片。
• 对人脸进行裁剪，并使用【皮肤修饰模型】对裁剪后的人脸进行处理。
• 使用【人脸识别模型】获取人脸的嵌入向量。
• 记录人脸的嵌入向量和角度，以及相关的图像路径和裁剪后的人脸图像。对人脸的嵌入向量进行比较和排序，根据中位值最近来筛选出参考人脸（用于人脸融合，俗称换脸）。
• 根据人脸的嵌入向量以根据中位值最近进行排序，选择前15个人脸作为处理的目标。
• 对选定的人脸进行图像增强处理，最多处理一半的人脸（避免全部处理反而效果不佳）。对增强后的人脸进行人脸边界框的修正。
• 使用语义分割模型进行显著性检测，生成人脸掩码。根据掩码提取人脸区域，将人脸与背景合成。
• 将处理后的图像保存到指定的输出目录，并将验证提示写入对应的文本文件。
• 将结果以JSON格式保存。

预处理核心就是规避由于用户传入错误的人脸图片（没有人脸的图片，其他人脸的图片等情况），照片格式不对、效果不佳、尺寸不合适、背景过于复杂等问题。这一块的逻辑是非常重要的。处理后的照片如下所示：

其中，*.txt文件中的prompt为：

easyphoto_face, easyphoto, 1person.

可以发现，如果原图质量不佳，处理后的训练图也会有一些问题。

推理过程（写真图片生成）代码解析

easyphoto_infer.py

其中最关键的是easyphoto_infer_forward函数，它实现了一个完整的图像处理流程，包括图像修复、颜色转移、背景恢复、皮肤修饰、肖像增强等操作，并可以计算生成图像与roop图像之间的FaceID相似度。核心步骤为：

• 一、检查和下载模型权重文件。
• 二、针对每个用户ID进行处理：
  • 检查用户ID的有效性。
  • 构建输入提示（input_prompt）字符串，其中包括用户ID、默认的正面提示（default_positive_prompt）、额外的提示信息（additional_prompt）等。
  • 获取用户ID对应的最佳图像（face_id_image）和roop图像（roop_image）。
  • 使用RetinaFace模型检测人脸，获取人脸框（face_id_retinaface_boxes）和关键点（face_id_retinaface_keypoints）。
  • 对模板图像进行预处理，包括裁剪和缩放。
  • 根据输入的参数，进行图像处理的操作：
    • 第一次扩散（first_diffusion）：
      • 使用输入图像和控制网络单元，进行图像修复（inpaint）。
      • 可选的颜色转移操作，将修复后的图像颜色转移到原始模板人脸区域。
    • 第二次扩散（second_diffusion，第二次生成将使用更高的分辨率的图片）：
      • 如果选择了面部融合操作，将roop图像与第一次扩散结果进行融合。
      • 可选的颜色转移操作，将修复后的图像颜色转移到原始模板人脸区域。
    • 如果选择了背景恢复操作，对整个图像进行第三次扩散（third_diffusion），恢复背景。
    • 进行皮肤修饰和肖像增强操作，使用对应的模型进行处理。
  • 将处理后的图像保存到输出列表中。
  • 如果选择了显示分数（display_score），计算第二次扩散生成图像与roop图像之间的FaceID相似度，并将结果保存到FaceID输出列表中。
• 三、返回处理结果和FaceID输出列表。

总结

本文详述了训练和推理两个过程的代码解析，输出了核心业务逻辑。希望能为大家理解EasyPhoto提供帮助。