Update app.py

app.py

@@ -5,6 +5,7 @@ import time
 import torch
 import shutil
 import argparse
+import subprocess
 import onnxruntime
 import numpy as np
 import gradio as gr
@@ -12,6 +13,7 @@ import concurrent.futures
 from tqdm import tqdm
 from moviepy.editor import VideoFileClip
 
+# Import spaces only after dependencies are pinned in requirements.txt
 import spaces
 
 from face_swapper import Inswapper, paste_to_whole
@@ -21,13 +23,14 @@ from face_enhancer import get_available_enhancer_names, load_face_enhancer_model
 from utils import trim_video, open_directory, split_list_by_lengths, merge_img_sequence_from_ref, create_image_grid
 
 parser = argparse.ArgumentParser(description="Free Face Swapper")
-parser.add_argument("--out_dir", default=os.getcwd())
-parser.add_argument("--batch_size", default=32)
-parser.add_argument("--cuda", action="store_true", default=False)
-parser.add_argument("--colab", action="store_true", default=False)
-user_args = parser.parse_args()
+parser.add_argument("--out_dir", help="Default Output directory", default=os.getcwd())
+parser.add_argument("--batch_size", help="Gpu batch size", default=32)
+parser.add_argument("--cuda", action="store_true", help="Enable cuda", default=False)
+parser.add_argument("--colab", action="store_true", help="Enable colab mode", default=False)
+user_args, _ = parser.parse_known_args()
 
 USE_COLAB = user_args.colab
+USE_CUDA = False
 DEF_OUTPUT_PATH = user_args.out_dir
 BATCH_SIZE = int(user_args.batch_size)
 
@@ -35,12 +38,10 @@ WORKSPACE = None
 OUTPUT_FILE = None
 PREVIEW = None
 STREAMER = None
-
 DETECT_CONDITION = "best detection"
 DETECT_SIZE = 640
 DETECT_THRESH = 0.6
 NUM_OF_SRC_SPECIFIC = 10
-
 MASK_INCLUDE = ["Skin", "R-Eyebrow", "L-Eyebrow", "L-Eye", "R-Eye", "Nose", "Mouth", "L-Lip", "U-Lip"]
 MASK_SOFT_KERNEL = 17
 MASK_SOFT_ITERATIONS = 10
@@ -51,13 +52,11 @@ FACE_SWAPPER = None
 FACE_ANALYSER = None
 FACE_ENHANCER = None
 FACE_PARSER = None
-
 FACE_ENHANCER_LIST = ["NONE"]
 FACE_ENHANCER_LIST.extend(get_available_enhancer_names())
 FACE_ENHANCER_LIST.extend(cv2_interpolations)
 
 PROVIDER = ["CPUExecutionProvider"]
-USE_CUDA = False
 device = "cpu"
 
 def empty_cache():
@@ -84,17 +83,39 @@ load_face_analyser_model()
 load_face_swapper_model()
 
 @spaces.GPU
-def gpu_process(input_type, image_path, video_path, directory_path, source_path, output_path, output_name,
-                keep_output_sequence, condition, age, distance, face_enhancer_name, enable_face_parser,
-                mask_includes, mask_soft_kernel, mask_soft_iterations, blur_amount, erode_amount, face_scale,
-                enable_laplacian_blend, crop_top, crop_bott, crop_left, crop_right, *specifics):
-
+def process(
+    input_type,
+    image_path,
+    video_path,
+    directory_path,
+    source_path,
+    output_path,
+    output_name,
+    keep_output_sequence,
+    condition,
+    age,
+    distance,
+    face_enhancer_name,
+    enable_face_parser,
+    mask_includes,
+    mask_soft_kernel,
+    mask_soft_iterations,
+    blur_amount,
+    erode_amount,
+    face_scale,
+    enable_laplacian_blend,
+    crop_top,
+    crop_bott,
+    crop_left,
+    crop_right,
+    *specifics,
+):
     global WORKSPACE, OUTPUT_FILE, PREVIEW, FACE_ENHANCER, FACE_PARSER
-
     start_time = time.time()
 
-    def make_status(msg):
-        return f"### \n ⌛ {msg}"
+    def finish_text():
+        mins, secs = divmod(time.time() - start_time, 60)
+        return f"✔️ Completed in {int(mins)} min {int(secs)} sec."
 
     includes = mask_regions_to_list(mask_includes)
     specifics = list(specifics)
@@ -125,7 +146,7 @@ def gpu_process(input_type, image_path, video_path, directory_path, source_path,
             source_data,
             swap_condition=condition,
             detect_condition=DETECT_CONDITION,
-            scale=face_scale
+            scale=face_scale,
         )
 
         preds = []
