Skip to content
/ DualStyleGAN-Mix Public
forked from williamyang1991/DualStyleGAN

[CVPR 2022] Pastiche Master: Exemplar-Based High-Resolution Portrait Style Transfer

License

View license
0 stars 256 forks Branches Tags Activity
Star
Notifications You must be signed in to change notification settings

SlieFamily/DualStyleGAN-Mix

BranchesTags
 
 

Repository files navigation

DualStyleGAN - Mix | 基于官方 PyTorch 项目的自定义修改

原始论文:Pastiche Master: Exemplar-Based High-Resolution Portrait Style Transfer
Shuai Yang, Liming Jiang, Ziwei Liu and Chen Change Loy
In CVPR 2022.

项目主页 | 论文 | 相关视频

摘要：
近期，在学习 StyleGAN 的过程中我们发现，StyleGAN 只通过有限数据的训练学习，就在人脸艺术肖像风格迁移（下称 人脸风格化 或 人像风格转移） 工作中表现得很出色。 在本文中，我们通过引入一种名为 DualStyleGAN 的较新颖的技术来探索更具挑战性的基于样本的高分辨率人脸风格化，可以灵活控制原始人脸和艺术肖像方面的双重风格。 不同于 StyleGAN，DualStyleGAN 分别通过内在风格路径和新的外在风格路径来表征图像内容和某个艺术画的风格（Style），以此来提供更自然的风格迁移方法。精心设计的外在风格路径使我们的模型能够分层调整颜色和复杂的结构风格，从而精确地贴合给出的风格图示。 此外，我们还引入了一种新颖的渐进式微调方案，使得即使在对网络架构进行上述修改的情况下，也可以将模型的生成空间平滑地转换到目标风格上。 实验证明， DualStyleGAN 在高质量的人像风格迁移和灵活的风格控制方面优于目前其他的先进的方法。

关键词：
High-Resolution (1024) | Training Data-Efficient (~200 Images) | Exemplar-Based Color and Structure Transfer

WebUI

可通过在Colab或 Hugging Face Spaces 运行 webui.py 开启可视化界面。

该 GUI 修改自官方 Web Demo: Hugging Face Spaces or Hugging Face Spaces

安装

克隆仓库:

git clone https://github.com/williamyang1991/DualStyleGAN.git
cd DualStyleGAN

依赖:

环境内的所有依赖包均在 environment/dualstylegan_env.yaml 中列出. 推荐使用本仓库的 Anaconda 进行安装:

conda env create -f ./environment/dualstylegan_env.yaml

原项目所使用的是 CUDA 10.1，因此对应Torch版本是 PyTorch 1.7.1 (在文件 dualstylegan_env.yaml 中的 Line 22, Line 25, Line 26 中有着对应版本号 ).

你可以通过https://pytorch.org/ 查询并下载安装与你使用的 CUDA 相符版本的 PyTorch

(1) 数据集准备

Cartoon, Caricature 以及 Anime 数据集 可在如下的官方渠道中进行下载.

我们也将提供构建新数据集的脚本.

Dataset Description
Cartoon 317 cartoon face images from Toonify.
Caricature 199 images from WebCaricature. Please refer to dataset preparation for more details.
Anime 174 images from Danbooru Portraits. Please refer to dataset preparation for more details.
Fantasy 137 fantasy face images generated by StableDiffusion.
Illustration 156 illustration face images generated by StableDiffusion.
Impasto 120 impasto face images generated by StableDiffusion.
Other styles Please refer to dataset preparation for the way of building new datasets.

(2) 风格化的流程

Inference Notebook

我们提供了 Jupyter notebook (路径为 ./notebooks/inference_playground.ipynb) 供用户体验 DualStyleGAN 的实际表现. 这个 notebook 笔记将会下载必要的预训练模型，然后以 ./data/ 中的图片为例进行风格迁移.

如果没有 GPU 供使用，可以参考CPU下运行，需要在 notebook 中设置 device = 'cpu'.

模型的预处理

预训练模型可以通过 Google Drive 或 百度云|Baidu Cloud (密码: cvpr) 进行下载.

