"""OmniAgent-RL: Native Active Perception as Reasoning for Omni-Modal Understanding.
This demo implements the agentic Observation-Thought-Action (OTA) loop described in the
OmniAgent paper. The model (Qwen2.5-Omni-7B fine-tuned with agentic RL) iteratively
requests frames, audio clips, or video clips from a video to answer a question.
Paper: https://huggingface.co/papers/2606.19341
Code: https://github.com/harryhsing/OmniAgent
"""
import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "False")
import spaces # MUST be first
import json
import math
import re
import shlex
import shutil
import subprocess
import tempfile
import time
import uuid
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import numpy as np
import torch
import gradio as gr
from transformers import AutoProcessor, Qwen2_5OmniForConditionalGeneration
from qwen_vl_utils import process_vision_info
from qwen_omni_utils import process_audio_info
# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
MODEL_ID = "harryhsing/OmniAgent-RL-7B"
MAX_STEPS_DEFAULT = 32
MAX_FRAMES = 60
MAX_AUDIO_LEN = 300.0
MAX_CLIP_LEN = 60.0
KEEP_RECENT_MEDIA = 1 # how many recent media turns to keep before compressing
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
.step-card {
border: 1px solid #e0e0e0; border-radius: 12px; padding: 14px; margin-bottom: 12px;
background: #fafafa;
}
.dark .step-card { background: #1a1a2e; border-color: #333; }
.step-header { font-weight: bold; font-size: 14px; margin-bottom: 8px; }
.step-think { color: #555; font-size: 13px; margin: 4px 0; }
.dark .step-think { color: #aaa; }
.step-action { font-family: monospace; font-size: 12px; color: #1a73e8; }
.dark .step-action { color: #64b5f6; }
.step-obs { font-size: 13px; color: #333; margin: 4px 0; }
.dark .step-obs { color: #ccc; }
"""
# ---------------------------------------------------------------------------
# System prompt (from OmniAgent's video_prompt.py)
# ---------------------------------------------------------------------------
SYSTEM_PROMPT = """You are the **Deep-Omni-Research Agent**, a specialized multi-modal analyst for temporal forensic investigation. Your goal is to solve complex queries by meticulously inspecting video and audio data through a step-by-step "Observe-Think-Action" loop.
============== GLOBAL OPERATING RULES ==============
- **META-Validation**: The first message provides "Video META" (duration, fps, has_audio). Validate every timestamp against these limits.
- **Audio Constraint**: If `has_audio` is false, the `get_audio` action is FORBIDDEN. Skip audio analysis and rely on visual cues only.
- **Media Persistence**: Once media is returned, it becomes a TEXT PLACEHOLDER in the next turn.
* **Frame Placeholder Example**: "Frames 10.00s-12.00s (num=5). Timestamps: [10.00s, 10.50s, 11.00s, 11.50s, 12.00s] [MEDIA OMITTED - Refer to your Observation]"
* **AUDIO/CLIP Placeholder Example**: "Audio 10.00s-20.00s [MEDIA OMITTED - Refer to your Observation]"
- **The "Memory" Requirement**: Your `observation` must be an exhaustive, high-fidelity log. Once media is omitted, you will "forget" any detail not recorded here.
- **Strategic Efficiency**: DO NOT request the exact same action and range twice. However, you are encouraged to re-inspect important ranges via different modalities (e.g., `get_clip` after `get_audio`) or higher density (Zooming in) to extract NEW forensic details.
- **Strict Fidelity**: You MUST use exact timestamps (including decimals) provided in environment labels (e.g., 481.84s). Never round or approximate numbers.
- **Evidence Traceability**: You MUST prefix findings with the **Full Evidence ID** (e.g., "[Frames 10.0s-12.0s (num=5)]") in both `observation` and `think` fields.
- **Environment Feedback**: Pay attention to `[ERROR]` and `[NOTICE]` (remaining steps). Adjust your strategy immediately.
========== STRATEGIC INSPECTION GUIDELINES ==========
1. **Visual Search (get_frames)**: (Max {max_frames} frames).
- **Scanning**: Use wide ranges (e.g., start=0, end=duration, num={max_frames}) to discover the overall timeline and identify key milestones or potential scene cuts.
- **Precision**: Use narrow windows (1-2s) with high `num` for micro-details (logos, text, fast motions, or subtle object state changes).
2. **Counting & Re-ID**: Assign approximate spatial locations [y, x] (0-100 scale; [0,0] is top-left) to each unique instance (e.g., "Person_A at [20, 45]") in your `observation`. This spatial ID prevents re-counting the same object across different frames/steps.
3. **Temporal Bisection**: Find 'start' and 'end' boundary frames where a state changes, then iteratively narrow the interval to locate the exact transition second or frame.
4. **Audio Analysis (get_audio)**: (Max {max_audio}s).
- **Verbatim Logging**: Identify speakers and transcribe speech near-verbatim. **CRITICAL**: Do not paraphrase or infer words to fit your hypothesis.
- **Acoustic Context**: Identify critical off-screen or background sounds (e.g., footsteps, sirens, clicks) that provide environmental clues for temporal reasoning.
5. **Multi-Modal Action Analysis (get_clip)**: (Max {max_clip}s).
- **Action & Temporal Dynamics**: Analyze the nature of movement (speed, direction, continuity) and precise sequencing to solve "Who moved first?" or "Was the motion deliberate?".
- **Process Logic**: Use when the continuous *process* of a state change (e.g., an object falling) is more critical than discrete start/end points.
- **Audio-Visual Synergy (Conditional)**: If `has_audio` is true, perform high-fidelity forensic matching (Sync, Active Speaker ID, Causality with time-lag).
====================== ACTIONS ======================
Exactly ONE action per turn in valid JSON:
1. {{"type": "get_frames", "start": float, "end": float, "num": int}}
2. {{"type": "get_audio", "start": float, "end": float}}
3. {{"type": "get_clip", "start": float, "end": float}}
4. {{"type": "answer", "content": "string"}}
- **MCQ**: Letter only (e.g., "A").
- **TR**: JSON array of one or more pairs, e.g., "[[10.5, 20.0], [35.0, 40.0]]".
- **NUM/SIZE**: A single number string, e.g., "10.3".
- **FF**: Detailed descriptive text.
============= STRICT EXECUTION PROTOCOL =============
- **Forensic Rigor**: Answering incorrectly is a failure. Rule out every possible distractor before concluding.
- **The Confidence Gatekeeper**: **You MUST include a numeric `confidence` field (0.0-1.0) as a top-level JSON key.** This represents your assessment of whether the evidence is sufficient to conclude.
- **The "0.9" Behavioral Rule**: You should only initiate the "answer" action when your `confidence` is >= 0.9. If it is lower, continue gathering evidence unless `[NOTICE]` indicates "FINAL STEP".
- **Evidence Contradiction**: In your `think` field, actively look for evidence that *disproves* your current leading hypothesis.
- **Deadline Management**: In "FINAL STEP", bypass the 0.9 threshold and provide your best-informed `answer` immediately.
=================== OUTPUT SCHEMA ===================
The response must contain **ONLY the JSON object itself**. Any text outside the curly braces ({{ }})—including thoughts, explanations, or markdown fences (```json)—is strictly forbidden and will result in system failure.
{{"observation": "[Clip 00.00s-00.00s] (T: 00.00s)[Obj_A at y,x] visual_detail. [Audio 00.0s-00.0s] exact_audio_log. [Key Fact]: forensic_finding.", "think": "Evidence Review: [Clip 00.00s-00.00s] confirms_or_contradicts [Frames 00.00s-00.00s (num=0)]. Gap Analysis: missing_or_ambiguous_details. Deduction: logical_path_to_action_or_answer.", "confidence": 0.0, "action": {{"type": "get_frames|get_audio|get_clip|answer", "start": 0.0, "end": 0.0, "num": 0, "content": ""}}}}
============= CRITICAL FORMATTING RULES =============
- **Physical Boundary**: Your entire response MUST start with '{{' and end with '}}' exactly.
- **The "One-Line" Mandate**: Your entire output MUST be ONE single line of text. NO newlines (\\n) allowed anywhere.
- **NO Markdown**: Output raw text ONLY. DO NOT use code blocks or wrappers.
""".format(max_frames=MAX_FRAMES, max_audio=MAX_AUDIO_LEN, max_clip=MAX_CLIP_LEN)
# ---------------------------------------------------------------------------
# Model loading (module scope — ZeroGPU rule)
# ---------------------------------------------------------------------------
print("[OmniAgent] Loading model…")
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
attn_implementation="sdpa",
trust_remote_code=True,
).to("cuda").eval()
# Disable the talker to save VRAM — we only use the thinker for text generation.
model.disable_talker = True
print("[OmniAgent] Model loaded.")
# ---------------------------------------------------------------------------
# Video helpers (ffmpeg)
# ---------------------------------------------------------------------------
def probe_video(video_path: str) -> Tuple[float, float, bool]:
"""Return (duration, fps, has_audio) via ffprobe."""
try:
r = subprocess.run(
["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
"-of", "default=nw=1", video_path],
capture_output=True, text=True, timeout=10,
)
duration = float(r.stdout.strip())
except Exception:
duration = 120.0
try:
r = subprocess.run(
["ffprobe", "-v", "quiet", "-show_entries", "stream=r_frame_rate",
"-select_streams", "v:0", "-of", "default=nw=1", video_path],
capture_output=True, text=True, timeout=10,
)
fps = eval(r.stdout.strip())
except Exception:
fps = 30.0
try:
r = subprocess.run(
["ffprobe", "-v", "quiet", "-show_entries", "stream=codec_type",
"-of", "default=nw=1", video_path],
capture_output=True, text=True, timeout=10,
)
has_audio = "audio" in r.stdout
except Exception:
has_audio = True
return duration, fps, has_audio
def extract_frame(video_path: str, ts: float, out_dir: str, step: int, duration: float) -> str:
"""Extract a single frame at timestamp *ts*."""
ts = max(0.0, ts)
if duration > 0:
ts = min(ts, duration - 0.2)
out = os.path.join(out_dir, f"step{step}_frame_{ts:.3f}.jpg")
cmd = (
f"ffmpeg -hide_banner -loglevel error -nostdin -y -threads 4 "
f"-ss {ts:.3f} -i {shlex.quote(video_path)} -frames:v 1 -q:v 2 {shlex.quote(out)}"
)
subprocess.run(cmd, shell=True, capture_output=True, timeout=30)
if not os.path.isfile(out) or os.path.getsize(out) == 0:
# fallback: seek from end
if duration > 0 and ts >= duration - 3:
for off in [0.5, 1.0, 2.0]:
if off < duration:
cmd2 = (
f"ffmpeg -hide_banner -loglevel error -nostdin -y -threads 4 "
f"-sseof -{off} -i {shlex.quote(video_path)} "
f"-frames:v 1 -q:v 2 {shlex.quote(out)}"
)
subprocess.run(cmd2, shell=True, capture_output=True, timeout=30)
if os.path.isfile(out) and os.path.getsize(out) > 0:
return out
return out
def extract_audio(video_path: str, start: float, end: float, out_dir: str, step: int, duration: float) -> str:
"""Extract audio segment [start, end] as wav."""
start = max(0.0, start)
end = min(end, duration - 0.2) if duration > 0 else end
dur = end - start
if dur <= 0:
raise ValueError("Invalid audio range")
out = os.path.join(out_dir, f"step{step}_audio_{start:.3f}_{end:.3f}.wav")
cmd = (
f"ffmpeg -hide_banner -loglevel error -nostdin -y -threads 4 "
f"-ss {start:.3f} -i {shlex.quote(video_path)} "
f"-ss 0 -t {dur:.3f} "
f"-map 0:a:0? -vn -ac 1 -ar 16000 -c:a pcm_s16le "
f"{shlex.quote(out)}"
)
subprocess.run(cmd, shell=True, capture_output=True, timeout=60)
return out
def extract_clip(video_path: str, start: float, end: float, out_dir: str, step: int, duration: float) -> str:
"""Extract a video clip [start, end] as mp4."""
start = max(0.0, start)
end = min(end, duration - 0.2) if duration > 0 else end
dur = end - start
if dur <= 0:
raise ValueError("Invalid clip range")
out = os.path.join(out_dir, f"step{step}_clip_{start:.3f}_{end:.3f}.mp4")
cmd = (
f"ffmpeg -hide_banner -loglevel error -nostdin -y -threads 4 "
f"-ss {start:.3f} -i {shlex.quote(video_path)} "
f"-ss 0 -t {dur:.3f} "
f"-map 0:v:0? -map 0:a:0? "
f"-c:v libx264 -pix_fmt yuv420p -preset superfast -crf 20 "
f"-movflags +faststart "
f"-c:a aac -b:a 128k -ar 48000 "
f"{shlex.quote(out)}"
)
subprocess.run(cmd, shell=True, capture_output=True, timeout=120)
return out
# ---------------------------------------------------------------------------
# Action parsing
# ---------------------------------------------------------------------------
def parse_response(raw: str) -> Optional[dict]:
"""Parse the model's JSON response. Returns None on failure."""
s = raw.strip()
# Strip code fences
if s.startswith("```") or s.endswith("```"):
s = s.strip("`").strip()
if s.startswith("json"):
s = s[4:].strip()
if not (s.startswith("{") and s.endswith("}")):
return None
try:
obj, end_pos = json.JSONDecoder().raw_decode(s)
except json.JSONDecodeError:
return None
if not isinstance(obj, dict):
return None
if "think" not in obj or "action" not in obj:
return None
return obj
# ---------------------------------------------------------------------------
# OTA environment
# ---------------------------------------------------------------------------
class OTAEnvironment:
"""Lightweight reimplementation of the OmniAgent video environment."""
def __init__(self, video_path: str, question: str, q_type: str,
options: Optional[List[str]], answer: str,
max_steps: int = MAX_STEPS_DEFAULT,
max_frames: int = MAX_FRAMES,
max_audio_len: float = MAX_AUDIO_LEN,
max_clip_len: float = MAX_CLIP_LEN):
self.video_path = video_path
self.question = question
self.q_type = q_type
self.options = options or []
self.answer = answer
self.max_steps = max_steps
self.max_frames = max_frames
self.max_audio_len = max_audio_len
self.max_clip_len = max_clip_len
self.duration, self.fps, self.has_audio = probe_video(video_path)
self.step_count = 0
self.done = False
self.history: List[dict] = []
self.temp_dir = tempfile.mkdtemp(prefix="omniagent_")
self.last_frames: List[str] = []
self.last_clip: Optional[str] = None
self.last_audio: Optional[str] = None
self.final_answer: str = ""
def build_initial_messages(self) -> List[dict]:
"""Build the initial system + user messages."""
def trunc(x, n=2):
if not isinstance(x, (int, float)):
return "unknown"
return f"{math.floor(x * 10**n) / 10**n:.{n}f}"
meta = (
f"Video META:\n- duration_seconds: {trunc(self.duration)}\n"
f"- fps: {trunc(self.fps)}\n- has_audio: {self.has_audio}\n\n"
)
if self.q_type == "MCQ":
opts = (
"\nOptions:\n" + "\n".join(self.options) +
"\nWhen answering, set action.content to ONE uppercase letter (A, B, C …)."
)
qtext = meta + "Question: " + self.question + opts
elif self.q_type == "TR":
guide = (
"\nWhen answering, set action.content to a JSON array "
"of timestamp pairs such as [[10.5, 20.0]]."
)
qtext = meta + "Question: " + self.question + guide
elif self.q_type == "FF":
guide = "\nWhen answering, set action.content to **your free-form answer text**."
qtext = meta + "Question: " + self.question + guide
elif self.q_type in ("NUM", "SIZE"):
guide = "\nWhen answering, set action.content to **ONE number**, e.g. 42 or 3.14159."
qtext = meta + "Question: " + self.question + guide
else:
qtext = meta + "Question: " + self.question
self.history = [
{"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]},
{"role": "user", "content": [{"type": "text", "text": qtext}]},
]
self._append_step_notice()
return self.history
def _append_step_notice(self):
"""Append a step notice to the last user message."""
remaining = self.max_steps - self.step_count
if remaining <= 1:
notice = "\n[NOTICE] FINAL STEP."
else:
notice = f"\n[NOTICE] Step {self.step_count + 1}/{self.max_steps}. {remaining - 1} steps remaining."
# Append to the last user message
if self.history and self.history[-1]["role"] == "user":
content = self.history[-1]["content"]
if isinstance(content, list) and content and content[-1].get("type") == "text":
content[-1]["text"] += notice
else:
content.append({"type": "text", "text": notice})
def _replace_old_media(self, keep_recent: int = KEEP_RECENT_MEDIA):
"""Compress older media in history to text placeholders."""
media_kept = 0
SUFFIX = "[MEDIA OMITTED - Refer to your Observation]"
for i in range(len(self.history) - 1, -1, -1):
msg = self.history[i]
if msg.get("role") != "user" or not isinstance(msg.get("content"), list):
continue
has_image = any(p.get("type") == "image" for p in msg["content"])
has_other = any(p.get("type") in ("video", "audio") for p in msg["content"])
if not (has_image or has_other):
continue
media_kept += 1
if media_kept <= keep_recent:
continue
# Compress
raw_header = "Media content"
if msg["content"] and msg["content"][0].get("type") == "text":
raw_header = msg["content"][0]["text"].strip()
if has_image:
all_ts = []
for p in msg["content"][1:]:
if p.get("type") == "text":
found = re.findall(r"(\d+(?:\.\d+)?)s", p.get("text", ""))
if found:
all_ts.extend(found)
ts_str = ", ".join([f"{float(x):.2f}s" for x in all_ts])
new_text = f"{raw_header} Timestamps: [{ts_str}] {SUFFIX}"
else:
new_text = f"{raw_header} {SUFFIX}"
msg["content"] = [{"type": "text", "text": new_text}]
def step(self, raw_response: str) -> Tuple[bool, str, List[str], Optional[str], Optional[str]]:
"""Process the model's response, execute the action, and update history.
Returns (done, action_type, frame_paths, clip_path, audio_path).
"""
self.history.append({
"role": "assistant",
"content": [{"type": "text", "text": raw_response}]
})
if self.done:
return True, "done", [], None, None
self.step_count += 1
parsed = parse_response(raw_response)
if parsed is None:
self.history.append({
"role": "user",
"content": [{"type": "text", "text": "[ERROR] Invalid JSON. Please output exactly ONE line of valid JSON matching the schema."}]
})
self._append_step_notice()
return False, "error", [], None, None
action = parsed.get("action", {})
atype = action.get("type", "")
frames, clip, audio = [], None, None
# Step limit check
if self.step_count > self.max_steps or (self.step_count == self.max_steps and atype != "answer"):
self.history.append({
"role": "user",
"content": [{"type": "text", "text": "[ERROR] Step limit reached. You must answer now."}]
})
# Force answer
self.done = True
self.final_answer = parsed.get("think", "")
return True, "forced_answer", [], None, None
if self.step_count == 1 and atype == "answer":
self.history.append({
"role": "user",
"content": [{"type": "text", "text": "[ERROR] You must gather evidence before answering. Use get_frames, get_audio, or get_clip first."}]
})
self._append_step_notice()
return False, "early_answer", [], None, None
if atype == "get_frames":
s = float(action.get("start", 0))
e = float(action.get("end", 0))
num = int(action.get("num", 0))
if s < 0 or e > self.duration or e <= s or num < 1 or num > self.max_frames:
self.history.append({
"role": "user",
"content": [{"type": "text", "text": f"[ERROR] Invalid get_frames params: start={s}, end={e}, num={num}. Duration={self.duration:.2f}s."}]
})
self._append_step_notice()
return False, "error", [], None, None
ts_list = np.linspace(s, e, num).tolist() if num > 1 else [s]
header = f"Frames {s:.2f}s-{e:.2f}s (num={num})."
parts = [{"type": "text", "text": header}]
for t in ts_list:
try:
img = extract_frame(self.video_path, t, self.temp_dir, self.step_count, self.duration)
if os.path.isfile(img) and os.path.getsize(img) > 0:
parts.append({"type": "text", "text": f"Frame {t:.2f}s:"})
parts.append({"type": "image", "image": img})
frames.append(img)
except Exception:
pass
self.history.append({"role": "user", "content": parts})
self.last_frames = frames
elif atype == "get_audio":
if not self.has_audio:
self.history.append({
"role": "user",
"content": [{"type": "text", "text": "[ERROR] This video has no audio stream."}]
})
self._append_step_notice()
return False, "error", [], None, None
s = float(action.get("start", 0))
e = float(action.get("end", 0))
if s < 0 or e > self.duration or e <= s:
self.history.append({
"role": "user",
"content": [{"type": "text", "text": f"[ERROR] Invalid audio range: {s}-{e}. Duration={self.duration:.2f}s."}]
})
self._append_step_notice()
return False, "error", [], None, None
try:
audio = extract_audio(self.video_path, s, e, self.temp_dir, self.step_count, self.duration)
if os.path.isfile(audio) and os.path.getsize(audio) > 0:
self.history.append({
"role": "user",
"content": [
{"type": "text", "text": f"Audio {s:.2f}s-{e:.2f}s"},
{"type": "audio", "audio": audio},
]
})
self.last_audio = audio
else:
self.history.append({
"role": "user",
"content": [{"type": "text", "text": f"[ERROR] Failed to extract audio {s:.2f}s-{e:.2f}s."}]
})
except Exception as ex:
self.history.append({
"role": "user",
"content": [{"type": "text", "text": f"[ERROR] Audio extraction failed: {ex}"}]
})
elif atype == "get_clip":
s = float(action.get("start", 0))
e = float(action.get("end", 0))
if s < 0 or e > self.duration or e <= s:
self.history.append({
"role": "user",
"content": [{"type": "text", "text": f"[ERROR] Invalid clip range: {s}-{e}. Duration={self.duration:.2f}s."}]
})
self._append_step_notice()
return False, "error", [], None, None
try:
clip = extract_clip(self.video_path, s, e, self.temp_dir, self.step_count, self.duration)
if os.path.isfile(clip) and os.path.getsize(clip) > 0:
self.history.append({
"role": "user",
"content": [
{"type": "text", "text": f"Clip {s:.2f}s-{e:.2f}s"},
{"type": "video", "video": clip},
]
})
self.last_clip = clip
else:
self.history.append({
"role": "user",
"content": [{"type": "text", "text": f"[ERROR] Failed to extract clip {s:.2f}s-{e:.2f}s."}]
})
except Exception as ex:
self.history.append({
"role": "user",
"content": [{"type": "text", "text": f"[ERROR] Clip extraction failed: {ex}"}]
})
elif atype == "answer":
content = action.get("content", "")
self.final_answer = content
self.done = True
self.history.append({
"role": "user",
"content": [{"type": "text", "text": f"[ANSWER] {content}"}]
})
return True, "answer", [], None, None
else:
self.history.append({
"role": "user",
"content": [{"type": "text", "text": f"[ERROR] Unknown action type: {atype}"}]
})
# Compress old media
self._replace_old_media()
# Append step notice
self._append_step_notice()
return self.done, atype, frames, clip, audio
def cleanup(self):
if self.temp_dir and os.path.exists(self.temp_dir):
shutil.rmtree(self.temp_dir, ignore_errors=True)
# ---------------------------------------------------------------------------
# Model generation
# ---------------------------------------------------------------------------
def generate_response(messages: List[dict], has_audio: bool) -> str:
"""Run the model to generate a single OTA response."""
# Build prompt text
prompt = processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
if isinstance(prompt, list):
prompt = prompt[0] if prompt else ""
# Process multi-modal inputs
imgs, vids, video_kwargs = process_vision_info(messages, return_video_kwargs=True)
audios = None
if has_audio:
try:
audios = process_audio_info(messages, use_audio_in_video=True)
except Exception:
audios = None
# Build processor kwargs
proc_kwargs = {"text": [prompt], "return_tensors": "pt"}
if imgs:
proc_kwargs["images"] = imgs
if vids:
proc_kwargs["videos"] = vids
if video_kwargs and vids:
for key, val in video_kwargs.items():
if val is None:
continue
if isinstance(val, list) and len(val) == 0:
continue
if isinstance(val, list) and len(val) == 1:
proc_kwargs[key] = val[0]
else:
proc_kwargs[key] = val
if audios:
proc_kwargs["audio"] = audios
proc_kwargs["use_audio_in_video"] = True
else:
proc_kwargs["use_audio_in_video"] = False
# Process inputs
inputs = processor(**proc_kwargs)
input_ids = inputs["input_ids"].to("cuda")
input_len = input_ids.shape[1]
# Move all tensors to cuda
gen_kwargs = {}
for key, val in inputs.items():
if key == "input_ids":
continue
if isinstance(val, torch.Tensor):
gen_kwargs[key] = val.to("cuda")
else:
gen_kwargs[key] = val
# Generate using the thinker only (no audio output)
with torch.no_grad():
output = model.generate(
input_ids=input_ids,
thinker_max_new_tokens=1024,
thinker_do_sample=True,
thinker_temperature=1.0,
thinker_top_p=0.95,
thinker_top_k=20,
generation_mode="text",
**gen_kwargs,
)
# Extract only the new tokens (response)
if isinstance(output, torch.Tensor):
new_tokens = output[0][input_len:]
text = processor.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
else:
# If it's a GenerationOutput or similar
seq = output.sequences if hasattr(output, "sequences") else output
new_tokens = seq[0][input_len:]
text = processor.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
return text
# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
BUILTIN_EXAMPLES = [
{
"video": "example_video_mcq.mp4",
"question": 'Who or what lauds "Immigrant Diaries" as "A SURE FIRE HIT", according to the video?',
"answer": "A",
"type": "MCQ",
"options": "A. Remote Goat.\nB. The New York Times.\nC. Variety.\nD. IndieWire.",
},
{
"video": "example_video_tr.mp4",
"question": 'What are all the time ranges corresponding to the text query: "A man with tousled dark hair and a beaded necklace thoughtfully shares his perspective, the subtle floral pattern of his light green shirt contrasting against the light-colored wall behind him as he speaks about challenges and roles."?',
"answer": "[51.72, 62.92]",
"type": "TR",
"options": "",
},
{
"video": "example_video_ff.mp4",
"question": "During the montage, what color was the horse that the boy in yellow is riding?",
"answer": "White",
"type": "FF",
"options": "",
},
]
def _step_html(step_num: int, action_type: str, think: str, observation: str,
confidence, action_detail: str, frames: List[str] = None) -> str:
"""Render an OTA step as an HTML card."""
colors = {
"get_frames": "#2196F3", "get_clip": "#4CAF50",
"get_audio": "#FF9800", "answer": "#F44336",
"error": "#9E9E9E", "early_answer": "#9E9E9E",
"forced_answer": "#F44336",
}
icons = {
"get_frames": "🖼️", "get_clip": "🎥", "get_audio": "🔊",
"answer": "✅", "error": "⚠️", "early_answer": "⚠️",
"forced_answer": "✅",
}
c = colors.get(action_type, "#607D8B")
icon = icons.get(action_type, "📌")
conf_text = f"{confidence:.3f}" if isinstance(confidence, (int, float)) else "N/A"
# Escape HTML
import html as html_mod
think_esc = html_mod.escape(think or "")[:500]
obs_esc = html_mod.escape(observation or "")[:500]
action_esc = html_mod.escape(action_detail or "")
frames_html = ""
if frames:
frame_items = []
for idx, fp in enumerate(frames[:6], 1):
from urllib.parse import quote
fp_escaped = quote(fp)
frame_items.append(
f'
'
)
frames_html = f'{" ".join(frame_items)}
'
return f"""
Observation: {obs_esc}
Think: {think_esc}
Action: {action_esc}
{frames_html}
"""
@spaces.GPU(duration=300)
def run_omniagent(video_path: str, question: str, q_type: str,
options_text: str, max_steps: int,
progress=gr.Progress()):
"""Run the OmniAgent OTA loop on a video and question.
Args:
video_path: Path to the input video file.
question: The question to answer about the video.
q_type: Question type — MCQ, TR, FF, NUM, or SIZE.
options_text: MCQ options (one per line, e.g. "A. Option\\nB. Option").
max_steps: Maximum number of agentic steps (default 12).
"""
if not video_path or not question:
yield "Please provide a video and a question.", "", "", gr.update(visible=False)
return
options_list = None
if q_type == "MCQ" and options_text:
options_list = [o.strip() for o in options_text.splitlines() if o.strip()]
steps_html = ""
final_answer = ""
status = "Initializing…"
env = OTAEnvironment(
video_path=video_path,
question=question,
q_type=q_type,
options=options_list,
answer="",
max_steps=min(max_steps, MAX_STEPS_DEFAULT),
)
try:
messages = env.build_initial_messages()
status = f"Running OTA loop (max {max_steps} steps)…"
yield steps_html, final_answer, status, gr.update(visible=False)
for step_num in range(1, max_steps + 1):
progress(step_num / max_steps, desc=f"Step {step_num}/{max_steps}")
try:
raw_response = generate_response(messages, env.has_audio)
except Exception as e:
raw_response = json.dumps({
"observation": f"Error: {e}",
"think": "Generation failed.",
"confidence": 0.0,
"action": {"type": "answer", "content": f"Error: {e}"}
})
# Parse for display
parsed = parse_response(raw_response) or {}
think = parsed.get("think", "")
observation = parsed.get("observation", "")
confidence = parsed.get("confidence")
action = parsed.get("action", {})
atype = action.get("type", "unknown")
action_detail = atype
if atype == "get_frames":
action_detail = f"get_frames(start={action.get('start')}, end={action.get('end')}, num={action.get('num')})"
elif atype in ("get_audio", "get_clip"):
action_detail = f"{atype}(start={action.get('start')}, end={action.get('end')})"
elif atype == "answer":
action_detail = f"answer: {action.get('content', '')}"
# Execute the action
done, exec_type, frames, clip, audio = env.step(raw_response)
messages = env.history
# Update display
step_html = _step_html(step_num, exec_type, think, observation,
confidence, action_detail, frames if exec_type == "get_frames" else None)
steps_html += step_html
if done:
final_answer = env.final_answer or parsed.get("action", {}).get("content", "")
status = f"Done — {step_num} step(s)."
yield steps_html, final_answer, status, gr.update(visible=True)
return
else:
yield steps_html, final_answer, status, gr.update(visible=False)
# Exhausted all steps
final_answer = env.final_answer or "(No answer produced within step limit.)"
status = f"Finished — {max_steps} steps (step limit reached)."
yield steps_html, final_answer, status, gr.update(visible=True)
except Exception as e:
import traceback
tb = traceback.format_exc()
steps_html += f'Error: {str(e)[:300]}
{tb[:1000]} '
yield steps_html, "", f"Error: {e}", gr.update(visible=True)
finally:
env.cleanup()
# ---------------------------------------------------------------------------
# UI layout
# ---------------------------------------------------------------------------
with gr.Blocks(title="OmniAgent-RL") as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown(
"# OmniAgent: Native Active Perception as Reasoning for Omni-Modal Understanding\n\n"
"An agentic audio-visual understanding model that iteratively requests frames, "
"audio, and clips to answer questions — powered by a Qwen2.5-Omni-7B "
"model fine-tuned with agentic SFT and agentic RL.\n\n"
"> 🤗 **Huge thanks to the Hugging Face team** for building this interactive demo based on our original UI! Please note that as an adapted version, its underlying logic is not 100% identical to our official release. For the exact standard and fully-featured experience, please use our **[official demo script](https://github.com/HarryHsing/OmniAgent/blob/main/demo/omniagent_demo_pro.py)**.\n\n"
"[Paper](https://huggingface.co/papers/2606.19341) · "
"[GitHub](https://github.com/harryhsing/OmniAgent) · "
"[Model](https://huggingface.co/harryhsing/OmniAgent-RL-7B)"
)
with gr.Row():
with gr.Column(scale=1):
video_input = gr.Video(label="Input Video", sources=["upload"])
question_input = gr.Textbox(
label="Question", placeholder="Ask a question about the video…",
lines=3,
)
with gr.Accordion("Question type & options", open=False):
q_type = gr.Radio(
["MCQ", "TR", "FF", "NUM", "SIZE"],
label="Question Type", value="MCQ",
info="MCQ: multiple choice · TR: temporal grounding · "
"FF: free-form · NUM/SIZE: numeric answer",
)
options_input = gr.Textbox(
label="MCQ Options (one per line)",
placeholder="A. Option one\nB. Option two\nC. Option three",
lines=4, visible=True,
)
with gr.Accordion("Advanced", open=False):
max_steps_slider = gr.Slider(
3, MAX_STEPS_DEFAULT, value=32, step=1,
label="Max Agentic Steps",
info="More steps = more thorough investigation (slower).",
)
run_btn = gr.Button("Run OmniAgent", variant="primary")
with gr.Column(scale=1):
status_box = gr.Textbox(label="Status", interactive=False)
final_answer_box = gr.Textbox(
label="Final Answer", interactive=False,
visible=False,
)
steps_output = gr.HTML(label="Agent Trace", value="")
gr.Examples(
examples=[
["example_video_mcq.mp4",
'Who or what lauds "Immigrant Diaries" as "A SURE FIRE HIT", according to the video?',
"MCQ",
"A. Remote Goat.\nB. The New York Times.\nC. Variety.\nD. IndieWire.",
32],
["example_video_tr.mp4",
'What are all the time ranges corresponding to the text query: "A man with tousled dark hair and a beaded necklace thoughtfully shares his perspective, the subtle floral pattern of his light green shirt contrasting against the light-colored wall behind him as he speaks about challenges and roles."?',
"TR",
"",
32],
["example_video_ff.mp4",
"During the montage, what color was the horse that the boy in yellow is riding?",
"FF",
"",
32],
],
inputs=[video_input, question_input, q_type, options_input, max_steps_slider],
outputs=[steps_output, final_answer_box, status_box, final_answer_box],
fn=run_omniagent,
cache_examples=False,
run_on_click=True,
)
run_btn.click(
fn=run_omniagent,
inputs=[video_input, question_input, q_type, options_input, max_steps_slider],
outputs=[steps_output, final_answer_box, status_box, final_answer_box],
)
if __name__ == "__main__":
demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)