如何在矩池云上运行 AI 图像编辑工具 DragGAN

发布时间 2023-06-27 18:24:11作者: 矩池云

5 月,DragGAN 横空出世,在开源代码尚未公布前,就在Github上斩获近 20000 Star,彼时,页面上只有效果图和一句“Code will be released in June”,然而这也足够带给人们无限期待。

在6月末,在若干非官方复现代码发布后,官方版DragGAN终于发布,我们也第一时间在矩池云上进行了尝试。

上传模型和数据

在租用机器前我们先上传需要的模型和数据到矩池云网盘,这样租用机器后就可以直接使用了。

DragGAN 提供了一个下载模型的脚本,但是里面的链接都是外网链接,服务器里直接运行下载可能很慢,所以我们需要提前本地下载好,然后上传到矩池云网盘再租用机器使用。

# 脚本地址 https://github.com/XingangPan/DragGAN/blob/main/scripts/download_model.sh
mkdir checkpoints
cd checkpoints

wget https://storage.googleapis.com/self-distilled-stylegan/lions_512_pytorch.pkl
mv lions_512_pytorch.pkl stylegan2_lions_512_pytorch.pkl

wget https://storage.googleapis.com/self-distilled-stylegan/dogs_1024_pytorch.pkl
mv dogs_1024_pytorch.pkl stylegan2_dogs_1024_pytorch.pkl

wget https://storage.googleapis.com/self-distilled-stylegan/horses_256_pytorch.pkl
mv horses_256_pytorch.pkl stylegan2_horses_256_pytorch.pkl

wget https://storage.googleapis.com/self-distilled-stylegan/elephants_512_pytorch.pkl
mv elephants_512_pytorch.pkl stylegan2_elephants_512_pytorch.pkl

wget https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhq-512x512.pkl
wget https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-afhqcat-512x512.pkl
wget http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-car-config-f.pkl
wget http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-cat-config-f.pkl

如果你本地下载也很慢,也可以通过云盘下载,谢谢网友分享:https://pan.quark.cn/s/d63e04863dd8

本地下载好后,把相关文件上传到矩池云网盘即可。

租用机器

本次复现使用的是 NVIDIA RTX A2000 单卡,镜像使用的是 Pytorch 2.0 ,选择好机器和镜像后我们先需要按下图指示开启 VNC 功能(Ubuntu 可视化桌面),然后高级选项自定义一个 7860 端口。

  • VNC 里我们可以运行启动 DragGAN 桌面软件
  • DragGAN Web应用使用的是 Gradio 开发,默认端口是 7860

租用成功后我们可以在租用页面看到机器的自定义端口、VNC、SSH、Jupyterlab等链接,矩池云官网有详细的教程介绍了如何使用这些链接连接服务器。

配置环境

clone 代码

经常修改的代码和数据集建议放矩池云网盘,Linux 系统实例中网盘对应机器中的 /mnt目录,永久存储,可以离线上传下载数据。

我们点击 Jupyterlab 链接即可直接访问服务器,新建一个 Terminal,输入以下指令 clone 项目。

# 进入网盘目录,我复现项目代码会归纳在一起 /mnt/example/,你改成自己目录即可
cd /mnt/example/
# clone 代码
git clone https://ghproxy.com/https://github.com/XingangPan/DragGAN

安装依赖

clone 好代码,我们进入项目目录下,安装 requirements.txt 里的依赖,按以下指令输入即可:

# 进入项目目录
cd DragGAN
# 安装依赖
pip install -r requirements.txt 

看到最后这个警告信息不用管,意思是不建议你在 root 权限下操作安装,实际是安装成功了的。

运行项目

环境配置好后就可以开始运行项目了,DragGAN 提供了两种运行方法 基于 Gradio 的 Web 应用和基于 imgui 的桌面程序。

配置模型权重

最开始我们已经下载上传了模型权重文件到网盘,所以这里我们只需要将相关文件复制到项目对应文件夹下即可。

以我为例,我的模型都上传在网盘的 example/models 下了,对应机器中的/mnt/example/models目录。

# 先进入项目目录 如果紧接着之前步骤,
# 你现在应该已经在 DragGAN 项目目录下了
# 如不在,请执行 cd 指令进入对应目录
# 如:cd /mnt/example/DragGAN
# 在项目目录下创建一个 checkpoints 目录
mkdir checkpoints
# 将网盘内模型复制到 checkpoints
cp /mnt/example/models/* checkpoints

注意,按照脚本里的,里面有四个文件需要改下文件名称,cd 进入 checkpoints 目录,然后执行对应指令即可。

cd checkpoints
# 有上传对应权重就执行对应mv指令改名字
# mv lions_512_pytorch.pkl stylegan2_lions_512_pytorch.pkl
# mv dogs_1024_pytorch.pkl stylegan2_dogs_1024_pytorch.pkl
# mv horses_256_pytorch.pkl stylegan2_horses_256_pytorch.pkl
# mv elephants_512_pytorch.pkl stylegan2_elephants_512_pytorch.pkl

运行基于 Gradio 的 Web 应用

Gradio web 应用默认是在启动在本地回环地址(127.0.0.1),这样就只能服务器里访问,没法通过公网地址访问,所以我们需要稍微改动下代码。

Jupyterlab 中打开项目文件夹下的visualizer_drag_gradio.py文件,移动到最后一行,在app.launch中添加server_name="0.0.0.0"即可解决前面说的问题。(如下图所示)

项目文件夹下输入下面指令即可启动服务。

python visualizer_drag_gradio.py 

项目启动成功后,你会看到项目默认启动在 7860 端口,我们直接访问租用页面的 7860 端口对应公网链接即可。

基于 imgui 的桌面程序

桌面程序需要在 VNC 里才能启动成功。复制 VNC 链接,按配置说明方法即可连接上 VNC 。

# 设置环境变量
export MESA_GL_VERSION_OVERRIDE=3.3
# 进入项目目录
cd /mnt/example/DragGAN/
# 启动程序
sh scripts/gui.sh

程序启动后即会弹出程序操作界面。

浅尝 Drag

方便下次使用,还可以在租用页面保存环境。