Upload 3 files

app.py

@@ -198,7 +198,9 @@ with gr.Blocks() as app:
                                 form_split_res, form_split_num):
         log.log = io.BytesIO()
         if form_inference_mode == "Square Image":
-            from efficient_inference_for_square_image import parse_args, main_process
+            from efficient_inference_for_square_image import parse_args, main_process, global_state
+            global_state[0] = 1
+
             opt = parse_args()
             opt.transform_mean = [.5, .5, .5]
             opt.transform_var = [.5, .5, .5]
@@ -219,7 +221,9 @@ with gr.Blocks() as app:
                 raise gr.Error("Patches too big. Try to reduce the `split_res`!")
 
         else:
-            from inference_for_arbitrary_resolution_image import parse_args, main_process
+            from inference_for_arbitrary_resolution_image import parse_args, main_process, global_state
+            global_state[0] = 1
+
             opt = parse_args()
             opt.transform_mean = [.5, .5, .5]
             opt.transform_var = [.5, .5, .5]
@@ -240,12 +244,20 @@ with gr.Blocks() as app:
                 raise gr.Error("Patches too big. Try to increase the `split_num`!")
 
 
-    form_start_btn.click(on_click_form_start_btn,
-                         inputs=[form_composite_image, form_mask_image, form_pretrained_dropdown, form_inference_mode,
-                                 form_split_res, form_split_num], outputs=[form_harmonized_image])
+    generate = form_start_btn.click(on_click_form_start_btn,
+                                    inputs=[form_composite_image, form_mask_image, form_pretrained_dropdown,
+                                            form_inference_mode,
+                                            form_split_res, form_split_num], outputs=[form_harmonized_image])
 
 
-    def on_click_form_reset_btn():
+    def on_click_form_reset_btn(form_inference_mode):
+        if form_inference_mode == "Square Image":
+            from efficient_inference_for_square_image import global_state
+            global_state[0] = 0
+        else:
+            from inference_for_arbitrary_resolution_image import global_state
+            global_state[0] = 0
+
         log.log = io.BytesIO()
         return gr.update(value=None), gr.update(value=None, interactive=True), gr.update(value=None,
                                                                                          interactive=False), gr.update(
@@ -253,13 +265,27 @@ with gr.Blocks() as app:
 
 
     form_reset_btn.click(on_click_form_reset_btn,
-                         inputs=None, outputs=[form_log, form_composite_image, form_mask_image, form_start_btn])
+                         inputs=[form_inference_mode],
+                         outputs=[form_log, form_composite_image, form_mask_image, form_start_btn], cancels=generate)
+
+
+    def on_click_form_stop(form_inference_mode):
+        if form_inference_mode == "Square Image":
+            from efficient_inference_for_square_image import global_state
+            global_state[0] = 0
+        else:
+            from inference_for_arbitrary_resolution_image import global_state
+            global_state[0] = 0
+
+        log.log = io.BytesIO()
+        return gr.update(value=None), gr.update(value=None, interactive=True), gr.update(value=None,
+                                                                                         interactive=False), gr.update(
+            interactive=False)
 
-    def on_click_form_stop():
-        gr.close_all()
 
     form_stop_btn.click(on_click_form_stop,
-                         inputs=None, outputs=None)
+                        inputs=[form_inference_mode],
+                        outputs=[form_log, form_composite_image, form_mask_image, form_start_btn], cancels=generate)
 
     gr.Markdown("""
         ## Quick Start

efficient_inference_for_square_image.py

@@ -24,6 +24,7 @@ from utils.misc import normalize
 
 import math
 
+global_state = [1]  # For Gradio Stop Button.
 
 class single_image_dataset(torch.utils.data.Dataset):
     def __init__(self, opt, composite_image=None, mask=None):
@@ -273,6 +274,10 @@ def inference(model, opt, composite_image=None, mask=None):
         fg_INR_coordinates = coordinate_map[1:]
 
         try:
+            if global_state[0] == 0:
+                print("Stop Harmonizing...!")
+                break
+
             if step == 0:  # This is for CUDA Kernel Warm-up, or the first inference step will be quite slow.
                 fg_content_bg_appearance_construct, _, lut_transform_image = model(
                     composite_image,
@@ -317,7 +322,8 @@ def inference(model, opt, composite_image=None, mask=None):
             init_img[start_points[id][0]:start_points[id][0] + singledataset.split_height_resolution,
             start_points[id][1]:start_points[id][1] + singledataset.split_width_resolution] = pred_harmonized_tmp
 
-    print(f'Inference time: {time_all}')
+    if opt.device == "cuda":
+        print(f'Inference time: {time_all}')
     if opt.save_path is not None:
         os.makedirs(opt.save_path, exist_ok=True)
         cv2.imwrite(os.path.join(opt.save_path, "pred_harmonized_image.jpg"), init_img)
@@ -329,7 +335,7 @@ def main_process(opt, composite_image=None, mask=None):
 
     model = build_model(opt).to(opt.device)
 
-    load_dict = torch.load(opt.pretrained, map_location='cpu')['model']
+    load_dict = torch.load(opt.pretrained)['model']
     for k in load_dict.keys():
         if k not in model.state_dict().keys():
             print(f"Skip {k}")

inference_for_arbitrary_resolution_image.py

@@ -24,6 +24,7 @@ from utils.misc import normalize
 
 import math
 
+global_state = [1]  # For Gradio Stop Button.
 
 class single_image_dataset(torch.utils.data.Dataset):
     def __init__(self, opt, composite_image=None, mask=None):
@@ -265,6 +266,10 @@ def inference(model, opt, composite_image=None, mask=None):
         fg_INR_coordinates = coordinate_map[1:]
 
         try:
+            if global_state[0] == 0:
+                print("Stop Harmonizing...!")
+                break
+
             if step == 0:  # This is for CUDA Kernel Warm-up, or the first inference step will be quite slow.
                 fg_content_bg_appearance_construct, _, lut_transform_image = model(
                     composite_image,
@@ -309,7 +314,8 @@ def inference(model, opt, composite_image=None, mask=None):
             init_img[start_points[id][0]:start_points[id][0] + singledataset.split_height_resolution,
             start_points[id][1]:start_points[id][1] + singledataset.split_width_resolution] = pred_harmonized_tmp
 
-    print(f'Inference time: {time_all}')
+    if opt.device == "cuda":
+        print(f'Inference time: {time_all}')
     if opt.save_path is not None:
         os.makedirs(opt.save_path, exist_ok=True)
         cv2.imwrite(os.path.join(opt.save_path, "pred_harmonized_image.jpg"), init_img)
@@ -321,7 +327,7 @@ def main_process(opt, composite_image=None, mask=None):
 
     model = build_model(opt).to(opt.device)
 
-    load_dict = torch.load(opt.pretrained, map_location='cpu')['model']
+    load_dict = torch.load(opt.pretrained)['model']
     for k in load_dict.keys():
         if k not in model.state_dict().keys():
             print(f"Skip {k}")