add pretrained model for stage II.

sczhou · Jul 23, 2023 · 8392d03 · mdsantoc · Feb 29, 2024 · mdsantoc
1 parent 4cc6d9d
commit 8392d03
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/docs/train.md b/docs/train.md
@@ -15,17 +15,19 @@
 - After VQGAN training, you can pre-calculate code sequence for the training dataset to speed up the later training stages:
   > python scripts/generate_latent_gt.py
 
-- If you don't require training your own VQGAN, you can find pre-trained VQGAN and the corresponding code sequence in the folder of Releases v0.1.0: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
+- If you don't require training your own VQGAN, you can find pre-trained VQGAN (`vqgan_code1024.pth`) and the corresponding code sequence (`latent_gt_code1024.pth`) in the folder of Releases v0.1.0: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
 
 ### 🚀 Stage II - CodeFormer (w=0)
 - Training Code Sequence Prediction Module:
   > python -m torch.distributed.launch --nproc_per_node=8 --master_port=4322 basicsr/train.py -opt options/CodeFormer_stage2.yml --launcher pytorch
 
+- Pre-trained CodeFormer of stage II (`codeformer_stage2.pth`) can be found in the folder of Releases v0.1.0: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
+
 ### 🛸 Stage III - CodeFormer (w=1)
 - Training Controllable Module:
   > python -m torch.distributed.launch --nproc_per_node=8 --master_port=4323 basicsr/train.py -opt options/CodeFormer_stage3.yml --launcher pytorch
 
-- Pre-trained CodeFormer can be found in the folder of Releases v0.1.0: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
+- Pre-trained CodeFormer (`codeformer.pth`) can be found in the folder of Releases v0.1.0: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
 
 ---
 

diff --git a/docs/train_CN.md b/docs/train_CN.md
@@ -15,17 +15,19 @@
 - 训练完VQGAN后，可以通过下面代码预先获得训练数据集的密码本序列，从而加速后面阶段的训练过程:
   > python scripts/generate_latent_gt.py
 
-- 如果你不需要训练自己的VQGAN，可以在Release v0.1.0文档中找到预训练的VQGAN和对应的密码本序列: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
+- 如果你不需要训练自己的VQGAN，可以在Release v0.1.0文档中找到预训练的VQGAN (`vqgan_code1024.pth`)和对应的密码本序列 (`latent_gt_code1024.pth`): https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
 
 ### 🚀 阶段 II - CodeFormer (w=0)
 - 训练密码本训练预测模块:
   > python -m torch.distributed.launch --nproc_per_node=8 --master_port=4322 basicsr/train.py -opt options/CodeFormer_stage2.yml --launcher pytorch
 
+- 预训练CodeFormer第二阶段模型 (`codeformer_stage2.pth`)可以在Releases v0.1.0文档里下载: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
+
 ### 🛸 阶段 III - CodeFormer (w=1)
 - 训练可调模块:
   > python -m torch.distributed.launch --nproc_per_node=8 --master_port=4323 basicsr/train.py -opt options/CodeFormer_stage3.yml --launcher pytorch
 
-- 预训练CodeFormer模型可以在Releases v0.1.0文档里下载: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
+- 预训练CodeFormer模型 (`codeformer.pth`)可以在Releases v0.1.0文档里下载: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
 
 ---