app.py

import gradio as gr
import os
import subprocess
import shutil

# ===== RIFE 모델 가중치 자동 다운로드 =====
if not os.path.exists("train_log"):
    os.makedirs("train_log", exist_ok=True)
    subprocess.run(["wget", "-O", "train_log/RIFE_HDv3.pkl", "https://huggingface.co/datasets/wnshi/RIFE/resolve/main/train_log/RIFE_HDv3.pkl"])
    subprocess.run(["wget", "-O", "train_log/RIFE_HDv2.pkl", "https://huggingface.co/datasets/wnshi/RIFE/resolve/main/train_log/RIFE_HDv2.pkl"])
    subprocess.run(["wget", "-O", "train_log/RIFE_HD.pkl", "https://huggingface.co/datasets/wnshi/RIFE/resolve/main/train_log/RIFE_HD.pkl"])

# ===== GitHub에서 RIFE 원본 코드 가져오기 =====
REPO_URL = "https://github.com/kimjongin0325-cyber/ECCV2022-RIFE.git"
REPO_DIR = "ECCV2022-RIFE"

if not os.path.exists(REPO_DIR):
    print("🔄 Cloning RIFE repository from GitHub...")
    os.system(f"git clone {REPO_URL}")

# ===== RIFE 실행 함수 =====
def process_video(video_file, scale=2.0, codec="OpenH264"):
    if video_file is None:
        raise gr.Error("Please upload a video file.")

    input_path = "input_video.mp4"
    output_path = "output_interp.mp4"

    # 업로드된 비디오를 저장
    shutil.copy(video_file, input_path)

    # ECCV2022-RIFE 실행 명령어
    cmd = (
        f"python {REPO_DIR}/inference_video.py "
        f"--exp=1 --scale={scale} "
        f"--video={input_path} --output={output_path}"
    )

    print("🚀 Running command:", cmd)
    result = subprocess.run(cmd, shell=True, text=True, capture_output=True)

    if result.returncode != 0:
        print("❌ Error:", result.stderr)
        raise gr.Error(f"RIFE execution failed:\n{result.stderr}")

    if not os.path.exists(output_path):
        raise gr.Error("Output video was not generated.")

    return output_path

# ===== Gradio UI 구성 =====
with gr.Blocks() as demo:
    gr.Markdown("## 🌀 RIFE Frame Interpolation (ECCV 2022 Version)")
    gr.Markdown("Upload a video and generate interpolated frames using RIFE (Kimjongin0325 Repository).")

    with gr.Row():
        with gr.Column():
            input_video = gr.Video(label="Upload a video")
            scale = gr.Slider(1.0, 4.0, value=2.0, step=0.5, label="Scale Factor")
            codec = gr.Radio(["OpenH264", "H265", "AV1"], value="OpenH264", label="Codec")
            submit_btn = gr.Button("🪄 Interpolate Frames")

        with gr.Column():
            output_video = gr.Video(label="Interpolated Result")

    submit_btn.click(
        fn=process_video,
        inputs=[input_video, scale, codec],
        outputs=output_video
    )

demo.launch()