Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 from refacer import Refacer
-import argparse
 import os
 import requests
 
@@ -25,65 +24,60 @@ 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)
+refacer = Refacer(force_cpu=False, colab_performance=False)
 
-num_faces = args.max_num_faces
+# Number of faces to handle
+num_faces = 5
 
 # Run function for refacing video
-def run(video_path, *vars):
-    origins = vars[:num_faces]
-    destinations = vars[num_faces:2*num_faces]
-    thresholds = vars[2*num_faces:]
+def run(video_path, *inputs):
+    origins = inputs[:num_faces]
+    destinations = inputs[num_faces:2*num_faces]
+    thresholds = inputs[2*num_faces:]
 
     faces = []
-    for k in range(num_faces):
-        if origins[k] is not None and destinations[k] is not None:
+    for i in range(num_faces):
+        if origins[i] is not None and destinations[i] is not None:
             faces.append({
-                'origin': origins[k],
-                'destination': destinations[k],
-                'threshold': thresholds[k]
+                "origin": origins[i],
+                "destination": destinations[i],
+                "threshold": thresholds[i]
             })
 
-    # Call refacer to process video
-    refaced_video = refacer.reface(video_path, faces)  # Directly returns the processed video as bytes
-    return refaced_video  # Return video content to display in Gradio UI
-
-# Prepare Gradio components
-origin = []
-destination = []
-thresholds = []
+    # Perform refacing
+    refaced_video_path = refacer.reface(video_path, faces)
+    return refaced_video_path  # Return the path of the refaced video
 
+# Define Gradio UI
 with gr.Blocks() as demo:
+    gr.Markdown("# Refacer - Replace Faces in Videos")
+
+    # Input and output components
     with gr.Row():
-        gr.Markdown("# Refacer - Replace Faces in Video")
-    with gr.Row():
-        video = gr.Video(label="Original Video", format="mp4")
-        video2 = gr.Video(label="Refaced Video", format="mp4", interactive=False)
+        input_video = gr.Video(label="Original Video", format="mp4")
+        output_video = gr.Video(label="Refaced Video", interactive=False, format="mp4")
 
+    # Face replacement tabs
+    origins = []
+    destinations = []
+    thresholds = []
     for i in range(num_faces):
         with gr.Tab(f"Face #{i+1}"):
-            with gr.Row():
-                origin.append(gr.Image(label="Face to replace"))
-                destination.append(gr.Image(label="Destination face"))
-            with gr.Row():
-                thresholds.append(gr.Slider(label="Threshold", minimum=0.0, maximum=1.0, value=0.2))
-    
+            origins.append(gr.Image(label="Origin Face"))
+            destinations.append(gr.Image(label="Destination Face"))
+            thresholds.append(gr.Slider(label="Threshold", minimum=0.0, maximum=1.0, value=0.2))
+
+    # Reface button
     with gr.Row():
-        button = gr.Button("Reface", variant="primary")
+        reface_button = gr.Button("Reface")
 
-    # Updated: Directly display the refaced video in the output
-    button.click(fn=run, inputs=[video] + origin + destination + thresholds, outputs=[video2])
+    # Connect the function to the Gradio UI
+    reface_button.click(
+        fn=run,
+        inputs=[input_video] + origins + destinations + thresholds,
+        outputs=output_video
+    )
 
 # 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.queue().launch(share=True)