app.py

import os

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

import spaces  # noqa: E402  MUST come before torch / transformers
import threading  # noqa: E402
import time  # noqa: E402

import decord  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
import gradio as gr  # noqa: E402
from torchvision import transforms  # noqa: E402
from transformers import (  # noqa: E402
    AutoProcessor,
    Qwen2_5_VLForConditionalGeneration,
    TextIteratorStreamer,
)

# ----------------------------------------------------------------------------
# Config
# ----------------------------------------------------------------------------
MODEL_ID = "Catalan258/VST-7B"

# Qwen2-VL patch(14) * merge(2) = 28
IMAGE_FACTOR = 28
FRAME_FACTOR = 2
# per-frame pixel budget (tokens ~ pixels / (28*28))
MIN_PIXELS = 128 * 28 * 28
MAX_PIXELS = 512 * 28 * 28
# total frames sampled across the whole clip
MAX_NUM_FRAMES = 64

SYSTEM_PROMPT = "You are a helpful assistant."
STREAM_THINK_PROMPT = "[System]\nYou are a Streaming Video Analyst.\n"

DEFAULT_QUESTION = "Describe what happens in this video and explain the sequence of events."

# ----------------------------------------------------------------------------
# Model — loaded once at module scope, eagerly moved to CUDA (ZeroGPU packs it)
# ----------------------------------------------------------------------------
processor = AutoProcessor.from_pretrained(MODEL_ID, max_pixels=MAX_PIXELS)
tokenizer = processor.tokenizer
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda").eval()


# ----------------------------------------------------------------------------
# Video helpers (port of the VST reference frame reader / spatial resize)
# ----------------------------------------------------------------------------
def _round_by_factor(number: int, factor: int) -> int:
    return round(number / factor) * factor


def smart_resize(height, width, factor=IMAGE_FACTOR, min_pixels=MIN_PIXELS, max_pixels=MAX_PIXELS):
    """Rescale so H,W are multiples of `factor` and total pixels within [min,max]."""
    h_bar = max(factor, _round_by_factor(height, factor))
    w_bar = max(factor, _round_by_factor(width, factor))
    if h_bar * w_bar > max_pixels:
        beta = (height * width / max_pixels) ** 0.5
        h_bar = max(factor, int(np.floor(height / beta / factor)) * factor)
        w_bar = max(factor, int(np.floor(width / beta / factor)) * factor)
    elif h_bar * w_bar < min_pixels:
        beta = (min_pixels / (height * width)) ** 0.5
        h_bar = int(np.ceil(height * beta / factor)) * factor
        w_bar = int(np.ceil(width * beta / factor)) * factor
    return h_bar, w_bar


