Update app.py

app.py

@@ -13,6 +13,7 @@ import onnxruntime
 import numpy as np
 import gradio as gr
 from tqdm import tqdm
+import concurrent.futures
 from moviepy.editor import VideoFileClip
 
 from nsfw_detector import get_nsfw_detector
@@ -262,17 +263,59 @@ def process(
         torch.cuda.empty_cache()
 
         split_preds = split_list_by_lengths(preds, num_faces_per_frame)
+        del preds
         split_aimgs = split_list_by_lengths(aimgs, num_faces_per_frame)
+        del aimgs
         split_matrs = split_list_by_lengths(matrs, num_faces_per_frame)
+        del matrs
 
         yield "### \n ⌛ Post-processing...", *ui_before()
-        for idx, frame_img in tqdm(enumerate(image_sequence), total=len(image_sequence), desc="Post-Processing"):
+        def process_frame(frame_idx, frame_img, split_preds, split_aimgs, split_matrs, enable_laplacian_blend, crop_top, crop_bott, crop_left, crop_right):
             whole_img_path = frame_img
             whole_img = cv2.imread(whole_img_path)
-            for p, a, m in zip(split_preds[idx], split_aimgs[idx], split_matrs[idx]):
-                whole_img = paste_to_whole(p, a, m, whole_img, laplacian_blend=enable_laplacian_blend, crop_mask=(crop_top,crop_bott,crop_left,crop_right))
+            for p, a, m in zip(split_preds[frame_idx], split_aimgs[frame_idx], split_matrs[frame_idx]):
+                whole_img = paste_to_whole(p, a, m, whole_img, laplacian_blend=enable_laplacian_blend, crop_mask=(crop_top, crop_bott, crop_left, crop_right))
             cv2.imwrite(whole_img_path, whole_img)
 
+        def optimize_processing(image_sequence, split_preds, split_aimgs, split_matrs, enable_laplacian_blend, crop_top, crop_bott, crop_left, crop_right):
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                futures = []
+                for idx, frame_img in enumerate(image_sequence):
+                    future = executor.submit(
+                        process_frame,
+                        idx,
+                        frame_img,
+                        split_preds,
+                        split_aimgs,
+                        split_matrs,
+                        enable_laplacian_blend,
+                        crop_top,
+                        crop_bott,
+                        crop_left,
+                        crop_right
+                    )
+                    futures.append(future)
+
+                for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="Post-Processing"):
+                    try:
+                        result = future.result()
+                    except Exception as e:
+                        print(f"An error occurred: {e}")
+
+        # Usage:
+        optimize_processing(
+            image_sequence,
+            split_preds,
+            split_aimgs,
+            split_matrs,
+            enable_laplacian_blend,
+            crop_top,
+            crop_bott,
+            crop_left,
+            crop_right
+        )
+
+
 
     ## ------------------------------ IMAGE ------------------------------
 
@@ -621,7 +664,7 @@ with gr.Blocks(css=css) as interface:
 
                 with gr.Group():
                     input_type = gr.Radio(
-                        ["Image", "Video"],
+                        ["Image", "Video", "Directory"],
                         label="Target Type",
                         value="Video",
                     )