Model Description
encoder Pixel2style2pixel encoder that embeds FFHQ images into StyleGAN2 Z+ latent code
encoder_wplus Original Pixel2style2pixel encoder that embeds FFHQ images into StyleGAN2 W+ latent code
cartoon DualStyleGAN and sampling models trained on Cartoon dataset, 317 (refined) extrinsic style codes
caricature DualStyleGAN and sampling models trained on Caricature dataset, 199 (refined) extrinsic style codes
anime DualStyleGAN and sampling models trained on Anime dataset, 174 (refined) extrinsic style codes
arcane DualStyleGAN and sampling models trained on Arcane dataset, 100 extrinsic style codes
comic DualStyleGAN and sampling models trained on Comic dataset, 101 extrinsic style codes
pixar DualStyleGAN and sampling models trained on Pixar dataset, 122 extrinsic style codes
slamdunk DualStyleGAN and sampling models trained on Slamdunk dataset, 120 extrinsic style codes
fantasy DualStyleGAN models trained on Fantasy dataset, 137 extrinsic style codes
illustration DualStyleGAN models trained on Illustration dataset, 156 extrinsic style codes
impasto DualStyleGAN models trained on Impasto dataset, 120 extrinsic style codes

模型的 checkpoints 遵循如下的目录结构:

checkpoint
|--encoder.pt                     % Pixel2style2pixel model
|--encoder_wplus.pt               % Pixel2style2pixel model (optional)
|--cartoon
    |--generator.pt               % DualStyleGAN model
    |--sampler.pt                 % The extrinsic style code sampling model
    |--exstyle_code.npy           % extrinsic style codes of Cartoon dataset
    |--refined_exstyle_code.npy   % refined extrinsic style codes of Cartoon dataset
|--caricature
    % the same files as in Cartoon
...

人脸风格化示例

对一张默认的人脸图利用 Cartoon Style 进行风格迁移:

python style_transfer.py

结果图 cartoon_transfer_53_081680.jpg 保存至 .\output\,

其中， 53 是 卡通数据集(Cartoon dataset)的风格ID, 081680 是 content face image 的名字. 额外保存相应的概览图 cartoon_transfer_53_081680_overview.jpg 以说明输入内容图像、编码内容图像、样式图像（* 样式图像只有在文件夹中才会显示）

结果如下：

可以通过 --style 和 --style_id 指定风格样式 (点击此处查看 , find the visual mapping between id and the style image here).

可以通过 --name 指定结果图文件的名字.

Specify the weight to adjust the degree of style with --weight.

下面展示的是论文摘要中的风格迁移结果以及对应的命令/脚本。

python style_transfer.py
python style_transfer.py --style cartoon --style_id 10
python style_transfer.py --style caricature --name caricature_transfer --style_id 0 --weight 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
python style_transfer.py --style caricature --name caricature_transfer --style_id 187 --weight 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
python style_transfer.py --style anime --name anime_transfer --style_id 17 --weight 0 0 0 0 0.75 0.75 0.75 1 1 1 1 1 1 1 1 1 1 1
python style_transfer.py --style anime --name anime_transfer --style_id 48 --weight 0 0 0 0 0.75 0.75 0.75 1 1 1 1 1 1 1 1 1 1 1

可以通过 --content 指定content image. 如果 content image 无法与 FFHQ 进行规格匹配, 可以添加上 --align_face.

如需保留content image 的颜色样式，可使用 --preserve_color 或将 --weight 的最后 11 个元素全置零。

python style_transfer.py --content ./data/content/unsplash-rDEOVtE7vOs.jpg --align_face --preserve_color \
       --style arcane --name arcane_transfer --style_id 13 \
       --weight 0.6 0.6 0.6 0.6 0.6 0.6 0.6 0.6 0.6 0.6 0.6 1 1 1 1 1 1 1

→

可通过 --wplus 以指定使用原始 pSp 编码器提取 W+ 固有样式代码，这可能会更好地保留 content image 的人脸特征。

