app.py

import io
import gradio as gr
from refacer import Refacer
import os
import requests
from huggingface_hub import HfApi

# Hugging Face URL to download the model
model_url = "https://huggingface.co/ofter/4x-UltraSharp/resolve/main/inswapper_128.onnx"
model_path = "/home/user/app/inswapper_128.onnx"  # Absolute path for the model in your environment

# Function to download the model if not exists
def download_model():
    if not os.path.exists(model_path):
        print("Downloading the inswapper_128.onnx model...")
        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:
            print(f"Error: Model download failed. Status code: {response.status_code}")
    else:
        print("Model already exists.")

# Download the model when the script runs
download_model()

# Initialize the Refacer class
refacer = Refacer(force_cpu=False, colab_performance=False)

# 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]
            })

    # Specify the output path for the refaced video
    output_dir = "/home/user/app/out"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    refaced_video_path = os.path.join(output_dir, "refaced_video.mp4")

    # Call refacer to process video and get refaced video path
    refacer.reface(video_path, faces, output_path=refaced_video_path)
    print(f"Refaced video can be found at {refaced_video_path}")
    
    # 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 = []

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")

    # Click event: Refacing the video and showing the refaced video in Gradio
    button.click(fn=run, inputs=[video] + origin + destination + thresholds, outputs=[video2])

# Function to upload the refaced video to Hugging Face Spaces
def upload_to_hf(video_path):
    api = HfApi()
    repo_id = "your-username/your-space-name"  # Replace with your actual repository name
    api.upload_file(
        path_or_fileobj=video_path,
        path_in_repo="out/refaced_video.mp4",
        repo_id=repo_id,
        repo_type="space"
    )
    print("Refaced video uploaded to Hugging Face Spaces.")

# Call the upload function after refacing is complete
upload_to_hf(refaced_video_path)

# Launch the Gradio app
demo.queue().launch(show_error=True, server_name="0.0.0.0", server_port=7860)