# app.py
import subprocess
from pathlib import Path
import gradio as gr

# --- Paths ---
DEMO_DIR = Path("demo")                         # Folder with input frames
FRAME1 = DEMO_DIR / "frame1.jpg"                # First image
FRAME2 = DEMO_DIR / "frame2.jpg"                # Second image
RETRO_GIF = DEMO_DIR / "demo.gif"               # Final result

# --- Interpolation function ---
def interpolate(img_a_path, img_b_path):
    """
    Run interpolation pipeline using existing shell scripts and commands.
    Expects two input image paths. Outputs interpolated GIF.
    """
    try:
        # Step 1: Run your inference script to generate intermediate frames
        subprocess.run([
            "python3", "img_inference.py",
            "--img", str(img_a_path), str(img_b_path)
        ], check=True)

        # Step 2: Generate optimized palette using ffmpeg
        subprocess.run([
            "ffmpeg", "-r", "14", "-f", "image2", "-i", "output/img%d.png",
            "-vf", "palettegen=stats_mode=single", "palette.png"
        ], check=True)

        # Step 3: Apply palette to produce final gif
        subprocess.run([
            "ffmpeg", "-r", "14", "-f", "image2", "-i", "output/img%d.png",
            "-i", "palette.png", "-lavfi", "paletteuse", str(RETRO_GIF)
        ], check=True)

        return str(RETRO_GIF)
    except subprocess.CalledProcessError:
        raise gr.Error("Interpolation failed. Please check script and inputs.")

# --- Demo reset: returns static frames ---
def reset_demo():
    return str(FRAME1), str(FRAME2)

# --- Gradio UI ---
with gr.Blocks(title="RIFE Image Interpolation") as demo:
    with gr.Tab("Demo"):
        gr.Markdown("### Interpolate between two images")
        img_a = gr.Image(type="filepath", value=str(FRAME1), label="Image A")
        img_b = gr.Image(type="filepath", value=str(FRAME2), label="Image B")
        run = gr.Button("Interpolate")
        reset = gr.Button("Reset")

        result = gr.Image(type="filepath", label="Interpolated GIF")
        run.click(interpolate, inputs=[img_a, img_b], outputs=result)
        reset.click(reset_demo, outputs=[img_a, img_b])

demo.launch()