Repository: open-mmlab/mmagic (language: python)

Issue:
File not found while compiling dcn module.
There are some missing files, such as `deform_conv_cuda_kernel.cu`. More details can be found here, https://github.com/open-mmlab/mmdetection/tree/master/mmdet/ops/dcn/src, copy this into src directory can be fixed.


Discussion / hints:
Thanks, we have uploaded this file. 

Fix patch:
diff --git a/configs/draggan/README.md b/configs/draggan/README.md
new file mode 100644
index 0000000000..1e3629e082
--- /dev/null
+++ b/configs/draggan/README.md
@@ -0,0 +1,86 @@
+# DragGAN (SIGGRAPH'2023)
+
+> [Drag Your GAN: Interactive Point-based Manipulation on the Generative Image Manifold](https://arxiv.org/pdf/2305.10973.pdf)
+
+> **Task**: DragGAN
+
+<!-- [ALGORITHM] -->
+
+## Abstract
+
+<!-- [ABSTRACT] -->
+
+Synthesizing visual content that meets users’ needs often requires flexible
+and precise controllability of the pose, shape, expression, and layout of the
+generated objects. Existing approaches gain controllability of generative
+adversarial networks (GANs) via manually annotated training data or a
+prior 3D model, which often lack flexibility, precision, and generality. In
+this work, we study a powerful yet much less explored way of controlling
+GANs, that is, to "drag" any points of the image to precisely reach target
+points in a user-interactive manner, as shown in Fig.1. To achieve this, we
+propose DragGAN, which consists of two main components: 1) a feature-based motion supervision that drives the handle point to move towards
+the target position, and 2) a new point tracking approach that leverages
+the discriminative generator features to keep localizing the position of the
+handle points. Through DragGAN, anyone can deform an image with precise
+control over where pixels go, thus manipulating the pose, shape, expression,and layout of diverse categories such as animals, cars, humans, landscapes,
+etc. As these manipulations are performed on the learned generative image
+manifold of a GAN, they tend to produce realistic outputs even for challenging scenarios such as hallucinating occluded content and deforming
+shapes that consistently follow the object’s rigidity. Both qualitative and
+quantitative comparisons demonstrate the advantage of DragGAN over prior
+approaches in the tasks of image manipulation and point tracking. We also
+showcase the manipulation of real images through GAN inversion.
+
+<!-- [IMAGE] -->
+
+<div align=center>
+<img src="https://github.com/open-mmlab/mmagic/assets/55343765/7c397bd0-fa07-48fe-8a7c-a4022907404b"/>
+</div>
+
+## Results and Models
+
+<table><tr>
+<b> Gradio Demo of DragGAN StyleGAN2-elephants-512 by MMagic </b>
+<td><img src="https://github.com/open-mmlab/mmagic/assets/55343765/08e9a687-0a6e-4d3f-94ec-22c46bd61819" border=0></td>
+<td><img src="https://github.com/open-mmlab/mmagic/assets/55343765/6fab1ccd-e190-4cd0-a8d5-0e843f65930b" border=0></td>
+</tr></table>
+
+|                        Model                         |      Dataset       |             Comment             | FID50k | Precision50k | Recall50k |                                 Download                                 |
+| :--------------------------------------------------: | :----------------: | :-----------------------------: | :----: | :----------: | :-------: | :----------------------------------------------------------------------: |
+|   [stylegan2_lion_512x512](./stylegan2_512x512.py)   |   Internet Lions   |     self-distilled StyleGAN     |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-Lions-internet) |
+| [stylegan2_elphants_512x512](./stylegan2_512x512.py) | Internet Elephants |     self-distilled StyleGAN     |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-elephants-internet) |
+|   [stylegan2_cats_512x512](./stylegan2_512x512.py)   |      Cat AFHQ      |     self-distilled StyleGAN     |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-cat-AFHQ) |
+|   [stylegan2_face_512x512](./stylegan2_512x512.py)   |        FFHQ        |     self-distilled StyleGAN     |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-FFHQ) |
+|  [stylegan2_horse_256x256](./stylegan2_256x256.py)   |     LSUN-Horse     |     self-distilled StyleGAN     |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-lsun-horses) |
+| [stylegan2_dogs_1024x1024](./stylegan2_1024x1024.py) |   Internet Dogs    |     self-distilled StyleGAN     |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-dogs-internet) |
+|   [stylegan2_car_512x512](./stylegan2_512x512.py)    |        Car         | transfer from official training |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-car-official) |
+|   [stylegan2_cat_256x256](./stylegan2_256x256.py)    |        Cat         | transfer from official training |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-cat-official) |
+
+## Demo
+
+To run DragGAN demo, please follow these two steps:
+
+First, put your checkpoint path in `./checkpoints`, *e.g.* `./checkpoints/stylegan2_lions_512_pytorch_mmagic.pth`. To be specific,
+
+```shell
+mkdir checkpoints
+cd checkpoints
+wget -O stylegan2_lions_512_pytorch_mmagic.pth https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-Lions-internet
+```
+
+Then, try on the script:
+
+```shell
+python demo/gradio_draggan.py
+```
+
+## Citation
+
+```latex
+@inproceedings{pan2023drag,
+  title={Drag your gan: Interactive point-based manipulation on the generative image manifold},
+  author={Pan, Xingang and Tewari, Ayush and Leimk{\"u}hler, Thomas and Liu, Lingjie and Meka, Abhimitra and Theobalt, Christian},
+  booktitle={ACM SIGGRAPH 2023 Conference Proceedings},
+  pages={1--11},
+  year={2023}
+}
+```
diff --git a/configs/draggan/README_zh-CN.md b/configs/draggan/README_zh-CN.md
new file mode 100644
index 0000000000..82cbdbdce9
--- /dev/null
+++ b/configs/draggan/README_zh-CN.md
@@ -0,0 +1,79 @@
+# DragGAN (SIGGRAPH'2023)
+
+> [Drag Your GAN: Interactive Point-based Manipulation on the Generative Image Manifold](https://arxiv.org/pdf/2305.10973.pdf)
+
+> **Task**: DragGAN
+
+<!-- [ALGORITHM] -->
+
+## 摘要
+
+<!-- [ABSTRACT] -->
+
+要合成满足用户需求的视觉内容，通常需要灵活、精确地控制人物的姿势、形状、表情和布局。
+和精确地控制生成对象的姿势、形状、表情和布局。
+生成对象。现有的方法是通过人工标注训练数据来获得生成对抗网络（GANs）的可控性。
+生成式对抗网络 (GAN) 的可控性是通过人工标注的训练数据或预先的 3D 模型来实现的。
+生成式对抗网络（GAN）的可控性是通过人工标注的训练数据或事先建立的三维模型来实现的，而这些方法往往缺乏灵活性、精确性和通用性。在中，我们研究了一种功能强大但探索较少的生成式对抗网络（GAN）控制方法。
+即 "拖动 "图像中的任意点，以用户互动的方式精确到达目标点。
+如图 1 所示。为此，我们
+DragGAN 由两个主要部分组成： 1) 基于特征的运动监督，驱动手柄点向目标位置移动；以及目标位置，以及 2) 一种新的点跟踪方法，利用的新点跟踪方法。
+手柄点的位置。通过 DragGAN，任何人都可以对图像进行变形，并精确控制像素的移动位置。
+通过 DragGAN，任何人都可以精确控制像素的位置，从而改变图像的姿态、形状、表情和布局，如动物、汽车、人类、风景等、
+等等。由于这些操作是在 GAN 的学习生成图像流形上进行的流形上进行，因此即使是在具有挑战性的情况下，它们也能产生逼真的输出，例如幻觉遮挡内容和变形形状始终遵循物体的刚性。定性和定量比较都表明，在图像处理和点跟踪任务中，DragGAN
+在图像处理和点跟踪任务中的优势。我们还展示了通过 GAN 反演对真实图像的操作。
+
+<!-- [IMAGE] -->
+
+<div align=center>
+<img src="https://github.com/open-mmlab/mmagic/assets/55343765/7c397bd0-fa07-48fe-8a7c-a4022907404b"/>
+</div>
+
+## 结果和模型
+
+<table><tr>
+<b> MMagic下实现的DragGAN（StyleGAN2-elephants-512） Gradio示例 </b>
+<td><img src="https://github.com/open-mmlab/mmagic/assets/55343765/08e9a687-0a6e-4d3f-94ec-22c46bd61819" border=0></td>
+<td><img src="https://github.com/open-mmlab/mmagic/assets/55343765/6fab1ccd-e190-4cd0-a8d5-0e843f65930b" border=0></td>
+</tr></table>
+
+|                         模型                         |       数据集       |      评论      | FID50k | Precision50k | Recall50k |                                       权重下载链接                                        |
+| :--------------------------------------------------: | :----------------: | :------------: | :----: | :----------: | :-------: | :---------------------------------------------------------------------------------------: |
+|   [stylegan2_lion_512x512](./stylegan2_512x512.py)   |   Internet Lions   | 自蒸馏StyleGAN |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-Lions-internet) |
+| [stylegan2_elphants_512x512](./stylegan2_512x512.py) | Internet Elephants | 自蒸馏StyleGAN |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-elephants-internet) |
+|   [stylegan2_cats_512x512](./stylegan2_512x512.py)   |      Cat AFHQ      | 自蒸馏StyleGAN |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-cat-AFHQ) |
+|   [stylegan2_face_512x512](./stylegan2_512x512.py)   |        FFHQ        | 自蒸馏StyleGAN |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-FFHQ) |
+|  [stylegan2_horse_256x256](./stylegan2_256x256.py)   |     LSUN-Horse     | 自蒸馏StyleGAN |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-lsun-horses) |
+| [stylegan2_dogs_1024x1024](./stylegan2_1024x1024.py) |   Internet Dogs    | 自蒸馏StyleGAN |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-dogs-internet) |
+|   [stylegan2_car_512x512](./stylegan2_512x512.py)    |        Car         |    官方训练    |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-car-official) |
+|   [stylegan2_cat_256x256](./stylegan2_256x256.py)    |        Cat         |    官方训练    |  0.0   |     0.0      |    0.0    | [model](https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-cat-official) |
+
+## 演示
+
+为了使用DragGAN演示, 请执行以下两步:
+
+首先，把模型文件放在 `./checkpoints`下, 比如 `./checkpoints/stylegan2_lions_512_pytorch_mmagic.pth`. 具体来说,
+
+```shell
+mkdir checkpoints
+cd checkpoints
+wget -O stylegan2_lions_512_pytorch_mmagic.pth https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-Lions-internet
+```
+
+然后，执行下面的脚本:
+
+```shell
+python demo/gradio_draggan.py
+```
+
+## 引用
+
+```latex
+@inproceedings{pan2023drag,
+  title={Drag your gan: Interactive point-based manipulation on the generative image manifold},
+  author={Pan, Xingang and Tewari, Ayush and Leimk{\"u}hler, Thomas and Liu, Lingjie and Meka, Abhimitra and Theobalt, Christian},
+  booktitle={ACM SIGGRAPH 2023 Conference Proceedings},
+  pages={1--11},
+  year={2023}
+}
+```
diff --git a/configs/draggan/metafile.yml b/configs/draggan/metafile.yml
new file mode 100644
index 0000000000..3e981f29d5
--- /dev/null
+++ b/configs/draggan/metafile.yml
@@ -0,0 +1,74 @@
+Collections:
+- Name: DragGAN
+  Paper:
+    Title: 'Drag Your GAN: Interactive Point-based Manipulation on the Generative
+      Image Manifold'
+    URL: https://arxiv.org/pdf/2305.10973.pdf
+  README: configs/draggan/README.md
+  Task:
+  - draggan
+  Year: 2023
+Models:
+- Config: configs/draggan/stylegan2_512x512.py
+  In Collection: DragGAN
+  Name: stylegan2_512x512
+  Results:
+  - Dataset: InternetLions
+    Metrics:
+      FID50k: 0.0
+      Precision50k: 0.0
+      Recall50k: 0.0
+    Task: DragGAN
+  - Dataset: InternetElephants
+    Metrics:
+      FID50k: 0.0
+      Precision50k: 0.0
+      Recall50k: 0.0
+    Task: DragGAN
+  - Dataset: CatAFHQ
+    Metrics:
+      FID50k: 0.0
+      Precision50k: 0.0
+      Recall50k: 0.0
+    Task: DragGAN
+  - Dataset: FFHQ
+    Metrics:
+      FID50k: 0.0
+      Precision50k: 0.0
+      Recall50k: 0.0
+    Task: DragGAN
+  - Dataset: Car
+    Metrics:
+      FID50k: 0.0
+      Precision50k: 0.0
+      Recall50k: 0.0
+    Task: DragGAN
+  Weights: https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-car-official
+- Config: configs/draggan/stylegan2_256x256.py
+  In Collection: DragGAN
+  Name: stylegan2_256x256
+  Results:
+  - Dataset: LSUN-Horse
+    Metrics:
+      FID50k: 0.0
+      Precision50k: 0.0
+      Recall50k: 0.0
+    Task: DragGAN
+  - Dataset: Cat
+    Metrics:
+      FID50k: 0.0
+      Precision50k: 0.0
+      Recall50k: 0.0
+    Task: DragGAN
+  Weights: https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-cat-official
+- Config: configs/draggan/stylegan2_1024x1024.py
+  In Collection: DragGAN
+  Name: stylegan2_1024x1024
+  Results:
+  - Dataset: InternetDogs
+    Metrics:
+      FID50k: 0.0
+      Precision50k: 0.0
+      Recall50k: 0.0
+    Task: DragGAN
+  Weights: https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-dogs-internet
diff --git a/configs/draggan/stylegan2_1024x1024.py b/configs/draggan/stylegan2_1024x1024.py
new file mode 100644
index 0000000000..2a2422adf4
--- /dev/null
+++ b/configs/draggan/stylegan2_1024x1024.py
@@ -0,0 +1,87 @@
+"""Config for the `config-f` setting in StyleGAN2."""
+
+_base_ = [
+    '../_base_/datasets/ffhq_flip.py', '../_base_/models/base_styleganv2.py',
+    '../_base_/gen_default_runtime.py'
+]
+
+# reg params
+d_reg_interval = 16
+g_reg_interval = 4
+
+g_reg_ratio = g_reg_interval / (g_reg_interval + 1)
+d_reg_ratio = d_reg_interval / (d_reg_interval + 1)
+
+ema_half_life = 10.  # G_smoothing_kimg
+
+model = dict(
+    generator=dict(
+        out_size=1024,
+        update_mean_latent_with_ema=True,
+        bgr2rgb=False,
+        fixed_noise=True),
+    discriminator=dict(in_size=1024),  # useless
+    ema_config=dict(
+        type='ExponentialMovingAverage',
+        interval=1,
+        momentum=1. - (0.5**(32. / (ema_half_life * 1000.)))),
+    loss_config=dict(
+        r1_loss_weight=10. / 2. * d_reg_interval,
+        r1_interval=d_reg_interval,
+        norm_mode='HWC',
+        g_reg_interval=g_reg_interval,
+        g_reg_weight=2. * g_reg_interval,
+        pl_batch_shrink=2))
+
+train_cfg = dict(max_iters=800002)
+
+optim_wrapper = dict(
+    generator=dict(
+        optimizer=dict(
+            type='Adam', lr=0.002 * g_reg_ratio, betas=(0,
+                                                        0.99**g_reg_ratio))),
+    discriminator=dict(
+        optimizer=dict(
+            type='Adam', lr=0.002 * d_reg_ratio, betas=(0,
+                                                        0.99**d_reg_ratio))))
+
+batch_size = 1
+data_root = './data/ffhq/images'
+
+train_dataloader = dict(
+    batch_size=batch_size, dataset=dict(data_root=data_root))
+
+val_dataloader = dict(batch_size=batch_size, dataset=dict(data_root=data_root))
+
+test_dataloader = dict(
+    batch_size=batch_size, dataset=dict(data_root=data_root))
+
+# VIS_HOOK
+custom_hooks = [
+    dict(
+        type='VisualizationHook',
+        interval=5000,
+        fixed_input=True,
+        vis_kwargs_list=dict(type='GAN', name='fake_img'))
+]
+
+# METRICS
+metrics = [
+    dict(
+        type='FrechetInceptionDistance',
+        prefix='FID-Full-50k',
+        fake_nums=50000,
+        inception_style='StyleGAN',
+        sample_model='ema'),
+    dict(type='PrecisionAndRecall', fake_nums=50000, prefix='PR-50K'),
+    dict(type='PerceptualPathLength', fake_nums=50000, prefix='ppl-w')
+]
+# NOTE: config for save multi best checkpoints
+# default_hooks = dict(
+#     checkpoint=dict(
+#         save_best=['FID-Full-50k/fid', 'IS-50k/is'],
+#         rule=['less', 'greater']))
+default_hooks = dict(checkpoint=dict(save_best='FID-Full-50k/fid'))
+
+val_evaluator = dict(metrics=metrics)
+test_evaluator = dict(metrics=metrics)
diff --git a/configs/draggan/stylegan2_256x256.py b/configs/draggan/stylegan2_256x256.py
new file mode 100644
index 0000000000..aa20fa4e00
--- /dev/null
+++ b/configs/draggan/stylegan2_256x256.py
@@ -0,0 +1,87 @@
+"""Config for the `config-f` setting in StyleGAN2."""
+
+_base_ = [
+    '../_base_/datasets/ffhq_flip.py', '../_base_/models/base_styleganv2.py',
+    '../_base_/gen_default_runtime.py'
+]
+
+# reg params
+d_reg_interval = 16
+g_reg_interval = 4
+
+g_reg_ratio = g_reg_interval / (g_reg_interval + 1)
+d_reg_ratio = d_reg_interval / (d_reg_interval + 1)
+
+ema_half_life = 10.  # G_smoothing_kimg
+
+model = dict(
+    generator=dict(
+        out_size=256,
+        update_mean_latent_with_ema=True,
+        bgr2rgb=False,
+        fixed_noise=True),
+    discriminator=dict(in_size=256),  # useless
+    ema_config=dict(
+        type='ExponentialMovingAverage',
+        interval=1,
+        momentum=1. - (0.5**(32. / (ema_half_life * 1000.)))),
+    loss_config=dict(
+        r1_loss_weight=10. / 2. * d_reg_interval,
+        r1_interval=d_reg_interval,
+        norm_mode='HWC',
+        g_reg_interval=g_reg_interval,
+        g_reg_weight=2. * g_reg_interval,
+        pl_batch_shrink=2))
+
+train_cfg = dict(max_iters=800002)
+
+optim_wrapper = dict(
+    generator=dict(
+        optimizer=dict(
+            type='Adam', lr=0.002 * g_reg_ratio, betas=(0,
+                                                        0.99**g_reg_ratio))),
+    discriminator=dict(
+        optimizer=dict(
+            type='Adam', lr=0.002 * d_reg_ratio, betas=(0,
+                                                        0.99**d_reg_ratio))))
+
+batch_size = 1
+data_root = './data/ffhq/images'
+
+train_dataloader = dict(
+    batch_size=batch_size, dataset=dict(data_root=data_root))
+
+val_dataloader = dict(batch_size=batch_size, dataset=dict(data_root=data_root))
+
+test_dataloader = dict(
+    batch_size=batch_size, dataset=dict(data_root=data_root))
+
+# VIS_HOOK
+custom_hooks = [
+    dict(
+        type='VisualizationHook',
+        interval=5000,
+        fixed_input=True,
+        vis_kwargs_list=dict(type='GAN', name='fake_img'))
+]
+
+# METRICS
+metrics = [
+    dict(
+        type='FrechetInceptionDistance',
+        prefix='FID-Full-50k',
+        fake_nums=50000,
+        inception_style='StyleGAN',
+        sample_model='ema'),
+    dict(type='PrecisionAndRecall', fake_nums=50000, prefix='PR-50K'),
+    dict(type='PerceptualPathLength', fake_nums=50000, prefix='ppl-w')
+]
+# NOTE: config for save multi best checkpoints
+# default_hooks = dict(
+#     checkpoint=dict(
+#         save_best=['FID-Full-50k/fid', 'IS-50k/is'],
+#         rule=['less', 'greater']))
+default_hooks = dict(checkpoint=dict(save_best='FID-Full-50k/fid'))
+
+val_evaluator = dict(metrics=metrics)
+test_evaluator = dict(metrics=metrics)
diff --git a/configs/draggan/stylegan2_512x512.py b/configs/draggan/stylegan2_512x512.py
new file mode 100644
index 0000000000..6630103e36
--- /dev/null
+++ b/configs/draggan/stylegan2_512x512.py
@@ -0,0 +1,87 @@
+"""Config for the `config-f` setting in StyleGAN2."""
+
+_base_ = [
+    '../_base_/datasets/ffhq_flip.py', '../_base_/models/base_styleganv2.py',
+    '../_base_/gen_default_runtime.py'
+]
+
+# reg params
+d_reg_interval = 16
+g_reg_interval = 4
+
+g_reg_ratio = g_reg_interval / (g_reg_interval + 1)
+d_reg_ratio = d_reg_interval / (d_reg_interval + 1)
+
+ema_half_life = 10.  # G_smoothing_kimg
+
+model = dict(
+    generator=dict(
+        out_size=512,
+        update_mean_latent_with_ema=True,
+        bgr2rgb=False,
+        fixed_noise=True),
+    discriminator=dict(in_size=512),  # useless
+    ema_config=dict(
+        type='ExponentialMovingAverage',
+        interval=1,
+        momentum=1. - (0.5**(32. / (ema_half_life * 1000.)))),
+    loss_config=dict(
+        r1_loss_weight=10. / 2. * d_reg_interval,
+        r1_interval=d_reg_interval,
+        norm_mode='HWC',
+        g_reg_interval=g_reg_interval,
+        g_reg_weight=2. * g_reg_interval,
+        pl_batch_shrink=2))
+
+train_cfg = dict(max_iters=800002)
+
+optim_wrapper = dict(
+    generator=dict(
+        optimizer=dict(
+            type='Adam', lr=0.002 * g_reg_ratio, betas=(0,
+                                                        0.99**g_reg_ratio))),
+    discriminator=dict(
+        optimizer=dict(
+            type='Adam', lr=0.002 * d_reg_ratio, betas=(0,
+                                                        0.99**d_reg_ratio))))
+
+batch_size = 1
+data_root = './data/ffhq/images'
+
+train_dataloader = dict(
+    batch_size=batch_size, dataset=dict(data_root=data_root))
+
+val_dataloader = dict(batch_size=batch_size, dataset=dict(data_root=data_root))
+
+test_dataloader = dict(
+    batch_size=batch_size, dataset=dict(data_root=data_root))
+
+# VIS_HOOK
+custom_hooks = [
+    dict(
+        type='VisualizationHook',
+        interval=5000,
+        fixed_input=True,
+        vis_kwargs_list=dict(type='GAN', name='fake_img'))
+]
+
+# METRICS
+metrics = [
+    dict(
+        type='FrechetInceptionDistance',
+        prefix='FID-Full-50k',
+        fake_nums=50000,
+        inception_style='StyleGAN',
+        sample_model='ema'),
+    dict(type='PrecisionAndRecall', fake_nums=50000, prefix='PR-50K'),
+    dict(type='PerceptualPathLength', fake_nums=50000, prefix='ppl-w')
+]
+# NOTE: config for save multi best checkpoints
+# default_hooks = dict(
+#     checkpoint=dict(
+#         save_best=['FID-Full-50k/fid', 'IS-50k/is'],
+#         rule=['less', 'greater']))
+default_hooks = dict(checkpoint=dict(save_best='FID-Full-50k/fid'))
+
+val_evaluator = dict(metrics=metrics)
+test_evaluator = dict(metrics=metrics)
diff --git a/configs/styleganv2/stylegan2_c2_8xb4_ffhq-1024x1024.py b/configs/styleganv2/stylegan2_c2_8xb4_ffhq-1024x1024.py
index f19330961c..a27c7ac3d9 100644
--- a/configs/styleganv2/stylegan2_c2_8xb4_ffhq-1024x1024.py
+++ b/configs/styleganv2/stylegan2_c2_8xb4_ffhq-1024x1024.py
@@ -29,6 +29,8 @@
         g_reg_weight=2. * g_reg_interval,
         pl_batch_shrink=2))
 
