Spaces:
Running on Zero
Running on Zero
| 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 | |
| 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) |