app.py

"""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'<img src="/gradio_api/file={fp_escaped}" '
                f'style="width:80px;height:60px;object-fit:cover;border-radius:4px;'
                f'border:1px solid #ddd;margin:2px;" />'
            )
        frames_html = f'<div style="margin-top:8px;">{"&nbsp;".join(frame_items)}</div>'

    return f"""
    <div class="step-card" style="border-left:4px solid {c};">
        <div class="step-header">
            <span style="font-size:18px;">{icon}</span>
            Step {step_num} — <span style="color:{c};text-transform:uppercase;font-size:11px;">{action_type}</span>
            <span style="float:right;font-size:11px;color:#888;">Confidence: {conf_text}</span>
        </div>
        <div class="step-obs"><strong>Observation:</strong> {obs_esc}</div>
        <div class="step-think"><strong>Think:</strong> {think_esc}</div>
        <div class="step-action"><strong>Action:</strong> {action_esc}</div>
        {frames_html}
    </div>
    """


@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'<div class="step-card" style="border-left:4px solid #F44336;"><b>Error:</b> {str(e)[:300]}<pre>{tb[:1000]}</pre></div>'
        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)