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)