@@ -146,7 +167,14 @@ def gpu_process(input_type, image_path, video_path, directory_path, source_path,
 
         if enable_face_parser:
             masks = []
-            for batch_mask in get_parsed_mask(FACE_PARSER, preds, classes=includes, device=device, batch_size=BATCH_SIZE, softness=int(mask_soft_iterations)):
+            for batch_mask in get_parsed_mask(
+                FACE_PARSER,
+                preds,
+                classes=includes,
+                device=device,
+                batch_size=BATCH_SIZE,
+                softness=int(mask_soft_iterations),
+            ):
                 masks.append(batch_mask)
                 empty_cache()
             masks = np.concatenate(masks, axis=0) if len(masks) >= 1 else masks
@@ -157,23 +185,27 @@ def gpu_process(input_type, image_path, video_path, directory_path, source_path,
         split_matrs = split_list_by_lengths(matrs, num_faces_per_frame)
         split_masks = split_list_by_lengths(masks, num_faces_per_frame)
 
-        def post_process(frame_idx, frame_img, split_preds, split_matrs, split_masks, enable_laplacian_blend, crop_mask, blur_amount, erode_amount):
+        def post_process(frame_idx, frame_img):
             whole_img = cv2.imread(frame_img)
             blend_method = "laplacian" if enable_laplacian_blend else "linear"
             for p, m, mask in zip(split_preds[frame_idx], split_matrs[frame_idx], split_masks[frame_idx]):
                 p = cv2.resize(p, (512, 512))
                 mask = cv2.resize(mask, (512, 512)) if mask is not None else None
                 m /= 0.25
-                whole_img = paste_to_whole(p, whole_img, m, mask=mask, crop_mask=crop_mask,
-                                           blend_method=blend_method, blur_amount=blur_amount, erode_amount=erode_amount)
+                whole_img = paste_to_whole(
+                    p,
+                    whole_img,
+                    m,
+                    mask=mask,
+                    crop_mask=crop_mask,
+                    blend_method=blend_method,
+                    blur_amount=blur_amount,
+                    erode_amount=erode_amount,
+                )
             cv2.imwrite(frame_img, whole_img)
 
         with concurrent.futures.ThreadPoolExecutor() as executor:
-            futures = [
-                executor.submit(post_process, idx, frame_img, split_preds, split_matrs, split_masks,
-                                enable_laplacian_blend, crop_mask, blur_amount, erode_amount)
-                for idx, frame_img in enumerate(image_sequence)
-            ]
+            futures = [executor.submit(post_process, idx, frame_img) for idx, frame_img in enumerate(image_sequence)]
             for _ in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="Pasting back"):
                 pass
 
@@ -184,7 +216,7 @@ def gpu_process(input_type, image_path, video_path, directory_path, source_path,
         OUTPUT_FILE = output_file
         WORKSPACE = output_path
         PREVIEW = cv2.imread(output_file)[:, :, ::-1]
-        return f"✔️ Completed in {int((time.time() - start_time) // 60)} min.", gr.update(visible=True, value=PREVIEW), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+        return finish_text(), gr.update(visible=True, value=PREVIEW), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
     if input_type == "Video":
         temp_path = os.path.join(output_path, output_name, "sequence")
@@ -213,7 +245,7 @@ def gpu_process(input_type, image_path, video_path, directory_path, source_path,
 
         WORKSPACE = output_path
         OUTPUT_FILE = output_video_path
-        return f"✔️ Completed in {int((time.time() - start_time) // 60)} min.", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True, value=OUTPUT_FILE)
+        return finish_text(), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True, value=OUTPUT_FILE)
 
     return "Unsupported input type", gr.update(), gr.update(), gr.update(), gr.update()
 
@@ -281,12 +313,35 @@ with gr.Blocks(css=css) as interface:
     gen_variable_txt = ",".join([f"src{i+1}" for i in range(NUM_OF_SRC_SPECIFIC)] + [f"trg{i+1}" for i in range(NUM_OF_SRC_SPECIFIC)])
     exec(f"src_specific_inputs = ({gen_variable_txt})")
 
-    swap_inputs = [input_type, image_input, video_input, gr.Text(), source_image_input, output_directory, output_name,
-                   keep_output_sequence, swap_option, age, gr.Number(value=0.6), face_enhancer_name, enable_face_parser_mask,
-                   mask_include, mask_soft_kernel, mask_soft_iterations, blur_amount, erode_amount, face_scale,
-                   enable_laplacian_blend, crop_top, crop_bott, crop_left, crop_right, *src_specific_inputs]
-
-    swap_button.click(fn=gpu_process, inputs=swap_inputs, outputs=[info, preview_image, output_directory_button, output_video_button, preview_video], show_progress=True)
+    swap_inputs = [
+        input_type,
+        image_input,
+        video_input,
+        gr.Text(),
+        source_image_input,
+        output_directory,
+        output_name,
+        keep_output_sequence,
+        swap_option,
+        age,
+        gr.Number(value=0.6),
+        face_enhancer_name,
+        enable_face_parser_mask,
+        mask_include,
+        mask_soft_kernel,
+        mask_soft_iterations,
+        blur_amount,
+        erode_amount,
+        face_scale,
+        enable_laplacian_blend,
+        crop_top,
+        crop_bott,
+        crop_left,
+        crop_right,
+        *src_specific_inputs,
+    ]
+
+    swap_button.click(fn=process, inputs=swap_inputs, outputs=[info, preview_image, output_directory_button, output_video_button, preview_video], show_progress=True)
     cancel_button.click(fn=stop_running, inputs=None, outputs=[info])
 
 if __name__ == "__main__":