Adding support for Consistency Models

re run circle check
open-mmlab · Oct 7, 2023 · 864b3aa · 864b3aa
1 parent 2c4c010
commit 864b3aa
Show file tree

Hide file tree

Showing 8 changed files with 68 additions and 54 deletions.
diff --git a/configs/consistency_models/README.md b/configs/consistency_models/README.md
@@ -18,30 +18,30 @@ Diffusion models have significantly advanced the fields of image, audio, and vid
 
 ## Pre-trained models
 
-|                                             Model                                             |  Dataset   | Download |
-| :-------------------------------------------------------------------------------------------: | :--------: | :------: |
-|       [onestep on ImageNet-64](./consistency_models_8xb256-imagenet1k-onestep-64x64.py)       | imagenet1k |    -     |
-|     [multistep on ImageNet-64](./consistency_models_8xb256-imagenet1k-multistep-64x64.py)     | imagenet1k |    -     |
-|   [onestep on LSUN Bedroom-256](./consistency_models_8xb32-LSUN-bedroom-onestep-256x256.py)   |    LSUN    |    -     |
-| [multistep on LSUN Bedroom-256](./consistency_models_8xb32-LSUN-bedroom-multistep-256x256.py) |    LSUN    |    -     |
-|       [onstep on LSUN Cat-256](./consistency_models_8xb32-LSUN-cat-onestep-256x256.py)        |    LSUN    |    -     |
-|     [multistep on LSUN Cat-256](./consistency_models_8xb32-LSUN-cat-multistep-256x256.py)     |    LSUN    |    -     |
+|                                             Model                                             |  Dataset   | Conditional | Download |
+| :-------------------------------------------------------------------------------------------: | :--------: | :---------: | :------: |
+|       [onestep on ImageNet-64](./consistency_models_8xb256-imagenet1k-onestep-64x64.py)       | imagenet1k |     yes     |    -     |
+|     [multistep on ImageNet-64](./consistency_models_8xb256-imagenet1k-multistep-64x64.py)     | imagenet1k |     yes     |    -     |
+|   [onestep on LSUN Bedroom-256](./consistency_models_8xb32-LSUN-bedroom-onestep-256x256.py)   |    LSUN    |     no      |    -     |
+| [multistep on LSUN Bedroom-256](./consistency_models_8xb32-LSUN-bedroom-multistep-256x256.py) |    LSUN    |     no      |    -     |
+|       [onstep on LSUN Cat-256](./consistency_models_8xb32-LSUN-cat-onestep-256x256.py)        |    LSUN    |     no      |    -     |
+|     [multistep on LSUN Cat-256](./consistency_models_8xb32-LSUN-cat-multistep-256x256.py)     |    LSUN    |     no      |    -     |
 
 You can also download checkpoints which is the main models in the paper to local machine and deliver the path to 'model_path' before infer.
 Here are the download links for each model checkpoint:
 