def read_video_frames(video_path, max_num_frames=MAX_NUM_FRAMES):
    """Uniformly sample up to `max_num_frames` frames. Returns (T,C,H,W) uint8 tensor,
    the timestamp of each frame, and total duration in seconds."""
    vr = decord.VideoReader(video_path, num_threads=2)
    video_fps = vr.get_avg_fps()
    total = len(vr)
    duration = total / max(video_fps, 1e-6)

    nframes = min(max_num_frames, total)
    # keep frame count a multiple of FRAME_FACTOR (Qwen video temporal merge)
    nframes = max(FRAME_FACTOR, (nframes // FRAME_FACTOR) * FRAME_FACTOR)
    idxs = np.linspace(0, total - 1, nframes).round().astype(int)
    frames = torch.from_numpy(vr.get_batch(idxs).asnumpy()).permute(0, 3, 1, 2)  # T,C,H,W
    pts = [float(i) / max(video_fps, 1e-6) for i in idxs]
    return frames, pts, duration


def spatial_resize(video):
    """Resize a (T,C,H,W) uint8 tensor to a Qwen-friendly resolution, return float."""
    _, _, h, w = video.shape
    rh, rw = smart_resize(h, w, factor=IMAGE_FACTOR, min_pixels=MIN_PIXELS, max_pixels=MAX_PIXELS)
    video = transforms.functional.resize(
        video,
        [rh, rw],
        interpolation=transforms.InterpolationMode.BICUBIC,
        antialias=True,
    )
    return video.float()


def split_into_chunks(frames, n_chunks):
    """Split (T,C,H,W) frames into n_chunks contiguous temporal segments."""
    total = len(frames)
    n_chunks = max(1, min(n_chunks, total // FRAME_FACTOR if total >= FRAME_FACTOR else 1))
    seg = max(1, total // n_chunks)
    chunks = []
    for j in range(n_chunks):
        start = j * seg
        end = total if j == n_chunks - 1 else min((j + 1) * seg, total)
        if start >= end:
            continue
        chunk = frames[start:end]
        # ensure even number of frames for the temporal merge
        if len(chunk) % FRAME_FACTOR != 0 and len(chunk) > 1:
            chunk = chunk[: (len(chunk) // FRAME_FACTOR) * FRAME_FACTOR]
        if len(chunk) == 0:
            continue
        chunks.append(chunk)
    return chunks


# ----------------------------------------------------------------------------
# Generation
# ----------------------------------------------------------------------------
def _generate_stream(cur_msg, video_chunk, max_new_tokens):
    """Run model.generate with a streamer in a thread; yield decoded text incrementally."""
    text = processor.apply_chat_template(cur_msg, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=text, videos=[video_chunk], padding=True, return_tensors="pt").to("cuda")

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )
    gen_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=False,
        num_beams=1,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
        use_cache=True,
    )
    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
    thread.start()
    acc = ""
    for token in streamer:
        acc += token
        yield acc
    thread.join()


def _estimate_duration(video, question, n_chunks, *args, **kwargs):
    try:
        n = int(n_chunks)
    except Exception:
        n = 3
    # Measured ~17-20s for short clips; budget for longer user videos with
    # full-length generations (512-token thoughts + 1024-token final answer)
    # plus cold-start GPU allocation headroom.
    return min(150, 40 + n * 18)


@spaces.GPU(duration=_estimate_duration)
def watch_and_think(video, question, n_chunks, progress=gr.Progress(track_tqdm=False)):
    """Watch a video while thinking: VST processes the clip chunk-by-chunk, emitting
    intermediate streaming thoughts, then delivers a grounded final answer.

    Args:
        video: path to the input video file.
        question: the question to answer about the video.
        n_chunks: how many streaming segments to split the video into.

    Yields:
        (streaming_thoughts_markdown, final_answer_markdown)
    """
    t0 = time.perf_counter()
    question = (question or DEFAULT_QUESTION).strip()
    n_chunks = max(1, int(n_chunks))

    yield "⏳ Reading and sampling video frames…", ""

    frames, pts, duration = read_video_frames(video)
    frames = spatial_resize(frames)
    chunks = split_into_chunks(frames, n_chunks)
    n_real = len(chunks)

    # timestamp boundaries per chunk
    seg_seconds = duration / n_real if n_real else duration

    # Build the running textual memory, exactly like the VST reference:
    # memory[0] is the streaming-analyst system preamble.
    textual_memory = [STREAM_THINK_PROMPT]
    thoughts_md = ""

    # ---- Streaming-thinking loop: reason over each chunk except the last ----
    for j, chunk in enumerate(chunks[:-1]) if n_real > 1 else []:
        start_t = j * seg_seconds
        end_t = (j + 1) * seg_seconds
        time_str = f"{start_t:.1f}-{end_t:.1f}s"

        header = f"\n\n### 🧠 Streaming thought {j + 1}/{n_real - 1} · watching `{time_str}`\n"
        thoughts_md += header
        yield thoughts_md, ""

        cur_msg = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": "".join(textual_memory)},
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": f"Time={time_str} "},
                    {"type": "video", "video": chunk},
                ],
            },
        ]
        partial = ""
        for partial in _generate_stream(cur_msg, chunk, max_new_tokens=512):
            yield thoughts_md + partial, ""
        thoughts_md += partial
        # accumulate into memory for subsequent rounds (mirrors reference)
        textual_memory.append(f"Time={time_str} " + partial + "\n")
        yield thoughts_md, ""

    # ---- Final answer over the last chunk + full accumulated memory ----
    final_header = "\n\n### ✅ Watching final segment & answering…\n"
    thoughts_md += final_header
    yield thoughts_md, "⏳ Generating final answer…"

    last_chunk = chunks[-1]
    last_start = (n_real - 1) * seg_seconds
    last_time = f"{last_start:.1f}-{duration:.1f}s"
    cur_msg = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": "".join(textual_memory)},
        {
            "role": "user",
            "content": [
                {"type": "text", "text": f"Time={last_time} "},
                {"type": "video", "video": last_chunk},
            ],
        },
        {"role": "user", "content": f"Time={duration:.1f}s " + question},
    ]

    final = ""
    for final in _generate_stream(cur_msg, last_chunk, max_new_tokens=1024):
        yield thoughts_md, final

    elapsed = time.perf_counter() - t0
    thoughts_md += f"\n\n---\n*Watched {duration:.1f}s of video across {n_real} streaming segment(s) in {elapsed:.1f}s.*"
    yield thoughts_md, final


# ----------------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1150px; 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(
            """
            # 🎬 Video Streaming Thinking (VST-7B)

            **VideoLLMs that watch and think *simultaneously*.** Instead of waiting for the
            whole video before reasoning, VST consumes the clip in segments, emitting
            **intermediate streaming thoughts** as it watches, then delivers a grounded
            final answer — amortizing test-time reasoning for real-time responsiveness.

            Built on [`Catalan258/VST-7B`](https://huggingface.co/Catalan258/VST-7B) ·
            [paper](https://huggingface.co/papers/2603.12262) ·
            [code](https://github.com/1ranGuan/VST) · ECCV 2026
            """
        )
        with gr.Row():
            with gr.Column(scale=1):
                video_in = gr.Video(label="Video", height=320)
                question_in = gr.Textbox(
                    label="Question",
                    value=DEFAULT_QUESTION,
                    lines=2,
                )
                with gr.Accordion("Advanced settings", open=False):
                    n_chunks_in = gr.Slider(
                        1, 5, value=3, step=1,
                        label="Streaming segments",
                        info="How many chunks to split the video into. More segments = more "
                        "intermediate 'thinking-while-watching' rounds before the final answer.",
                    )
                run_btn = gr.Button("Watch & Think", variant="primary")
            with gr.Column(scale=1):
                thoughts_out = gr.Markdown(label="Streaming thoughts", value="")
                answer_out = gr.Markdown(label="Final answer", value="")

        gr.Examples(
            examples=[
                ["examples/chef_wok_flames.mp4", "What is the person doing, and what technique are they using?", 3],
                ["examples/basketball_dribble.mp4", "Describe the sequence of actions performed with the ball.", 3],
                ["examples/couple_dancing.mp4", "What activity is shown and how does it progress over time?", 2],
                ["examples/dog_chasing_ball.mp4", "What is the dog doing throughout the video?", 3],
            ],
            inputs=[video_in, question_in, n_chunks_in],
            outputs=[thoughts_out, answer_out],
            fn=watch_and_think,
            cache_examples=False,
        )

    run_btn.click(
        watch_and_think,
        inputs=[video_in, question_in, n_chunks_in],
        outputs=[thoughts_out, answer_out],
        api_name="watch_and_think",
    )


if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)