ZhanshenzhuiandSy03
48ae412a2f
Signed-off-by: Zhanshenzhui <99891002+zyforsure@users.noreply.github.com> Co-authored-by: Sy03 <1370724210@qq.com>
260 lines
9.3 KiB
Python
260 lines
9.3 KiB
Python
#!/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()
|