-- EDM on ImageNet-64: [edm_imagenet64_ema.pt](https://openaipublic.blob.core.windows.net/consistency/edm_imagenet64_ema.pt)
-- CD on ImageNet-64 with l2 metric: [cd_imagenet64_l2.pt](https://openaipublic.blob.core.windows.net/consistency/cd_imagenet64_l2.pt)
-- CD on ImageNet-64 with LPIPS metric: [cd_imagenet64_lpips.pt](https://openaipublic.blob.core.windows.net/consistency/cd_imagenet64_lpips.pt)
-- CT on ImageNet-64: [ct_imagenet64.pt](https://openaipublic.blob.core.windows.net/consistency/ct_imagenet64.pt)
-- EDM on LSUN Bedroom-256: [edm_bedroom256_ema.pt](https://openaipublic.blob.core.windows.net/consistency/edm_bedroom256_ema.pt)
-- CD on LSUN Bedroom-256 with l2 metric: [cd_bedroom256_l2.pt](https://openaipublic.blob.core.windows.net/consistency/cd_bedroom256_l2.pt)
-- CD on LSUN Bedroom-256 with LPIPS metric: [cd_bedroom256_lpips.pt](https://openaipublic.blob.core.windows.net/consistency/cd_bedroom256_lpips.pt)
-- CT on LSUN Bedroom-256: [ct_bedroom256.pt](https://openaipublic.blob.core.windows.net/consistency/ct_bedroom256.pt)
-- EDM on LSUN Cat-256: [edm_cat256_ema.pt](https://openaipublic.blob.core.windows.net/consistency/edm_cat256_ema.pt)
-- CD on LSUN Cat-256 with l2 metric: [cd_cat256_l2.pt](https://openaipublic.blob.core.windows.net/consistency/cd_cat256_l2.pt)
-- CD on LSUN Cat-256 with LPIPS metric: [cd_cat256_lpips.pt](https://openaipublic.blob.core.windows.net/consistency/cd_cat256_lpips.pt)
-- CT on LSUN Cat-256: [ct_cat256.pt](https://openaipublic.blob.core.windows.net/consistency/ct_cat256.pt)
+- EDM on ImageNet-64: [edm_imagenet64_ema.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/edm_imagenet64_ema.pt)
+- CD on ImageNet-64 with l2 metric: [cd_imagenet64_l2.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/cd_imagenet64_l2.pt)
+- CD on ImageNet-64 with LPIPS metric: [cd_imagenet64_lpips.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/cd_imagenet64_lpips.pt)
+- CT on ImageNet-64: [ct_imagenet64.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/ct_imagenet64.pt)
+- EDM on LSUN Bedroom-256: [edm_bedroom256_ema.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/edm_bedroom256_ema.pt)
+- CD on LSUN Bedroom-256 with l2 metric: [cd_bedroom256_l2.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/cd_bedroom256_l2.pt)
+- CD on LSUN Bedroom-256 with LPIPS metric: [cd_bedroom256_lpips.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/cd_bedroom256_lpips.pt)
+- CT on LSUN Bedroom-256: [ct_bedroom256.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/ct_bedroom256.pt)
+- EDM on LSUN Cat-256: [edm_cat256_ema.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/edm_cat256_ema.pt)
+- CD on LSUN Cat-256 with l2 metric: [cd_cat256_l2.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/cd_cat256_l2.pt)
+- CD on LSUN Cat-256 with LPIPS metric: [cd_cat256_lpips.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/cd_cat256_lpips.pt)
+- CT on LSUN Cat-256: [ct_cat256.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/ct_cat256.pt)
 
 ## quick start
 
@@ -64,6 +64,13 @@ python demo\mmagic_inference_demo.py \
     --model-name consistency_models \
     --model-config configs/consistency_models/consistency_models_8xb256-imagenet1k-multistep-64x64.py \
     --result-out-dir demo_consistency_model.jpg
+
+# conditional
+python demo\mmagic_inference_demo.py \
+    --model-name consistency_models \
+    --model-config configs/consistency_models/consistency_models_8xb256-imagenet1k-onestep-64x64.py \
+    --label 145 \
+    --result-out-dir demo_consistency_model.jpg
 ```
 
 </details>

diff --git a/configs/consistency_models/README_zh-CN.md b/configs/consistency_models/README_zh-CN.md
@@ -18,30 +18,30 @@
 
 ## 预训练模型
 
-|                                             Model                                             |  Dataset   | Download |
-| :-------------------------------------------------------------------------------------------: | :--------: | :------: |
-|       [onestep on ImageNet-64](./consistency_models_8xb256-imagenet1k-onestep-64x64.py)       | imagenet1k |    -     |
-|     [multistep on ImageNet-64](./consistency_models_8xb256-imagenet1k-multistep-64x64.py)     | imagenet1k |    -     |
-|   [onestep on LSUN Bedroom-256](./consistency_models_8xb32-LSUN-bedroom-onestep-256x256.py)   |    LSUN    |    -     |
-| [multistep on LSUN Bedroom-256](./consistency_models_8xb32-LSUN-bedroom-multistep-256x256.py) |    LSUN    |    -     |
-|       [onstep on LSUN Cat-256](./consistency_models_8xb32-LSUN-cat-onestep-256x256.py)        |    LSUN    |    -     |
-|     [multistep on LSUN Cat-256](./consistency_models_8xb32-LSUN-cat-multistep-256x256.py)     |    LSUN    |    -     |
+|                                             Model                                             |  Dataset   | Conditional | Download |
+| :-------------------------------------------------------------------------------------------: | :--------: | :---------: | :------: |
+|       [onestep on ImageNet-64](./consistency_models_8xb256-imagenet1k-onestep-64x64.py)       | imagenet1k |     yes     |    -     |
+|     [multistep on ImageNet-64](./consistency_models_8xb256-imagenet1k-multistep-64x64.py)     | imagenet1k |     yes     |    -     |
+|   [onestep on LSUN Bedroom-256](./consistency_models_8xb32-LSUN-bedroom-onestep-256x256.py)   |    LSUN    |     no      |    -     |
+| [multistep on LSUN Bedroom-256](./consistency_models_8xb32-LSUN-bedroom-multistep-256x256.py) |    LSUN    |     no      |    -     |
+|       [onstep on LSUN Cat-256](./consistency_models_8xb32-LSUN-cat-onestep-256x256.py)        |    LSUN    |     no      |    -     |
+|     [multistep on LSUN Cat-256](./consistency_models_8xb32-LSUN-cat-multistep-256x256.py)     |    LSUN    |     no      |    -     |
 
 你也可以在进行推理前先把论文中主要模型的权重下载到本地的机器上并将权重路径传给'model_path'。
 以下是每个模型权重的下载链接：
 
-- EDM on ImageNet-64: [edm_imagenet64_ema.pt](https://openaipublic.blob.core.windows.net/consistency/edm_imagenet64_ema.pt)
-- CD on ImageNet-64 with l2 metric: [cd_imagenet64_l2.pt](https://openaipublic.blob.core.windows.net/consistency/cd_imagenet64_l2.pt)
-- CD on ImageNet-64 with LPIPS metric: [cd_imagenet64_lpips.pt](https://openaipublic.blob.core.windows.net/consistency/cd_imagenet64_lpips.pt)
-- CT on ImageNet-64: [ct_imagenet64.pt](https://openaipublic.blob.core.windows.net/consistency/ct_imagenet64.pt)
-- EDM on LSUN Bedroom-256: [edm_bedroom256_ema.pt](https://openaipublic.blob.core.windows.net/consistency/edm_bedroom256_ema.pt)
-- CD on LSUN Bedroom-256 with l2 metric: [cd_bedroom256_l2.pt](https://openaipublic.blob.core.windows.net/consistency/cd_bedroom256_l2.pt)
-- CD on LSUN Bedroom-256 with LPIPS metric: [cd_bedroom256_lpips.pt](https://openaipublic.blob.core.windows.net/consistency/cd_bedroom256_lpips.pt)
-- CT on LSUN Bedroom-256: [ct_bedroom256.pt](https://openaipublic.blob.core.windows.net/consistency/ct_bedroom256.pt)
-- EDM on LSUN Cat-256: [edm_cat256_ema.pt](https://openaipublic.blob.core.windows.net/consistency/edm_cat256_ema.pt)
-- CD on LSUN Cat-256 with l2 metric: [cd_cat256_l2.pt](https://openaipublic.blob.core.windows.net/consistency/cd_cat256_l2.pt)
-- CD on LSUN Cat-256 with LPIPS metric: [cd_cat256_lpips.pt](https://openaipublic.blob.core.windows.net/consistency/cd_cat256_lpips.pt)
-- CT on LSUN Cat-256: [ct_cat256.pt](https://openaipublic.blob.core.windows.net/consistency/ct_cat256.pt)
+- EDM on ImageNet-64: [edm_imagenet64_ema.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/edm_imagenet64_ema.pt)
+- CD on ImageNet-64 with l2 metric: [cd_imagenet64_l2.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/cd_imagenet64_l2.pt)
+- CD on ImageNet-64 with LPIPS metric: [cd_imagenet64_lpips.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/cd_imagenet64_lpips.pt)
+- CT on ImageNet-64: [ct_imagenet64.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/ct_imagenet64.pt)
+- EDM on LSUN Bedroom-256: [edm_bedroom256_ema.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/edm_bedroom256_ema.pt)
+- CD on LSUN Bedroom-256 with l2 metric: [cd_bedroom256_l2.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/cd_bedroom256_l2.pt)
+- CD on LSUN Bedroom-256 with LPIPS metric: [cd_bedroom256_lpips.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/cd_bedroom256_lpips.pt)
+- CT on LSUN Bedroom-256: [ct_bedroom256.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/ct_bedroom256.pt)
+- EDM on LSUN Cat-256: [edm_cat256_ema.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/edm_cat256_ema.pt)
+- CD on LSUN Cat-256 with l2 metric: [cd_cat256_l2.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/cd_cat256_l2.pt)
+- CD on LSUN Cat-256 with LPIPS metric: [cd_cat256_lpips.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/cd_cat256_lpips.pt)
+- CT on LSUN Cat-256: [ct_cat256.pt](https://download.openxlab.org.cn/models/xiaomile/consistency_models/weight/ct_cat256.pt)
 -
 
 ## 快速开始
@@ -54,17 +54,24 @@
 您可以使用以下命令来使用该模型进行推理：
 
 ```shell
-# onestep
+# 一步生成
 python demo\mmagic_inference_demo.py \
     --model-name consistency_models \
     --model-config configs/consistency_models/consistency_models_8xb256-imagenet1k-onestep-64x64.py \
     --result-out-dir demo_consistency_model.jpg
 
-# multistep
+# 多步生成
 python demo\mmagic_inference_demo.py \
     --model-name consistency_models \
     --model-config configs/consistency_models/consistency_models_8xb256-imagenet1k-multistep-64x64.py \
     --result-out-dir demo_consistency_model.jpg
+
+# 条件控制生成
+python demo\mmagic_inference_demo.py \
+    --model-name consistency_models \
+    --model-config configs/consistency_models/consistency_models_8xb256-imagenet1k-onestep-64x64.py \
+    --label 145 \
+    --result-out-dir demo_consistency_model.jpg
 ```
 
 </details>

diff --git a/configs/consistency_models/consistency_models_8xb256-imagenet1k-multistep-64x64.py b/configs/consistency_models/consistency_models_8xb256-imagenet1k-multistep-64x64.py
@@ -35,8 +35,8 @@
     generator='determ',
     image_size=64,
     learn_sigma=False,
-    model_path='https://openaipublic.blob.core.windows.net/'
-    'consistency/cd_imagenet64_l2.pt',
+    model_path='https://download.openxlab.org.cn/models/xiaomile/'
+    'consistency_models/weight/cd_imagenet64_l2.pt',
     num_classes=1000,
     sampler='multistep',
     seed=42,

diff --git a/configs/consistency_models/consistency_models_8xb256-imagenet1k-onestep-64x64.py b/configs/consistency_models/consistency_models_8xb256-imagenet1k-onestep-64x64.py
@@ -35,8 +35,8 @@
     generator='determ',
     image_size=64,
     learn_sigma=False,
-    model_path='https://openaipublic.blob.core.windows.net/'
-    'consistency/cd_imagenet64_l2.pt',
+    model_path='https://download.openxlab.org.cn/models/xiaomile/'
+    'consistency_models/weight/cd_imagenet64_l2.pt',
     num_classes=1000,
     sampler='onestep',
     seed=42,

diff --git a/configs/consistency_models/consistency_models_8xb32-LSUN-bedroom-multistep-256x256.py b/configs/consistency_models/consistency_models_8xb32-LSUN-bedroom-multistep-256x256.py
@@ -35,8 +35,8 @@
     generator='determ-indiv',
     image_size=256,
     learn_sigma=False,
-    model_path='https://openaipublic.blob.core.windows.net'
-    '/consistency/ct_bedroom256.pt',
+    model_path='https://download.openxlab.org.cn/models/xiaomile/'
+    'consistency_models/weight/ct_bedroom256.pt',
     num_classes=1000,
     sampler='multistep',
     seed=42,

diff --git a/configs/consistency_models/consistency_models_8xb32-LSUN-bedroom-onestep-256x256.py b/configs/consistency_models/consistency_models_8xb32-LSUN-bedroom-onestep-256x256.py
@@ -35,8 +35,8 @@
     generator='determ-indiv',
     image_size=256,
     learn_sigma=False,
-    model_path=
-    'https://openaipublic.blob.core.windows.net/consistency/ct_bedroom256.pt',
+    model_path='https://download.openxlab.org.cn/models/xiaomile/'
+    'consistency_models/weight/ct_bedroom256.pt',
     num_classes=1000,
     sampler='onestep',
     seed=42,

diff --git a/configs/consistency_models/consistency_models_8xb32-LSUN-cat-multistep-256x256.py b/configs/consistency_models/consistency_models_8xb32-LSUN-cat-multistep-256x256.py
@@ -35,8 +35,8 @@
     generator='determ-indiv',
     image_size=256,
     learn_sigma=False,
-    model_path=
-    'https://openaipublic.blob.core.windows.net/consistency/ct_cat256.pt',
+    model_path='https://download.openxlab.org.cn/models/xiaomile/'
+    'consistency_models/weight/ct_cat256.pt',
     num_classes=1000,
     sampler='multistep',
     seed=42,

diff --git a/configs/consistency_models/consistency_models_8xb32-LSUN-cat-onestep-256x256.py b/configs/consistency_models/consistency_models_8xb32-LSUN-cat-onestep-256x256.py
@@ -35,8 +35,8 @@
     generator='determ-indiv',
     image_size=256,
     learn_sigma=False,
-    model_path=
-    'https://openaipublic.blob.core.windows.net/consistency/ct_cat256.pt',
+    model_path='https://download.openxlab.org.cn/models/xiaomile/'
+    'consistency_models/weight/ct_cat256.pt',
     num_classes=1000,
     sampler='onestep',
     seed=42,