Upload 9 files

LICENSE.txt

@@ -0,0 +1,97 @@
+Copyright (c) 2021, NVIDIA Corporation & affiliates. All rights reserved.
+
+
+NVIDIA Source Code License for StyleGAN3
+
+
+=======================================================================
+
+1. Definitions
+
+"Licensor" means any person or entity that distributes its Work.
+
+"Software" means the original work of authorship made available under
+this License.
+
+"Work" means the Software and any additions to or derivative works of
+the Software that are made available under this License.
+
+The terms "reproduce," "reproduction," "derivative works," and
+"distribution" have the meaning as provided under U.S. copyright law;
+provided, however, that for the purposes of this License, derivative
+works shall not include works that remain separable from, or merely
+link (or bind by name) to the interfaces of, the Work.
+
+Works, including the Software, are "made available" under this License
+by including in or with the Work either (a) a copyright notice
+referencing the applicability of this License to the Work, or (b) a
+copy of this License.
+
+2. License Grants
+
+    2.1 Copyright Grant. Subject to the terms and conditions of this
+    License, each Licensor grants to you a perpetual, worldwide,
+    non-exclusive, royalty-free, copyright license to reproduce,
+    prepare derivative works of, publicly display, publicly perform,
+    sublicense and distribute its Work and any resulting derivative
+    works in any form.
+
+3. Limitations
+
+    3.1 Redistribution. You may reproduce or distribute the Work only
+    if (a) you do so under this License, (b) you include a complete
+    copy of this License with your distribution, and (c) you retain
+    without modification any copyright, patent, trademark, or
+    attribution notices that are present in the Work.
+
+    3.2 Derivative Works. You may specify that additional or different
+    terms apply to the use, reproduction, and distribution of your
+    derivative works of the Work ("Your Terms") only if (a) Your Terms
+    provide that the use limitation in Section 3.3 applies to your
+    derivative works, and (b) you identify the specific derivative
+    works that are subject to Your Terms. Notwithstanding Your Terms,
+    this License (including the redistribution requirements in Section
+    3.1) will continue to apply to the Work itself.
+
+    3.3 Use Limitation. The Work and any derivative works thereof only
+    may be used or intended for use non-commercially. Notwithstanding
+    the foregoing, NVIDIA and its affiliates may use the Work and any
+    derivative works commercially. As used herein, "non-commercially"
+    means for research or evaluation purposes only.
+
+    3.4 Patent Claims. If you bring or threaten to bring a patent claim
+    against any Licensor (including any claim, cross-claim or
+    counterclaim in a lawsuit) to enforce any patents that you allege
+    are infringed by any Work, then your rights under this License from
+    such Licensor (including the grant in Section 2.1) will terminate
+    immediately.
+
+    3.5 Trademarks. This License does not grant any rights to use any
+    Licensor’s or its affiliates’ names, logos, or trademarks, except
+    as necessary to reproduce the notices described in this License.
+
+    3.6 Termination. If you violate any term of this License, then your
+    rights under this License (including the grant in Section 2.1) will
+    terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
+NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
+THIS LICENSE.
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
+THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
+SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
+INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
+OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
+(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
+LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
+COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
+THE POSSIBILITY OF SUCH DAMAGES.
+
+=======================================================================

environment.yml

@@ -0,0 +1,24 @@
+name: stylegan3
+channels:
+  - pytorch
+  - nvidia
+dependencies:
+  - python >= 3.8
+  - pip
+  - numpy>=1.20
+  - click>=8.0
+  - pillow=8.3.1
+  - scipy=1.7.1
+  - pytorch=1.9.1
+  - cudatoolkit=11.1
+  - requests=2.26.0
+  - tqdm=4.62.2
+  - ninja=1.10.2
+  - matplotlib=3.4.2
+  - imageio=2.9.0
+  - pip:
+    - imgui==1.3.0
+    - glfw==2.2.0
+    - pyopengl==3.1.5
+    - imageio-ffmpeg==0.4.3
+    - pyspng

gen_images.py

@@ -0,0 +1,149 @@
+# 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.
+
+"""Generate images using pretrained network pickle."""
+
+import os
+import re
+from typing import List, Optional, Tuple, Union
+
+import click
+import dnnlib
+import numpy as np
+import PIL.Image
+import torch
+
+import legacy
+
+#----------------------------------------------------------------------------
+
+def parse_range(s: Union[str, List]) -> List[int]:
+    '''Parse a comma separated list of numbers or ranges and return a list of ints.
+
+    Example: '1,2,5-10' returns [1, 2, 5, 6, 7]
+    '''
+    if isinstance(s, list): return s
+    ranges = []
+    range_re = re.compile(r'^(\d+)-(\d+)$')
+    for p in s.split(','):
+        m = range_re.match(p)
+        if m:
+            ranges.extend(range(int(m.group(1)), int(m.group(2))+1))
+        else:
+            ranges.append(int(p))
+    return ranges
+
+#----------------------------------------------------------------------------
+
+def parse_vec2(s: Union[str, Tuple[float, float]]) -> Tuple[float, float]:
+    '''Parse a floating point 2-vector of syntax 'a,b'.
+
+    Example:
+        '0,1' returns (0,1)
+    '''
+    if isinstance(s, tuple): return s
+    parts = s.split(',')
+    if len(parts) == 2:
+        return (float(parts[0]), float(parts[1]))
+    raise ValueError(f'cannot parse 2-vector {s}')
+
+#----------------------------------------------------------------------------
+
+def make_transform(translate: Tuple[float,float], angle: float):
+    m = np.eye(3)
+    s = np.sin(angle/360.0*np.pi*2)
+    c = np.cos(angle/360.0*np.pi*2)
+    m[0][0] = c
+    m[0][1] = s
+    m[0][2] = translate[0]
+    m[1][0] = -s
+    m[1][1] = c
+    m[1][2] = translate[1]
+    return m
+
+#----------------------------------------------------------------------------
+
+@click.command()
+@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
+@click.option('--seeds', type=parse_range, help='List of random seeds (e.g., \'0,1,4-6\')', required=True)
+@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=1, show_default=True)
+@click.option('--class', 'class_idx', type=int, help='Class label (unconditional if not specified)')
+@click.option('--noise-mode', help='Noise mode', type=click.Choice(['const', 'random', 'none']), default='const', show_default=True)
+@click.option('--translate', help='Translate XY-coordinate (e.g. \'0.3,1\')', type=parse_vec2, default='0,0', show_default=True, metavar='VEC2')
+@click.option('--rotate', help='Rotation angle in degrees', type=float, default=0, show_default=True, metavar='ANGLE')
+@click.option('--outdir', help='Where to save the output images', type=str, required=True, metavar='DIR')
+def generate_images(
+    network_pkl: str,
+    seeds: List[int],
+    truncation_psi: float,
+    noise_mode: str,
+    outdir: str,
+    translate: Tuple[float,float],
+    rotate: float,
+    class_idx: Optional[int]
+):
+    """Generate images using pretrained network pickle.
+
+    Examples:
+
+    \b
+    # Generate an image using pre-trained AFHQv2 model ("Ours" in Figure 1, left).
+    python gen_images.py --outdir=out --trunc=1 --seeds=2 \\
+        --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl
+
+    \b
+    # Generate uncurated images with truncation using the MetFaces-U dataset
+    python gen_images.py --outdir=out --trunc=0.7 --seeds=600-605 \\
+        --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-t-metfacesu-1024x1024.pkl
+    """
+
+    print('Loading networks from "%s"...' % network_pkl)
+    device = torch.device('cuda')
+    with dnnlib.util.open_url(network_pkl) as f:
+        G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore
+        # import pickle
+        # G = legacy.load_network_pkl(f)
+        # output = open('checkpoints/stylegan2-car-config-f-pt.pkl', 'wb')
+        # pickle.dump(G, output)
+
+    os.makedirs(outdir, exist_ok=True)
+
+    # Labels.
+    label = torch.zeros([1, G.c_dim], device=device)
+    if G.c_dim != 0:
+        if class_idx is None:
+            raise click.ClickException('Must specify class label with --class when using a conditional network')
+        label[:, class_idx] = 1
+    else:
+        if class_idx is not None:
+            print ('warn: --class=lbl ignored when running on an unconditional network')
+
+    # Generate images.
+    for seed_idx, seed in enumerate(seeds):
+        print('Generating image for seed %d (%d/%d) ...' % (seed, seed_idx, len(seeds)))
+        z = torch.from_numpy(np.random.RandomState(seed).randn(1, G.z_dim)).to(device)
+
+        # Construct an inverse rotation/translation matrix and pass to the generator.  The
+        # generator expects this matrix as an inverse to avoid potentially failing numerical
+        # operations in the network.
+        if hasattr(G.synthesis, 'input'):
+            m = make_transform(translate, rotate)
+            m = np.linalg.inv(m)
+            G.synthesis.input.transform.copy_(torch.from_numpy(m))
+
+        img = G(z, label, truncation_psi=truncation_psi, noise_mode=noise_mode)
+        img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+        PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(f'{outdir}/seed{seed:04d}.png')
+
+
+#----------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    generate_images() # pylint: disable=no-value-for-parameter
+
+#----------------------------------------------------------------------------

