app.py

import gradio as gr
from refacer import Refacer
import os
import requests

# Model download URL and path
model_url = "https://huggingface.co/ofter/4x-UltraSharp/resolve/main/inswapper_128.onnx"
model_path = "./inswapper_128.onnx"

# Function to download the model
def download_model():
    if not os.path.exists(model_path):
        print("Downloading inswapper_128.onnx...")
        response = requests.get(model_url)
        if response.status_code == 200:
            with open(model_path, 'wb') as f:
                f.write(response.content)
            print("Model downloaded successfully!")
        else:
            raise Exception(f"Failed to download the model. Status code: {response.status_code}")
    else:
        print("Model already exists.")

# Download the model
download_model()

# Initialize the Refacer class
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
if __name__ == "__main__":
    ui = build_ui()
    ui.queue().launch(server_name="0.0.0.0", server_port=7860, debug=True)