app.py

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / transformers
import time

import cv2
import gradio as gr
import numpy as np
import torch
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

MODEL_ID = "MINT-SJTU/RoboFAC-7B"
FRAME_INTERVAL = 30  # sample one keyframe per second at the dataset's 30 Hz
FRAME_SIZE = 512  # frames are resized to 512x512, matching the paper's eval pipeline
MAX_FRAMES = 32  # safety cap so a very long upload cannot blow the context

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
model.eval()


def extract_keyframes(video_path: str) -> list[np.ndarray]:
    """Decode a video and return one keyframe per second, resized to 512x512 RGB.

    This mirrors the authors' reference pipeline (evaluation/eval_utils.py in
    mint-sjtu/robofac): cv2 decodes the stream, every FRAME_INTERVAL-th frame is
    kept, resized to (512, 512), and JPEG-encoded before being sent to the model.
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError(f"Could not open video: {video_path}")
    frames = []
    count = 0
    ok, frame = cap.read()
    while ok:
        if count % FRAME_INTERVAL == 0:
            frame = cv2.resize(frame, (FRAME_SIZE, FRAME_SIZE))
            # cv2 decodes BGR; imencode expects BGR input, so the reference pipeline
            # produces correctly-colored JPEGs. We convert to RGB for the processor.
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(frame)
            if len(frames) >= MAX_FRAMES:
                break
        count += 1
        ok, frame = cap.read()
    cap.release()
    if not frames:
        raise ValueError(
            "No frames could be decoded from the video. The file may be "
            "corrupt, empty, or in a codec this environment cannot decode "
            "(e.g. AV1). Please try an H.264-encoded MP4."
        )
    return frames


@spaces.GPU(duration=60)
def analyze(
    video: str | None,
    question: str,
    max_new_tokens: int = 512,
    progress=gr.Progress(),
) -> str:
    """Analyze a robot-manipulation video with RoboFAC-7B and answer a question.

    Args:
        video: filepath of an uploaded robot-manipulation video.
        question: the question to ask about the video (task identification,
            failure detection/explanation, correction guidance, ...).
        max_new_tokens: maximum number of new tokens to generate.
        progress: Gradio progress bar (injected automatically).
    Returns:
        The model's textual analysis/answer.
    """
    if not question or not question.strip():
        raise gr.Error("Please enter a question about the video.")
    if not video:
        raise gr.Error("Please upload a video first.")

    t0 = time.perf_counter()
    frames = extract_keyframes(video)
    progress((0.2, 0.9), desc=f"Extracted {len(frames)} keyframes")

    # Reference prompt format: all frames as image blocks + the question in a
    # single user message (no answer-format suffixes for RoboFAC).
    conversation = [
        {
            "role": "user",
            "content": [
                *[{"type": "image", "image": f} for f in frames],
                {"type": "text", "text": question.strip()},
            ],
        }
    ]
    inputs = processor.apply_chat_template(
        conversation,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to("cuda")

    with torch.inference_mode():
        generated = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=False,
        )

    generated = generated[:, inputs["input_ids"].shape[1]:]
    answer = processor.batch_decode(generated, skip_special_tokens=True)[0].strip()
    elapsed = time.perf_counter() - t0
    return f"{answer}\n\n_({len(frames)} keyframes · {elapsed:.1f}s on GPU)_"


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🤖 RoboFAC-7B — Robotic Failure Analysis & Correction

            Upload a video of a robot arm performing a manipulation task and ask
            anything about it: what the task is, whether it succeeded, why it
            failed, or how to correct the failure. The model samples one
            keyframe per second from the video and answers in text.

            **Model**: [`MINT-SJTU/RoboFAC-7B`](https://huggingface.co/MINT-SJTU/RoboFAC-7B) ·
            **Paper**: [RoboFAC](https://huggingface.co/papers/2505.12224) ·
            **Code**: [mint-sjtu/robofac](https://github.com/mint-sjtu/robofac)
            """
        )
        with gr.Row():
            with gr.Column(scale=1):
                video_in = gr.Video(label="Robot manipulation video")
                question_in = gr.Textbox(
                    label="Question",
                    placeholder="e.g. Did the robot finish the task? If not, why did it fail?",
                    lines=2,
                )
                run_btn = gr.Button("Analyze", variant="primary")
                with gr.Accordion("Advanced", open=False):
                    max_new_tokens = gr.Slider(
                        minimum=64, maximum=2048, value=512, step=64,
                        label="Max new tokens",
                    )
            with gr.Column(scale=1):
                answer_out = gr.Textbox(
                    label="Analysis",
                    lines=14,
                    buttons=["copy"],
                )

        run_btn.click(
            analyze,
            inputs=[video_in, question_in, max_new_tokens],
            outputs=answer_out,
            api_name="analyze",
        )

        gr.Examples(
            examples=[
                ["examples/sim_microwave_spoon.mp4",
                 "Please describe the task the robot is performing in the video."],
                ["examples/sim_microwave_spoon.mp4",
                 "Based on the video, provide a detailed explanation of why the robotic arm failed to complete the task."],
                ["examples/realworld_insert_cylinder.mp4",
                 "In the video, the robotic arm executes a task, was it successful?"],
                ["examples/realworld_insert_cylinder.mp4",
                 "Based on the video, an error happened while the robot was executing a task, give detailed low-level instructions to correct the issue and allow the task to be finished."],
            ],
            inputs=[video_in, question_in],
            outputs=answer_out,
            fn=analyze,
            cache_examples=True,
            cache_mode="lazy",
            label="Try the authors' test questions (from the RoboFAC dataset)",
        )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)