legacy.py

@@ -0,0 +1,323 @@
+# 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.
+
+"""Converting legacy network pickle into the new format."""
+
+import click
+import pickle
+import re
+import copy
+import numpy as np
+import torch
+import dnnlib
+from torch_utils import misc
+
+#----------------------------------------------------------------------------
+
+def load_network_pkl(f, force_fp16=False):
+    data = _LegacyUnpickler(f).load()
+
+    # Legacy TensorFlow pickle => convert.
+    if isinstance(data, tuple) and len(data) == 3 and all(isinstance(net, _TFNetworkStub) for net in data):
+        tf_G, tf_D, tf_Gs = data
+        G = convert_tf_generator(tf_G)
+        D = convert_tf_discriminator(tf_D)
+        G_ema = convert_tf_generator(tf_Gs)
+        data = dict(G=G, D=D, G_ema=G_ema)
+
+    # Add missing fields.
+    if 'training_set_kwargs' not in data:
+        data['training_set_kwargs'] = None
+    if 'augment_pipe' not in data:
+        data['augment_pipe'] = None
+
+    # Validate contents.
+    assert isinstance(data['G'], torch.nn.Module)
+    assert isinstance(data['D'], torch.nn.Module)
+    assert isinstance(data['G_ema'], torch.nn.Module)
+    assert isinstance(data['training_set_kwargs'], (dict, type(None)))
+    assert isinstance(data['augment_pipe'], (torch.nn.Module, type(None)))
+
+    # Force FP16.
+    if force_fp16:
+        for key in ['G', 'D', 'G_ema']:
+            old = data[key]
+            kwargs = copy.deepcopy(old.init_kwargs)
+            fp16_kwargs = kwargs.get('synthesis_kwargs', kwargs)
+            fp16_kwargs.num_fp16_res = 4
+            fp16_kwargs.conv_clamp = 256
+            if kwargs != old.init_kwargs:
+                new = type(old)(**kwargs).eval().requires_grad_(False)
+                misc.copy_params_and_buffers(old, new, require_all=True)
+                data[key] = new
+    return data
+
+#----------------------------------------------------------------------------
+
+class _TFNetworkStub(dnnlib.EasyDict):
+    pass
+
+class _LegacyUnpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if module == 'dnnlib.tflib.network' and name == 'Network':
+            return _TFNetworkStub
+        return super().find_class(module, name)
+
+#----------------------------------------------------------------------------
+
+def _collect_tf_params(tf_net):
+    # pylint: disable=protected-access
+    tf_params = dict()
+    def recurse(prefix, tf_net):
+        for name, value in tf_net.variables:
+            tf_params[prefix + name] = value
+        for name, comp in tf_net.components.items():
+            recurse(prefix + name + '/', comp)
+    recurse('', tf_net)
+    return tf_params
+
+#----------------------------------------------------------------------------
+
+def _populate_module_params(module, *patterns):
+    for name, tensor in misc.named_params_and_buffers(module):
+        found = False
+        value = None
+        for pattern, value_fn in zip(patterns[0::2], patterns[1::2]):
+            match = re.fullmatch(pattern, name)
+            if match:
+                found = True
+                if value_fn is not None:
+                    value = value_fn(*match.groups())
+                break
+        try:
+            assert found
+            if value is not None:
+                tensor.copy_(torch.from_numpy(np.array(value)))
+        except:
+            print(name, list(tensor.shape))
+            raise
+
+#----------------------------------------------------------------------------
+
+def convert_tf_generator(tf_G):
+    if tf_G.version < 4:
+        raise ValueError('TensorFlow pickle version too low')
+
+    # Collect kwargs.
+    tf_kwargs = tf_G.static_kwargs
+    known_kwargs = set()
+    def kwarg(tf_name, default=None, none=None):
+        known_kwargs.add(tf_name)
+        val = tf_kwargs.get(tf_name, default)
+        return val if val is not None else none
+
+    # Convert kwargs.
+    from training import networks_stylegan2
+    network_class = networks_stylegan2.Generator
+    kwargs = dnnlib.EasyDict(
+        z_dim               = kwarg('latent_size',          512),
+        c_dim               = kwarg('label_size',           0),
+        w_dim               = kwarg('dlatent_size',         512),
+        img_resolution      = kwarg('resolution',           1024),
+        img_channels        = kwarg('num_channels',         3),
+        channel_base        = kwarg('fmap_base',            16384) * 2,
+        channel_max         = kwarg('fmap_max',             512),
+        num_fp16_res        = kwarg('num_fp16_res',         0),
+        conv_clamp          = kwarg('conv_clamp',           None),
+        architecture        = kwarg('architecture',         'skip'),
+        resample_filter     = kwarg('resample_kernel',      [1,3,3,1]),
+        use_noise           = kwarg('use_noise',            True),
+        activation          = kwarg('nonlinearity',         'lrelu'),
+        mapping_kwargs      = dnnlib.EasyDict(
+            num_layers      = kwarg('mapping_layers',       8),
+            embed_features  = kwarg('label_fmaps',          None),
+            layer_features  = kwarg('mapping_fmaps',        None),
+            activation      = kwarg('mapping_nonlinearity', 'lrelu'),
+            lr_multiplier   = kwarg('mapping_lrmul',        0.01),
+            w_avg_beta      = kwarg('w_avg_beta',           0.995,  none=1),
+        ),
+    )
+
+    # Check for unknown kwargs.
+    kwarg('truncation_psi')
+    kwarg('truncation_cutoff')
+    kwarg('style_mixing_prob')
+    kwarg('structure')
+    kwarg('conditioning')
+    kwarg('fused_modconv')
+    unknown_kwargs = list(set(tf_kwargs.keys()) - known_kwargs)
+    if len(unknown_kwargs) > 0:
+        raise ValueError('Unknown TensorFlow kwarg', unknown_kwargs[0])
+
+    # Collect params.
+    tf_params = _collect_tf_params(tf_G)
+    for name, value in list(tf_params.items()):
+        match = re.fullmatch(r'ToRGB_lod(\d+)/(.*)', name)
+        if match:
+            r = kwargs.img_resolution // (2 ** int(match.group(1)))
+            tf_params[f'{r}x{r}/ToRGB/{match.group(2)}'] = value
+            kwargs.synthesis.kwargs.architecture = 'orig'
+    #for name, value in tf_params.items(): print(f'{name:<50s}{list(value.shape)}')
+
+    # Convert params.
+    G = network_class(**kwargs).eval().requires_grad_(False)
+    # pylint: disable=unnecessary-lambda
+    # pylint: disable=f-string-without-interpolation
+    _populate_module_params(G,
+        r'mapping\.w_avg',                                  lambda:     tf_params[f'dlatent_avg'],
+        r'mapping\.embed\.weight',                          lambda:     tf_params[f'mapping/LabelEmbed/weight'].transpose(),
+        r'mapping\.embed\.bias',                            lambda:     tf_params[f'mapping/LabelEmbed/bias'],
+        r'mapping\.fc(\d+)\.weight',                        lambda i:   tf_params[f'mapping/Dense{i}/weight'].transpose(),
+        r'mapping\.fc(\d+)\.bias',                          lambda i:   tf_params[f'mapping/Dense{i}/bias'],
+        r'synthesis\.b4\.const',                            lambda:     tf_params[f'synthesis/4x4/Const/const'][0],
+        r'synthesis\.b4\.conv1\.weight',                    lambda:     tf_params[f'synthesis/4x4/Conv/weight'].transpose(3, 2, 0, 1),
+        r'synthesis\.b4\.conv1\.bias',                      lambda:     tf_params[f'synthesis/4x4/Conv/bias'],
+        r'synthesis\.b4\.conv1\.noise_const',               lambda:     tf_params[f'synthesis/noise0'][0, 0],
+        r'synthesis\.b4\.conv1\.noise_strength',            lambda:     tf_params[f'synthesis/4x4/Conv/noise_strength'],
+        r'synthesis\.b4\.conv1\.affine\.weight',            lambda:     tf_params[f'synthesis/4x4/Conv/mod_weight'].transpose(),
+        r'synthesis\.b4\.conv1\.affine\.bias',              lambda:     tf_params[f'synthesis/4x4/Conv/mod_bias'] + 1,
+        r'synthesis\.b(\d+)\.conv0\.weight',                lambda r:   tf_params[f'synthesis/{r}x{r}/Conv0_up/weight'][::-1, ::-1].transpose(3, 2, 0, 1),
+        r'synthesis\.b(\d+)\.conv0\.bias',                  lambda r:   tf_params[f'synthesis/{r}x{r}/Conv0_up/bias'],
+        r'synthesis\.b(\d+)\.conv0\.noise_const',           lambda r:   tf_params[f'synthesis/noise{int(np.log2(int(r)))*2-5}'][0, 0],
+        r'synthesis\.b(\d+)\.conv0\.noise_strength',        lambda r:   tf_params[f'synthesis/{r}x{r}/Conv0_up/noise_strength'],
+        r'synthesis\.b(\d+)\.conv0\.affine\.weight',        lambda r:   tf_params[f'synthesis/{r}x{r}/Conv0_up/mod_weight'].transpose(),
+        r'synthesis\.b(\d+)\.conv0\.affine\.bias',          lambda r:   tf_params[f'synthesis/{r}x{r}/Conv0_up/mod_bias'] + 1,
+        r'synthesis\.b(\d+)\.conv1\.weight',                lambda r:   tf_params[f'synthesis/{r}x{r}/Conv1/weight'].transpose(3, 2, 0, 1),
+        r'synthesis\.b(\d+)\.conv1\.bias',                  lambda r:   tf_params[f'synthesis/{r}x{r}/Conv1/bias'],
+        r'synthesis\.b(\d+)\.conv1\.noise_const',           lambda r:   tf_params[f'synthesis/noise{int(np.log2(int(r)))*2-4}'][0, 0],
+        r'synthesis\.b(\d+)\.conv1\.noise_strength',        lambda r:   tf_params[f'synthesis/{r}x{r}/Conv1/noise_strength'],
+        r'synthesis\.b(\d+)\.conv1\.affine\.weight',        lambda r:   tf_params[f'synthesis/{r}x{r}/Conv1/mod_weight'].transpose(),
+        r'synthesis\.b(\d+)\.conv1\.affine\.bias',          lambda r:   tf_params[f'synthesis/{r}x{r}/Conv1/mod_bias'] + 1,
+        r'synthesis\.b(\d+)\.torgb\.weight',                lambda r:   tf_params[f'synthesis/{r}x{r}/ToRGB/weight'].transpose(3, 2, 0, 1),
+        r'synthesis\.b(\d+)\.torgb\.bias',                  lambda r:   tf_params[f'synthesis/{r}x{r}/ToRGB/bias'],
+        r'synthesis\.b(\d+)\.torgb\.affine\.weight',        lambda r:   tf_params[f'synthesis/{r}x{r}/ToRGB/mod_weight'].transpose(),
+        r'synthesis\.b(\d+)\.torgb\.affine\.bias',          lambda r:   tf_params[f'synthesis/{r}x{r}/ToRGB/mod_bias'] + 1,
+        r'synthesis\.b(\d+)\.skip\.weight',                 lambda r:   tf_params[f'synthesis/{r}x{r}/Skip/weight'][::-1, ::-1].transpose(3, 2, 0, 1),
+        r'.*\.resample_filter',                             None,
+        r'.*\.act_filter',                                  None,
+    )
+    return G
+
+#----------------------------------------------------------------------------
+
+def convert_tf_discriminator(tf_D):
+    if tf_D.version < 4:
+        raise ValueError('TensorFlow pickle version too low')
+
+    # Collect kwargs.
+    tf_kwargs = tf_D.static_kwargs
+    known_kwargs = set()
+    def kwarg(tf_name, default=None):
+        known_kwargs.add(tf_name)
+        return tf_kwargs.get(tf_name, default)
+
+    # Convert kwargs.
+    kwargs = dnnlib.EasyDict(
+        c_dim                   = kwarg('label_size',           0),
+        img_resolution          = kwarg('resolution',           1024),
+        img_channels            = kwarg('num_channels',         3),
+        architecture            = kwarg('architecture',         'resnet'),
+        channel_base            = kwarg('fmap_base',            16384) * 2,
+        channel_max             = kwarg('fmap_max',             512),
+        num_fp16_res            = kwarg('num_fp16_res',         0),
+        conv_clamp              = kwarg('conv_clamp',           None),
+        cmap_dim                = kwarg('mapping_fmaps',        None),
+        block_kwargs = dnnlib.EasyDict(
+            activation          = kwarg('nonlinearity',         'lrelu'),
+            resample_filter     = kwarg('resample_kernel',      [1,3,3,1]),
+            freeze_layers       = kwarg('freeze_layers',        0),
+        ),
+        mapping_kwargs = dnnlib.EasyDict(
+            num_layers          = kwarg('mapping_layers',       0),
+            embed_features      = kwarg('mapping_fmaps',        None),
+            layer_features      = kwarg('mapping_fmaps',        None),
+            activation          = kwarg('nonlinearity',         'lrelu'),
+            lr_multiplier       = kwarg('mapping_lrmul',        0.1),
+        ),
+        epilogue_kwargs = dnnlib.EasyDict(
+            mbstd_group_size    = kwarg('mbstd_group_size',     None),
+            mbstd_num_channels  = kwarg('mbstd_num_features',   1),
+            activation          = kwarg('nonlinearity',         'lrelu'),
+        ),
+    )
+
+    # Check for unknown kwargs.
+    kwarg('structure')
+    kwarg('conditioning')
+    unknown_kwargs = list(set(tf_kwargs.keys()) - known_kwargs)
+    if len(unknown_kwargs) > 0:
+        raise ValueError('Unknown TensorFlow kwarg', unknown_kwargs[0])
+
+    # Collect params.
+    tf_params = _collect_tf_params(tf_D)
+    for name, value in list(tf_params.items()):
+        match = re.fullmatch(r'FromRGB_lod(\d+)/(.*)', name)
+        if match:
+            r = kwargs.img_resolution // (2 ** int(match.group(1)))
+            tf_params[f'{r}x{r}/FromRGB/{match.group(2)}'] = value
+            kwargs.architecture = 'orig'
+    #for name, value in tf_params.items(): print(f'{name:<50s}{list(value.shape)}')
+
+    # Convert params.
+    from training import networks_stylegan2
+    D = networks_stylegan2.Discriminator(**kwargs).eval().requires_grad_(False)
+    # pylint: disable=unnecessary-lambda
+    # pylint: disable=f-string-without-interpolation
+    _populate_module_params(D,
+        r'b(\d+)\.fromrgb\.weight',     lambda r:       tf_params[f'{r}x{r}/FromRGB/weight'].transpose(3, 2, 0, 1),
+        r'b(\d+)\.fromrgb\.bias',       lambda r:       tf_params[f'{r}x{r}/FromRGB/bias'],
+        r'b(\d+)\.conv(\d+)\.weight',   lambda r, i:    tf_params[f'{r}x{r}/Conv{i}{["","_down"][int(i)]}/weight'].transpose(3, 2, 0, 1),
+        r'b(\d+)\.conv(\d+)\.bias',     lambda r, i:    tf_params[f'{r}x{r}/Conv{i}{["","_down"][int(i)]}/bias'],
+        r'b(\d+)\.skip\.weight',        lambda r:       tf_params[f'{r}x{r}/Skip/weight'].transpose(3, 2, 0, 1),
+        r'mapping\.embed\.weight',      lambda:         tf_params[f'LabelEmbed/weight'].transpose(),
+        r'mapping\.embed\.bias',        lambda:         tf_params[f'LabelEmbed/bias'],
+        r'mapping\.fc(\d+)\.weight',    lambda i:       tf_params[f'Mapping{i}/weight'].transpose(),
+        r'mapping\.fc(\d+)\.bias',      lambda i:       tf_params[f'Mapping{i}/bias'],
+        r'b4\.conv\.weight',            lambda:         tf_params[f'4x4/Conv/weight'].transpose(3, 2, 0, 1),
+        r'b4\.conv\.bias',              lambda:         tf_params[f'4x4/Conv/bias'],
+        r'b4\.fc\.weight',              lambda:         tf_params[f'4x4/Dense0/weight'].transpose(),
+        r'b4\.fc\.bias',                lambda:         tf_params[f'4x4/Dense0/bias'],
+        r'b4\.out\.weight',             lambda:         tf_params[f'Output/weight'].transpose(),
+        r'b4\.out\.bias',               lambda:         tf_params[f'Output/bias'],
+        r'.*\.resample_filter',         None,
+    )
+    return D
+
+#----------------------------------------------------------------------------
+
+@click.command()
+@click.option('--source', help='Input pickle', required=True, metavar='PATH')
+@click.option('--dest', help='Output pickle', required=True, metavar='PATH')
+@click.option('--force-fp16', help='Force the networks to use FP16', type=bool, default=False, metavar='BOOL', show_default=True)
+def convert_network_pickle(source, dest, force_fp16):
+    """Convert legacy network pickle into the native PyTorch format.
+
+    The tool is able to load the main network configurations exported using the TensorFlow version of StyleGAN2 or StyleGAN2-ADA.
+    It does not support e.g. StyleGAN2-ADA comparison methods, StyleGAN2 configs A-D, or StyleGAN1 networks.
+
+    Example:
+
+    \b
+    python legacy.py \\
+        --source=https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-cat-config-f.pkl \\
+        --dest=stylegan2-cat-config-f.pkl
+    """
+    print(f'Loading "{source}"...')
+    with dnnlib.util.open_url(source) as f:
+        data = load_network_pkl(f, force_fp16=force_fp16)
+    print(f'Saving "{dest}"...')
+    with open(dest, 'wb') as f:
+        pickle.dump(data, f)
+    print('Done.')
+
+#----------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    convert_network_pickle() # pylint: disable=no-value-for-parameter
+
+#----------------------------------------------------------------------------

