Update visualizer_drag_gradio.py

visualizer_drag_gradio.py

@@ -15,28 +15,31 @@ 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 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)
+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')
+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', '')
+IS_SPACE = "DragGan/DragGan" in os.environ.get("SPACE_ID", "")
+device = "cuda" if torch.cuda.is_available() else "cpu"
 TIMEOUT = 80
 
 
@@ -54,23 +57,24 @@ def clear_state(global_state, target=None):
     2. set global_state['mask'] as full-one mask.
     """
     if target is None:
-        target = ['point', 'mask']
+        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:
+    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!')
+        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.
+    """This function is called only once when 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
@@ -83,43 +87,48 @@ def init_images(global_state):
     else:
         state = global_state
 
-    state['renderer'].init_network(
-        state['generator_params'],  # res
-        valid_checkpoints_dict[state['pretrained_weight']],  # pkl
-        state['params']['seed'],  # w0_seed,
+    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,
+        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["params"]["lr"],  # lr,
+    )
+
+    state["renderer"]._render_drag_impl(
+        state["generator_params"], is_drag=False, to_pil=True
     )
 
-    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)
+    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():
+    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
+        global_state["images"]["image_show"] = image_draw
     return image_draw
 
 
@@ -131,34 +140,31 @@ def preprocess_mask_info(global_state, image):
         2.2 global_state is add_mask:
     """
     if isinstance(image, dict):
-        last_mask = get_valid_mask(image['mask'])
+        last_mask = get_valid_mask(image["mask"])
     else:
         last_mask = None
-    mask = global_state['mask']
+    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']
+    editing_mode = global_state["editing_state"]
 
     if last_mask is None:
         return global_state
 
-    if editing_mode == 'remove_mask':
+    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':
+        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.')
+        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.')
+        print(f"Last editing_state is {editing_mode}, do nothing to mask.")
 
-    global_state['mask'] = updated_mask
-    # global_state['last_mask'] = None  # clear buffer
+    global_state["mask"] = updated_mask
     return global_state
 
 
@@ -171,10 +177,12 @@ def print_memory_usage():
         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")
+            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()
+        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")
@@ -183,21 +191,24 @@ def print_memory_usage():
 # 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"]
+    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')
+    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 checkpoint file:")
 print(valid_checkpoints_dict)
 
-init_pkl = 'stylegan_human_v2_512'
+init_pkl = "stylegan_human_v2_512"
 
 with gr.Blocks() as app:
-    gr.Markdown("""
+    gr.Markdown(
+        """
 # DragGAN - Drag Your GAN
 ## Interactive Point-based Manipulation on the Generative Image Manifold
 ### Unofficial Gradio Demo
@@ -209,58 +220,49 @@ with gr.Blocks() as app:
 
 * 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
-    })
+"""
+    )
+
+    global_state = gr.State(
+        {
+            "images": {},
+            "temporal_params": {},
+            "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)
+                        gr.Markdown(value="Pickle", show_label=False)
 
                     with gr.Column(scale=4, min_width=10):
                         form_pretrained_dropdown = gr.Dropdown(
@@ -272,14 +274,14 @@ with gr.Blocks() as app:
                 # Latent
                 with gr.Row():
                     with gr.Column(scale=1, min_width=10):
-                        gr.Markdown(value='Latent', show_label=False)
+                        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,
+                            minimum=0,
+                            maximum=2**32 - 1,
                             step=1,
-                            value=global_state.value['params']['seed'],
+                            value=global_state.value["params"]["seed"],
                             interactive=True,
                             # randomize=True,
                             label="Seed",
@@ -287,60 +289,66 @@ with gr.Blocks() as app:
                         form_lr_number = gr.Number(
                             value=global_state.value["params"]["lr"],
                             interactive=True,
-                            label="Step Size")
+                            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'],
+                                    ["w", "w+"],
+                                    value=global_state.value["params"][
+                                        "latent_space"
+                                    ],
                                     interactive=True,
-                                    label='Latent space to optimize',
+                                    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)
+                        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')
+                                enable_add_points = gr.Button("Add Points")
                             with gr.Column(scale=1, min_width=10):
-                                undo_points = gr.Button('Reset Points')
+                                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)
+                        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)
+                        gr.Markdown(value="Mask", show_label=False)
                     with gr.Column(scale=4, min_width=10):
