Proact-VL-Washing/mage_demo/hug_mage.py

"""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)