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

# 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(video_path, *vars):
    origins = vars[:5]
    destinations = vars[5:10]
    thresholds = vars[10:]

    faces = []
    for k in range(5):
        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 refaced video path
    refaced_video_path = refacer.reface(video_path, faces)
    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
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")

    origins, destinations, thresholds = [], [], []
    for i in range(5):
        with gr.Tab(f"Face #{i+1}"):
            with gr.Row():
                origins.append(gr.Image(label="Face to replace"))
                destinations.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] + origins + destinations + thresholds, outputs=[video2])

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