import gradio as gr
import os
import shutil
from base64 import b64encode


# Clone the Real-ESRGAN repository (Hugging Face Spaces already runs in a Unix-like environment)
os.system("git clone https://github.com/xinntao/Real-ESRGAN.git")
os.chdir("Real-ESRGAN")
os.system("pip install basicsr facexlib gfpgan ffmpeg-python ipython")
os.system("pip install -r requirements.txt")
os.system("python setup.py develop")

with open('/usr/local/lib/python3.10/site-packages/basicsr/data/degradations.py', 'r') as file:
    file_data = file.read()
    file_data = file_data.replace('from torchvision.transforms.functional_tensor import rgb_to_grayscale', 'from torchvision.transforms.functional import rgb_to_grayscale')
with open('/usr/local/lib/python3.10/site-packages/basicsr/data/degradations.py', 'w') as file:
    file.write(file_data)

from IPython.display import HTML

# Set up folders
upload_folder = 'upload'
result_folder = 'results'

if os.path.isdir(upload_folder):
    shutil.rmtree(upload_folder)
if os.path.isdir(result_folder):
    shutil.rmtree(result_folder)
os.mkdir(upload_folder)
os.mkdir(result_folder)

# Function to perform super-resolution
def super_resolve_video(video_file):
    video_filename = os.path.basename(video_file)
    video_path = os.path.join(upload_folder, video_filename)
    shutil.move(video_file, video_path)
    
    # Run the Real-ESRGAN model
    os.system(f"python inference_realesrgan_video.py -i {video_path} -n realesr-animevideov3 -s 2 --suffix outx2")
    
    output_video_path = video_path.replace(".mp4", "_outx2.mp4")
    return video_path, output_video_path

# Function to display video
def show_video(video_path, video_width=600):
    with open(video_path, "rb") as video_file:
        video_url = f"data:video/mp4;base64,{b64encode(video_file.read()).decode()}"
    return f"""<video width="{video_width}" controls><source src="{video_url}"></video>"""

# Gradio interface
def gradio_interface(video):
    original_video_path, enhanced_video_path = super_resolve_video(video)
    return show_video(original_video_path, video_width=600), show_video(enhanced_video_path, video_width=600)

iface = gr.Interface(
    fn=gradio_interface, 
    inputs=gr.Video(label="Upload Video"), 
    outputs=[gr.HTML(label="Original Video"), gr.HTML(label="Enhanced Video")],
    title="Real-ESRGAN Video Super-Resolution",
    description="Upload a video to enhance its resolution using Real-ESRGAN."
)

iface.launch()