Update app.py

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
-from refacer import Refacer
-import argparse
+import cv2
+import multiprocessing
 import os
 import requests
+from refacer import Refacer
 
 # Hugging Face URL to download the model
 model_url = "https://huggingface.co/ofter/4x-UltraSharp/resolve/main/inswapper_128.onnx"
@@ -25,42 +26,67 @@ def download_model():
 # Download the model when the script runs
 download_model()
 
-# Argument parser
-parser = argparse.ArgumentParser(description='Refacer')
-parser.add_argument("--max_num_faces", type=int, help="Max number of faces on UI", default=5)
-parser.add_argument("--force_cpu", help="Force CPU mode", default=False, action="store_true")
-parser.add_argument("--share_gradio", help="Share Gradio", default=False, action="store_true")
-parser.add_argument("--server_name", type=str, help="Server IP address", default="127.0.0.1")
-parser.add_argument("--server_port", type=int, help="Server port", default=7860)
-parser.add_argument("--colab_performance", help="Use in colab for better performance", default=False, action="store_true")
-args = parser.parse_args()
-
-# Initialize the Refacer class
-refacer = Refacer(force_cpu=args.force_cpu, colab_performance=args.colab_performance)
-
-num_faces = args.max_num_faces
-
-# Run function for refacing video
-def run(*vars):
-    video_path = vars[0]
-    origins = vars[1:(num_faces+1)]
-    destinations = vars[(num_faces+1):(num_faces*2)+1]
-    thresholds = vars[(num_faces*2)+1:]
-
-    faces = []
-    for k in range(0, num_faces):
-        if origins[k] is not None and destinations[k] is not None:
-            faces.append({
-                'origin': origins[k],
-                'destination': destinations[k],
-                'threshold': thresholds[k]
-            })
-
-    # Call refacer to process video and get file path
-    refaced_video_path = refacer.reface(video_path, faces)  # refaced video path
+# Initialize Refacer class (force CPU mode)
+refacer = Refacer(force_cpu=True)
+
+# Dummy function to simulate frame-level processing
+def process_frame(frame, origin_face, destination_face, threshold):
+    # Simulate face swapping or any processing needed
+    result_frame = refacer.reface(frame, [{
+        'origin': origin_face,
+        'destination': destination_face,
+        'threshold': threshold
+    }])
+    return result_frame
+
+# Function to process the video in parallel using multiprocessing
+def process_video(video_path, origins, destinations, thresholds, max_processes=2):
+    cap = cv2.VideoCapture(video_path)
+    frames = []
+
+    # Read all frames from the video
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frames.append(frame)
+
+    cap.release()
+
+    # Parallel processing of frames with limited processes (for CPU optimization)
+    with multiprocessing.Pool(processes=max_processes) as pool:
+        processed_frames = pool.starmap(process_frame, [
+            (frame, origins[min(i, len(origins) - 1)], destinations[min(i, len(destinations) - 1)], thresholds[min(i, len(thresholds) - 1)])
+            for i, frame in enumerate(frames)
+        ])
+
+    # Saving the processed frames back into a video
+    output_video_path = "processed_video.mp4"
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Compression using mp4 codec
+    out = cv2.VideoWriter(output_video_path, fourcc, 30.0, (640, 360))  # Reduce resolution to speed up processing
+
+    for frame in processed_frames:
+        out.write(frame)
+
+    out.release()
+    return output_video_path
+
+# Gradio Interface function
+def run(video_path, *vars):
+    # Split the inputs into origins, destinations, and thresholds based on num_faces
+    num_faces = 5  # You can adjust this based on your UI
+    origins = vars[:num_faces]
+    destinations = vars[num_faces:2*num_faces]
+    thresholds = vars[2*num_faces:]
+
+    # Ensure there are no index errors by limiting the number of inputs
+    if len(origins) != num_faces or len(destinations) != num_faces or len(thresholds) != num_faces:
+        return "Please provide input for all faces."
+
+    refaced_video_path = process_video(video_path, origins, destinations, thresholds)
     print(f"Refaced video can be found at {refaced_video_path}")
 
-    return refaced_video_path  # Return the file path to show in Gradio output
+    return refaced_video_path
 
 # Prepare Gradio components
 origin = []
@@ -71,10 +97,10 @@ with gr.Blocks() as demo:
     with gr.Row():
         gr.Markdown("# Refacer")
     with gr.Row():
-        video = gr.Video(label="Original video", format="mp4")
-        video2 = gr.Video(label="Refaced video", interactive=False, format="mp4")
+        video_input = gr.Video(label="Original video", format="mp4")
+        video_output = gr.Video(label="Refaced video", interactive=False, format="mp4")
 
-    for i in range(0, num_faces):
+    for i in range(5):  # Set max faces to 5
         with gr.Tab(f"Face #{i+1}"):
             with gr.Row():
                 origin.append(gr.Image(label="Face to replace"))
@@ -85,7 +111,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         button = gr.Button("Reface", variant="primary")
 
-    button.click(fn=run, inputs=[video] + origin + destination + thresholds, outputs=[video2])
+    button.click(fn=run, inputs=[video_input] + origin + destination + thresholds, outputs=[video_output])
 
 # Launch the Gradio app
-demo.queue().launch(show_error=True, share=args.share_gradio, server_name="0.0.0.0", server_port=args.server_port)
+demo.launch(show_error=True, server_name="0.0.0.0", server_port=7860)