[model]Add experimental Mage-VL full-duplex adapter and offline integration (#6537)

Signed-off-by: Zhanshenzhui <99891002+zyforsure@users.noreply.github.com>
Co-authored-by: Sy03 <1370724210@qq.com>

examples/offline_inference/mage_vl/end2end.py

@@ -0,0 +1,259 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Run Mage-VL offline end-to-end inference through the upstream Mage scripts.
+
+Mage-VL currently ships its custom Mage-ViT and codec processors as
+Transformers remote code in the Microsoft Mage repository. This example keeps
+the vLLM-Omni entrypoint stable while delegating the checkpoint-specific forward
+path to that upstream implementation.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+import time
+from importlib.util import find_spec
+from pathlib import Path
+
+DEFAULT_MODEL = "microsoft/Mage-VL"
+DEFAULT_IMAGE_QUESTION = "Describe this image in detail."
+DEFAULT_VIDEO_QUESTION = "Describe this video."
+
+
+def _default_mage_repo() -> Path | None:
+    value = os.getenv("MAGE_REPO_DIR")
+    if value:
+        return Path(value)
+    sibling = Path(__file__).resolve().parents[3].parent / "Mage"
+    if sibling.exists():
+        return sibling
+    return None
+
+
+def _resolve_mage_repo(value: str | None) -> Path:
+    repo = Path(value) if value else _default_mage_repo()
+    if repo is None:
+        raise ValueError("Set --mage-repo or MAGE_REPO_DIR to a checkout of https://github.com/microsoft/Mage.")
+    repo = repo.expanduser().resolve()
+    if not (repo / "mage_vl" / "inference_base.py").is_file():
+        raise FileNotFoundError(f"Missing Mage-VL scripts under {repo}. Expected mage_vl/inference_base.py.")
+    return repo
+
+
+def _default_media_path(mage_repo: Path, task: str) -> Path:
+    if task == "image":
+        return mage_repo / "mage_vl" / "assets" / "examples" / "dog.jpg"
+    return mage_repo / "mage_vl" / "assets" / "examples" / "soccer-broadcast.mp4"
+
+
+def _resolve_media_path(path: str | None, mage_repo: Path, task: str) -> Path:
+    resolved = Path(path).expanduser().resolve() if path else _default_media_path(mage_repo, task)
+    if not resolved.is_file():
+        raise FileNotFoundError(f"Input media file not found: {resolved}")
+    return resolved
+
+
+def build_command(args: argparse.Namespace, mage_repo: Path, media_path: Path) -> list[str]:
+    if args.task == "streaming":
+        return [
+            sys.executable,
+            str(mage_repo / "mage_vl" / "inference_streaming.py"),
+            "--video",
+            str(media_path),
+            "--checkpoint",
+            args.model,
+            "--video_backend",
+            args.video_backend,
+            "--num_frames",
+            str(args.num_frames),
+            "--cur_fps",
+            str(args.cur_fps),
+            "--segment_sec",
+            str(args.segment_sec),
+            "--max_new_tokens",
+            str(args.max_new_tokens),
+            "--max_segments",
+            str(args.max_segments),
+            "--gate_threshold",
+            str(args.gate_threshold),
+            "--device",
+            args.device,
+            "--attn_impl",
+            args.attn_impl,
+        ]
+
+    cmd = [
+        sys.executable,
+        str(mage_repo / "mage_vl" / "inference_base.py"),
+        "--mode",
+        "offline",
+        "--question",
+        args.question or (DEFAULT_IMAGE_QUESTION if args.task == "image" else DEFAULT_VIDEO_QUESTION),
+        "--model",
+        args.model,
+        "--max-new-tokens",
+        str(args.max_new_tokens),
+    ]
+    if args.task == "image":
+        cmd.extend(["--image", str(media_path)])
+        return cmd
+
+    cmd.extend(
+        [
+            "--video",
+            str(media_path),
+            "--video-backend",
+            args.video_backend,
+            "--num-frames",
+            str(args.num_frames),
+            "--max-pixels",
+            str(args.max_pixels),
+        ]
+    )
+    if args.video_backend == "codec":
+        cmd.extend(["--codec-engine", args.codec_engine])
+    return cmd
+
+
+def check_runtime_dependencies(task: str) -> None:
+    required = ["torch"]
+    if task == "streaming":
+        required.append("mamba_ssm")
+    missing = [name for name in required if find_spec(name) is None]
+    if not missing:
+        return
+
+    install_hint = (
+        "Install Mage-VL runtime dependencies in this environment first:\n"
+        "  pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121\n"
+        "  grep -v '^mamba-ssm' /path/to/Mage/mage_vl/requirements.txt > /tmp/mage_vl_requirements_no_mamba.txt\n"
+        "  pip install -r /tmp/mage_vl_requirements_no_mamba.txt\n"
+        "  pip install 'causal-conv1d>=1.4.0' --no-build-isolation\n"
+        "  pip install 'mamba-ssm>=2.2' --no-build-isolation"
+    )
+    raise RuntimeError(f"Missing Python packages: {', '.join(missing)}.\n{install_hint}")
+
+
+def run_command(cmd: list[str], cwd: Path, timeout: int) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        cmd,
+        cwd=str(cwd),
+        input="y\n",
+        text=True,
+        capture_output=True,
+        timeout=timeout,
+        check=False,
+    )
+
+
+def write_outputs(
+    *,
+    args: argparse.Namespace,
+    mage_repo: Path,
+    media_path: Path,
+    cmd: list[str],
+    result: subprocess.CompletedProcess[str] | None,
+    elapsed_s: float | None,
+) -> Path:
+    output_dir = Path(args.output_dir).expanduser().resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    summary = {
+        "task": args.task,
+        "model": args.model,
+        "mage_repo": str(mage_repo),
+        "media_path": str(media_path),
+        "command": cmd,
+        "dry_run": args.dry_run,
+        "elapsed_s": elapsed_s,
+        "returncode": None if result is None else result.returncode,
+    }
+    if result is not None:
+        stdout = result.stdout.strip()
+        stderr = result.stderr.strip()
+        summary["stdout"] = stdout
+        summary["stderr"] = stderr
+        (output_dir / "stdout.txt").write_text(result.stdout, encoding="utf-8")
+        (output_dir / "stderr.txt").write_text(result.stderr, encoding="utf-8")
+    summary_path = output_dir / "summary.json"
+    summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+    return summary_path
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Mage-VL offline end-to-end inference example.")
+    parser.add_argument(
+        "--mage-repo",
+        default=None,
+        help="Path to a Microsoft Mage checkout. Defaults to MAGE_REPO_DIR.",
+    )
+    parser.add_argument("--model", default=DEFAULT_MODEL, help="HF model id or local checkpoint path.")
+    parser.add_argument("--task", choices=("image", "video", "streaming"), default="image")
+    parser.add_argument("--image-path", default=None, help="Image input. Defaults to the Mage sample image.")
+    parser.add_argument("--video-path", default=None, help="Video input. Defaults to the Mage sample video.")
+    parser.add_argument("--question", default=None, help="Question for image/video offline QA.")
+    parser.add_argument("--video-backend", choices=("frames", "codec"), default="frames")
+    parser.add_argument("--codec-engine", choices=("traditional", "neural"), default="traditional")
+    parser.add_argument("--num-frames", type=int, default=32)
+    parser.add_argument("--max-pixels", type=int, default=150000)
+    parser.add_argument("--max-new-tokens", type=int, default=256)
+    parser.add_argument("--cur-fps", type=float, default=2.0)
+    parser.add_argument("--segment-sec", type=float, default=8.0)
+    parser.add_argument("--max-segments", type=int, default=0)
+    parser.add_argument("--gate-threshold", type=float, default=0.5)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--attn-impl", choices=("flash_attention_2", "sdpa", "eager"), default="flash_attention_2")
+    parser.add_argument("--timeout", type=int, default=1800)
+    parser.add_argument("--output-dir", default="outputs/mage_vl")
+    parser.add_argument("--dry-run", action="store_true", help="Build the command and summary without executing it.")
+    args = parser.parse_args()
+    if args.num_frames <= 0:
+        parser.error("--num-frames must be positive")
+    if args.segment_sec <= 0:
+        parser.error("--segment-sec must be positive")
+    if args.task == "image" and args.video_path:
+        parser.error("--video-path cannot be used with --task image")
+    if args.task in {"video", "streaming"} and args.image_path:
+        parser.error("--image-path cannot be used with video tasks")
+    return args
+
+
+def main() -> None:
+    args = parse_args()
+    mage_repo = _resolve_mage_repo(args.mage_repo)
+    media_arg = args.image_path if args.task == "image" else args.video_path
+    media_path = _resolve_media_path(media_arg, mage_repo, args.task)
+    cmd = build_command(args, mage_repo, media_path)
+
+    result = None
+    elapsed_s = None
+    if not args.dry_run:
+        check_runtime_dependencies(args.task)
+        start = time.perf_counter()
+        result = run_command(cmd, cwd=mage_repo, timeout=args.timeout)
+        elapsed_s = time.perf_counter() - start
+
+    summary_path = write_outputs(
+        args=args,
+        mage_repo=mage_repo,
+        media_path=media_path,
+        cmd=cmd,
+        result=result,
+        elapsed_s=elapsed_s,
+    )
+    print(f"summary saved to {summary_path}")
+    if result is not None:
+        if result.stdout:
+            print(result.stdout, end="")
+        if result.returncode != 0:
+            if result.stderr:
+                print(result.stderr, file=sys.stderr, end="")
+            raise SystemExit(result.returncode)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/mage_vl/duplex_client.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import base64
+import json
+import time
+from pathlib import Path
+
+import websockets
+
+
+async def run(args: argparse.Namespace) -> None:
+    headers = {"Authorization": f"Bearer {args.auth_token}"} if args.auth_token else None
+    started = time.perf_counter()
+    async with websockets.connect(
+        f"{args.url}?session_id={args.session_id}",
+        additional_headers=headers,
+        max_size=args.max_message_bytes,
+        open_timeout=args.timeout,
+    ) as websocket:
+        created = json.loads(await asyncio.wait_for(websocket.recv(), args.timeout))
+        print(json.dumps(created, ensure_ascii=False))
+        await websocket.send(
+            json.dumps(
+                {
+                    "type": "input.append",
+                    "modality": "text",
+                    "data": args.prompt,
+                }
+            )
+        )
+        encoded = base64.b64encode(Path(args.video).read_bytes()).decode("ascii")
+        await websocket.send(
+            json.dumps(
+                {
+                    "type": "input.append",
+                    "modality": "video",
+                    "data": {
+                        "kind": "video/mp4",
+                        "segment_id": args.segment_id,
+                        "pts_ms": args.pts_ms,
+                        "duration_ms": args.duration_ms,
+                        "data": {"video_base64": encoded},
+                    },
+                }
+            )
+        )
+        while True:
+            event = json.loads(await asyncio.wait_for(websocket.recv(), args.timeout))
+            print(json.dumps(event, ensure_ascii=False))
+            if event.get("type") in {"response.done", "error"}:
+                break
+        await websocket.send(json.dumps({"type": "close"}))
+    print(f"elapsed_s={time.perf_counter() - started:.3f}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Send one Mage-VL video window over the duplex WebSocket.")
+    parser.add_argument("--url", default="ws://127.0.0.1:8090/v1/mage-vl/duplex")
+    parser.add_argument("--video", required=True)
+    parser.add_argument("--prompt", default="Describe this video segment in detail.")
+    parser.add_argument("--session-id", default="mage-vl-example")
+    parser.add_argument("--segment-id", default="segment-0")
+    parser.add_argument("--pts-ms", type=int, default=0)
+    parser.add_argument("--duration-ms", type=int, default=4000)
+    parser.add_argument("--auth-token")
+    parser.add_argument("--timeout", type=float, default=120.0)
+    parser.add_argument("--max-message-bytes", type=int, default=16 * 1024 * 1024)
+    args = parser.parse_args()
+    asyncio.run(run(args))
+
+
+if __name__ == "__main__":
+    main()