+extra_parameters = dict(num_batches=1, sample_model='orig')
+
 train_cfg = dict(max_iters=800002)
 
 optim_wrapper = dict(
diff --git a/demo/README.md b/demo/README.md
index 5c36b969df..0634543170 100644
--- a/demo/README.md
+++ b/demo/README.md
@@ -38,6 +38,10 @@ Table of contents:
 
 [3. Other demos](#3-other-demos)
 
+&#8195;    [3.1. Gradio demo](#31-gradio-demo)
+
+&#8195;      [3.1.1. DragGAN](#311-draggan)
+
 ## 1. Download sample images or videos
 
 We prepared some images and videos for you to run demo with. After MMagic is well installed, you could use demos in this folder to infer these data.
@@ -283,3 +287,23 @@ python demo/mmagic_inference_demo.py \
     --model-name eg3d \
     --result-out-dir ./resources/output/eg3d-output
 ```
+
+## 3. Other demos
+
+## 3.1 gradio demo
+
+#### 3.1.1 DragGAN
+
+First, put your checkpoint path in `./checkpoints`, *e.g.* `./checkpoints/stylegan2_lions_512_pytorch_mmagic.pth`. For example,
+
+```shell
+mkdir checkpoints
+cd checkpoints
+wget -O stylegan2_lions_512_pytorch_mmagic.pth https://download.openxlab.org.cn/models/qsun1/DragGAN-StyleGAN2-checkpoint/weight//StyleGAN2-Lions-internet
+```
+
+Then, try on the script:
+
+```shell
+python demo/gradio_draggan.py
+```
diff --git a/demo/gradio_draggan.py b/demo/gradio_draggan.py
new file mode 100644
index 0000000000..d30b8ce186
--- /dev/null
+++ b/demo/gradio_draggan.py
@@ -0,0 +1,901 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import os
+import os.path as osp
+from argparse import ArgumentParser
+from functools import partial
+
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+
+from demo.utils.gradio_utils import (EasyDict, ImageMask, draw_mask_on_image,
+                                     draw_points_on_image,
+                                     get_latest_points_pair, get_valid_mask,
+                                     on_change_single_global_state)
+from demo.utils.renderer import Renderer, add_watermark_np
+
+parser = ArgumentParser()
+parser.add_argument('--share', action='store_true', default='True')
+parser.add_argument('--cache-dir', type=str, default='./checkpoints')
+parser.add_argument(
+    '--listen',
+    action='store_true',
+    help=''
+    'launch gradio with 0.0.0.0 as server name, \
+        allowing to respond to network requests',
+)
+args = parser.parse_args()
+cache_dir = args.cache_dir
+
+
+def get_default_gr_state(device='cuda',
+                         init_pkl='stylegan2_lions_512_pytorch_mmagic.pth'):
+    global_state = gr.State({
+        'images': {
+            # image_orig: the original image, change with seed/model is changed
+            # image_raw: image with mask and points, change during optimization
+            # image_show: image showed on screen
+        },
+        'temporal_params': {
+            # stop
+        },
+        'mask':
+        None,  # mask for visualization, 1 for editing and 0 for unchange
+        'last_mask': None,  # last edited mask
+        'show_mask': True,  # add button
+        'generator_params': EasyDict(),
+        'params': {
+            'seed': 0,
+            'motion_lambda': 20,
+            'r1_in_pixels': 3,
+            'r2_in_pixels': 12,
+            'magnitude_direction_in_pixels': 1.0,
+            'latent_space': 'w+',
+            'trunc_psi': 0.7,
+            'trunc_cutoff': None,
+            'lr': 0.001,
+        },
+        'device': device,
+        'draw_interval': 1,
+        'renderer': Renderer(disable_timing=True),
+        'points': {},
+        'curr_point': None,
+        'curr_type_point': 'start',
+        'editing_state': 'add_points',
+        'pretrained_weight': init_pkl
+    })
+    return global_state
+
+
+def reverse_point_pairs(points):
+    new_points = []
+    for p in points:
+        new_points.append([p[1], p[0]])
+    return new_points
+
+
+def clear_state(global_state, target=None):
+    """Clear target history state from global_state If target is not defined,
+    points and mask will be both removed.
+
+    1. set global_state['points'] as empty dict
+    2. set global_state['mask'] as full-one mask.
+    """
+    if target is None:
+        target = ['point', 'mask']
+    if not isinstance(target, list):
+        target = [target]
+    if 'point' in target:
+        global_state['points'] = dict()
+        print('Clear Points State!')
+    if 'mask' in target:
+        image_raw = global_state['images']['image_raw']
+        global_state['mask'] = np.ones((image_raw.size[1], image_raw.size[0]),
+                                       dtype=np.uint8)
+        print('Clear mask State!')
+
+    return global_state
+
+
+def init_images(global_state, cache_dir='./checkpoints'):
+    """This function is called only ones with Gradio App is started.
+
+    0. pre-process global_state, unpack value from global_state of need
+    1. Re-init renderer
+    2. run `renderer._render_drag_impl` with `is_drag=False` to generate
+       new image
+    3. Assign images to global state and re-generate mask
+    """
+
+    if isinstance(global_state, gr.State):
+        state = global_state.value
+    else:
+        state = global_state
+
+    state['renderer'].init_network(
+        state['generator_params'],  # res
+        os.path.join(cache_dir, state['pretrained_weight']
+                     ),  # StyleGAN2 checkpoint path from Official DragGAN
+        state['params']['seed'],  # w0_seed,
+        None,  # w_load
+        state['params']['latent_space'] == 'w+',  # w_plus
+        'const',
+        state['params']['trunc_psi'],  # trunc_psi,
+        state['params']['trunc_cutoff'],  # trunc_cutoff,
+        None,  # input_transform
+        state['params']['lr']  # lr,
+    )
+    state['renderer']._render_drag_impl(
+        state['generator_params'], is_drag=False, to_pil=True)
+
+    init_image = state['generator_params'].image
+    state['images']['image_orig'] = init_image
+    state['images']['image_raw'] = init_image
+    state['images']['image_show'] = Image.fromarray(
+        add_watermark_np(np.array(init_image)))
+    state['mask'] = np.ones((init_image.size[1], init_image.size[0]),
+                            dtype=np.uint8)
+    return global_state
+
+
+def update_image_draw(image, points, mask, show_mask, global_state=None):
+    image_draw = draw_points_on_image(image, points)
+    if show_mask and mask is not None and not (mask == 0).all() and not (
+            mask == 1).all():
+        image_draw = draw_mask_on_image(image_draw, mask)
+
+    image_draw = Image.fromarray(add_watermark_np(np.array(image_draw)))
+    if global_state is not None:
+        global_state['images']['image_show'] = image_draw
+    return image_draw
+
+
+def preprocess_mask_info(global_state, image):
+    """Function to handle mask information.
+
+    1. last_mask is None: Do not need to change mask, return mask
+    2. last_mask is not None:
+        2.1 global_state is remove_mask:
+        2.2 global_state is add_mask:
+    """
+
+    if isinstance(image, dict):
+        last_mask = get_valid_mask(image['mask'])
+    else:
+        last_mask = None
+    mask = global_state['mask']
+    # mask in global state is a placeholder with all 1.
+    if (mask == 1).all():
+        mask = last_mask
+
+    # last_mask = global_state['last_mask']
+    editing_mode = global_state['editing_state']
+
+    if last_mask is None:
+        return global_state
+
+    if editing_mode == 'remove_mask':
+        updated_mask = np.clip(mask - last_mask, 0, 1)
+        print(f'Last editing_state is {editing_mode}, do remove.')
+    elif editing_mode == 'add_mask':
+        updated_mask = np.clip(mask + last_mask, 0, 1)
+        print(f'Last editing_state is {editing_mode}, do add.')
+    else:
+        updated_mask = mask
+        print(f'Last editing_state is {editing_mode}, ' 'do nothing to mask.')
+
+    global_state['mask'] = updated_mask
+    # global_state['last_mask'] = None  # clear buffer
+    return global_state
+
+
+valid_checkpoints_dict = {
+    f.split('/')[-1]: osp.join(cache_dir, f)
+    for f in os.listdir(cache_dir)
+    if (f.endswith('pth') and osp.exists(osp.join(cache_dir, f)))
+}
+print(f'File under cache_dir ({cache_dir}):')
+print(os.listdir(cache_dir))
+print('Valid checkpoint file:')
+print(valid_checkpoints_dict)
+
+if __name__ == '__main__':
+    init_pkl = 'stylegan2_lions_512_pytorch_mmagic.pth'
+    device = 'cuda'
+
+    with gr.Blocks() as app:
+
+        # renderer = Renderer()
+        global_state = get_default_gr_state(device, init_pkl)
+
+        # init image
+        global_state = init_images(global_state)
+
+        with gr.Row():
+
+            with gr.Row():
+
+                # Left --> tools
+                with gr.Column(scale=3):
+
+                    # Pickle
+                    with gr.Row():
+
+                        with gr.Column(scale=1, min_width=10):
+                            gr.Markdown(value='Pickle', show_label=False)
+
+                        with gr.Column(scale=4, min_width=10):
+                            form_pretrained_dropdown = gr.Dropdown(
+                                choices=list(valid_checkpoints_dict.keys()),
+                                label='Pretrained Model',
+                                value=init_pkl,
+                            )
+
+                    # Latent
+                    with gr.Row():
+                        with gr.Column(scale=1, min_width=10):
+                            gr.Markdown(value='Latent', show_label=False)
+
+                        with gr.Column(scale=4, min_width=10):
+                            form_seed_number = gr.Number(
+                                value=global_state.value['params']['seed'],
+                                interactive=True,
+                                label='Seed',
+                            )
+                            form_lr_number = gr.Number(
+                                value=global_state.value['params']['lr'],
+                                interactive=True,
+                                label='Step Size')
+
+                            with gr.Row():
+                                with gr.Column(scale=2, min_width=10):
+                                    form_reset_image = gr.Button('Reset Image')
+                                with gr.Column(scale=3, min_width=10):
+                                    form_latent_space = gr.Radio(
+                                        ['w', 'w+'],
+                                        value=global_state.value['params']
+                                        ['latent_space'],
+                                        interactive=True,
+                                        label='Latent space to optimize',
+                                        show_label=False,
+                                    )
+
+                    # Drag
+                    with gr.Row():
+                        with gr.Column(scale=1, min_width=10):
+                            gr.Markdown(value='Drag', show_label=False)
+                        with gr.Column(scale=4, min_width=10):
+                            with gr.Row():
+                                with gr.Column(scale=1, min_width=10):
+                                    enable_add_points = gr.Button('Add Points')
+                                with gr.Column(scale=1, min_width=10):
+                                    undo_points = gr.Button('Reset Points')
+                            with gr.Row():
+                                with gr.Column(scale=1, min_width=10):
+                                    form_start_btn = gr.Button('Start')
+                                with gr.Column(scale=1, min_width=10):
+                                    form_stop_btn = gr.Button('Stop')
+
+                            form_steps_number = gr.Number(
+                                value=0, label='Steps', interactive=False)
+
+                    # Mask
+                    with gr.Row():
+                        with gr.Column(scale=1, min_width=10):
+                            gr.Markdown(value='Mask', show_label=False)
+                        with gr.Column(scale=4, min_width=10):
+                            enable_add_mask = gr.Button('Edit Flexible Area')
+                            with gr.Row():
+                                with gr.Column(scale=1, min_width=10):
+                                    form_reset_mask_btn = gr.Button(
+                                        'Reset mask')
+                                with gr.Column(scale=1, min_width=10):
+                                    show_mask = gr.Checkbox(
+                                        label='Show Mask',
+                                        value=global_state.value['show_mask'],
+                                        show_label=False)
+
+                            with gr.Row():
+                                form_lambda_number = gr.Number(
+                                    value=global_state.value['params']
+                                    ['motion_lambda'],
+                                    interactive=True,
+                                    label='Lambda',
+                                )
+
+                    form_draw_interval_number = gr.Number(
+                        value=global_state.value['draw_interval'],
+                        label='Draw Interval (steps)',
+                        interactive=True,
+                        visible=False)
+
+                # Right --> Image
+                with gr.Column(scale=8):
+                    form_image = ImageMask(
+                        value=global_state.value['images']['image_show'],
+                        brush_radius=20).style(
+                            width=768,
+                            height=768)  # NOTE: hard image size code here.
+        gr.Markdown("""
+            ## Quick Start
+
+            1. Select desired `Pretrained Model` and adjust `Seed` to generate
+            an initial image.
+            2. Click on image to add control points.
+            3. Click `Start` and enjoy it!
+
+            ## Advance Usage
+
+            1. Change `Step Size` to adjust learning rate in drag optimization.
+            2. Select `w` or `w+` to change latent space to optimize:
+            * Optimize on `w` space may cause greater influence to the image.
+            * Optimize on `w+` space may work slower than `w`, but usually
+            achieve better results.
+            * Note that changing the latent space will reset the image, points
+            and mask (this has the same effect as `Reset Image` button).
+            3. Click `Edit Flexible Area` to create a mask and constrain the
+            unmasked region to remain unchanged.
+            """)
+        gr.HTML("""
+            <style>
+                .container {
+                    position: absolute;
+                    height: 50px;
+                    text-align: center;
+                    line-height: 50px;
+                    width: 100%;
+                }
+            </style>
+            <div class="container">
+            Gradio demo supported by
+            <img src="https://avatars.githubusercontent.com/u/10245193?s=
+            200&v=4"\
+                height="20" width="20" style="display:inline;">
+            <a href="https://github.com/open-mmlab/mmagic">OpenMMLab MMagic</a>
+            </div>
+            """)
+
+        # Network & latents tab listeners
+        def on_change_pretrained_dropdown(pretrained_value, global_state):
+            """Function to handle model change.
+
+            1. Set pretrained value to global_state
+            2. Re-init images and clear all states
+            """
+
+            global_state['pretrained_weight'] = pretrained_value
+            init_images(global_state)
+            clear_state(global_state)
+
+            return global_state, global_state['images']['image_show']
+
+        form_pretrained_dropdown.change(
+            on_change_pretrained_dropdown,
+            inputs=[form_pretrained_dropdown, global_state],
+            outputs=[global_state, form_image],
+        )
+
+        def on_click_reset_image(global_state):
+            """Reset image to the original one and clear all states
+            1. Re-init images
+            2. Clear all states
+            """
+
+            init_images(global_state)
+            clear_state(global_state)
+
+            return global_state, global_state['images']['image_show']
+
+        form_reset_image.click(
+            on_click_reset_image,
+            inputs=[global_state],
+            outputs=[global_state, form_image],
+        )
+
+        # Update parameters
+        def on_change_update_image_seed(seed, global_state):
+            """Function to handle generation seed change.
+
+            1. Set seed to global_state
+            2. Re-init images and clear all states
+            """
+
+            global_state['params']['seed'] = int(seed)
+            init_images(global_state)
+            clear_state(global_state)
+
+            return global_state, global_state['images']['image_show']
+
+        form_seed_number.change(
+            on_change_update_image_seed,
+            inputs=[form_seed_number, global_state],
+            outputs=[global_state, form_image],
+        )
+
+        def on_click_latent_space(latent_space, global_state):
+            """Function to reset latent space to optimize.
+
+            NOTE: this function we reset the image and all controls
+            1. Set latent-space to global_state
+            2. Re-init images and clear all state
+            """
+
+            global_state['params']['latent_space'] = latent_space
+            init_images(global_state)
+            clear_state(global_state)
+
+            return global_state, global_state['images']['image_show']
+
+        form_latent_space.change(
+            on_click_latent_space,
+            inputs=[form_latent_space, global_state],
+            outputs=[global_state, form_image])
+
+        # ==== Params
+        form_lambda_number.change(
+            partial(on_change_single_global_state,
+                    ['params', 'motion_lambda']),
+            inputs=[form_lambda_number, global_state],
+            outputs=[global_state],
+        )
+
+        def on_change_lr(lr, global_state):
+            if lr == 0:
+                print('lr is 0, do nothing.')
+                return global_state
+            else:
+                global_state['params']['lr'] = lr
+                renderer = global_state['renderer']
+                renderer.update_lr(lr)
+                print('New optimizer: ')
+                print(renderer.w_optim)
+            return global_state
+
+        form_lr_number.change(
+            on_change_lr,
+            inputs=[form_lr_number, global_state],
+            outputs=[global_state],
+        )
+
+        def on_click_start(global_state, image):
+            p_in_pixels = []
+            t_in_pixels = []
+            valid_points = []
+
+            # handle of start drag in mask editing mode
+            global_state = preprocess_mask_info(global_state, image)
+
+            # Prepare the points for the inference
+            if len(global_state['points']) == 0:
+                # yield on_click_start_wo_points(global_state, image)
+                image_raw = global_state['images']['image_raw']
+                update_image_draw(
+                    image_raw,
+                    global_state['points'],
+                    global_state['mask'],
+                    global_state['show_mask'],
+                    global_state,
+                )
+
+                yield (
+                    global_state,
+                    0,
+                    global_state['images']['image_show'],
+                    # gr.File.update(visible=False),
+                    gr.Button.update(interactive=True),
+                    gr.Button.update(interactive=True),
+                    gr.Button.update(interactive=True),
+                    gr.Button.update(interactive=True),
+                    gr.Button.update(interactive=True),
+                    # latent space
+                    gr.Radio.update(interactive=True),
+                    gr.Button.update(interactive=True),
+                    # NOTE: disable stop button
+                    gr.Button.update(interactive=False),
+
+                    # update other comps
+                    gr.Dropdown.update(interactive=True),
+                    gr.Number.update(interactive=True),
+                    gr.Number.update(interactive=True),
+                    gr.Button.update(interactive=True),
+                    gr.Button.update(interactive=True),
+                    gr.Checkbox.update(interactive=True),
+                    # gr.Number.update(interactive=True),
+                    gr.Number.update(interactive=True),
+                )
+            else:
+
+                # Transform the points into torch tensors
+                for key_point, point in global_state['points'].items():
+                    try:
+                        p_start = point.get('start_temp', point['start'])
+                        p_end = point['target']
+
+                        if p_start is None or p_end is None:
+                            continue
+
+                    except KeyError:
+                        continue
+
+                    p_in_pixels.append(p_start)
+                    t_in_pixels.append(p_end)
+                    valid_points.append(key_point)
+
+                mask = torch.tensor(global_state['mask']).float()
+                drag_mask = 1 - mask
+
+                renderer: Renderer = global_state['renderer']
+                global_state['temporal_params']['stop'] = False
+                global_state['editing_state'] = 'running'
+
+                # reverse points order
+                p_to_opt = reverse_point_pairs(p_in_pixels)
+                t_to_opt = reverse_point_pairs(t_in_pixels)
+                print(p_to_opt)
+                print('Running with:')
+                print(f'    Source: {p_in_pixels}')
+                print(f'    Target: {t_in_pixels}')
+                step_idx = 0
+                while True:
+                    if global_state['temporal_params']['stop']:
+                        break
+
+                    # do drage here!
+                    renderer._render_drag_impl(
+                        global_state['generator_params'],
+                        p_to_opt,  # point
+                        t_to_opt,  # target
+                        drag_mask,  # mask,
+                        global_state['params']['motion_lambda'],  # lambda_mask
+                        reg=0,
+                        feature_idx=5,  # NOTE: do not support change for now
+                        r1=global_state['params']['r1_in_pixels'],  # r1
+                        r2=global_state['params']['r2_in_pixels'],  # r2
+                        # random_seed     = 0,
+                        # noise_mode      = 'const',
+                        trunc_psi=global_state['params']['trunc_psi'],
+                        # force_fp32      = False,
+                        # layer_name      = None,
+                        # sel_channels    = 3,
+                        # base_channel    = 0,
+                        # img_scale_db    = 0,
+                        # img_normalize   = False,
+                        # untransform     = False,
+                        is_drag=True,
+                        to_pil=True)
+
+                    if step_idx % global_state['draw_interval'] == 0:
+                        print('Current Source:')
+                        for key_point, p_i, t_i in zip(valid_points, p_to_opt,
+                                                       t_to_opt):
+                            global_state['points'][key_point]['start_temp'] = [
+                                p_i[1],
+                                p_i[0],
+                            ]
+                            global_state['points'][key_point]['target'] = [
+                                t_i[1],
+                                t_i[0],
+                            ]
+                            start_temp = global_state['points'][key_point][
+                                'start_temp']
+                            print(f'    {start_temp}')
+
+                        image_result = global_state['generator_params'][
+                            'image']
+                        image_draw = update_image_draw(  # noqa
+                            image_result,
+                            global_state['points'],
+                            global_state['mask'],
+                            global_state['show_mask'],
+                            global_state,
+                        )
+                        global_state['images']['image_raw'] = image_result
+
+                    yield (
+                        global_state,
+                        step_idx,
+                        global_state['images']['image_show'],
+                        # gr.File.update(visible=False),
+                        gr.Button.update(interactive=False),
+                        gr.Button.update(interactive=False),
+                        gr.Button.update(interactive=False),
+                        gr.Button.update(interactive=False),
+                        gr.Button.update(interactive=False),
+                        # latent space
+                        gr.Radio.update(interactive=False),
+                        gr.Button.update(interactive=False),
+                        # enable stop button in loop
+                        gr.Button.update(interactive=True),
+
+                        # update other comps
+                        gr.Dropdown.update(interactive=False),
+                        gr.Number.update(interactive=False),
+                        gr.Number.update(interactive=False),
+                        gr.Button.update(interactive=False),
+                        gr.Button.update(interactive=False),
+                        gr.Checkbox.update(interactive=False),
+                        # gr.Number.update(interactive=False),
+                        gr.Number.update(interactive=False),
+                    )
+
+                    # increate step
+                    step_idx += 1
+
+                image_result = global_state['generator_params']['image']
+                global_state['images']['image_raw'] = image_result
+                # image_draw = update_image_draw(image_result,
+                #                                global_state['points'],
+                #                                global_state['mask'],
+                #                                global_state['show_mask'],
+                #                                global_state)
+
+                # fp = NamedTemporaryFile(suffix=".png", delete=False)
+                # image_result.save(fp, "PNG")
+
+                global_state['editing_state'] = 'add_points'
+
+                yield (
+                    global_state,
+                    0,  # reset step to 0 after stop.
+                    global_state['images']['image_show'],
+                    # gr.File.update(visible=True, value=fp.name),
+                    gr.Button.update(interactive=True),
+                    gr.Button.update(interactive=True),
+                    gr.Button.update(interactive=True),
+                    gr.Button.update(interactive=True),
+                    gr.Button.update(interactive=True),
+                    # latent space
+                    gr.Radio.update(interactive=True),
+                    gr.Button.update(interactive=True),
+                    # NOTE: disable stop button with loop finish
+                    gr.Button.update(interactive=False),
+
+                    # update other comps
+                    gr.Dropdown.update(interactive=True),
+                    gr.Number.update(interactive=True),
+                    gr.Number.update(interactive=True),
+                    gr.Checkbox.update(interactive=True),
+                    gr.Number.update(interactive=True),
+                )
+
+        form_start_btn.click(
+            on_click_start,
+            inputs=[global_state, form_image],
+            outputs=[
+                global_state,
+                form_steps_number,
+                form_image,
+                # form_download_result_file,
+                # >>> buttons
+                form_reset_image,
+                enable_add_points,
+                enable_add_mask,
+                undo_points,
+                form_reset_mask_btn,
+                form_latent_space,
+                form_start_btn,
+                form_stop_btn,
+                # <<< buttonm
+                # >>> inputs comps
+                form_pretrained_dropdown,
+                form_seed_number,
+                form_lr_number,
+                show_mask,
+                form_lambda_number,
+            ],
+        )
+
+        def on_click_stop(global_state):
+            """Function to handle stop button is clicked.
+
+            1. send a stop signal by
+            set global_state["temporal_params"]["stop"] as True
+            2. Disable Stop button
+            """
+            global_state['temporal_params']['stop'] = True
+
+            return global_state, gr.Button.update(interactive=False)
+
+        form_stop_btn.click(
+            on_click_stop,
+            inputs=[global_state],
+            outputs=[global_state, form_stop_btn])
+
+        form_draw_interval_number.change(
+            partial(
+                on_change_single_global_state,
+                'draw_interval',
+                map_transform=lambda x: int(x),
+            ),
+            inputs=[form_draw_interval_number, global_state],
+            outputs=[global_state],
+        )
+
+        def on_click_remove_point(global_state):
+            choice = global_state['curr_point']
+            del global_state['points'][choice]
+
+            choices = list(global_state['points'].keys())
+
+            if len(choices) > 0:
+                global_state['curr_point'] = choices[0]
+
+            return (
+                gr.Dropdown.update(choices=choices, value=choices[0]),
+                global_state,
+            )
+
+        # Mask
+        def on_click_reset_mask(global_state):
+            global_state['mask'] = np.ones(
+                (
+                    global_state['images']['image_raw'].size[1],
+                    global_state['images']['image_raw'].size[0],
+                ),
+                dtype=np.uint8,
+            )
+            image_draw = update_image_draw(global_state['images']['image_raw'],
+                                           global_state['points'],
+                                           global_state['mask'],
+                                           global_state['show_mask'],
+                                           global_state)
+            return global_state, image_draw
+
+        form_reset_mask_btn.click(
+            on_click_reset_mask,
+            inputs=[global_state],
+            outputs=[global_state, form_image],
+        )
+
+        # Image
+        def on_click_enable_draw(global_state, image):
+            """Function to start add mask mode.
+
+            1. Preprocess mask info from last state
+            2. Change editing state to add_mask
+            3. Set curr image with points and mask
+            """
+            global_state = preprocess_mask_info(global_state, image)
+            global_state['editing_state'] = 'add_mask'
+            image_raw = global_state['images']['image_raw']
+            image_draw = update_image_draw(image_raw, global_state['points'],
+                                           global_state['mask'], True,
+                                           global_state)
+            return (global_state,
+                    gr.Image.update(value=image_draw, interactive=True))
+
+        def on_click_remove_draw(global_state, image):
+            """Function to start remove mask mode.
+
+            1. Preprocess mask info from last state
+            2. Change editing state to remove_mask
+            3. Set curr image with points and mask
+            """
+            global_state = preprocess_mask_info(global_state, image)
+            global_state['edinting_state'] = 'remove_mask'
+            image_raw = global_state['images']['image_raw']
+            image_draw = update_image_draw(image_raw, global_state['points'],
+                                           global_state['mask'], True,
+                                           global_state)
+            return (global_state,
+                    gr.Image.update(value=image_draw, interactive=True))
+
+        enable_add_mask.click(
+            on_click_enable_draw,
+            inputs=[global_state, form_image],
+            outputs=[
+                global_state,
+                form_image,
+            ])
+
+        def on_click_add_point(global_state, image: dict):
+            """Function switch from add mask mode to add points mode.
+
+            1. Updaste mask buffer if need
+            2. Change global_state['editing_state'] to 'add_points'
+            3. Set current image with mask
+            """
+
+            global_state = preprocess_mask_info(global_state, image)
+            global_state['editing_state'] = 'add_points'
+            mask = global_state['mask']
+            image_raw = global_state['images']['image_raw']
+            image_draw = update_image_draw(image_raw, global_state['points'],
+                                           mask, global_state['show_mask'],
+                                           global_state)
+
+            return (global_state,
+                    gr.Image.update(value=image_draw, interactive=False))
+
+        enable_add_points.click(
+            on_click_add_point,
+            inputs=[global_state, form_image],
+            outputs=[global_state, form_image])
+
+        def on_click_image(global_state, evt: gr.SelectData):
+            """This function only support click for point selection."""
+            xy = evt.index
+            if global_state['editing_state'] != 'add_points':
+                print(f'In {global_state["editing_state"]} state. '
+                      'Do not add points.')
+
+                return global_state, global_state['images']['image_show']
+
+            points = global_state['points']
+
+            point_idx = get_latest_points_pair(points)
+            if point_idx is None:
+                points[0] = {'start': xy, 'target': None}
+                print(f'Click Image - Start - {xy}')
+            elif points[point_idx].get('target', None) is None:
+                points[point_idx]['target'] = xy
+                print(f'Click Image - Target - {xy}')
+            else:
+                points[point_idx + 1] = {'start': xy, 'target': None}
+                print(f'Click Image - Start - {xy}')
+
+            image_raw = global_state['images']['image_raw']
+            image_draw = update_image_draw(
+                image_raw,
+                global_state['points'],
+                global_state['mask'],
+                global_state['show_mask'],
+                global_state,
+            )
+
+            return global_state, image_draw
+
+        form_image.select(
+            on_click_image,
+            inputs=[global_state],
+            outputs=[global_state, form_image],
+        )
+
+        def on_click_clear_points(global_state):
+            """Function to handle clear all control points
+            1. clear global_state['points'] (clear_state)
+            2. re-init network
+            2. re-draw image
+            """
+            clear_state(global_state, target='point')
+
+            renderer: Renderer = global_state['renderer']
+            renderer.feat_refs = None
+
+            image_raw = global_state['images']['image_raw']
+            image_draw = update_image_draw(image_raw, {}, global_state['mask'],
+                                           global_state['show_mask'],
+                                           global_state)
+            return global_state, image_draw
+
+        undo_points.click(
+            on_click_clear_points,
+            inputs=[global_state],
+            outputs=[global_state, form_image])
+
+        def on_click_show_mask(global_state, show_mask):
+            """Function to control whether show mask on image."""
+            global_state['show_mask'] = show_mask
+
+            image_raw = global_state['images']['image_raw']
+            image_draw = update_image_draw(
+                image_raw,
+                global_state['points'],
+                global_state['mask'],
+                global_state['show_mask'],
+                global_state,
+            )
+            return global_state, image_draw
+
+        show_mask.change(
+            on_click_show_mask,
+            inputs=[global_state, show_mask],
+            outputs=[global_state, form_image],
+        )
+
+    gr.close_all()
+    app.queue(concurrency_count=3, max_size=20)
+    app.launch(
+        share=args.share,
+        server_name='0.0.0.0' if args.listen else '127.0.0.1')
diff --git a/demo/utils/gradio_utils.py b/demo/utils/gradio_utils.py
new file mode 100644
index 0000000000..0d4c02078c
--- /dev/null
+++ b/demo/utils/gradio_utils.py
@@ -0,0 +1,168 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+# yapf: disable
+import gradio as gr
+import numpy as np
+from PIL import Image, ImageDraw
+
+
+class EasyDict(dict):
+    """Convenience class that behaves like a dict but allows access with the
+    attribute syntax."""
+
+    def __getattr__(self, name: str):
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(name)
+
+    def __setattr__(self, name: str, value) -> None:
+        self[name] = value
+
+    def __delattr__(self, name: str) -> None:
+        del self[name]
+
+
+class ImageMask(gr.components.Image):
+    """
+    Sets: source="canvas", tool="sketch"
+    """
+
+    is_template = True
+
+    def __init__(self, **kwargs):
+        super().__init__(
+            source='upload', tool='sketch', interactive=False, **kwargs)
+
+    def preprocess(self, x):
+        if x is None:
+            return x
+        if self.tool == 'sketch' and self.source in ['upload', 'webcam'
+                                                     ] and type(x) != dict:
+            decode_image = gr.processing_utils.decode_base64_to_image(x)
+            width, height = decode_image.size
+            mask = np.ones((height, width, 4), dtype=np.uint8)
+            mask[..., -1] = 255
+            mask = self.postprocess(mask)
+            x = {'image': x, 'mask': mask}
+        return super().preprocess(x)
+
+
+def get_valid_mask(mask: np.ndarray):
+    """Convert mask from gr.Image(0 to 255, RGBA) to binary mask."""
+    if mask.ndim == 3:
+        mask_pil = Image.fromarray(mask).convert('L')
+        mask = np.array(mask_pil)
+    if mask.max() == 255:
+        mask = mask / 255
+    return mask
+
+
+def draw_points_on_image(image,
+                         points,
+                         curr_point=None,
+                         highlight_all=True,
+                         radius_scale=0.01):
+    overlay_rgba = Image.new('RGBA', image.size, 0)
+    overlay_draw = ImageDraw.Draw(overlay_rgba)
+    for point_key, point in points.items():
+        if ((curr_point is not None and curr_point == point_key)
+                or highlight_all):
+            p_color = (255, 0, 0)
+            t_color = (0, 0, 255)
+
+        else:
+            p_color = (255, 0, 0, 35)
+            t_color = (0, 0, 255, 35)
+
+        rad_draw = int(image.size[0] * radius_scale)
+
+        p_start = point.get('start_temp', point['start'])
+        p_target = point['target']
+
+        if p_start is not None and p_target is not None:
+            p_draw = int(p_start[0]), int(p_start[1])
+            t_draw = int(p_target[0]), int(p_target[1])
+
+            overlay_draw.line(
+                (p_draw[0], p_draw[1], t_draw[0], t_draw[1]),
+                fill=(255, 255, 0),
+                width=2,
+            )
+
+        if p_start is not None:
+            p_draw = int(p_start[0]), int(p_start[1])
+            overlay_draw.ellipse(
+                (
+                    p_draw[0] - rad_draw,
+                    p_draw[1] - rad_draw,
+                    p_draw[0] + rad_draw,
+                    p_draw[1] + rad_draw,
+                ),
+                fill=p_color,
+            )
+
+            if curr_point is not None and curr_point == point_key:
+                overlay_draw.text(p_draw, 'p', align='center', fill=(0, 0, 0))
+
+        if p_target is not None:
+            t_draw = int(p_target[0]), int(p_target[1])
+            overlay_draw.ellipse(
+                (
+                    t_draw[0] - rad_draw,
+                    t_draw[1] - rad_draw,
+                    t_draw[0] + rad_draw,
+                    t_draw[1] + rad_draw,
+                ),
+                fill=t_color,
+            )
+
+            if curr_point is not None and curr_point == point_key:
+                overlay_draw.text(t_draw, 't', align='center', fill=(0, 0, 0))
+
+    return Image.alpha_composite(image.convert('RGBA'),
+                                 overlay_rgba).convert('RGB')
+
+
+def draw_mask_on_image(image, mask):
+    im_mask = np.uint8(mask * 255)
+    im_mask_rgba = np.concatenate(
+        (
+            np.tile(im_mask[..., None], [1, 1, 3]),
+            45 * np.ones(
+                (im_mask.shape[0], im_mask.shape[1], 1), dtype=np.uint8),
+        ),
+        axis=-1,
+    )
+    im_mask_rgba = Image.fromarray(im_mask_rgba).convert('RGBA')
+
+    return Image.alpha_composite(image.convert('RGBA'),
+                                 im_mask_rgba).convert('RGB')
+
+
+def on_change_single_global_state(keys,
+                                  value,
+                                  global_state,
+                                  map_transform=None):
+    if map_transform is not None:
+        value = map_transform(value)
+
+    curr_state = global_state
+    if isinstance(keys, str):
+        last_key = keys
+
+    else:
+        for k in keys[:-1]:
+            curr_state = curr_state[k]
+
+        last_key = keys[-1]
+
+    curr_state[last_key] = value
+    return global_state
+
+
+def get_latest_points_pair(points_dict):
+    if not points_dict:
+        return None
+    point_idx = list(points_dict.keys())
+    latest_point_idx = max(point_idx)
+    return latest_point_idx
diff --git a/demo/utils/renderer.py b/demo/utils/renderer.py
new file mode 100644
index 0000000000..9feaa56adb
--- /dev/null
+++ b/demo/utils/renderer.py
@@ -0,0 +1,453 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import sys
+import traceback
+
+import matplotlib.cm
+import matplotlib.font_manager
+import numpy as np
+import torch
+import torch.fft
+import torch.nn.functional as F
+from PIL import Image, ImageDraw, ImageFont
+
+from demo.utils.gradio_utils import EasyDict
+from mmagic.apis import MMagicInferencer
+
+
+class CapturedException(Exception):
+
+    def __init__(self, msg=None):
+        if msg is None:
+            _type, value, _traceback = sys.exc_info()
+            assert value is not None
+            if isinstance(value, CapturedException):
+                msg = str(value)
+            else:
+                msg = traceback.format_exc()
+        assert isinstance(msg, str)
+        super().__init__(msg)
+
+
+'''
+----------------------------------------------------------------------------
+'''
+
+
+class CaptureSuccess(Exception):
+
+    def __init__(self, out):
+        super().__init__()
+        self.out = out
+
+
+'''
+----------------------------------------------------------------------------
+'''
+
+
+def add_watermark_np(input_image_array, watermark_text='AI Generated'):
+    image = Image.fromarray(np.uint8(input_image_array)).convert('RGBA')
+
+    # Initialize text image
+    txt = Image.new('RGBA', image.size, (255, 255, 255, 0))
+    ttf_pth = matplotlib.font_manager.findSystemFonts(
+        fontpaths=None, fontext='ttf')[-1]
+
+    font = ImageFont.truetype(ttf_pth, round(25 / 512 * image.size[0]))
+    d = ImageDraw.Draw(txt)
+
+    text_width, text_height = font.getsize(watermark_text)
+    text_position = (image.size[0] - text_width - 10,
+                     image.size[1] - text_height - 10)
+    text_color = (
+        255, 255, 255, 128
+    )  # white color with the alpha channel set to semi-transparent
+
+    # Draw the text onto the text canvas
+    d.text(text_position, watermark_text, font=font, fill=text_color)
+
+    # Combine the image with the watermark
+    watermarked = Image.alpha_composite(image, txt)
+    watermarked_array = np.array(watermarked)
+    return watermarked_array
+
+
+class Renderer:
+
+    def __init__(self, disable_timing=False):
+        self._device = torch.device('cuda' if torch.cuda.is_available(
+        ) else 'mps' if torch.backends.mps.is_available() else 'cpu')
+        self._dtype = torch.float32 \
+            if self._device.type == 'mps' else torch.float64
+        self._pkl_data = dict()  # {pkl: dict | CapturedException, ...}
+        self._networks = dict()  # {cache_key: torch.nn.Module, ...}
+        self._pinned_bufs = dict()  # {(shape, dtype): torch.Tensor, ...}
+        self._cmaps = dict()  # {name: torch.Tensor, ...}
+        self._is_timing = False
+        if not disable_timing:
+            self._start_event = torch.cuda.Event(enable_timing=True)
+            self._end_event = torch.cuda.Event(enable_timing=True)
+        self._disable_timing = disable_timing
+        self._net_layers = dict()  # {cache_key: [dnnlib.EasyDict, ...], ...}
+
+    def render(self, **args):
+        if self._disable_timing:
+            self._is_timing = False
+        else:
+            self._start_event.record(torch.cuda.current_stream(self._device))
+            self._is_timing = True
+        res = EasyDict()
+        try:
+            init_net = False
+            if not hasattr(self, 'G'):
+                init_net = True
+            if hasattr(self, 'pkl'):
+                if self.pkl != args['pkl']:
+                    init_net = True
+            if hasattr(self, 'w_load'):
+                if self.w_load is not args['w_load']:
+                    init_net = True
+            if hasattr(self, 'w0_seed'):
+                if self.w0_seed != args['w0_seed']:
+                    init_net = True
+            if hasattr(self, 'w_plus'):
+                if self.w_plus != args['w_plus']:
+                    init_net = True
+            if args['reset_w']:
+                init_net = True
+            res.init_net = init_net
+            if init_net:
+                self.init_network(res, **args)
+            self._render_drag_impl(res, **args)
+        except CapturedException:
+            res.error = CapturedException()
+        if not self._disable_timing:
+            self._end_event.record(torch.cuda.current_stream(self._device))
+        if 'image' in res:
+            res.image = self.to_cpu(res.image).detach().numpy()
+            res.image = add_watermark_np(res.image, 'AI Generated')
+        if 'stats' in res:
+            res.stats = self.to_cpu(res.stats).detach().numpy()
+        if 'error' in res:
+            res.error = str(res.error)
+        # if 'stop' in res and res.stop:
+
+        if self._is_timing and not self._disable_timing:
+            self._end_event.synchronize()
+            res.render_time = self._start_event.elapsed_time(
+                self._end_event) * 1e-3
+            self._is_timing = False
+        return res
+
+    def get_network(self, pkl, key, **tweak_kwargs):
+        data = self._pkl_data.get(pkl, None)
+        if data is None:
+            print(f'Loading "{pkl}"... ', end='', flush=True)
+            try:
+                print('Done.')
+            except CapturedException:
+                data = CapturedException()
+                print('Failed!')
+            self._pkl_data[pkl] = data
+            self._ignore_timing()
+        if isinstance(data, CapturedException):
+            raise data
+
+        orig_net = data[key]
+        cache_key = (orig_net, self._device,
+                     tuple(sorted(tweak_kwargs.items())))
+        net = self._networks.get(cache_key, None)
+        if net is None:
+            try:
+                if 'stylegan2' in pkl:
+                    from training.networks_stylegan2 import Generator
+                else:
+                    raise NameError('Cannot infer model type from pkl name!')
+
+                print(data[key].init_args)
+                print(data[key].init_kwargs)
+                if 'stylegan_human' in pkl:
+                    net = Generator(
+                        *data[key].init_args,
+                        **data[key].init_kwargs,
+                        square=False,
+                        padding=True)
+                else:
+                    net = Generator(*data[key].init_args,
+                                    **data[key].init_kwargs)
+                net.load_state_dict(data[key].state_dict())
+                net.to(self._device)
+            except CapturedException:
+                net = CapturedException()
+            self._networks[cache_key] = net
+            self._ignore_timing()
+        if isinstance(net, CapturedException):
+            raise net
+        return net
+
+    def _get_pinned_buf(self, ref):
+        key = (tuple(ref.shape), ref.dtype)
+        buf = self._pinned_bufs.get(key, None)
+        if buf is None:
+            buf = torch.empty(ref.shape, dtype=ref.dtype).pin_memory()
+            self._pinned_bufs[key] = buf
+        return buf
+
+    def to_device(self, buf):
+        return self._get_pinned_buf(buf).copy_(buf).to(self._device)
+
+    def to_cpu(self, buf):
+        return self._get_pinned_buf(buf).copy_(buf).clone()
+
+    def _ignore_timing(self):
+        self._is_timing = False
+
+    def _apply_cmap(self, x, name='viridis'):
+        cmap = self._cmaps.get(name, None)
+        if cmap is None:
+            cmap = matplotlib.cm.get_cmap(name)
+            cmap = cmap(np.linspace(0, 1, num=1024), bytes=True)[:, :3]
+            cmap = self.to_device(torch.from_numpy(cmap))
+            self._cmaps[name] = cmap
+        hi = cmap.shape[0] - 1
+        x = (x * hi + 0.5).clamp(0, hi).to(torch.int64)
+        x = torch.nn.functional.embedding(x, cmap)
+        return x
+
+    def init_network(self,
+                     res,
+                     ckpt_pth=None,
+                     w0_seed=0,
+                     w_load=None,
+                     w_plus=True,
+                     noise_mode='const',
+                     trunc_psi=0.7,
+                     trunc_cutoff=None,
+                     input_transform=None,
+                     lr=0.001,
+                     **kwargs):
+        # Dig up network details.
+        if '256' in ckpt_pth:
+            editor = MMagicInferencer(
+                'draggan',
+                model_setting=1,
+                model_ckpt=ckpt_pth,
+            )
+        elif '512' in ckpt_pth:
+            editor = MMagicInferencer(
+                'draggan',
+                model_setting=0,
+                model_ckpt=ckpt_pth,
+            )
+        elif '1024' in ckpt_pth:
+            editor = MMagicInferencer(
+                'draggan',
+                model_setting=2,
+                model_ckpt=ckpt_pth,
+            )
+        else:
+            raise NotImplementedError
+        self.editor = editor
+        # Generate random latents.
+        self.w0_seed = w0_seed
+        self.w_load = w_load
+
+        if self.w_load is None:
+            # Generate random latents.
+            z = torch.from_numpy(np.random.RandomState(w0_seed).randn(
+                1, 512)).to(torch.float32).requires_grad_(True)
+
+            # Run mapping network.
+            sample_kwargs = {
+                'truncation': trunc_psi,
+                'return_noise': True,
+                'return_features': True
+            }
+            extra_parameters = {
+                'sample_kwargs': sample_kwargs,
+                'num_batches': 1,
+                'noise': z,
+                'sample_model': 'ema',
+                'infer_with_grad': True
+            }
+            results = self.editor.infer(extra_parameters=extra_parameters)
+
+            w = results[0]['latent'].unsqueeze(0)  # [1, 16, n_dim]
+        else:
+            w = self.w_load.clone().to(self._device)
+
+        self.w0 = w.detach().clone()
+        self.w_plus = w_plus
+        if w_plus:
+            self.w = w.detach()
+        else:
+            self.w = w[:, 0, :].detach()
+        self.w.requires_grad = True
+        self.w_optim = torch.optim.Adam([self.w], lr=lr)
+
+        self.feat_refs = None
+        self.points0_pt = None
+
+    def update_lr(self, lr):
+
+        del self.w_optim
+        self.w_optim = torch.optim.Adam([self.w], lr=lr)
+        print(f'Rebuild optimizer with lr: {lr}')
+        print('    Remain feat_refs and points0_pt')
+
+    def _render_drag_impl(self,
+                          res,
+                          points=[],
+                          targets=[],
+                          mask=None,
+                          lambda_mask=10,
+                          reg=0,
+                          feature_idx=5,
+                          r1=3,
+                          r2=12,
+                          random_seed=0,
+                          noise_mode='const',
+                          trunc_psi=0.7,
+                          force_fp32=False,
+                          layer_name=None,
+                          sel_channels=3,
+                          base_channel=0,
+                          img_scale_db=0,
+                          img_normalize=False,
+                          untransform=False,
+                          is_drag=False,
+                          reset=False,
+                          to_pil=False,
+                          **kwargs):
+        ws = self.w
+        if ws.dim() == 2:
+            ws = ws.unsqueeze(1).repeat(1, 6, 1)
+        ws = torch.cat([ws[:, :6, :], self.w0[:, 6:, :]],
+                       dim=1)  # [1, 16, n_dim]
+        if hasattr(self, 'points'):
+            if len(points) != len(self.points):
+                reset = True
+        if reset:
+            self.feat_refs = None
+            self.points0_pt = None
+        self.points = points
+
+        # Run synthesis network.
+        sample_kwargs = {
+            'truncation': 1,
+            'return_noise': True,
+            'return_features': True,
+            'input_is_latent': True
+        }
+        extra_parameters = {
+            'sample_kwargs': sample_kwargs,
+            'num_batches': 1,
+            'noise': ws[0][:1],
+            'sample_model': 'ema',
+            'infer_with_grad': True
+        }
+
+        results = self.editor.infer(extra_parameters=extra_parameters)
+        img = results[0]['fake_img'].unsqueeze(
+            0) / 127.5 - 1.  # requires_grad==True
+        feat_5 = results[0]['feats'].unsqueeze(0)  # requires_grad==True
+        h, w = img.shape[-1], img.shape[-2]
+
+        if is_drag:
+            X = torch.linspace(0, h, h)
+            Y = torch.linspace(0, w, w)
+            xx, yy = torch.meshgrid(X, Y)
+            feat_resize = F.interpolate(feat_5, [h, w], mode='bilinear')
+            if self.feat_refs is None:
+                self.feat0_resize = F.interpolate(
+                    feat_5.detach(), [h, w], mode='bilinear')
+                self.feat_refs = []
+                for point in points:
+                    py, px = round(point[0]), round(point[1])
+                    self.feat_refs.append(self.feat0_resize[:, :, py, px])
+                self.points0_pt = torch.Tensor(points).unsqueeze(0).to(
+                    self._device)  # 1, N, 2
+
+            # Point tracking with feature matching
+            with torch.no_grad():
+                for j, point in enumerate(points):
+                    r = round(r2 / 512 * h)
+                    up = max(point[0] - r, 0)
+                    down = min(point[0] + r + 1, h)
+                    left = max(point[1] - r, 0)
+                    right = min(point[1] + r + 1, w)
+                    feat_patch = feat_resize[:, :, up:down, left:right]
+                    L2 = torch.linalg.norm(
+                        feat_patch - self.feat_refs[j].reshape(1, -1, 1, 1),
+                        dim=1)
+                    _, idx = torch.min(L2.view(1, -1), -1)
+                    width = right - left
+                    point = [
+                        idx.item() // width + up,
+                        idx.item() % width + left
+                    ]
+                    points[j] = point
+
+            res.points = [[point[0], point[1]] for point in points]
+
+            # Motion supervision
+            loss_motion = 0
+            res.stop = True
+            for j, point in enumerate(points):
+                direction = torch.Tensor(
+                    [targets[j][1] - point[1], targets[j][0] - point[0]])
+                if torch.linalg.norm(direction) > max(2 / 512 * h, 2):
+                    res.stop = False
+                if torch.linalg.norm(direction) > 1:
+                    distance = ((xx.to(self._device) - point[0])**2 +
+                                (yy.to(self._device) - point[1])**2)**0.5
+                    relis, reljs = torch.where(distance < round(r1 / 512 * h))
+                    direction = direction / (
+                        torch.linalg.norm(direction) + 1e-7)
+                    gridh = (relis - direction[1]) / (h - 1) * 2 - 1
+                    gridw = (reljs - direction[0]) / (w - 1) * 2 - 1
+                    grid = torch.stack([gridw, gridh],
+                                       dim=-1).unsqueeze(0).unsqueeze(0)
+                    target = F.grid_sample(
+                        feat_resize.float(), grid,
+                        align_corners=True).squeeze(2)
+                    loss_motion += F.l1_loss(feat_resize[:, :, relis, reljs],
+                                             target.detach())
+
+            loss = loss_motion
+            if mask is not None:
+                if mask.min() == 0 and mask.max() == 1:
+                    mask_usq = mask.to(self._device).unsqueeze(0).unsqueeze(0)
+                    loss_fix = F.l1_loss(feat_resize * mask_usq,
+                                         self.feat0_resize * mask_usq)
+                    loss += lambda_mask * loss_fix
+
+            loss += reg * F.l1_loss(ws, self.w0)  # latent code regularization
+            if not res.stop:
+                self.w_optim.zero_grad()
+                print(loss)
+                loss.backward()
+                self.w_optim.step()
+
+        # Scale and convert to uint8.
+        img = img[0]
+        if img_normalize:
+            img = img / img.norm(
+                float('inf'), dim=[1, 2], keepdim=True).clip(1e-8, 1e8)
+        img = img * (10**(img_scale_db / 20))
+        img = (img * 127.5 + 128).clamp(0,
+                                        255).to(torch.uint8).permute(1, 2, 0)
+        if to_pil:
+            from PIL import Image
+            img = img.cpu().numpy()
+            img = Image.fromarray(img)
+        res.image = img
+        res.w = ws.detach().cpu().numpy()
diff --git a/mmagic/apis/inferencers/__init__.py b/mmagic/apis/inferencers/__init__.py
index 4827156031..b175a0c297 100644
--- a/mmagic/apis/inferencers/__init__.py
+++ b/mmagic/apis/inferencers/__init__.py
@@ -53,7 +53,7 @@ def __init__(self,
         elif self.task in ['colorization', 'Colorization']:
             self.inferencer = ColorizationInferencer(
                 config, ckpt, device, extra_parameters, seed=seed)
-        elif self.task in ['unconditional', 'Unconditional GANs']:
+        elif self.task in ['unconditional', 'Unconditional GANs', 'DragGAN']:
             self.inferencer = UnconditionalInferencer(
                 config, ckpt, device, extra_parameters, seed=seed)
         elif self.task in ['matting', 'Matting']:
diff --git a/mmagic/apis/inferencers/base_mmagic_inferencer.py b/mmagic/apis/inferencers/base_mmagic_inferencer.py
index ed05faf93f..99bfca2680 100644
--- a/mmagic/apis/inferencers/base_mmagic_inferencer.py
+++ b/mmagic/apis/inferencers/base_mmagic_inferencer.py
@@ -121,7 +121,6 @@ def _dispatch_kwargs(self, **kwargs) -> Tuple[Dict, Dict, Dict, Dict]:
 
         return results
 
-    @torch.no_grad()
     def __call__(self, **kwargs) -> Union[Dict, List[Dict]]:
         """Call the inferencer.
 
@@ -131,6 +130,25 @@ def __call__(self, **kwargs) -> Union[Dict, List[Dict]]:
         Returns:
             Union[Dict, List[Dict]]: Results of inference pipeline.
         """
+        if 'extra_parameters' in kwargs.keys():
+            if 'infer_with_grad' in kwargs['extra_parameters'].keys():
+                if kwargs['extra_parameters']['infer_with_grad']:
+                    results = self.base_call(**kwargs)
+        else:
+            with torch.no_grad():
+                results = self.base_call(**kwargs)
+        return results
+
+    def base_call(self, **kwargs) -> Union[Dict, List[Dict]]:
+        """Call the inferencer.
+
+        Args:
+            kwargs: Keyword arguments for the inferencer.
+
+        Returns:
+            Union[Dict, List[Dict]]: Results of inference pipeline.
+        """
+
         self._update_extra_parameters(**kwargs)
 
         params = self._dispatch_kwargs(**kwargs)
diff --git a/mmagic/apis/inferencers/unconditional_inferencer.py b/mmagic/apis/inferencers/unconditional_inferencer.py
index a6aca652cd..11960c0a92 100644
--- a/mmagic/apis/inferencers/unconditional_inferencer.py
+++ b/mmagic/apis/inferencers/unconditional_inferencer.py
@@ -20,7 +20,8 @@ class UnconditionalInferencer(BaseMMagicInferencer):
         visualize=['result_out_dir'],
         postprocess=[])
 
-    extra_parameters = dict(num_batches=4, sample_model='orig')
+    extra_parameters = dict(
+        num_batches=4, sample_model='orig', sample_kwargs=None, noise=None)
 
     def preprocess(self) -> Dict:
         """Process the inputs into a model-feedable format.
@@ -30,8 +31,14 @@ def preprocess(self) -> Dict:
         """
         num_batches = self.extra_parameters['num_batches']
         sample_model = self.extra_parameters['sample_model']
+        noise = self.extra_parameters['noise']
+        sample_kwargs = self.extra_parameters['sample_kwargs']
 
-        results = dict(num_batches=num_batches, sample_model=sample_model)
+        results = dict(
+            num_batches=num_batches,
+            sample_model=sample_model,
+            sample_kwargs=sample_kwargs,
+            noise=noise)
 
         return results
 
@@ -83,4 +90,8 @@ def _pred2dict(self, data_sample: DataSample) -> Dict:
         result = {}
         result['fake_img'] = data_sample.fake_img.data.cpu()
         result['noise'] = data_sample.noise.data.cpu()
+        if hasattr(data_sample, 'latent'):
+            result['latent'] = data_sample.latent
+        if hasattr(data_sample, 'feats'):
+            result['feats'] = data_sample.feats
         return result
diff --git a/mmagic/apis/mmagic_inferencer.py b/mmagic/apis/mmagic_inferencer.py
index 3855b20305..634d5e5718 100644
--- a/mmagic/apis/mmagic_inferencer.py
+++ b/mmagic/apis/mmagic_inferencer.py
@@ -35,7 +35,7 @@ class MMagicInferencer:
         >>> editor = MMagicInferencer(model_name='pix2pix')
         >>> editor.infer(img='./test.jpg', result_out_dir='./pix2pix_res.jpg')
 
-        >>> # see demo/mmagic_inference_tutorial.ipynb for more examples
+        >>> # see demo/mmediting_inference_tutorial.ipynb for more examples
     """
     # unsupported now
     # singan, liif
@@ -115,7 +115,10 @@ class MMagicInferencer:
         'eg3d',
 
         # diffusers inferencer
-        'controlnet_animation'
+        'controlnet_animation',
+
+        # draggan
+        'draggan'
     ]
 
     inference_supported_models_cfg = {}
diff --git a/mmagic/models/base_models/base_gan.py b/mmagic/models/base_models/base_gan.py
index d7d471ab30..725737773c 100644
--- a/mmagic/models/base_models/base_gan.py
+++ b/mmagic/models/base_models/base_gan.py
@@ -349,7 +349,25 @@ def forward(self,
                 generator = self.generator_ema
             else:
                 generator = self.generator
-            outputs = generator(noise, return_noise=False, **sample_kwargs)
+            if sample_kwargs:
+                if 'return_noise' in sample_kwargs.keys():
+                    outputs = generator(noise, **sample_kwargs)
+                else:
+                    outputs = generator(
+                        noise, return_noise=False,
+                        **sample_kwargs)  # yapf: disable
+            else:
+                sample_kwargs = {}
+                outputs = generator(
+                    noise, return_noise=False,
+                    **sample_kwargs)  # no need to be False all time
+            if isinstance(outputs, dict):
+                if 'latent' in outputs.keys():
+                    latent = outputs['latent']
+                if 'feats' in outputs.keys():
+                    feats = outputs['feats']
+                outputs = outputs['fake_img']
+
             outputs = self.data_preprocessor.destruct(outputs, data_samples)
 
             gen_sample = DataSample()
@@ -359,6 +377,10 @@ def forward(self,
                 gen_sample.gt_img = inputs['img']
             gen_sample.fake_img = outputs
             gen_sample.noise = noise
+            if 'latent' in locals():
+                gen_sample.latent = latent
+            if 'feats' in locals():
+                gen_sample.feats = feats
             gen_sample.sample_kwargs = deepcopy(sample_kwargs)
             gen_sample.sample_model = sample_model
             batch_sample_list = gen_sample.split(allow_nonseq_value=True, )
diff --git a/mmagic/models/data_preprocessors/data_preprocessor.py b/mmagic/models/data_preprocessors/data_preprocessor.py
index 02ec643a8b..bdea5f044b 100644
--- a/mmagic/models/data_preprocessors/data_preprocessor.py
+++ b/mmagic/models/data_preprocessors/data_preprocessor.py
@@ -142,11 +142,19 @@ def cast_data(self, data: CastData) -> CastData:
     def _parse_channel_index(inputs) -> int:
         """Parse channel index of inputs."""
         channel_index_mapping = {2: 1, 3: 0, 4: 1, 5: 2}
-        assert inputs.ndim in channel_index_mapping, (
-            'Only support (H*W, C), (C, H, W), (N, C, H, W) or '
-            '(N, t, C, H, W) inputs. But received '
-            f'\'({inputs.shape})\'.')
-        channel_index = channel_index_mapping[inputs.ndim]
+        if isinstance(inputs, dict):
+            ndim = inputs['fake_img'].ndim
+            assert ndim in channel_index_mapping, (
+                'Only support (H*W, C), (C, H, W), (N, C, H, W) or '
+                '(N, t, C, H, W) inputs. But received '
+                f'\'({inputs.shape})\'.')
+            channel_index = channel_index_mapping[ndim]
+        else:
+            assert inputs.ndim in channel_index_mapping, (
+                'Only support (H*W, C), (C, H, W), (N, C, H, W) or '
+                '(N, t, C, H, W) inputs. But received '
+                f'\'({inputs.shape})\'.')
+            channel_index = channel_index_mapping[inputs.ndim]
 
         return channel_index
 
@@ -155,7 +163,10 @@ def _parse_channel_order(self,
                              inputs: Tensor,
                              data_sample: Optional[DataSample] = None) -> str:
         channel_index = self._parse_channel_index(inputs)
-        num_color_channels = inputs.shape[channel_index]
+        if isinstance(inputs, dict):
+            num_color_channels = inputs['fake_img'].shape[channel_index]
+        else:
+            num_color_channels = inputs.shape[channel_index]
 
         # data sample is None, attempt to infer from input tensor
         if data_sample is None:
diff --git a/mmagic/models/editors/stylegan1/stylegan1_modules.py b/mmagic/models/editors/stylegan1/stylegan1_modules.py
index 272b734d93..04c1483a45 100644
--- a/mmagic/models/editors/stylegan1/stylegan1_modules.py
+++ b/mmagic/models/editors/stylegan1/stylegan1_modules.py
@@ -102,11 +102,14 @@ class NoiseInjection(BaseModule):
     Args:
         noise_weight_init (float, optional): Initialization weight for noise
             injection. Defaults to ``0.``.
+        fixed_noise (bool, optional): Whether to inject a fixed noise. Defaults
+        to ``False``.
     """
 
-    def __init__(self, noise_weight_init=0.):
+    def __init__(self, noise_weight_init=0., fixed_noise=False):
         super().__init__()
         self.weight = nn.Parameter(torch.zeros(1).fill_(noise_weight_init))
+        self.fixed_noise = fixed_noise
 
     def forward(self, image, noise=None, return_noise=False):
         """Forward Function.
@@ -124,6 +127,10 @@ def forward(self, image, noise=None, return_noise=False):
         if noise is None:
             batch, _, height, width = image.shape
             noise = image.new_empty(batch, 1, height, width).normal_()
+            if self.fixed_noise:
+                torch.manual_seed(1024)
+                noise = torch.randn(batch, 1, height, width).cuda()
+
         noise = noise.to(image.dtype)
         if return_noise:
             return image + self.weight.to(image.dtype) * noise, noise
diff --git a/mmagic/models/editors/stylegan2/stylegan2_generator.py b/mmagic/models/editors/stylegan2/stylegan2_generator.py
index 9456b33595..961856ea11 100644
--- a/mmagic/models/editors/stylegan2/stylegan2_generator.py
+++ b/mmagic/models/editors/stylegan2/stylegan2_generator.py
@@ -120,7 +120,8 @@ def __init__(self,
                  num_fp16_scales=0,
                  fp16_enabled=False,
                  bgr2rgb=False,
-                 pretrained=None):
+                 pretrained=None,
+                 fixed_noise=False):
         super().__init__()
         self.out_size = out_size
         self.style_channels = style_channels
@@ -189,7 +190,8 @@ def __init__(self,
             kernel_size=3,
             style_channels=style_channels,
             blur_kernel=blur_kernel,
-            fp16_enabled=fp16_enabled)
+            fp16_enabled=fp16_enabled,
+            fixed_noise=fixed_noise)
         self.to_rgb1 = ModulatedToRGB(
             self.channels[4],
             style_channels,
@@ -348,7 +350,10 @@ def forward(self,
                 injected_noise=None,
                 add_noise=True,
                 randomize_noise=True,
-                update_ws=False):
+                update_ws=False,
+                return_features=False,
+                feat_idx=5,
+                return_latent_only=False):
         """Forward function.
 
         This function has been integrated with the truncation trick. Please
@@ -393,6 +398,7 @@ def forward(self,
             torch.Tensor | dict: Generated image tensor or dictionary \
                 containing more data.
         """
+        # device = styles.device
         input_dim = self.style_channels if input_is_latent else self.noise_size
         # receive noise and conduct sanity check.
         if isinstance(styles, torch.Tensor):
@@ -503,6 +509,10 @@ def forward(self,
 
             latent = torch.cat([latent, latent2], 1)
 
+        if return_latent_only:
+            return latent
+
+        feats = []
         with autocast(enabled=self.fp16_enabled):
             # 4x4 stage
             out = self.constant_input(latent)
@@ -513,6 +523,7 @@ def forward(self,
                 latent[:, 0],
                 noise=injected_noise[0],
                 add_noise=add_noise)
+            feats.append(out)
             skip = self.to_rgb1(out, latent[:, 1])
 
             _index = 1
@@ -528,17 +539,26 @@ def forward(self,
                     latent[:, _index + 1],
                     noise=noise2,
                     add_noise=add_noise)
+                feats.append(out)
                 skip = to_rgb(out, latent[:, _index + 2], skip)
                 _index += 2
 
         # make sure the output image is torch.float32 to avoid RunTime Error
         # in other modules
         img = skip.to(torch.float32)
-
         if self.bgr2rgb:
             img = torch.flip(img, dims=[1])
 
         if return_latents or return_noise:
+            if return_features:
+                output_dict = dict(
+                    fake_img=img,
+                    latent=latent,
+                    inject_index=inject_index,
+                    noise_batch=noise_batch,
+                    feats=feats[feat_idx])
+                return output_dict
+
             output_dict = dict(
                 fake_img=img,
                 latent=latent,
diff --git a/mmagic/models/editors/stylegan2/stylegan2_modules.py b/mmagic/models/editors/stylegan2/stylegan2_modules.py
index 76dd83be24..66c9b5398f 100644
--- a/mmagic/models/editors/stylegan2/stylegan2_modules.py
+++ b/mmagic/models/editors/stylegan2/stylegan2_modules.py
@@ -226,7 +226,6 @@ def __init__(
 
     def forward(self, x, style, input_gain=None):
         n, c, h, w = x.shape
-
         weight = self.weight
         # Pre-normalize inputs to avoid FP16 overflow.
         # if x.dtype == torch.float16 and self.demodulate:
@@ -238,7 +237,6 @@ def forward(self, x, style, input_gain=None):
             )  # max_Ikk
             style = style / style.norm(
                 float('inf'), dim=1, keepdim=True)  # max_I
-
         with autocast(enabled=self.fp16_enabled):
             # process style code
             style = self.style_modulation(style).view(n, 1, c, 1,
@@ -329,7 +327,8 @@ def __init__(self,
                  style_mod_cfg=dict(bias_init=1.),
                  style_bias=0.,
                  fp16_enabled=False,
-                 conv_clamp=256):
+                 conv_clamp=256,
+                 fixed_noise=False):
         super().__init__()
 
         # add support for fp16
@@ -348,7 +347,7 @@ def __init__(self,
             style_bias=style_bias,
             fp16_enabled=fp16_enabled)
 
-        self.noise_injector = NoiseInjection()
+        self.noise_injector = NoiseInjection(fixed_noise=fixed_noise)
         self.activate = _FusedBiasLeakyReLU(out_channels)
 
     def forward(self,
@@ -384,7 +383,6 @@ def forward(self,
 
             # TODO: FP16 in activate layers
             out = self.activate(out)
-
             if self.fp16_enabled:
                 out = torch.clamp(
                     out, min=-self.conv_clamp, max=self.conv_clamp)
diff --git a/mmagic/structures/data_sample.py b/mmagic/structures/data_sample.py
index dcd4e20ce2..5d925ec95a 100644
--- a/mmagic/structures/data_sample.py
+++ b/mmagic/structures/data_sample.py
@@ -191,7 +191,10 @@ class DataSample(BaseDataElement):
         'pred_img': 'pred_img',
         'ori_trimap': 'ori_trimap',
         # For text to images
-        'prompt': 'prompt'
+        'prompt': 'prompt',
+        # For StyleGAN
+        'latent': 'latent',
+        'feats': 'feats'
     }
 
     def set_predefined_data(self, data: dict) -> None:
diff --git a/model-index.yml b/model-index.yml
index e04c95e9f2..76e0476520 100644
--- a/model-index.yml
+++ b/model-index.yml
@@ -14,6 +14,7 @@ Import:
 - configs/dic/metafile.yml
 - configs/dim/metafile.yml
 - configs/disco_diffusion/metafile.yml
+- configs/draggan/metafile.yml
 - configs/dreambooth/metafile.yml
 - configs/edsr/metafile.yml
 - configs/edvr/metafile.yml


Test patch:
diff --git a/tests/test_apis/test_inferencers/test_unconditional_inferencer.py b/tests/test_apis/test_inferencers/test_unconditional_inferencer.py
index 4b18d0416f..5eb7d5f27d 100644
--- a/tests/test_apis/test_inferencers/test_unconditional_inferencer.py
+++ b/tests/test_apis/test_inferencers/test_unconditional_inferencer.py
@@ -1,6 +1,8 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import os.path as osp
 
+import torch
+
 from mmagic.apis.inferencers.unconditional_inferencer import \
     UnconditionalInferencer
 from mmagic.utils import register_all_modules
@@ -27,3 +29,24 @@ def test_unconditional_inferencer():
     inference_result = inferencer_instance(result_out_dir=result_out_dir)
     result_img = inference_result[1]
     assert result_img.detach().numpy().shape == (1, 3, 256, 256)
+
+    # test additional input for extra parameters
+    noise = torch.randn(1, 512)
+    sample_kwargs = sample_kwargs = {
+        'return_noise': True,
+        'input_is_latent': False
+    }
+
+    inferencer_instance = \
+        UnconditionalInferencer(cfg,
+                                None,
+                                extra_parameters={
+                                    'num_batches': 1,
+                                    'sample_model': 'orig',
+                                    'noise': noise,
+                                    'sample_kwargs': sample_kwargs})
+    # test no result_out_dir
+    inferencer_instance()
+    inference_result = inferencer_instance(result_out_dir=result_out_dir)
+    result_img = inference_result[1]
+    assert result_img.detach().numpy().shape == (1, 3, 256, 256)
diff --git a/tests/test_models/test_base_models/test_base_gan.py b/tests/test_models/test_base_models/test_base_gan.py
index 82c398fde6..7523106c2b 100644
--- a/tests/test_models/test_base_models/test_base_gan.py
+++ b/tests/test_models/test_base_models/test_base_gan.py
@@ -278,6 +278,14 @@ def test_forward(self):
         for idx in range(4):
             self.assertEqual(outputs[idx].fake_img.shape, (3, 8, 8))
 
+        # test additional sample kwargs
+        sample_kwargs = dict(return_noise=True)
+        inputs = dict(noise=torch.randn(4, 10), sample_kwargs=sample_kwargs)
+        outputs = gan(inputs)
+        self.assertEqual(len(outputs), 4)
+        for idx in range(4):
+            self.assertEqual(outputs[idx].fake_img.shape, (3, 8, 8))
+
         # test when data samples is not None
         inputs = dict(num_batches=3, sample_model='ema/orig')
         data_samples = [DataSample(id=1), DataSample(id=2), DataSample(id=3)]
diff --git a/tests/test_models/test_data_preprocessors/test_data_preprocessor.py b/tests/test_models/test_data_preprocessors/test_data_preprocessor.py
index 9d06700a53..ce7e2523f6 100644
--- a/tests/test_models/test_data_preprocessors/test_data_preprocessor.py
+++ b/tests/test_models/test_data_preprocessors/test_data_preprocessor.py
@@ -54,6 +54,10 @@ def test_parse_channel_index(self):
             data_preprocessor._parse_channel_index(
                 torch.rand(1, 2, 10, 3, 5, 5, 5))
 
+        # test dict input
+        inputs = dict(fake_img=torch.rand(2, 3, 5, 5))
+        self.assertEqual(data_preprocessor._parse_channel_index(inputs), 1)
+
     def test_parse_channel_order(self):
         data_preprocessor = DataPreprocessor()
         parse_fn = data_preprocessor._parse_channel_order
@@ -62,6 +66,10 @@ def test_parse_channel_order(self):
         self.assertEqual(parse_fn('img', torch.rand(1, 3, 3, 4)), 'BGR')
         self.assertEqual(parse_fn('img', torch.rand(1, 1, 3, 4)), 'single')
 
+        # test dict input
+        inputs = dict(fake_img=torch.rand(1, 3, 3, 4))
+        self.assertEqual(parse_fn('img', inputs), 'BGR')
+
         # test data sample is not None
         data_sample = DataSample()
 
diff --git a/tests/test_models/test_editors/test_stylegan2/test_stylegan2_generator.py b/tests/test_models/test_editors/test_stylegan2/test_stylegan2_generator.py
index b9b22f0ca3..89b16c41b7 100644
--- a/tests/test_models/test_editors/test_stylegan2/test_stylegan2_generator.py
+++ b/tests/test_models/test_editors/test_stylegan2/test_stylegan2_generator.py
@@ -132,6 +132,29 @@ def test_stylegan2_g_cpu(self):
         # should not be update in test
         assert (mean_latent_test == mean_latent).all()
 
+        # test return features
+        g.eval()
+        res = g(
+            None,
+            num_batches=1,
+            injected_noise=None,
+            randomize_noise=False,
+            return_noise=True,
+            return_features=True,
+            feat_idx=1)
+        assert res['feats'].shape == (1, 512, 8, 8)
+        assert res['latent'].shape == (1, 10, 16)
+
+        # test return latent only
+        g.eval()
+        res = g(
+            None,
+            num_batches=1,
+            injected_noise=None,
+            randomize_noise=False,
+            return_latent_only=True)
+        assert res.shape == (1, 10, 16)
+
         g.train()
         res = g(
             None,
diff --git a/tests/test_structures/test_data_sample.py b/tests/test_structures/test_data_sample.py
index 95be641124..7fe3f3e036 100644
--- a/tests/test_structures/test_data_sample.py
+++ b/tests/test_structures/test_data_sample.py
@@ -62,6 +62,8 @@ def test_set_prefined_data(self):
         img_path, gt_path, merged_path = 'aaa', 'bbb', 'ccc'
         gt_channel_order, gt_color_type = 'rgb', 'color'
         prompt = 'prompt'
+        latent = torch.randn(1, 16, 512)
+        feats = torch.randn(64, 256, 256)
 
         data = dict(
             gt=gt,
@@ -76,7 +78,9 @@ def test_set_prefined_data(self):
             merged_path=merged_path,
             gt_channel_order=gt_channel_order,
             gt_color_type=gt_color_type,
-            prompt=prompt)
+            prompt=prompt,
+            latent=latent,
+            feats=feats)
         data_sample = DataSample()
         data_sample.set_predefined_data(data)
 
@@ -94,6 +98,8 @@ def test_set_prefined_data(self):
         self._check_in_and_same(data_sample, 'gt_color_type', gt_color_type,
                                 True)
         self._check_in_and_same(data_sample, 'prompt', prompt, False)
+        self._check_in_and_same(data_sample, 'latent', latent)
+        self._check_in_and_same(data_sample, 'feats', feats)
         # check gt label
         data_sample.gt_label.data = gt_label