import gradio as gr
import torch
from safetensors.torch import load_file

from model import ImprovedEfficientViT
from inference import predict_vedio

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = ImprovedEfficientViT()
state_dict = load_file("model.safetensors")
model.load_state_dict(state_dict)
model.to(device)
model.eval()

def detect(video):
    result = predict_vedio(video, model)

    if result["class"] == 1:
        return "FAKE VIDEO DETECTED"
    else:
        return "REAL VIDEO"

demo = gr.Interface(
    fn=detect,
    inputs=gr.Video(),
    outputs="text",
    title="Deepfake Detector"
)

demo.launch()