前言

随着 Hugging Face 的爆火，平台上出现了越来越多的有趣的模型和数据集，目前仅模型数量就高达 4 万 5 千多个。

这些模型有一个有趣的特点，在云平台上跑的好好的，但是一旦想在本地跑起来就得各种“费劲”折腾，项目关联的 GitHub 中总是能看到用户反馈：这个模型和代码，我本地跑不起来，运行环境和调用代码搞起来太麻烦了。

在日常的工作和学习中，我们也会经常遇到类似上面 Hugging Face 的情况：许多模型在“云上”跑的好好的，但是一到本地就跑不起来了，这或许是因为“操作系统环境、设备 CPU 架构（x86/ ARM）差异”、或许是因为“Python 运行时版本过高或过低”、或许是因为“某个 PIP 安装的软件包版本不对”、“冗长的示例代码中写死了一堆东西”…

用 Docker 容器搭配 Towhee，制作模型的一键运行环境。以腾讯 ARC 实验室开源的 GFPGAN 模型为例。

制作 PyTorch 模型使用的通用 Docker 基础镜像

通过容器解决或避免以下的情况

想要避免不同项目之间的环境造成干扰（污染）
想要确保项目依赖清晰，任何人都能够在任何设备上复现结果
想要复现模型的时间成本更低一些，不喜欢折腾 80% 重复的模型调优之外的工作内容（尤其是环境、基础配置）

考虑到模型可能需要在 x86 和 ARM 两类设备上运行，推荐使用 miniconda3 这个基于 debian 内置了 conda 工具包的基础镜像。

Dockerfile

代码块

language	text

FROM continuumio/miniconda3:4.11.0
LABEL maintainer=waringid@gmail.com

ARG USE_MIRROR=1

RUN if [ "$USE_MIRROR" = true ] ; then sed -i -e "s/deb.debian.org/mirrors.tuna.tsinghua.edu.cn/" /etc/apt/sources.list && sed -i -e "s/security.debian.org/mirrors.tuna.tsinghua.edu.cn/" /etc/apt/sources.list; fi

RUN apt update && \
    apt install -y libgl1-mesa-glx && \
    apt autoremove -y

RUN if [ "$USE_MIRROR" = true ] ; then pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/; fi

RUN conda install -y pytorch

RUN pip3 install --upgrade torch==1.9.0 torchvision==0.10.0

RUN pip install towhee

代码块

language	shell

docker build -t waringid/

制作具体模型的应用镜像

下载所需的模型文件

考虑到国内网络下载 Hugging Face 和 GitHub 模型比较慢，还容易出现网络中断。我推荐大家在做应用模型构建的时候，可以考虑提前进行依赖模型的下载，在构建镜像的过程中，将模型放置到合适的目录位置即可。模型下载可以参考 D004-使用 llama.cpp 转换模型程序的模型下载方法，当然也可以单独下载。

在 GFPGAN 项目中一共依赖俩模型文件，一个是 https://github.com/xinntao/facexlib 项目中基于 ResNet50 的人脸检测模型，另一个是用于图片修复的 GFPGAN 对抗网络模型，也就是传统意义上的“主角”。

第一个模型文件 detection_Resnet50_Final.pth ，可以在 https://github.com/xinntao/facexlib/releases/tag/v0.1.0 中获取；第二个模型则需要我们根据自己的设备状况，来做具体选择：

如果你需要使用 CPU 来跑模型，可以在 https://github.com/TencentARC/GFPGAN/releases/tag/v0.2.0 中下载 GFPGANCleanv1-NoCE-C2.pth；或者在https://github.com/TencentARC/GFPGAN/releases/tag/v1.3.0中下载 GFPGANv1.3.pth，这类模型可以完成黑白人像图片的处理。
如果你可以使用 GPU 来跑模型，可以在 https://github.com/TencentARC/GFPGAN/releases/tag/v0.1.0 中下载 GFPGANv1.pth；或者在 https://share.weiyun.com/ShYoCCoc 中进行模型文件下载，这类模型可以处理带颜色的人像图片。
除了 GitHub 之外，我们也可以选择直接从 Hugging Face 下载模型（只是可选版本不像上面那么多）：https://huggingface.co/TencentARC/GFPGANv1/tree/main

python 编写模型调用程序

在 GFPGAN 项目中可以找到官方的模型使用示例：https://github.com/TencentARC/GFPGAN/blob/master/inference_gfpgan.py。

