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)