Update app.py

app.py

@@ -1,10 +1,8 @@
+import io
 import gradio as gr
 from refacer import Refacer
-import argparse
 import os
 import requests
-import tempfile
-import shutil
 
 # Hugging Face URL to download the model
 model_url = "https://huggingface.co/ofter/4x-UltraSharp/resolve/main/inswapper_128.onnx"
@@ -27,20 +25,8 @@ 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
+refacer = Refacer(force_cpu=False, colab_performance=False)
 
 # Run function for refacing video
 def run(*vars):
@@ -61,11 +47,19 @@ def run(*vars):
     # Call refacer to process video and get refaced video path
     refaced_video_path = refacer.reface(video_path, faces)  # Get refaced video path
     print(f"Refaced video can be found at {refaced_video_path}")
-
-    # Directly return the path to the Gradio UI without using ffmpeg or temp files
-    return refaced_video_path  # Gradio will handle the video display
+    
+    # Convert the output video to memory buffer
+    video_buffer = io.BytesIO()
+    with open(refaced_video_path, "rb") as f:
+        video_buffer.write(f.read())
+    
+    # Rewind the buffer to the beginning
+    video_buffer.seek(0)
+    
+    return video_buffer  # Gradio will handle the video display
 
 # Prepare Gradio components
+num_faces = 5
 origin = []
 destination = []
 thresholds = []
@@ -92,4 +86,4 @@ with gr.Blocks() as demo:
     button.click(fn=run, inputs=[video] + origin + destination + thresholds, outputs=[video2])
 
 # 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(show_error=True, server_name="0.0.0.0", server_port=7860)