代码块

language	python

import os
import cv2
import gradio as gr
import torch
from basicsr.archs.srvgg_arch import SRVGGNetCompact
from basicsr.archs.rrdbnet_arch import RRDBNet
from gfpgan.utils import GFPGANer
from realesrgan.utils import RealESRGANer

def set_realesrgan(version):
    half=True if torch.cuda.is_available() else False
    if version == 'v3':
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
        return RealESRGANer(scale=4, model_path='model/realesr-general-x4v3.pth', model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
    else:
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
        return RealESRGANer(scale=2, model_path='model/RealESRGAN_x2plus.pth', model=model, tile=400, tile_pad=10, pre_pad=0, half=half)

def set_face_enhancer(upsampler, version):
    if version == 'RestoreFormer':
        return GFPGANer(model_path='model/RestoreFormer.pth', upscale=2, arch='RestoreFormer', channel_multiplier=2, bg_upsampler=upsampler)
    else:
        if version == 'CPU':
            model_path='model/GFPGANCleanv1-NoCE-C2.pth'
        elif version == 'v1':
            arch = 'original'
            channel_multiplier = 1
            model_path='model/GFPGANv1.pth'
        elif version == 'v1.2':
            arch = 'clean'
            channel_multiplier = 2
            model_path='model/GFPGANv1.2.pth'
        elif version == 'v1.3':
            arch = 'clean'
            channel_multiplier = 2
            model_path='model/GFPGANv1.3.pth'
        elif version == 'v1.4':
            arch = 'clean'
            channel_multiplier = 2
            model_path='model/GFPGANv1.4.pth'

    return GFPGANer(model_path=model_path, upscale=2, arch=arch, channel_multiplier=channel_multiplier, bg_upsampler=upsampler)

os.makedirs('output', exist_ok=True)

def inference(img, realesrgan_version, gfpgan_version, scale):
    if scale > 4:
        scale = 4
    elif scale < 0:
        scale = 1

    try:
        extension = os.path.splitext(os.path.basename(str(img)))[1]
        img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
        if len(img.shape) == 3 and img.shape[2] == 4:
            img_mode = 'RGBA'
        elif len(img.shape) == 2:  # for gray inputs
            img_mode = None
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        else:
            img_mode = None

        h, w = img.shape[0:2]
        if h > 3500 or w > 3500:
            print('too large size')
            return None, None
        
        if h < 300:
            img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)

        upsampler = set_realesrgan(realesrgan_version)
        face_enhancer = set_face_enhancer(upsampler, gfpgan_version)

        try:
            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True, weight=None)
        except RuntimeError as error:
            print('Error', error)

        try:
            if scale != 2:
                interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
                h, w = img.shape[0:2]
                output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
        except Exception as error:
            print('wrong scale input.', error)
        if img_mode == 'RGBA':  # RGBA images should be saved in png format
            extension = 'png'
        else:
            extension = 'jpg'
        save_path = f'output/out.{extension}'
        cv2.imwrite(save_path, output)

        output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
        return output, save_path
    except Exception as error:
        print('global exception', error)
        return None, None

description = r"""Gradio demo for <a href='https://github.com/TencentARC/GFPGAN' target='_blank'>GFPGAN: Towards Real-World Blind Face Restoration with Generative Facial Prior</a>.<br>
这个应用可以被用来修复你的**老照片**或者改善**AI生成的人脸**。<br>
你只需要上传你的图片，就可以使用它。<br>
如果你觉得这个项目帮助到了你，不妨为 <a href='https://github.com/TencentARC/GFPGAN' target='_blank'>它</a> 点个Star吧 :-D <br>
"""

radio_realesrgan_version = gr.Radio(['v3'], type="value", value='v3', label='RealESR GAN version')
radio_gfpgan_version = gr.Radio(['CPU', 'v1.2', 'v1.3', 'v1.4', 'RestoreFormer'], type="value", value='v1.4', label='version')
if os.environ.get('BASICSR_JIT') is not None:
    radio_realesrgan_version = gr.Radio(['v2', 'v3'], type="value", value='v3', label='RealESR GAN version')
    radio_gfpgan_version = gr.Radio(['CPU', 'v1', 'v1.2', 'v1.3', 'v1.4', 'RestoreFormer'], type="value", value='v1.4', label='version')