评价: 我们在 Z+/W+ 空间中训练好的 pSp 编码器(encoder) 并不能完美地编码出 the content image. 如果希望风格化之后的结果与 content image 更加一致，可以使用潜在优化来更好地适应内容图像或使用其他 StyleGAN 编码器(如 williamyang1991#11 和 williamyang1991#29 中的讨论).

可通过 python style_transfer.py -h 查询更多选项.

生成艺术画

首先，生成随机卡通人脸图像 (结果保存在 ./output/ folder):

python generate.py

可通过 --style 指定 Style， --name 指定保存的图像文件名:

python generate.py --style arcane --name arcane_generate

Specify the weight to adjust the degree of style with --weight.

Keep the intrinsic style code, extrinsic color code or extrinsic structure code fixed using --fix_content, --fix_color and --fix_structure, respectively.

python generate.py --style caricature --name caricature_generate --weight 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 --fix_content

可通过 python generate.py -h 查询更多选项.


(3) 训练 DualStyleGAN

下载支撑模型到 ./checkpoint/ 目录下:

Model Description
stylegan2-ffhq-config-f.pt StyleGAN model trained on FFHQ taken from rosinality.
model_ir_se50.pth Pretrained IR-SE50 model taken from TreB1eN for ID loss.

面部风格提取

Step 1:数据准备. 将数据集放入 ./data/DATASET_NAME/images/train/目录. 首先生成 lmdb 数据集:

python ./model/stylegan/prepare_data.py --out LMDB_PATH 
                                        --n_worker N_WORKER 
                                        --size SIZE1,SIZE2,SIZE3,... 
                                        DATASET_PATH

例如, 下载 317 张 Cartoon 图像 保存于 ./data/cartoon/images/train/ 中，然后运行：

python ./model/stylegan/prepare_data.py --out ./data/cartoon/lmdb/ --n_worker 4 --size 1024 ./data/cartoon/images/

其中的 1024 即是 Cartoon 数据集中每张图像的长宽。

Step 2: 微调 StyleGAN. Fine-tune StyleGAN in distributed settings:

python -m torch.distributed.launch --nproc_per_node=N_GPU 
                                   --master_port=PORT finetune_stylegan.py 
                                   --batch BATCH_SIZE \
                                   --ckpt FFHQ_MODEL_PATH 
                                   --iter ITERATIONS 
                                   --style DATASET_NAME 
                                   --augment LMDB_PATH

以 cartoon dataset 为例, 运行 (建议将 batch size of 设置为 8*4=32)：

python -m torch.distributed.launch --nproc_per_node=8 --master_port=8765 finetune_stylegan.py --iter 600
                          --batch 4 --ckpt ./checkpoint/stylegan2-ffhq-config-f.pt --style cartoon
                          --augment ./data/cartoon/lmdb/

fine-tuned model 模型将会 放入 ./checkpoint/cartoon/finetune-000600.pt. 中间结果或数据放入 ./log/cartoon/.

Step 3: 风格提取.

python destylize.py --model_name FINETUNED_MODEL_NAME 
                    --batch BATCH_SIZE 
                    --iter ITERATIONS DATASET_NAME

以 cartoon dataset 为例, 运行:

python destylize.py --model_name finetune-000600.pt --batch 1 --iter 300 cartoon

内外部风格(Style)将会分别存放于 ./checkpoint/cartoon/instyle_code.npy 和 ./checkpoint/cartoon/exstyle_code.npy之中, 中间结果或数据放入./log/cartoon/destylization/.

为了快速进行去风格化处理, 可以设置 --batch 的参数值为较大的数，比如 16.

对于与真实面孔严重不同的 Style，可设置 --truncation 为较小的值，比如 0.5，以使结果更逼真（它使 DualStyleGAN 能够学习更大的结构变形）。

渐进微调

Stage 1 & 2: 在 FFHQ 中对 DualStyleGAN 进行预训练 .

我们提供已训练好的预训练模型 generator-pretrain.pt . 在 Google Drive 或 百度云|Baidu Cloud (密码: cvpr) 中均可下载.

该模型通过如下命令获得:

python -m torch.distributed.launch --nproc_per_node=1 --master_port=8765 pretrain_dualstylegan.py --iter 3000
                          --batch 4 ./data/ffhq/lmdb/

其中 ./data/ffhq/lmdb/ 内包含通过 ./model/stylegan/prepare_data.py 从 FFHQ 数据集创建的 lmdb 数据.

Stage 3: 在目标风格域中对 DualStyleGAN 微调.

在分布式设置中微调 DualStyleGAN:

python -m torch.distributed.launch --nproc_per_node=N_GPU 
                                   --master_port=PORT finetune_dualstylegan.py 
                                   --iter ITERATIONS \ 
                                   --batch BATCH_SIZE 
                                   --ckpt PRETRAINED_MODEL_PATH 
                                   --augment DATASET_NAME

损失项的权重通过 --style_loss ($λ_{FM}$), --CX_loss ($λ_{CX}$), --perc_loss ($λ_{perc}$), --id_loss ($λ_{id}$) and --L2_reg_loss ($λ_{reg}$) 设置.

建议针对每个不同风格的数据集，逐步调整 $λ_{id}$ 和 $λ_{reg}$ 以达到理想的效果.

可通过 python finetune_dualstylegan.py -h 查询更多选项.

以 cartoon dataset 为例, 运行(multi-GPU 最大支持 batch size 为 8*4=32, 可获得较快的训练):

python -m torch.distributed.launch --nproc_per_node=8 --master_port=8765 finetune_dualstylegan.py --iter 1500 --batch 4 --ckpt ./checkpoint/generator-pretrain.pt 
--style_loss 0.25 --CX_loss 0.25 --perc_loss 1 --id_loss 1 --L2_reg_loss 0.015 --augment cartoon

微调模型存入 ./checkpoint/cartoon/generator-ITER.pt 重，其中的 ITER 取 001000, 001100, ..., 001500.

中间结果或数据放入 ./log/cartoon/. 较大的 ITER 有着更强的 cartoon styles 但却可能伴随着原图风格的损失。因此可能需要人工从 1000-1500 中选取更加平衡的结果. 在我们的论文里，我们选取的值为 1400.

（可选）潜在优化和采样

完善外部风格 |extrinsic style code. 优化颜色和结构风格以更好地适应Sample 图像。

python refine_exstyle.py --lr_color COLOR_LEARNING_RATE --lr_structure STRUCTURE_LEARNING_RATE DATASET_NAME

默认情况下, 此命令会 加载存放在 ./checkpoint/DATASET_NAME/ 中的 instyle_code.npy, exstyle_code.npy 和 generator.pt. 可使用 --instyle_path, --exstyle_path, --ckpt 指定其他保存的 Style 或模型.

以 cartoon dataset 为例, 运行:

python refine_exstyle.py --lr_color 0.1 --lr_structure 0.005 --ckpt ./checkpoint/cartoon/generator-001400.pt cartoon

完善得到的更精致的外部风格特征保存在 ./checkpoint/DATASET_NAME/refined_exstyle_code.npy中.

可通过调整参数lr_color 和 lr_structure 以获得更好的效果.

训练采样网络. 训练采样网络以将单位高斯噪声映射到外部风格特征的分布:

python train_sampler.py DATASET_NAME

默认情况下, 上述命令将从 ./checkpoint/DATASET_NAME/ 中加载 refined_exstyle_code.npy 或 exstyle_code.npy 文件 .

可使用 --exstyle_path 指定其他已保存的外部风格特征. 已保存的模型存入 ./checkpoint/DATASET_NAME/sampler.pt.


(4) 结果展示

Exemplar-based cartoon style trasnfer

cartoon.mp4

Exemplar-based caricature style trasnfer

caricature.mp4

Exemplar-based anime style trasnfer

anime.mp4

Other styles

Combine DualStyleGAN with State-of-the-Art Diffusion model

We use StableDiffusion to generate face images of the specified style of famous artists. Trained with these images, DualStyleGAN is able to pastiche these famous artists and generates appealing results.

引用本文

如果你认为本项目内容对你的研究有帮助的话，还请引用:

@inproceedings{yang2022Pastiche,
  title={Pastiche Master: Exemplar-Based High-Resolution Portrait Style Transfer},
  author={Yang, Shuai and Jiang, Liming and Liu, Ziwei and Loy, Chen Change},
  booktitle={CVPR},
  year={2022}
}

致谢

本项目代码主要基于 stylegan2-pytorch 和 pixel2style2pixel 实现.

About

[CVPR 2022] Pastiche Master: Exemplar-Based High-Resolution Portrait Style Transfer

Resources

Readme

License

View license
Activity

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

No releases published

Packages

No packages published

Languages

  • Jupyter Notebook 93.2%
  • Python 6.0%
  • Other 0.8%