import streamlit as st
import os
import subprocess
from PIL import Image
from io import BytesIO

# Function to run the provided script and return the path of the generated video
def run_script(image_path, video_path):
    output_video_path = "/content/roop/swapped.mp4"
    script = """
    cd /content/roop

    !git clone https://github.com/s0md3v/roop.git
    %cd roop

    !wget https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx -O inswapper_128.onnx
    !mkdir models
    !mv inswapper_128.onnx ./models

    !python run.py --target {} --output-video-quality 80 --source {} -o {} --execution-provider cuda --frame-processor face_swapper face_enhancer
    """.format(video_path, image_path, output_video_path)
    
    subprocess.run(script, shell=True, cwd="/content")

    # Return the path of the generated video
    return output_video_path

# Streamlit app
def main():
    st.title("Face Swapper App")

    # File upload
    st.sidebar.header("Upload Files")
    image_file = st.sidebar.file_uploader("Upload Image", type=["jpg", "png", "jpeg"])
    video_file = st.sidebar.file_uploader("Upload Video", type=["mp4"])

    if st.sidebar.button("Swap Faces"):
        if image_file is not None and video_file is not None:
            # Save uploaded files
            image_path = "/content/uploaded_image.jpg"
            video_path = "/content/uploaded_video.mp4"

            with open(image_path, "wb") as f:
                f.write(image_file.read())

            with open(video_path, "wb") as f:
                f.write(video_file.read())

            # Run the script and get the path of the generated video
            generated_video_path = run_script(image_path, video_path)

            # Display the swapped video
            video_bytes = open(generated_video_path, "rb").read()
            st.video(video_bytes)

if __name__ == "__main__":
    main()