app = gr.Interface(
    inference, [
        gr.Image(type="filepath", label="Input"),
        radio_realesrgan_version,
        radio_gfpgan_version,
        gr.Number(label="Rescaling factor", value=2),
    ], [
        gr.Image(type="numpy", label="Output (The whole image)"),
        gr.File(label="Download the output image")
    ],
    title="GFPGAN: 实用的人脸修复算法",
    description=description,
    article="<p style='text-align: center'>written by: <a href='https://wiki.waringid.me' target='_blank'></a></p>",    # examples=[['AI-generate.jpg', 'v1.4', 2], ['lincoln.jpg', 'v1.4', 2], ['Blake_Lively.jpg', 'v1.4', 2], ['10045.png', 'v1.4', 2]]
)
app.queue()
app.launch(server_name="0.0.0.0")

将上面的内容保存为 app.py 后续备用。

简易版本 app.py

以该版本为主

代码块

language	python

from gfpgan import GFPGANer
import towhee

@towhee.register
class GFPGANerOp:

    def __init__(self,
                 model_path='/GFPGAN.pth',
                 upscale=2,
                 arch='clean',
                 channel_multiplier=2,
                 bg_upsampler=None) -> None:
        self._restorer = GFPGANer(model_path, upscale, arch, channel_multiplier, bg_upsampler)

    def __call__(self, img):
        cropped_faces, restored_faces, restored_img = self._restorer.enhance(
            img, has_aligned=False, only_center_face=False, paste_back=True)

        return restored_faces[0][:, :, ::-1]

(
    towhee.glob['path']('*.jpg') 
        .image_load['path', 'img']() 
        .GFPGANerOp['img','face']() 
        .show(formatter=dict(img='image', face='image'))

Dockerfile

将下载好的模型文件和新的 Dockerfile 文件放置于相同目录之后，通过整合 Dockerfile 的内容，完成项目依赖的安装，并将模型放置在容器内合适的目录位置。

代码块

language	text

FROM waringid/docker-pytorch-playground

RUN pip install gfpgan==1.3.8 realesrgan==0.3.0 facexlib==0.3.0 gradio==3.39.0
COPY detection_Resnet50_Final.pth /opt/conda/lib/python3.9/site-packages/facexlib/weights/detection_Resnet50_Final.pth

# 尺寸大一些的模型文件，可以选择使用挂载的方式
# 而不在此处直接 COPY 到容器内部
COPY GFPGANCleanv1-NoCE-C2.pth /GFPGAN.pth

COPY app.py /entrypoint.py
WORKDIR /data
RUN pip install IPython pandas
RUN sed -i -e "s/display(HTML(table))/with open('result.html', 'w') as file:\n            file.write(HTML(table).data)/" /opt/conda/lib/python3.9/site-packages/towhee/datacollection/mixins/display.py
CMD ["python3", "/entrypoint.py"]

代码块

language	shell

docker build -t pytorch-playground-gfpgan -f Dockerfile .

Image Added

模型应用镜像的使用

如果上一步已经下载了模型文件，并将模型文件打包到了镜像中，那么我们只需要下载一些黑白或者彩色的包含人像的图片（根据模型来选择），将它们放在一个目录中（比如 data目录），然后执行一行命令就能够完成模型的调用啦

代码块

language	shell

docker run --rm -it -v `pwd`/model/GFPGANCleanv1-NoCE-C2.pth:/GFPGAN.pth -v `pwd`/data:/data pytorch-playground-gfpgan

如果在上文构建应用模型镜像时，没有选择将 GFPGAN 模型打包到镜像中，那么我们就需要使用文件挂载的方式，来运行模型了。为了项目结构的清晰，可以在项目中创建 model 的目录，来存放上文中提到的模型文件。请忽略 Dockerfile 和 app.py 这几个文件。

Image Added

目录

空间快捷链接

页面树结构

版本比较

旧版本 2

新版本 3

标识

前言

制作 PyTorch 模型使用的通用 Docker 基础镜像

Dockerfile

制作具体模型的应用镜像

下载所需的模型文件

python 编写模型调用程序

简易版本 app.py

Dockerfile

模型应用镜像的使用

空间快捷链接

页面树结构

页面历史记录

版本比较

旧版本 2

新版本 3

标识

前言

制作 PyTorch 模型使用的通用 Docker 基础镜像

Dockerfile

制作具体模型的应用镜像

下载所需的模型文件

python 编写模型调用程序

简易版本 app.py

Dockerfile

模型应用镜像的使用