Update app.py

app.py

@@ -1,20 +1,17 @@
 import os
-import cv2
-import glob
 import time
-import torch
 import shutil
 import argparse
-import subprocess
-import onnxruntime
-import numpy as np
-import gradio as gr
 import concurrent.futures
-from tqdm import tqdm
-from moviepy.editor import VideoFileClip
 
-# Import spaces only after dependencies are pinned in requirements.txt
+import cv2
+import numpy as np
+import torch
+import gradio as gr
 import spaces
+import onnxruntime
+from moviepy.editor import VideoFileClip
+from tqdm import tqdm
 
 from face_swapper import Inswapper, paste_to_whole
 from face_analyser import detect_conditions, get_analysed_data, swap_options_list
@@ -23,25 +20,26 @@ 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", 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()
+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)
+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)
+USE_COLAB = args.colab
+DEF_OUTPUT_PATH = args.out_dir
+BATCH_SIZE = int(args.batch_size)
 
 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
@@ -52,6 +50,7 @@ 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)
@@ -156,13 +155,10 @@ def process(
             matrs.extend(batch_matr)
             empty_cache()
 
-        generated_len = len(preds)
-
         if face_enhancer_name != "NONE" and FACE_ENHANCER is not None:
             enhancer_model, enhancer_model_runner = FACE_ENHANCER
             for idx, pred in enumerate(preds):
-                pred = enhancer_model_runner(pred, enhancer_model)
-                preds[idx] = cv2.resize(pred, (512, 512))
+                preds[idx] = cv2.resize(enhancer_model_runner(pred, enhancer_model), (512, 512))
             empty_cache()
 
         if enable_face_parser:
@@ -179,7 +175,7 @@ def process(
                 empty_cache()
             masks = np.concatenate(masks, axis=0) if len(masks) >= 1 else masks
         else:
-            masks = [None] * generated_len
+            masks = [None] * len(preds)
 
         split_preds = split_list_by_lengths(preds, num_faces_per_frame)
         split_matrs = split_list_by_lengths(matrs, num_faces_per_frame)
@@ -265,32 +261,25 @@ with gr.Blocks(css=css) as interface:
         with gr.Column(scale=0.4):
             swap_option = gr.Dropdown(swap_options_list, value=swap_options_list[0], interactive=True, show_label=False)
             age = gr.Number(value=25, interactive=True, visible=False)
-
             detect_condition_dropdown = gr.Dropdown(detect_conditions, label="Condition", value=DETECT_CONDITION)
             detection_size = gr.Number(label="Detection Size", value=DETECT_SIZE)
             detection_threshold = gr.Number(label="Detection Threshold", value=DETECT_THRESH)
-
             output_directory = gr.Text(label="Output Directory", value=DEF_OUTPUT_PATH)
             output_name = gr.Text(label="Output Name", value="Result")
             keep_output_sequence = gr.Checkbox(label="Keep output sequence", value=False)
-
             face_scale = gr.Slider(label="Face Scale", minimum=0, maximum=2, value=1)
             face_enhancer_name = gr.Dropdown(FACE_ENHANCER_LIST, label="Face Enhancer", value="NONE")
-
             enable_face_parser_mask = gr.Checkbox(label="Enable Face Parsing", value=False)
             mask_include = gr.Dropdown(mask_regions.keys(), value=MASK_INCLUDE, multiselect=True, label="Include")
             mask_soft_kernel = gr.Number(label="Soft Erode Kernel", value=MASK_SOFT_KERNEL, visible=False)
             mask_soft_iterations = gr.Number(label="Soft Erode Iterations", value=MASK_SOFT_ITERATIONS)
-
             crop_top = gr.Slider(label="Top", minimum=0, maximum=511, value=0, step=1)
             crop_bott = gr.Slider(label="Bottom", minimum=0, maximum=511, value=511, step=1)
             crop_left = gr.Slider(label="Left", minimum=0, maximum=511, value=0, step=1)
             crop_right = gr.Slider(label="Right", minimum=0, maximum=511, value=511, step=1)
-
             erode_amount = gr.Slider(label="Mask Erode", minimum=0, maximum=1, value=MASK_ERODE_AMOUNT, step=0.05)
             blur_amount = gr.Slider(label="Mask Blur", minimum=0, maximum=1, value=MASK_BLUR_AMOUNT, step=0.05)
             enable_laplacian_blend = gr.Checkbox(label="Laplacian Blending", value=True)
-
             source_image_input = gr.Image(label="Source face", type="filepath", interactive=True)
             input_type = gr.Radio(["Image", "Video"], label="Target Type", value="Image")
             image_input = gr.Image(label="Target Image", interactive=True, type="filepath")
@@ -309,40 +298,18 @@ with gr.Blocks(css=css) as interface:
     for i in range(NUM_OF_SRC_SPECIFIC):
         exec(f"src{i+1} = gr.Image(interactive=True, type='numpy', label='Source Face {i+1}')")
         exec(f"trg{i+1} = gr.Image(interactive=True, type='numpy', label='Specific Face {i+1}')")
-
-    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})")
+    exec("src_specific_inputs = (" + ",".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)]) + ")")
 
     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,
+        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__":
-    interface.queue(concurrency_count=1, max_size=10).launch(server_name="0.0.0.0", share=False)
+    interface.queue().launch(server_name="0.0.0.0", share=False)