"""Mage-VL demo — codec-native image / video understanding + proactive streaming. Preprocessing mirrors the reference `inference.py` (model repo) and `inference_streaming.py` (github.com/microsoft/Mage/tree/main/mage_vl): * image -> processor(text=..., images=[PIL]) * video / frames -> processor(text=..., videos=[path], num_frames=N) * video / codec (HEVC) -> processor(..., video_backend="codec", codec_config={"engine": "hevc", "target_canvas": N, "patch": 16}, max_pixels=150000) * streaming -> ffmpeg segment cuts -> codec inputs per segment -> model.streammind_gate_forward_segments(...) -> generate() only on segments above the gate threshold """ import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") os.environ.setdefault("ONLINE_CODEC_CACHE_DIR", "/tmp/mage_vl_codec_cache") os.environ.setdefault("STREAMMIND_CLIP_CACHE", "/tmp/mage_vl_clips") os.makedirs(os.environ["ONLINE_CODEC_CACHE_DIR"], exist_ok=True) os.makedirs(os.environ["STREAMMIND_CLIP_CACHE"], exist_ok=True) import importlib import subprocess import time import traceback import uuid import warnings from pathlib import Path import spaces # noqa: F401 (must precede torch) import torch import gradio as gr from PIL import Image from transformers import AutoModelForCausalLM, AutoProcessor MODEL_ID = "microsoft/Mage-VL" FRAMES_LABEL = "Uniform frame sampling" CODEC_LABEL = "Codec-native (HEVC)" # Reference defaults (inference.py) CODEC_MAX_PIXELS = 150_000 # pixel budget per codec canvas FRAME_MAX_PIXELS = 589_824 # 768x768 cap per sampled frame (guards 4K uploads) MAX_VIDEO_SECONDS = 150.0 # longer uploads are trimmed to keep GPU time sane STREAM_PROMPT = ( "Please describe the video content in detail based on the provided information." ) # --------------------------------------------------------------------------- model processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, trust_remote_code=True, dtype=torch.bfloat16, attn_implementation="sdpa", ).to("cuda").eval() # The trust_remote_code modules live in a package next to the processor class. _REMOTE_PKG = type(processor).__module__.rsplit(".", 1)[0] codec_mod = importlib.import_module(_REMOTE_PKG + ".codec_video_processing_mage_vl") video_mod = importlib.import_module(_REMOTE_PKG + ".video_processing_mage_vl") CodecConfig = codec_mod.CodecConfig MageVLVideoProcessor = video_mod.MageVLVideoProcessor ProcessorCls = type(processor) print(f"[init] model loaded, remote package = {_REMOTE_PKG}", flush=True) # ----------------------------------------------------- StreamMind cognition gate # `streammind_gate.py` in the model repo imports # `mamba_ssm.models.mixer_seq_simple.create_block`; the real mamba-ssm ships CUDA # extensions with no wheel for this runtime, so `./mamba_ssm/` in this Space is a # pure-PyTorch stand-in with identical parameter names/shapes (see its docstring). # With that in place the released gate code + checkpoint load unmodified. GATE_ERROR = "" try: _t0 = time.perf_counter() model.model._load_streammind_gate() print(f"[gate] loaded in {time.perf_counter() - _t0:.1f}s", flush=True) except Exception as exc: # surfaced in the streaming tab, never kills the boot GATE_ERROR = f"{type(exc).__name__}: {exc}" print("[gate] FAILED:\n" + traceback.format_exc(), flush=True) # --------------------------------------------------------------------------- utils def _ffprobe_duration(video_path): try: out = subprocess.run( ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=nw=1:nk=1", str(video_path)], check=True, capture_output=True, timeout=60, text=True, ) return float(out.stdout.strip()) except Exception: return 0.0 def _probe(video_path): """(duration_seconds, fps, frame_count) via OpenCV, ffprobe as fallback.""" import cv2 cap = cv2.VideoCapture(video_path) try: fps = float(cap.get(cv2.CAP_PROP_FPS) or 0.0) frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0) finally: cap.release() if fps <= 0: fps = 30.0 duration = frames / fps if frames > 0 else 0.0 if duration <= 0: duration = _ffprobe_duration(video_path) frames = int(duration * fps) return duration, fps, frames def _trim_if_long(video_path, limit=MAX_VIDEO_SECONDS): """Keep the first `limit` seconds so a long upload can't blow the GPU budget.""" duration, _, _ = _probe(video_path) if duration <= limit + 1.0: return video_path, duration, "" out = Path("/tmp/mage_vl_clips") / f"trim_{uuid.uuid4().hex}.mp4" subprocess.run( ["ffmpeg", "-y", "-loglevel", "error", "-t", f"{limit:.3f}", "-i", video_path, "-c", "copy", "-avoid_negative_ts", "make_zero", str(out)], check=True, capture_output=True, timeout=300, ) note = (f"\n\n> Input was {duration:.0f}s long — only the first {limit:.0f}s " f"were used.") return str(out), min(duration, limit), note def _codec_cfg(target_canvas, max_pixels=CODEC_MAX_PIXELS): """Rebuild exactly the CodecConfig MageVLProcessor would build, so `process_codec_video` and the processor share one disk cache entry.""" override = {"engine": "hevc", "target_canvas": int(target_canvas), "patch": 16} kwargs = dict(processor._codec_config_defaults) merged_dcvc = dict(kwargs.get("dcvc") or {}) merged_dcvc.update(override.get("dcvc") or {}) kwargs.update(override) if merged_dcvc: kwargs["dcvc"] = merged_dcvc kwargs["max_pixels"] = int(max_pixels) return CodecConfig(**kwargs), override def _frames_processor(max_pixels): """Processor clone whose video processor honours a per-frame pixel budget.""" ip = processor.image_processor vp = MageVLVideoProcessor( image_processor=ip, min_pixels=getattr(ip, "min_pixels", 3136), max_pixels=int(max_pixels), patch_size=getattr(ip, "patch_size", 16), spatial_merge_size=getattr(ip, "merge_size", 2), ) return ProcessorCls( image_processor=ip, tokenizer=processor.tokenizer, video_processor=vp, chat_template=processor.chat_template, codec_config=processor._codec_config_defaults, ) def _prompt(kind, question): messages = [{"role": "user", "content": [{"type": kind}, {"type": "text", "text": question}]}] return processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) def _to_cuda(inputs): out = {} for key, value in inputs.items(): if not hasattr(value, "to"): continue out[key] = value.to("cuda") if key == "pixel_values": out[key] = out[key].to(model.dtype) return out def _visual_tokens(inputs): grid = inputs["image_grid_thw"] return int((grid[:, 0] * grid[:, 1] * grid[:, 2]).sum().item()) // 4 @torch.inference_mode() def _generate(inputs, max_new_tokens): out = model.generate(**inputs, max_new_tokens=int(max_new_tokens), do_sample=False) new = out[0, inputs["input_ids"].shape[1]:] return processor.tokenizer.decode(new, skip_special_tokens=True).strip() # ------------------------------------------------------------------------- handlers def _image_duration(image=None, question=None, max_new_tokens=384, *args, **kwargs): # measured on ZeroGPU: ~13 s for a 2 MP image at 1024 new tokens return int(16 + float(max_new_tokens) * 0.013) @spaces.GPU(duration=_image_duration) def run_image(image, question, max_new_tokens=384): """Answer a question about a single image with Mage-VL. Args: image: the input image. question: what to ask about it. max_new_tokens: generation length cap. """ if image is None: raise gr.Error("Please provide an image.") question = (question or "Describe this image in detail.").strip() if isinstance(image, str): image = Image.open(image) image = image.convert("RGB") inputs = processor(text=[_prompt("image", question)], images=[image], return_tensors="pt") tokens = _visual_tokens(inputs) started = time.perf_counter() answer = _generate(_to_cuda(inputs), max_new_tokens) elapsed = time.perf_counter() - started stats = (f"**{image.width}×{image.height}** · **{tokens:,}** visual tokens · " f"**{elapsed:.1f}s** on ZeroGPU") return answer, stats def _video_duration(video=None, question=None, backend=CODEC_LABEL, num_frames=32, max_new_tokens=384, *args, **kwargs): # measured on ZeroGPU (30 s clip, 64 canvases / 64 frames, 1024 new tokens): # codec 3.2 s prep + 13.3 s generate, frames 4.5 s prep + 16.0 s generate. # The base covers codec analysis on longer uploads (trimmed to 150 s). return int(30 + float(max_new_tokens) * 0.020) @spaces.GPU(duration=_video_duration) def run_video(video, question, backend=CODEC_LABEL, num_frames=32, max_new_tokens=384): """Answer a question about a video, with codec-native or frame-sampled tokens. Args: video: the input video file. question: what to ask about it. backend: "Codec-native (HEVC)" or "Uniform frame sampling". num_frames: codec canvases (codec backend) or sampled frames (frames backend). max_new_tokens: generation length cap. """ if not video: raise gr.Error("Please provide a video.") question = (question or "Describe this video.").strip() num_frames = max(4, int(num_frames) // 4 * 4) # target_canvas % images_per_group == 0 yield "", [], "Preparing video…" path, duration, note = _trim_if_long(video) _, fps, frame_count = _probe(path) text = _prompt("video", question) prep_start = time.perf_counter() if str(backend).startswith("Codec"): yield "", [], "Running codec analysis (cv-preinfer)…" cfg, override = _codec_cfg(num_frames) try: with warnings.catch_warnings(): warnings.simplefilter("ignore") # process_codec_video is disk-cached, so running it here for the # canvas gallery costs nothing extra: the processor call below # hits the same cache entry. payload = codec_mod.process_codec_video(path, cfg) canvases, _, _ = codec_mod.drop_padding_canvases( payload["images"], payload["src_positions"] ) inputs = processor( text=[text], videos=[path], video_backend="codec", max_pixels=CODEC_MAX_PIXELS, codec_config=override, return_tensors="pt", padding=True, ) except Exception as exc: print("[codec] " + traceback.format_exc(), flush=True) raise gr.Error( "Codec analysis failed for this file " f"({type(exc).__name__}: {exc}). The codec path needs a video the " "HEVC/H.264 bit-cost reader can parse — try a standard .mp4, or " 'switch the tokenizer to "Uniform frame sampling".' ) from exc gallery = [(im, f"canvas {i + 1}") for i, im in enumerate(canvases)] sampled = min(cfg.num_sampled_frames(), max(1, frame_count)) detail = (f"**codec-native (HEVC)** · {sampled} frames analysed → " f"**{len(canvases)} canvases** · {CODEC_MAX_PIXELS:,} px/canvas budget") else: yield "", [], "Sampling frames…" proc = _frames_processor(FRAME_MAX_PIXELS) inputs = proc(text=[text], videos=[path], num_frames=int(num_frames), return_tensors="pt", padding=True) frames, _, _ = video_mod.extract_video_frames_to_pil( video_path=path, max_frames=384, patch_size=16, min_pixels=3136, max_pixels=FRAME_MAX_PIXELS, resize_frames=True, fixed_num_frames=int(num_frames), target_fps=None, ) gallery = [(im, f"frame {i + 1}") for i, im in enumerate(frames)] detail = (f"**uniform frame sampling** · **{len(frames)} frames** · " f"{FRAME_MAX_PIXELS:,} px/frame budget") prep = time.perf_counter() - prep_start tokens = _visual_tokens(inputs) prompt_tokens = int(inputs["input_ids"].shape[1]) yield "", gallery, f"{detail} · **{tokens:,}** visual tokens — generating…" gen_start = time.perf_counter() answer = _generate(_to_cuda(inputs), max_new_tokens) gen = time.perf_counter() - gen_start stats = (f"{detail}\n\n" f"**{tokens:,}** visual tokens ({prompt_tokens:,} prompt tokens) · " f"video {duration:.1f}s @ {fps:.1f} fps · " f"preprocess **{prep:.1f}s** · generate **{gen:.1f}s**{note}") yield answer, gallery, stats def _extract_subclip(src, start, duration, out_path): if out_path.exists() and out_path.stat().st_size > 0: return out_path tmp = out_path.with_name(f".tmp_{out_path.name}") subprocess.run( ["ffmpeg", "-y", "-loglevel", "error", "-ss", f"{max(0.0, start):.3f}", "-t", f"{duration:.3f}", "-i", str(src), "-c", "copy", "-avoid_negative_ts", "make_zero", str(tmp)], check=True, capture_output=True, timeout=180, ) tmp.replace(out_path) return out_path def _stream_duration(video=None, segment_sec=8.0, gate_threshold=0.5, max_segments=4, max_new_tokens=96, *args, **kwargs): # measured on ZeroGPU: 54 s for 8 segments all speaking at 256 new tokens # (~1.3 s codec analysis per segment + ~5 s generation per spoken segment). return int(14 + float(max_segments) * (2.5 + float(max_new_tokens) * 0.021)) @spaces.GPU(duration=_stream_duration) def run_streaming(video, segment_sec=8.0, gate_threshold=0.5, max_segments=4, max_new_tokens=96): """Event-gated streaming commentary: the model only speaks on salient segments. Args: video: the input video file. segment_sec: length of each rolling codec segment, in seconds. gate_threshold: speak when the cognition gate's p(response) is above this. max_segments: how many leading segments to process. max_new_tokens: generation length cap per spoken segment. """ if not video: raise gr.Error("Please provide a video.") if GATE_ERROR: raise gr.Error(f"The StreamMind gate could not be loaded: {GATE_ERROR}") segment_sec = float(segment_sec) max_segments = int(max_segments) yield "Cutting segments…" path, duration, note = _trim_if_long(video) clip_dir = Path("/tmp/mage_vl_clips") clip_dir.mkdir(parents=True, exist_ok=True) stem = Path(path).stem text = _prompt("video", STREAM_PROMPT) cfg, override = _codec_cfg(32) segments, spans = [], [] start = 0.0 index = 0 while start < duration - 0.2 and index < max_segments: end = min(duration, start + segment_sec) clip = clip_dir / f"{stem}__t{int(start * 100):08d}__d{int((end - start) * 100):08d}.mp4" yield (f"Codec analysis for segment {index + 1} " f"(t={start:.1f}–{end:.1f}s)…") try: _extract_subclip(path, start, end - start, clip) with warnings.catch_warnings(): warnings.simplefilter("ignore") inputs = processor( text=[text], videos=[str(clip)], video_backend="codec", max_pixels=CODEC_MAX_PIXELS, codec_config=override, return_tensors="pt", padding=False, ) except Exception as exc: print(f"[stream] segment {start:.1f}-{end:.1f} unusable: {exc}", flush=True) start, index = end, index + 1 continue segments.append(_to_cuda(inputs)) spans.append((start, end)) start, index = end, index + 1 if not segments: yield "No usable segments — try a longer clip or a smaller segment length." return yield f"Scoring {len(segments)} segments with the cognition gate…" visual = [{k: seg[k] for k in ("pixel_values", "image_grid_thw", "patch_positions") if k in seg} for seg in segments] with torch.inference_mode(): logits = model.model.streammind_gate_forward_segments(visual)[0] lengths = [int(seg["image_grid_thw"][:, 0].sum()) for seg in visual] boundaries = torch.tensor(lengths, device=logits.device).cumsum(0) - 1 probs = torch.softmax(logits[boundaries].float(), dim=-1)[:, 1].tolist() lines = [] for (span, inputs, prob) in zip(spans, segments, probs): head = f"**t = {span[0]:.1f}–{span[1]:.1f}s** · p(respond) = `{prob:.2f}`" if prob < float(gate_threshold): lines.append(f"{head} → 🔇 *silence*") yield "\n\n".join(lines) continue lines.append(f"{head} → 🗣️ *generating…*") yield "\n\n".join(lines) try: said = _generate(inputs, max_new_tokens) except Exception as exc: said = f"(generation failed: {type(exc).__name__}: {exc})" lines[-1] = f"{head} → 🗣️ {said}" yield "\n\n".join(lines) yield "\n\n".join(lines) + note # ------------------------------------------------------------------------------ UI CSS = """ #col-container { max-width: 1180px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ INTRO = """# Mage-VL — codec-native streaming multimodal model [`microsoft/Mage-VL`](https://huggingface.co/microsoft/Mage-VL) is a 4B VLM whose visual encoder is trained **from scratch** to tokenize video the way a codec compresses it: keep every anchor (I) frame patch, keep only the predicted (P) frame patches where the codec spends bits, pack the survivors into **canvases**. That cuts visual tokens by >75% versus uniform frame sampling. A lightweight cognition gate on top turns the same weights into a proactive streaming model. [Project page](https://microsoft.github.io/Mage) · [GitHub](https://github.com/microsoft/Mage) · [Tech report](https://github.com/microsoft/Mage/blob/main/assets/mage_vl_tech_report.pdf) """ with gr.Blocks(title="Mage-VL") as demo: with gr.Column(elem_id="col-container"): gr.Markdown(INTRO) with gr.Tabs(): # ------------------------------------------------------------- image with gr.Tab("Image"): with gr.Row(): with gr.Column(): img_in = gr.Image(label="Image", type="pil", height=360) img_q = gr.Textbox( label="Question", lines=2, value="Describe this image in detail.", ) img_btn = gr.Button("Ask", variant="primary") with gr.Column(): img_out = gr.Textbox(label="Answer", lines=16, buttons=["copy"]) img_stats = gr.Markdown() with gr.Accordion("Advanced", open=False): img_tokens = gr.Slider(64, 1024, value=384, step=32, label="Max new tokens") gr.Examples( examples=[ ["examples/dog.jpg", "Describe this image in detail."], ["examples/dog.jpg", "Describe the rug behind the dog and the patterns woven into it."], ], inputs=[img_in, img_q], outputs=[img_out, img_stats], fn=run_image, cache_examples=True, cache_mode="lazy", ) img_btn.click(run_image, inputs=[img_in, img_q, img_tokens], outputs=[img_out, img_stats], api_name="ask_image") # ------------------------------------------------------------- video with gr.Tab("Video"): gr.Markdown( "Switch **Tokenizer** to compare the codec-native path against " "uniform frame sampling — the gallery shows exactly what the " "vision encoder sees, and the stats line reports the visual-token " "count for each path." ) with gr.Row(): with gr.Column(): vid_in = gr.Video(label="Video", height=320) vid_q = gr.Textbox(label="Question", lines=2, value="Describe this video.") vid_backend = gr.Radio( [CODEC_LABEL, FRAMES_LABEL], value=CODEC_LABEL, label="Tokenizer", ) vid_btn = gr.Button("Ask", variant="primary") with gr.Column(): vid_out = gr.Textbox(label="Answer", lines=14, buttons=["copy"]) vid_stats = gr.Markdown() vid_gallery = gr.Gallery( label="What the vision encoder sees (codec canvases / sampled frames)", columns=8, height=260, object_fit="contain", show_label=True, ) with gr.Accordion("Advanced", open=False): vid_frames = gr.Slider( 8, 64, value=32, step=4, label="Codec canvases (codec) / sampled frames (frames)", ) vid_tokens = gr.Slider(64, 1024, value=384, step=32, label="Max new tokens") gr.Examples( examples=[ ["examples/soccer-broadcast.mp4", "Describe this video.", CODEC_LABEL], ["examples/soccer-broadcast.mp4", "What is the score shown on the scoreboard, and which teams are playing?", CODEC_LABEL], ["examples/soccer-broadcast.mp4", "Describe this video.", FRAMES_LABEL], ], inputs=[vid_in, vid_q, vid_backend], outputs=[vid_out, vid_gallery, vid_stats], fn=run_video, cache_examples=True, cache_mode="lazy", ) vid_btn.click( run_video, inputs=[vid_in, vid_q, vid_backend, vid_frames, vid_tokens], outputs=[vid_out, vid_gallery, vid_stats], api_name="ask_video", ) # --------------------------------------------------------- streaming with gr.Tab("Proactive streaming"): gr.Markdown( "The video is split into non-overlapping codec segments. The " "**cognition gate** (System 1) scores every segment and stays " "silent on routine content; the full VLM (System 2) only " "generates commentary where `p(respond)` crosses the threshold." ) with gr.Row(): with gr.Column(scale=1): st_in = gr.Video(label="Video", height=300) st_btn = gr.Button("Run streaming", variant="primary") with gr.Accordion("Advanced", open=False): st_seg = gr.Slider(4.0, 12.0, value=8.0, step=1.0, label="Segment length (s)") st_thr = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Gate threshold") st_max = gr.Slider(2, 8, value=4, step=1, label="Max segments") st_tokens = gr.Slider(32, 256, value=96, step=16, label="Max new tokens per response") with gr.Column(scale=2): st_out = gr.Markdown( value="_Timeline will appear here._", label="Gated timeline", ) gr.Examples( examples=[["examples/soccer-broadcast.mp4"]], inputs=[st_in], outputs=[st_out], fn=run_streaming, cache_examples=True, cache_mode="lazy", ) st_btn.click( run_streaming, inputs=[st_in, st_seg, st_thr, st_max, st_tokens], outputs=[st_out], api_name="run_streaming", ) gr.Markdown( "Preprocessing follows the reference `inference.py` and " "`inference_streaming.py` from the model repo / " "[Mage GitHub](https://github.com/microsoft/Mage/tree/main/mage_vl). " "The codec path shells out to `cv-preinfer` (`codec-video-prep`) for " "HEVC bit-cost readiness; the DCVC-RT neural-codec engine is not " "wired up in this Space." ) if __name__ == "__main__": demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, show_error=True)