import os
from pathlib import Path

import gradio as gr
import spaces
from tribev2.demo_utils import TribeModel, download_file
from tribev2.plotting import PlotBrain

os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = "0" # For whisperx that is run through uvx with an old huggingface_hub version

model = TribeModel.from_pretrained("facebook/tribev2")
plotter = PlotBrain(mesh="fsaverage5")

@spaces.GPU(duration=480)
def predict(video_url: str):
    video_path = Path('cache') / video_url.split('/')[-1]
    download_file(video_url, video_path)
    df = model.get_events_dataframe(video_path=video_path)
    preds, segments = model.predict(events=df)
    return "ok"

demo = gr.Interface(
    predict,
    gr.Text(value="https://download.blender.org/durian/trailer/sintel_trailer-480p.mp4"),
    gr.Text(),
)
demo.launch()