pre-requirements.txt

@@ -0,0 +1 @@
+pip

requirements.txt

@@ -0,0 +1,8 @@
+torch
+torchvision
+Ninja
+gradio
+huggingface_hub
+hf_transfer
+Pillow==9.5.0
+psutil

visualizer_drag.py

@@ -0,0 +1,404 @@
+# 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 click
+import os
+
+import multiprocessing
+import numpy as np
+import torch
+import imgui
+import dnnlib
+from gui_utils import imgui_window
+from gui_utils import imgui_utils
+from gui_utils import gl_utils
+from gui_utils import text_utils
+from viz import renderer
+from viz import pickle_widget
+from viz import latent_widget
+from viz import drag_widget
+from viz import capture_widget
+
+#----------------------------------------------------------------------------
+
+class Visualizer(imgui_window.ImguiWindow):
+    def __init__(self, capture_dir=None):
+        super().__init__(title='DragGAN', window_width=3840, window_height=2160)
+
+        # Internals.
+        self._last_error_print  = None
+        self._async_renderer    = AsyncRenderer()
+        self._defer_rendering   = 0
+        self._tex_img           = None
+        self._tex_obj           = None
+        self._mask_obj          = None
+        self._image_area        = None
+        self._status            = dnnlib.EasyDict()
+
+        # Widget interface.
+        self.args               = dnnlib.EasyDict()
+        self.result             = dnnlib.EasyDict()
+        self.pane_w             = 0
+        self.label_w            = 0
+        self.button_w           = 0
+        self.image_w            = 0
+        self.image_h            = 0
+
+        # Widgets.
+        self.pickle_widget      = pickle_widget.PickleWidget(self)
+        self.latent_widget      = latent_widget.LatentWidget(self)
+        self.drag_widget        = drag_widget.DragWidget(self)
+        self.capture_widget     = capture_widget.CaptureWidget(self)
+
+        if capture_dir is not None:
+            self.capture_widget.path = capture_dir
+
+        # Initialize window.
+        self.set_position(0, 0)
+        self._adjust_font_size()
+        self.skip_frame() # Layout may change after first frame.
+
+    def close(self):
+        super().close()
+        if self._async_renderer is not None:
+            self._async_renderer.close()
+            self._async_renderer = None
+
+    def add_recent_pickle(self, pkl, ignore_errors=False):
+        self.pickle_widget.add_recent(pkl, ignore_errors=ignore_errors)
+
+    def load_pickle(self, pkl, ignore_errors=False):
+        self.pickle_widget.load(pkl, ignore_errors=ignore_errors)
+
+    def print_error(self, error):
+        error = str(error)
+        if error != self._last_error_print:
+            print('\n' + error + '\n')
+            self._last_error_print = error
+
+    def defer_rendering(self, num_frames=1):
+        self._defer_rendering = max(self._defer_rendering, num_frames)
+
+    def clear_result(self):
+        self._async_renderer.clear_result()
+
+    def set_async(self, is_async):
+        if is_async != self._async_renderer.is_async:
+            self._async_renderer.set_async(is_async)
+            self.clear_result()
+            if 'image' in self.result:
+                self.result.message = 'Switching rendering process...'
+                self.defer_rendering()
+
+    def _adjust_font_size(self):
+        old = self.font_size
+        self.set_font_size(min(self.content_width / 120, self.content_height / 60))
+        if self.font_size != old:
+            self.skip_frame() # Layout changed.
+
+    def check_update_mask(self, **args):
+        update_mask = False
+        if 'pkl' in self._status:
+            if self._status.pkl != args['pkl']:
+                update_mask = True
+        self._status.pkl = args['pkl']
+        if 'w0_seed' in self._status:
+            if self._status.w0_seed != args['w0_seed']:
+                update_mask = True
+        self._status.w0_seed = args['w0_seed']
+        return update_mask
+
+    def capture_image_frame(self):
+        self.capture_next_frame()
+        captured_frame = self.pop_captured_frame()
+        captured_image = None
+        if captured_frame is not None:
+            x1, y1, w, h = self._image_area
+            captured_image = captured_frame[y1:y1+h, x1:x1+w, :]
+        return captured_image
+
+    def get_drag_info(self):
+        seed = self.latent_widget.seed
+        points = self.drag_widget.points
+        targets = self.drag_widget.targets
+        mask = self.drag_widget.mask
+        w = self._async_renderer._renderer_obj.w
+        return seed, points, targets, mask, w
+
+    def draw_frame(self):
+        self.begin_frame()
+        self.args = dnnlib.EasyDict()
+        self.pane_w = self.font_size * 18
+        self.button_w = self.font_size * 5
+        self.label_w = round(self.font_size * 4.5)
+
+        # Detect mouse dragging in the result area.
+        if self._image_area is not None:
+            if not hasattr(self.drag_widget, 'width'):
+                self.drag_widget.init_mask(self.image_w, self.image_h)
+            clicked, down, img_x, img_y = imgui_utils.click_hidden_window(
+                '##image_area', self._image_area[0], self._image_area[1], self._image_area[2], self._image_area[3], self.image_w, self.image_h)
+            self.drag_widget.action(clicked, down, img_x, img_y)
+
+        # Begin control pane.
+        imgui.set_next_window_position(0, 0)
+        imgui.set_next_window_size(self.pane_w, self.content_height)
+        imgui.begin('##control_pane', closable=False, flags=(imgui.WINDOW_NO_TITLE_BAR | imgui.WINDOW_NO_RESIZE | imgui.WINDOW_NO_MOVE))
+
+        # Widgets.
+        expanded, _visible = imgui_utils.collapsing_header('Network & latent', default=True)
+        self.pickle_widget(expanded)
+        self.latent_widget(expanded)
+        expanded, _visible = imgui_utils.collapsing_header('Drag', default=True)
+        self.drag_widget(expanded)
+        expanded, _visible = imgui_utils.collapsing_header('Capture', default=True)
+        self.capture_widget(expanded)
+
+        # Render.
+        if self.is_skipping_frames():
+            pass
+        elif self._defer_rendering > 0:
+            self._defer_rendering -= 1
+        elif self.args.pkl is not None:
+            self._async_renderer.set_args(**self.args)
+            result = self._async_renderer.get_result()
+            if result is not None:
+                self.result = result        
+                if 'stop' in self.result and self.result.stop:
+                    self.drag_widget.stop_drag()
+                if 'points' in self.result:
+                    self.drag_widget.set_points(self.result.points)
+                if 'init_net' in self.result:
+                    if self.result.init_net:
+                        self.drag_widget.reset_point()
+
+        if self.check_update_mask(**self.args):
+            h, w, _ = self.result.image.shape
+            self.drag_widget.init_mask(w, h)
+
+        # Display.
+        max_w = self.content_width - self.pane_w
+        max_h = self.content_height
+        pos = np.array([self.pane_w + max_w / 2, max_h / 2])
+        if 'image' in self.result:
+            if self._tex_img is not self.result.image:
+                self._tex_img = self.result.image
+                if self._tex_obj is None or not self._tex_obj.is_compatible(image=self._tex_img):
+                    self._tex_obj = gl_utils.Texture(image=self._tex_img, bilinear=False, mipmap=False)
+                else:
+                    self._tex_obj.update(self._tex_img)
+                self.image_h, self.image_w = self._tex_obj.height, self._tex_obj.width
+            zoom = min(max_w / self._tex_obj.width, max_h / self._tex_obj.height)
+            zoom = np.floor(zoom) if zoom >= 1 else zoom
+            self._tex_obj.draw(pos=pos, zoom=zoom, align=0.5, rint=True)
+            if self.drag_widget.show_mask and hasattr(self.drag_widget, 'mask'):
+                mask = ((1-self.drag_widget.mask.unsqueeze(-1)) * 255).to(torch.uint8)
+                if self._mask_obj is None or not self._mask_obj.is_compatible(image=self._tex_img):
+                    self._mask_obj = gl_utils.Texture(image=mask, bilinear=False, mipmap=False)
+                else:
+                    self._mask_obj.update(mask)
+                self._mask_obj.draw(pos=pos, zoom=zoom, align=0.5, rint=True, alpha=0.15)
+
+            if self.drag_widget.mode in ['flexible', 'fixed']:
+                posx, posy = imgui.get_mouse_pos()
+                if posx >= self.pane_w:
+                    pos_c = np.array([posx, posy])
+                    gl_utils.draw_circle(center=pos_c, radius=self.drag_widget.r_mask * zoom, alpha=0.5)
+            
+            rescale = self._tex_obj.width / 512 * zoom
+            
+            for point in self.drag_widget.targets:
+                pos_x = self.pane_w + max_w / 2 + (point[1] - self.image_w//2) * zoom
+                pos_y = max_h / 2 + (point[0] - self.image_h//2) * zoom
+                gl_utils.draw_circle(center=np.array([pos_x, pos_y]), color=[0,0,1], radius=9 * rescale)
+            
+            for point in self.drag_widget.points:
+                pos_x = self.pane_w + max_w / 2 + (point[1] - self.image_w//2) * zoom
+                pos_y = max_h / 2 + (point[0] - self.image_h//2) * zoom
+                gl_utils.draw_circle(center=np.array([pos_x, pos_y]), color=[1,0,0], radius=9 * rescale)
+
+            for point, target in zip(self.drag_widget.points, self.drag_widget.targets):
+                t_x = self.pane_w + max_w / 2 + (target[1] - self.image_w//2) * zoom
+                t_y = max_h / 2 + (target[0] - self.image_h//2) * zoom
+
+                p_x = self.pane_w + max_w / 2 + (point[1] - self.image_w//2) * zoom
+                p_y = max_h / 2 + (point[0] - self.image_h//2) * zoom
+
+                gl_utils.draw_arrow(p_x, p_y, t_x, t_y, l=8 * rescale, width = 3 * rescale)
+
+            imshow_w = int(self._tex_obj.width * zoom)
+            imshow_h = int(self._tex_obj.height * zoom)
+            self._image_area = [int(self.pane_w + max_w / 2 - imshow_w / 2), int(max_h / 2 - imshow_h / 2), imshow_w, imshow_h]
+        if 'error' in self.result:
+            self.print_error(self.result.error)
+            if 'message' not in self.result:
+                self.result.message = str(self.result.error)
+        if 'message' in self.result:
+            tex = text_utils.get_texture(self.result.message, size=self.font_size, max_width=max_w, max_height=max_h, outline=2)
+            tex.draw(pos=pos, align=0.5, rint=True, color=1)
+
+        # End frame.
+        self._adjust_font_size()
+        imgui.end()
+        self.end_frame()
+
+#----------------------------------------------------------------------------
+
+class AsyncRenderer:
+    def __init__(self):
+        self._closed        = False
+        self._is_async      = False
+        self._cur_args      = None
+        self._cur_result    = None
+        self._cur_stamp     = 0
+        self._renderer_obj  = None
+        self._args_queue    = None
+        self._result_queue  = None
+        self._process       = None
+
+    def close(self):
+        self._closed = True
+        self._renderer_obj = None
+        if self._process is not None:
+            self._process.terminate()
+        self._process = None
+        self._args_queue = None
+        self._result_queue = None
+
+    @property
+    def is_async(self):
+        return self._is_async
+
+    def set_async(self, is_async):
+        self._is_async = is_async
+
+    def set_args(self, **args):
+        assert not self._closed
+        args2 = args.copy()
+        args_mask = args2.pop('mask')
+        if self._cur_args:
+            _cur_args = self._cur_args.copy()
+            cur_args_mask = _cur_args.pop('mask')
+        else:
+            _cur_args = self._cur_args
+        # if args != self._cur_args:
+        if args2 != _cur_args:
+            if self._is_async:
+                self._set_args_async(**args)
+            else:
+                self._set_args_sync(**args)
+            self._cur_args = args
+
+    def _set_args_async(self, **args):
+        if self._process is None:
+            self._args_queue = multiprocessing.Queue()
+            self._result_queue = multiprocessing.Queue()
+            try:
+                multiprocessing.set_start_method('spawn')
+            except RuntimeError:
+                pass
+            self._process = multiprocessing.Process(target=self._process_fn, args=(self._args_queue, self._result_queue), daemon=True)
+            self._process.start()
+        self._args_queue.put([args, self._cur_stamp])
+
+    def _set_args_sync(self, **args):
+        if self._renderer_obj is None:
+            self._renderer_obj = renderer.Renderer()
+        self._cur_result = self._renderer_obj.render(**args)
+
+    def get_result(self):
+        assert not self._closed
+        if self._result_queue is not None:
+            while self._result_queue.qsize() > 0:
+                result, stamp = self._result_queue.get()
+                if stamp == self._cur_stamp:
+                    self._cur_result = result
+        return self._cur_result
+
+    def clear_result(self):
+        assert not self._closed
+        self._cur_args = None
+        self._cur_result = None
+        self._cur_stamp += 1
+
+    @staticmethod
+    def _process_fn(args_queue, result_queue):
+        renderer_obj = renderer.Renderer()
+        cur_args = None
+        cur_stamp = None
+        while True:
+            args, stamp = args_queue.get()
+            while args_queue.qsize() > 0:
+                args, stamp = args_queue.get()
+            if args != cur_args or stamp != cur_stamp:
+                result = renderer_obj.render(**args)
+                if 'error' in result:
+                    result.error = renderer.CapturedException(result.error)
+                result_queue.put([result, stamp])
+                cur_args = args
+                cur_stamp = stamp
+
+#----------------------------------------------------------------------------
+
+@click.command()
+@click.argument('pkls', metavar='PATH', nargs=-1)
+@click.option('--capture-dir', help='Where to save screenshot captures', metavar='PATH', default=None)
+@click.option('--browse-dir', help='Specify model path for the \'Browse...\' button', metavar='PATH')
+def main(
+    pkls,
+    capture_dir,
+    browse_dir
+):
+    """Interactive model visualizer.
+
+    Optional PATH argument can be used specify which .pkl file to load.
+    """
+    viz = Visualizer(capture_dir=capture_dir)
+
+    if browse_dir is not None:
+        viz.pickle_widget.search_dirs = [browse_dir]
+
+    # List pickles.
+    if len(pkls) > 0:
+        for pkl in pkls:
+            viz.add_recent_pickle(pkl)
+        viz.load_pickle(pkls[0])
+    else:
+        pretrained = [
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-afhqcat-512x512.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-afhqdog-512x512.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-afhqv2-512x512.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-afhqwild-512x512.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-brecahad-512x512.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-celebahq-256x256.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-cifar10-32x32.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhq-1024x1024.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhq-256x256.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhq-512x512.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhqu-1024x1024.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhqu-256x256.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-lsundog-256x256.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-metfaces-1024x1024.pkl',
+            'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-metfacesu-1024x1024.pkl'
+        ]
+
+        # Populate recent pickles list with pretrained model URLs.
+        for url in pretrained:
+            viz.add_recent_pickle(url)
+
+    # Run.
+    while not viz.should_close():
+        viz.draw_frame()
+    viz.close()
+
+#----------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    main()
+
+#----------------------------------------------------------------------------

visualizer_drag_gradio.py

@@ -0,0 +1,940 @@
+# https://huggingface.co/DragGan/DragGan-Models
+# https://arxiv.org/abs/2305.10973
+import os
+import os.path as osp
+from argparse import ArgumentParser
+from functools import partial
+from pathlib import Path
+import time
+
+import psutil
+
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+
+import dnnlib
+from gradio_utils import (ImageMask, draw_mask_on_image, draw_points_on_image,
+                          get_latest_points_pair, get_valid_mask,
+                          on_change_single_global_state)
+from viz.renderer import Renderer, add_watermark_np
+
+
+# download models from Hugging Face hub
+from huggingface_hub import snapshot_download
+
+model_dir = Path('./checkpoints')
+snapshot_download('DragGan/DragGan-Models',
+                  repo_type='model', local_dir=model_dir)
+
+parser = ArgumentParser()
+parser.add_argument('--share', action='store_true')
+parser.add_argument('--cache-dir', type=str, default='./checkpoints')
+args = parser.parse_args()
+
+cache_dir = args.cache_dir
+
+device = 'cuda'
+IS_SPACE = "DragGan/DragGan" in os.environ.get('SPACE_ID', '')
+TIMEOUT = 80
+
+
+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):
+    """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
+        valid_checkpoints_dict[state['pretrained_weight']],  # pkl
+        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
+
+
+def print_memory_usage():
+    # Print system memory usage
+    print(f"System memory usage: {psutil.virtual_memory().percent}%")
+
+    # Print GPU memory usage
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+        print(f"GPU memory usage: {torch.cuda.memory_allocated() / 1e9} GB")
+        print(
+            f"Max GPU memory usage: {torch.cuda.max_memory_allocated() / 1e9} GB")
+        device_properties = torch.cuda.get_device_properties(device)
+        available_memory = device_properties.total_memory - \
+            torch.cuda.max_memory_allocated()
+        print(f"Available GPU memory: {available_memory / 1e9} GB")
+    else:
+        print("No GPU available")
+
+
+# filter large models running on SPACES
+allowed_checkpoints = []  # all checkpoints
+if IS_SPACE:
+    allowed_checkpoints = ["stylegan_human_v2_512.pkl",
+                           "stylegan2_dogs_1024_pytorch.pkl"]
+
+valid_checkpoints_dict = {
+    f.name.split('.')[0]: str(f)
+    for f in Path(cache_dir).glob('*.pkl')
+    if f.name in allowed_checkpoints or not IS_SPACE
+}
+print('Valid checkpoint file:')
+print(valid_checkpoints_dict)
+
+init_pkl = 'stylegan_human_v2_512'
+
+with gr.Blocks() as app:
+    gr.Markdown("""
+# DragGAN - Drag Your GAN
+## Interactive Point-based Manipulation on the Generative Image Manifold
+### Unofficial Gradio Demo
+
+**Due to high demand, only one model can be run at a time, or you can duplicate the space and run your own copy.**
+
+<a href="https://huggingface.co/spaces/radames/DragGan?duplicate=true" style="display: inline-block;margin-top: .5em;margin-right: .25em;" target="_blank">
+<img style="margin-bottom: 0em;display: inline;margin-top: -.25em;" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a> for no queue on your own hardware.</p>
+
+* Official Repo: [XingangPan](https://github.com/XingangPan/DragGAN)
+* Gradio Demo by: [LeoXing1996](https://github.com/LeoXing1996) © [OpenMMLab MMagic](https://github.com/open-mmlab/mmagic)
+""")
+
+    # renderer = Renderer()
+    global_state = gr.State({
+        "images": {
+            # image_orig: the original image, change with seed/model is changed
+            # image_raw: image with mask and points, change durning 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": dnnlib.EasyDict(),
+        "params": {
+            "seed": int(np.random.randint(0, 2**32 - 1)),
+            "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
+    })
+
+    # 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.Slider(
+                            mininium=0,
+                            maximum=2**32-1,
+                            step=1,
+                            value=global_state.value['params']['seed'],
+                            interactive=True,
+                            # randomize=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],
+        queue=True,
+    )
+
+    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],
+        queue=False,
+    )
+
+    # 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],
+        queue=False,
+    )
+
+    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('Running with:')
+            print(f'    Source: {p_in_pixels}')
+            print(f'    Target: {t_in_pixels}')
+            step_idx = 0
+            last_time = time.time()
+            while True:
+                print_memory_usage()
+                # add a TIMEOUT break
+                print(f'Running time: {time.time() - last_time}')
+                if IS_SPACE and time.time() - last_time > TIMEOUT:
+                    print('Timeout break!')
+                    break
+                if global_state["temporal_params"]["stop"] or global_state['generator_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(
+                        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],
+                        queue=False)
+
+    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],
+        queue=False,
+    )
+
+    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,
+                          ],
+                          queue=False)
+
+    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],
+                            queue=False)
+
+    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],
+        queue=False,
+    )
+
+    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],
+                      queue=False)
+
+    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],
+        queue=False,
+    )
+
+print("SHAReD: Start app", parser.parse_args())
+gr.close_all()
+app.queue(concurrency_count=1, max_size=200, api_open=False)
+app.launch(share=args.share, show_api=False)