-                        enable_add_mask = gr.Button('Edit Flexible Area')
+                        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)
+                                    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"],
+                                value=global_state.value["params"][
+                                    "motion_lambda"
+                                ],
                                 interactive=True,
                                 label="Lambda",
                             )
@@ -349,16 +357,21 @@ with gr.Blocks() as app:
                     value=global_state.value["draw_interval"],
                     label="Draw Interval (steps)",
                     interactive=True,
-                    visible=False)
+                    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("""
+                    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
@@ -377,10 +390,10 @@ with gr.Blocks() as app:
           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("""
+        """
+    )
+    gr.HTML(
+        """
         <style>
             .container {
                 position: absolute;
@@ -395,8 +408,8 @@ with gr.Blocks() as app:
         <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.
@@ -404,11 +417,11 @@ with gr.Blocks() as app:
         2. Re-init images and clear all states
         """
 
-        global_state['pretrained_weight'] = pretrained_value
+        global_state["pretrained_weight"] = pretrained_value
         init_images(global_state)
         clear_state(global_state)
 
-        return global_state, global_state["images"]['image_show']
+        return global_state, global_state["images"]["image_show"]
 
     form_pretrained_dropdown.change(
         on_change_pretrained_dropdown,
@@ -426,7 +439,7 @@ with gr.Blocks() as app:
         init_images(global_state)
         clear_state(global_state)
 
-        return global_state, global_state['images']['image_show']
+        return global_state, global_state["images"]["image_show"]
 
     form_reset_image.click(
         on_click_reset_image,
@@ -446,7 +459,7 @@ with gr.Blocks() as app:
         init_images(global_state)
         clear_state(global_state)
 
-        return global_state, global_state['images']['image_show']
+        return global_state, global_state["images"]["image_show"]
 
     form_seed_number.change(
         on_change_update_image_seed,
@@ -461,15 +474,17 @@ with gr.Blocks() as app:
         2. Re-init images and clear all state
         """
 
-        global_state['params']['latent_space'] = latent_space
+        global_state["params"]["latent_space"] = latent_space
         init_images(global_state)
         clear_state(global_state)
 
-        return global_state, global_state['images']['image_show']
+        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])
+    form_latent_space.change(
+        on_click_latent_space,
+        inputs=[form_latent_space, global_state],
+        outputs=[global_state, form_image],
+    )
 
     # ==== Params
     form_lambda_number.change(
@@ -480,13 +495,13 @@ with gr.Blocks() as app:
 
     def on_change_lr(lr, global_state):
         if lr == 0:
-            print('lr is 0, do nothing.')
+            print("lr is 0, do nothing.")
             return global_state
         else:
             global_state["params"]["lr"] = lr
-            renderer = global_state['renderer']
+            renderer = global_state["renderer"]
             renderer.update_lr(lr)
-            print('New optimizer: ')
+            print("New optimizer: ")
             print(renderer.w_optim)
         return global_state
 
@@ -507,45 +522,36 @@ with gr.Blocks() as app:
 
         # 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']
+            image_raw = global_state["images"]["image_raw"]
             update_image_draw(
                 image_raw,
-                global_state['points'],
-                global_state['mask'],
-                global_state['show_mask'],
+                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),
+                global_state["images"]["image_show"],
                 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"])
@@ -561,59 +567,51 @@ with gr.Blocks() as app:
                 t_in_pixels.append(p_end)
                 valid_points.append(key_point)
 
-            mask = torch.tensor(global_state['mask']).float()
+            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'
+            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}')
+            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}')
+                print(f"Running time: {time.time() - last_time}")
                 if IS_SPACE and time.time() - last_time > TIMEOUT:
