Update app.py

app.py

@@ -1,11 +1,9 @@
 import gradio as gr
 from refacer import Refacer
-import argparse
-import tempfile  # Temporary file handling
 import os
 import requests
 
-# Hugging Face URL to download the model
+# Model download URL and path
 model_url = "https://huggingface.co/ofter/4x-UltraSharp/resolve/main/inswapper_128.onnx"
 model_path = "./inswapper_128.onnx"
 
@@ -23,74 +21,72 @@ def download_model():
     else:
         print("Model already exists.")
 
-# Download the model when the script runs
+# Download the model
 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
-
-    # Save the output to a temporary file for Gradio UI
-    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
-        temp_file_path = temp_file.name
-        os.rename(refaced_video_path, temp_file_path)
-
-    return temp_file_path  # Return the file path to show in Gradio output
-
-# Prepare Gradio components
-origin = []
-destination = []
-thresholds = []
-
-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")
-
-    for i in range(0, 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))
-    
-    with gr.Row():
-        button = gr.Button("Reface", variant="primary")
-
-    button.click(fn=run, inputs=[video] + origin + destination + thresholds, outputs=[video2])
+refacer = Refacer(force_cpu=True)  # Use CPU for simplicity
+
+# Function to process the video
+def reface_video(video_path, *faces):
+    try:
+        # Prepare face data
+        face_data = [
+            {
+                "origin": face["origin"],
+                "destination": face["destination"],
+                "threshold": face.get("threshold", 0.2),
+            }
+            for face in faces if face.get("origin") and face.get("destination")
+        ]
+
+        # Process the video
+        print("Processing video...")
+        refaced_video_path = refacer.reface(video_path, face_data)
+
+        # Return only the refaced video
+        return refaced_video_path
+    except Exception as e:
+        return f"Error processing video: {e}"
+
+# Build Gradio UI
+def build_ui():
+    num_faces = 5  # Maximum number of faces
+
+    with gr.Blocks() as demo:
+        with gr.Row():
+            gr.Markdown("# Refacer: AI-Powered Video Face Replacement")
+
+        with gr.Row():
+            input_video = gr.Video(label="Upload a video", type="filepath")
+            output_video = gr.Video(label="Refaced video", interactive=False)
+
+        # Create inputs for multiple faces
+        faces = []
+        for i in range(num_faces):
+            with gr.Tab(f"Face #{i+1}"):
+                with gr.Row():
+                    origin = gr.Image(label="Origin Face", type="filepath")
+                    destination = gr.Image(label="Destination Face", type="filepath")
+                    threshold = gr.Slider(
+                        label="Threshold", minimum=0.0, maximum=1.0, value=0.2
+                    )
+                faces.append({"origin": origin, "destination": destination, "threshold": threshold})
+
+        with gr.Row():
+            submit_button = gr.Button("Reface Video")
+
+        # Connect inputs and outputs
+        inputs = [input_video] + [face["origin"] for face in faces] + [
+            face["destination"] for face in faces
+        ] + [face["threshold"] for face in faces]
+        submit_button.click(
+            fn=reface_video, inputs=inputs, outputs=[output_video]
+        )
+
+    return demo
 
 # 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)
+if __name__ == "__main__":
+    ui = build_ui()
+    ui.queue().launch(server_name="0.0.0.0", server_port=7860, debug=True)