Skip to content

Latest commit

 

History

History
95 lines (63 loc) · 2.03 KB

File metadata and controls

95 lines (63 loc) · 2.03 KB

基于扩散模型的Jittor草图生成风景

注:第一行是原图,第二行是语义分割图,第三行是我们的模型根据语义图生成的结果。

简介

本项目包含了第二届计图挑战赛计图 - 语义图生成风景比赛的代码实现。我们采用隐空间扩散模型生成风景图,模型分为两阶段:一阶段训练autoencoder,压缩和重建风景图;二阶段训练扩散模型,随机生成噪声,和语义图拼接在一起,映射到一阶段隐空间中进行扩散,最后使用一阶段的decoder对扩散结果进行解码生成风景图。

成绩

  • A榜:10/154
  • B榜:12/31
  • B榜+大众投票:10/31

安装

本项目可在 2 张 rtx titan 上运行。

运行环境

  • ubuntu 20.04 LTS
  • python >= 3.7
  • jittor >= 1.3.0

安装依赖

执行以下命令安装 python 依赖

pip install -r requirements.txt

预训练模型

下载预训练模型 download link,放到ckpts/下。

数据准备

DATA_DIR/
|-- train_val/
|   |-- 1.jpg
|   |
|   |-- train.txt
|   |-- val.txt
|
|-- test/
|   |-- 1.jpg
|   |
|   |-- test.txt
|
|-- test_B/
|   |-- 1.jpg
|   |
|   |-- test_B.txt

在configs文件夹下的配置文件中修改数据集路径。

训练

模型分为两阶段:vqgan和ldm,分开训练

  • 单卡训练可运行以下命令:
python train_vqgan.py

python train_ldm.py 
  • 多卡训练可以运行以下命令:
mpirun --allow-run-as-root -np 2 python train_vqgan.py

mpirun --allow-run-as-root -np 2 python train_ldm.py 

推理

生成测试集上的结果可以运行以下命令:

python inference_jittor.py

致谢

此项目基于论文 High-Resolution Image Synthesis with Latent Diffusion Models 实现,代码参考 latent-diffusion