本篇文章聊聊 Stable Diffusion WebUI 中的核心組件，強壯的人臉圖像面部畫面修復模型 CodeFormer 相關(guān)的事情。
寫在前面
在 Stable Diffusion WebUI 項目中，源碼 modules 目錄中，有一個有趣的目錄叫做 CodeFormer，它就是本文的主角啦。
CodeFormer 是一個很棒的開源項目sczhou/CodeFormer，被應用在許多項目中，它的論文（arxiv.org/abs/2206.11253）在 2022 年被 “經(jīng)信息處理系統(tǒng)大會”（NeurIPS）接收后，自 2022 年 6 月代碼開始放出至今的一年出頭的時間里，Star 數(shù)量迅速升到了接近萬星的水平，足見開源社區(qū)的認可程度。
從“點贊”者的地區(qū)分布來看，在國內(nèi)的支持者占了項目近一半的人數(shù)。
在展開代碼走讀之前，先玩一下有助于對項目的理解。和往常一樣，我將項目封裝成了 Docker 容器，完整的項目，我上傳到了 GitHub soulteary/docker-codeformer，自取的時候別忘記“一鍵三連”。
下面進入熱身階段。
CodeFormer 相關(guān)的前置知識
CodeFormer 是一個基于 Transformer 的預測網(wǎng)絡，利用 code prediction 根據(jù)上下文來優(yōu)化人臉圖像，能夠在畫面非常模糊、甚至有損壞的情況下，修復出接近原始的、極高質(zhì)量的圖像畫面。
項目核心的外部依賴有三個：
ultralytics/yolov5，是目前最受歡迎的目標檢測開源項目，在 CodeFormer 中，作者使用了項目中的非常少的一部分代碼實現(xiàn)，主要功能為 face_detector.py 人臉檢測模塊。相關(guān)代碼位于項目 facelib/detection/yolov5face。
xinntao/facexlib，提供了當前開源人臉相關(guān) STOA 的方法的工具庫。項目使用了其中的 detection 、 parsing 和 utils 三個模塊，并進行了一些修改和調(diào)整。相關(guān)代碼位于項目根目錄的 facelib。
XPixelGroup/BasicSR，開源圖像和視頻恢復工具箱，能夠提供超分辨率、去噪、去模糊等能力，項目包含了非常多的網(wǎng)絡： EDSR、RCAN、SRResNet、SRGAN、ESRGAN、EDVR、BasicVSR、SwinIR、ECBSR 等等，并且支持 StyleGAN2、DFDNet。相關(guān)代碼位于項目根目錄的 basicsr。
項目使用的 BasicSR 項目并非原始項目團隊發(fā)布的版本，而是經(jīng)過修改的，目前未在發(fā)布頁面提供的 1.3.2 版本。
相關(guān)細節(jié)在本文“代碼解讀部分”，感興趣可以跳轉(zhuǎn)瀏覽。
在 Stable Diffusion WebUI 中的使用
在 Stable Diffusion 圖片生成過程中，它并不直接參與圖片生成工作，而是在圖片繪制完畢之后，在“后處理”階段，進行面部細節(jié)恢復操作，這個后處理過程在 Stable Diffusion WebUI 的 process_images_inner 過程中。
因為本文主角是 CodeFormer，所以，我們就先不過多展開不相關(guān)的細節(jié)啦。有關(guān)于 WebUI 和 CodeFormer 相關(guān)需要注意的部分，在本文下面的章節(jié)中會聊。
準備工作
準備工作部分，我們還是只需要做兩個工作：準備模型文件和模型運行環(huán)境。
關(guān)于模型運行環(huán)境，可以參考之前的文章《基于 Docker 的深度學習環(huán)境：入門篇》，如果你是 Windows 環(huán)境的用戶，可以參考這篇《基于 Docker 的深度學習環(huán)境：Windows 篇》。
如果你不熟悉如何在 Docker 環(huán)境中使用 GPU，建議仔細閱讀。考慮篇幅問題，本文就不贅述相關(guān)的話題啦。
只要你安裝好 Docker 環(huán)境，配置好能夠在 Docker 容器中調(diào)用顯卡的基礎環(huán)境，就可以進行下一步啦。
快速封裝一個 CodeFormer Docker 容器應用
從 Docker CodeFormer 項目下載代碼，并進入項目目錄：
git clone https://github.com/soulteary/docker-codeformer.git
cd docker-codeformer
執(zhí)行項目中的鏡像構(gòu)建工具：
scripts/build.sh
耐心等待鏡像構(gòu)建完畢：
# bash scripts/build.sh
[+] Building 0.1s (13/13) FINISHED
=> [internal] load build definition from Dockerfile 0.0s
=> => transferring dockerfile: 449B 0.0s
=> [internal] load .dockerignore 0.0s
=> => transferring context: 2B 0.0s
=> [internal] load metadata for nvcr.io/nvidia/pytorch:23.04-py3 0.0s
=> [internal] load build context 0.0s
=> => transferring context: 387B 0.0s
=> [1/8] FROM nvcr.io/nvidia/pytorch:23.04-py3 0.0s
=> CACHED [2/8] RUN pip install gradio==3.39.0 lpips==0.1.4 0.0s
=> CACHED [3/8] WORKDIR /app 0.0s
=> CACHED [4/8] RUN git clone https://github.com/sczhou/CodeFormer.git && cd CodeFormer && git checkout 8392d0334956108ab53d9439c4b9fc9c4af0d66d 0.0s
=> CACHED [5/8] WORKDIR /app/CodeFormer/ 0.0s
=> CACHED [6/8] COPY assets ./assets 0.0s
=> CACHED [7/8] COPY src/*.py ./ 0.0s
=> CACHED [8/8] RUN python code-fix.py 0.0s
=> exporting to image 0.0s
=> => exporting layers 0.0s
=> => writing image sha256:58709f7b295be0a1c32c578e2897f5efa771ce75c19976718d812e7b55d7794d 0.0s
=> => naming to docker.io/soulteary/docker-codeformer
因為項目鎖定了 Python 3.8，所以我們暫時只能使用 nvidia/pytorch:23.04-py3 來作為基礎鏡像。
在完成基礎鏡像構(gòu)建之后，可以從網(wǎng)盤下載 weights.zip (如果地址失效，請前往項目 issue 反饋)。模型應用運行需要的所有模型都在這里了，下載完畢后，解壓縮模型壓縮包，將 CodeFormer、facelib、realesrgan 三個目錄放置到 weights 目錄中，完整的項目結(jié)構(gòu)這樣的：
.
├── LICENSE
├── README.md
├── assets
│ └── image
├── docker
│ └── Dockerfile
├── scripts
│ └── build.sh
├── src
│ ├── app.py
│ └── code-fix.py
└── weights
├── CodeFormer
├── facelib
└── realesrgan
準備好模型文件之后，使用下面的命令啟動模型應用：
docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --rm -it -v `pwd`/weights/:/app/CodeFormer/weights -p 7860:7860 soulteary/docker-codeformer
稍等片刻，我們將看到類似下面的日志：
Running on local URL: http://0.0.0.0:7860
To create a public link, set `share=True` in `launch()`.
接著，我們就可以打開瀏覽器訪問 http://localhost:7860 或者 http://你的IP地址:7860 來試試看啦。
完整的代碼和 Docker 封裝邏輯，都在 soulteary/docker-codeformer 里，因為接下來要聊 CodeFormer 的邏輯，所以我們就不展開啦。
顯卡資源使用
CodeFormer 不是我們之前使用的大模型，所以在顯卡資源使用上輕量了不少，一般情況只需要 2G 左右的資源，處理過程中會稍微高一些，但也還在 2G 出頭的水平：
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.125.06 Driver Version: 525.125.06 CUDA Version: 12.0 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 NVIDIA GeForce ... Off | 00000000:01:00.0 Off | Off |
| 32% 40C P2 64W / 450W | 2080MiB / 24564MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| 0 N/A N/A 1429 G /usr/lib/xorg/Xorg 167MiB |
| 0 N/A N/A 1621 G /usr/bin/gnome-shell 16MiB |
| 0 N/A N/A 5090 C python 1892MiB |
+-----------------------------------------------------------------------------+
圖片處理簡單測試對比
我分別選擇了三種不同風格，都包含人像，但是原始圖片像素、模糊燥點都比較多的圖片做了三個測試，能夠看到效果還是非常驚艷的。
CodeFormer 代碼執(zhí)行邏輯
CodeFormer 代碼執(zhí)行邏輯非常簡單：加載模型，使用模型處理圖片，獲得處理結(jié)果。
加載相關(guān)模型
在正式進行 CodeFormer 進行圖片處理流程前，需要先調(diào)用模型創(chuàng)建三個模型實例。
項目使用 RealESRGAN 創(chuàng)建了一個“增強器”，精簡并整理相關(guān)代碼后，代碼實現(xiàn)如下：
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.realesrgan_utils import RealESRGANer
# set enhancer with RealESRGAN
def set_realesrgan():
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
upsampler = RealESRGANer(
model_path="CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth",
scale=2, model=model, tile=400, tile_pad=40, pre_pad=0,
)
return upsampler
upsampler = set_realesrgan()
使用 BasicSR 注冊了一個名為 CodeFormer 的新網(wǎng)絡實例，用于后續(xù)處理圖片，大概代碼實現(xiàn)如下：
import torch
from basicsr.utils.registry import ARCH_REGISTRY
def set_codeformer():
codeformer_net = ARCH_REGISTRY.get("CodeFormer")(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=["32", "64", "128", "256"]).to("cuda")
codeformer_net.load_state_dict(torch.load("CodeFormer/weights/CodeFormer/codeformer.pth")["params_ema"])
codeformer_net.eval()
return codeformer_net
codeformer_net = set_codeformer()
在推理過程中，有許多邊邊角角的功能需要處理，比如判斷圖片是否是灰度圖片，對齊人臉的特征點等等操作，所以，還需要加載 retinaface_resnet50 模型：
def get_face_helper(upscale):
face_helper = FaceRestoreHelper(
det_model="retinaface_resnet50", upscale=upscale,
face_size=512, crop_ratio=(1, 1), save_ext="png", use_parse=True, device="cuda"
)
return face_helper
“工具人（模型）”準備齊了，就可以開始人臉圖片的修復和增強處理了。
圖片處理流程
第一步，使用模型讀取圖片，然后解析其中的人臉，并標記和進行人臉對齊：
face_helper.read_image(img)
# get face landmarks for each face
num_det_faces = face_helper.get_face_landmarks_5(
only_center_face=only_center_face, resize=640, eye_dist_threshold=5
)
print(f'\tdetect {num_det_faces} faces')
# align and warp each face
face_helper.align_warp_face()
第二步，依次處理上一步模型識別出的所有人臉：
for idx, cropped_face in enumerate(face_helper.cropped_faces):
對每一張人臉，使用上文中初始化好的 CodeFormer 網(wǎng)絡來進行處理，并將處理后的圖片進行暫存：
from torchvision.transforms.functional import normalize
from basicsr.utils import img2tensor, tensor2img
# prepare data
cropped_face_t = img2tensor(
cropped_face / 255.0, bgr2rgb=True, float32=True
)
normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
try:
with torch.no_grad():
output = codeformer_net(
cropped_face_t, w=codeformer_fidelity, adain=True
)[0]
restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
del output
torch.cuda.empty_cache()
except RuntimeError as error:
print(f"Failed inference for CodeFormer: {error}")
restored_face = tensor2img(
cropped_face_t, rgb2bgr=True, min_max=(-1, 1)
)
restored_face = restored_face.astype("uint8")
face_helper.add_restored_face(restored_face)
考慮到顯存資源有限，在處理過程中，程序還對處理后的臨時內(nèi)容進行了清理，也是很環(huán)保了：
del output
torch.cuda.empty_cache()
最后一步，判斷是否有尚未人臉對齊的圖片，如果有，使用 face_helper 將修復后的圖片“復制粘貼”到原圖中：
if not has_aligned:
# upsample the background
if bg_upsampler is not None:
# Now only support RealESRGAN for upsampling background
bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
else:
bg_img = None
face_helper.get_inverse_affine(None)
# paste each restored face to the input image
if face_upsample and face_upsampler is not None:
restored_img = face_helper.paste_faces_to_input_image(
upsample_img=bg_img,
draw_box=draw_box,
face_upsampler=face_upsampler,
)
else:
restored_img = face_helper.paste_faces_to_input_image(
upsample_img=bg_img, draw_box=draw_box
)
上面的一切都搞定后，將圖片進行保存就大功告成了。
Stable Diffusion WebUI 中的調(diào)用邏輯
在 WebUI 程序入口 webui.py 程序中，能夠看到 CodeFormer 在程序初始化時進行了模型的加載：
def initialize():
...
modules.sd_models.setup_model()
startup_timer.record("setup SD model")
codeformer.setup_model(cmd_opts.codeformer_models_path)
startup_timer.record("setup codeformer")
gfpgan.setup_model(cmd_opts.gfpgan_models_path)
startup_timer.record("setup gfpgan")
...
除了默認的位于項目根目錄下的 CodeFormer 的目錄外，我們可以通過手動指定 --codeformer-models-path 參數(shù)，來改變程序加載模型的位置。
雖然程序在啟動過程中，會調(diào)用modules/launch_utils.py#L271程序中的 prepare_environment 來準備組件代碼：
def prepare_environment():
codeformer_repo = os.environ.get('CODEFORMER_REPO', 'https://github.com/sczhou/CodeFormer.git')
codeformer_commit_hash = os.environ.get('CODEFORMER_COMMIT_HASH', "c5b4593074ba6214284d6acd5f1719b6c5d739af")
git_clone(codeformer_repo, repo_dir('CodeFormer'), "CodeFormer", codeformer_commit_hash)
if not is_installed("lpips"):
run_pip(f"install -r \"{os.path.join(repo_dir('CodeFormer'), 'requirements.txt')}\"", "requirements for CodeFormer")
相關(guān)提交記錄在十個月前，為保障程序可運行，并鎖定了代碼版本。所以如果我們想使用到最新的程序代碼，還需要自己進行代碼合并和更新。
在 modules/codeformer_model.py 程序中，作者重新實現(xiàn)了一個結(jié)構(gòu)更清晰的 CodeFormer 處理流程，只有一百多行，去掉邏輯，只看架子的話：
import modules.face_restoration
import modules.shared
from modules import shared, errors
codeformer = None
def setup_model(dirname):
class FaceRestorerCodeFormer(modules.face_restoration.FaceRestoration):
def name(self):
return "CodeFormer"
def __init__(self, dirname):
def create_models(self):
return net, face_helper
def send_model_to(self, device):
def restore(self, np_image, w=None):
return restored_img
global codeformer
codeformer = FaceRestorerCodeFormer(dirname)
shared.face_restorers.append(codeformer)
結(jié)構(gòu)非常清晰，包含了初始化模型網(wǎng)絡，將模型發(fā)送到設備（比如 cpu、cuda、mps 等等），進行圖片修復。
WebUI 中 CodeFormer 關(guān)鍵實現(xiàn)保存在 modules/codeformer，包含了兩個程序 codeformer_arch.py 和 vqgan_arch.py，這兩個文件來自 CodeFormer 項目，就不展開了。
實際調(diào)用 CodeFormer 的邏輯在 modules/postprocessing.py 和 scripts/postprocessing_codeformer.py。
后者是借助 modules/scripts_auto_postprocessing.py 程序中的 ScriptPostprocessingForMainUI 函數(shù)來調(diào)用的，也算是一種有趣的解耦方案了吧。
在 WebUI 中支持兩種面部修復方案，CodeFormer 和 GFP GAN，可以根據(jù)用戶的喜好來選擇：
def apply_face_restore(p, opt, x):
opt = opt.lower()
if opt == 'codeformer':
is_active = True
p.face_restoration_model = 'CodeFormer'
elif opt == 'gfpgan':
is_active = True
p.face_restoration_model = 'GFPGAN'
else:
is_active = opt in ('true', 'yes', 'y', '1')
p.restore_faces = is_active
Stable Diffusion WebUI 中 CodeFormer 的額外注意事項
簡單來說，當 CodeFormer 模型加載失敗的時候，WebUI 使用會有異常。但在 WebUI 初始化時，我們得不到任何錯誤提醒。
在 modules/codeformer_model.py 程序中，雖然代碼處理流程清晰，但也寫了一個坑：
import os
import modules.face_restoration
import modules.shared
from modules import errors
def setup_model(dirname):
os.makedirs(model_path, exist_ok=True)
path = modules.paths.paths.get("CodeFormer", None)
if path is None:
return
try:
...
except Exception:
errors.report("Error setting up CodeFormer", exc_info=True)
如果模型初始化失敗，程序會直接 return，沒有任何報錯。但是實際使用的過程中，WebUI 是需要這個組件的，而這個組件初始化成功的前提，除了設備資源足夠初始化網(wǎng)絡模型之外，還需要能夠成功下載模型文件。
為了避免網(wǎng)絡問題，導致模型下載失敗，我們可以將模型文件提前下載完畢，放置在 WebUI 模型讀取路徑中。
最后
這個項目在開源社區(qū)無疑是成功的項目之一，它能取得成功離不開許許多多的它基于的開源項目，從本文開頭介紹前置知識和 CodeFormer 代碼執(zhí)行邏輯就不難看出來：每一個開源項目都站在了其他開源項目的肩上，然后讓項目走的更遠。
甚至，非代碼之外的項目，也對這個項目的出現(xiàn)提供了非常多的助力，包括 Nvidia Lab 推出的開源的高質(zhì)量人臉數(shù)據(jù)集：Flickr-Faces-HQ Dataset (FFHQ)。
本篇文章就先寫到這里吧，下一篇文章再見。
–EOF
我們有一個小小的折騰群，里面聚集了一些喜歡折騰、彼此坦誠相待的小伙伴。
我們在里面會一起聊聊軟硬件、HomeLab、編程上的一些問題，也會在群里不定期的分享一些技術(shù)資料。
喜歡折騰的小伙伴，歡迎閱讀下面的內(nèi)容，掃碼添加好友。
關(guān)于“交友”的一些建議和看法
添加好友時，請備注實名和公司或?qū)W校、注明來源和目的，珍惜彼此的時間 ??
蘇洋：關(guān)于折騰群入群的那些事
本文使用「署名 4.0 國際 (CC BY 4.0)」許可協(xié)議，歡迎轉(zhuǎn)載、或重新修改使用，但需要注明來源。 署名 4.0 國際 (CC BY 4.0)
本文作者: 蘇洋文章來源：http://www.zghlxwxcb.cn/news/detail-626572.html
創(chuàng)建時間: 2023年08月02日
統(tǒng)計字數(shù): 13203字
閱讀時間: 27分鐘閱讀
本文鏈接: https://soulteary.com/2023/08/02/stable-diffusion-hardcore-survival-guide-codeformer-in-webui.html文章來源地址http://www.zghlxwxcb.cn/news/detail-626572.html
到了這里，關(guān)于Stable Diffusion 硬核生存指南：WebUI 中的 CodeFormer的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！