-                    print('Timeout break!')
+                    print("Timeout break!")
                     break
-                if global_state["temporal_params"]["stop"] or global_state['generator_params']["stop"]:
+                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
+                    global_state["generator_params"],
+                    p_to_opt,
+                    t_to_opt,
+                    drag_mask,
+                    global_state["params"]["motion_lambda"],
                     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,
+                    feature_idx=5,
+                    r1=global_state["params"]["r1_in_pixels"],
+                    r2=global_state["params"]["r2_in_pixels"],
+                    trunc_psi=global_state["params"]["trunc_psi"],
                     is_drag=True,
-                    to_pil=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):
+                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],
@@ -623,79 +621,67 @@ with gr.Blocks() as app:
                             t_i[0],
                         ]
                         start_temp = global_state["points"][key_point][
-                            "start_temp"]
-                        print(f'    {start_temp}')
+                            "start_temp"
+                        ]
+                        print(f"    {start_temp}")
 
-                    image_result = global_state['generator_params']['image']
+                    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["points"],
+                        global_state["mask"],
+                        global_state["show_mask"],
                         global_state,
                     )
-                    global_state['images']['image_raw'] = image_result
+                    global_state["images"]["image_raw"] = image_result
 
                 yield (
                     global_state,
                     step_idx,
-                    global_state['images']['image_show'],
-                    # gr.File.update(visible=False),
+                    global_state["images"]["image_show"],
                     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")
+            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,
+            )
 
-            global_state['editing_state'] = 'add_points'
+            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),
+                0,
+                global_state["images"]["image_show"],
                 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),
@@ -710,8 +696,6 @@ with gr.Blocks() as app:
             global_state,
             form_steps_number,
             form_image,
-            # form_download_result_file,
-            # >>> buttons
             form_reset_image,
             enable_add_points,
             enable_add_mask,
@@ -720,8 +704,6 @@ with gr.Blocks() as app:
             form_latent_space,
             form_start_btn,
             form_stop_btn,
-            # <<< buttonm
-            # >>> inputs comps
             form_pretrained_dropdown,
             form_seed_number,
             form_lr_number,
@@ -739,10 +721,12 @@ with gr.Blocks() as app:
 
         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_stop_btn.click(
+        on_click_stop,
+        inputs=[global_state],
+        outputs=[global_state, form_stop_btn],
+        queue=False,
+    )
 
     form_draw_interval_number.change(
         partial(
@@ -771,17 +755,20 @@ with gr.Blocks() as app:
 
     # Mask
     def on_click_reset_mask(global_state):
-        global_state['mask'] = np.ones(
+        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)
+        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(
@@ -798,13 +785,16 @@ with gr.Blocks() as app:
         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))
+        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.
@@ -813,73 +803,89 @@ with gr.Blocks() as app:
         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)
+        global_state["editing_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
+        1. Update 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)
+        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))
+        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)
+    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.')
+        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']
+            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}')
+            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}')
+            points[point_idx + 1] = {"start": xy, "target": None}
+            print(f"Click Image - Start - {xy}")
 
-        image_raw = global_state['images']['image_raw']
+        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["points"],
+            global_state["mask"],
+            global_state["show_mask"],
             global_state,
         )
 
@@ -898,31 +904,38 @@ with gr.Blocks() as app:
         2. re-init network
         2. re-draw image
         """
-        clear_state(global_state, target='point')
+        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)
+        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)
+    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):
+    def on_click_show_mask(global_state, show_mask_val):
         """Function to control whether show mask on image."""
-        global_state['show_mask'] = show_mask
+        global_state["show_mask"] = show_mask_val
 
-        image_raw = global_state['images']['image_raw']
+        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["points"],
+            global_state["mask"],
+            global_state["show_mask"],
             global_state,
         )
         return global_state, image_draw
@@ -937,4 +950,4 @@ with gr.Blocks() as app:
 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)
+app.launch(share=args.share, show_api=False)