Signed-off-by: Sy03 <1370724210@qq.com> Co-authored-by: shimei.jmh <shimei.jmh@taobao.com> Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
This commit is contained in:
14 files changed
+3597
No files matched your search
@@ -0,0 +1,208 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Example WebSocket client for the /v1/video/chat/stream endpoint.
|
||||
|
||||
Sends video frames from a local file (or generates synthetic ones), submits a
|
||||
query, and prints the streamed text response.
|
||||
|
||||
Requirements:
|
||||
pip install websockets pillow
|
||||
|
||||
Usage:
|
||||
# With a video file (requires opencv-python):
|
||||
python streaming_video_client.py --video my_clip.mp4 \\
|
||||
--query "What is happening in this video?"
|
||||
|
||||
# Synthetic frames (no extra deps):
|
||||
python streaming_video_client.py \\
|
||||
--query "Describe what you see." \\
|
||||
--synthetic-frames 10
|
||||
|
||||
# With audio (Phase 3):
|
||||
python streaming_video_client.py --video my_clip.mp4 \\
|
||||
--audio my_audio.pcm \\
|
||||
--query "What is the person saying and doing?"
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
|
||||
try:
|
||||
import websockets
|
||||
except ImportError:
|
||||
print("Please install websockets: pip install websockets")
|
||||
sys.exit(1)
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def _generate_synthetic_frame(index: int, width: int = 320, height: int = 240) -> bytes:
|
||||
"""Generate a simple synthetic JPEG frame with a colour gradient."""
|
||||
r = (index * 37) % 256
|
||||
g = (index * 73) % 256
|
||||
b = (index * 113) % 256
|
||||
img = Image.new("RGB", (width, height), (r, g, b))
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="JPEG", quality=80)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def _load_video_frames(path: str, max_frames: int = 64, fps: int = 2) -> list[bytes]:
|
||||
"""Extract frames from a video file using OpenCV."""
|
||||
try:
|
||||
import cv2
|
||||
except ImportError:
|
||||
print("opencv-python is required to read video files: pip install opencv-python")
|
||||
sys.exit(1)
|
||||
|
||||
cap = cv2.VideoCapture(path)
|
||||
if not cap.isOpened():
|
||||
print(f"Cannot open video: {path}")
|
||||
sys.exit(1)
|
||||
|
||||
video_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||
frame_interval = max(1, int(video_fps / fps))
|
||||
|
||||
frames: list[bytes] = []
|
||||
idx = 0
|
||||
while len(frames) < max_frames:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
if idx % frame_interval == 0:
|
||||
_, buf = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
|
||||
frames.append(buf.tobytes())
|
||||
idx += 1
|
||||
|
||||
cap.release()
|
||||
print(f"Loaded {len(frames)} frames from {path} (interval={frame_interval})")
|
||||
return frames
|
||||
|
||||
|
||||
async def run(args: argparse.Namespace) -> None:
|
||||
uri = f"ws://{args.host}:{args.port}/v1/video/chat/stream"
|
||||
|
||||
# Prepare frames
|
||||
if args.video:
|
||||
frames = _load_video_frames(args.video, max_frames=args.max_frames, fps=args.fps)
|
||||
else:
|
||||
frames = [_generate_synthetic_frame(i) for i in range(args.synthetic_frames)]
|
||||
print(f"Generated {len(frames)} synthetic frames")
|
||||
|
||||
# Prepare audio (optional, Phase 3)
|
||||
audio_data: bytes | None = None
|
||||
if args.audio:
|
||||
with open(args.audio, "rb") as f:
|
||||
audio_data = f.read()
|
||||
print(f"Loaded audio: {len(audio_data)} bytes")
|
||||
|
||||
async with websockets.connect(uri, max_size=16 * 1024 * 1024) as ws:
|
||||
# 1. Send session.config
|
||||
config = {
|
||||
"type": "session.config",
|
||||
"model": args.model,
|
||||
"modalities": ["text", "audio"] if audio_data else ["text"],
|
||||
"max_frames": args.max_frames,
|
||||
"num_frames": args.num_sample_frames,
|
||||
"enable_frame_filter": args.evs,
|
||||
"frame_filter_threshold": args.evs_threshold,
|
||||
"use_audio_in_video": bool(audio_data),
|
||||
}
|
||||
await ws.send(json.dumps(config))
|
||||
print(f"Sent session.config: model={args.model} evs={args.evs}")
|
||||
|
||||
# 2. Send frames
|
||||
for i, frame in enumerate(frames):
|
||||
msg = {
|
||||
"type": "video.frame",
|
||||
"data": base64.b64encode(frame).decode(),
|
||||
}
|
||||
await ws.send(json.dumps(msg))
|
||||
if (i + 1) % 10 == 0:
|
||||
print(f" Sent {i + 1}/{len(frames)} frames")
|
||||
print(f"Sent all {len(frames)} frames")
|
||||
|
||||
# 3. Send audio chunks (Phase 3)
|
||||
if audio_data:
|
||||
chunk_size = 16000 * 2 # 1 second of 16 kHz 16-bit PCM
|
||||
for offset in range(0, len(audio_data), chunk_size):
|
||||
chunk = audio_data[offset : offset + chunk_size]
|
||||
msg = {
|
||||
"type": "audio.chunk",
|
||||
"data": base64.b64encode(chunk).decode(),
|
||||
}
|
||||
await ws.send(json.dumps(msg))
|
||||
print(f"Sent audio in {(len(audio_data) + chunk_size - 1) // chunk_size} chunks")
|
||||
|
||||
# 4. Send query, then immediately send video.done so the server
|
||||
# knows the session is complete (avoids deadlock where client
|
||||
# waits for session.done while server waits for video.done).
|
||||
await ws.send(json.dumps({"type": "video.query", "text": args.query}))
|
||||
print(f"\nQuery: {args.query}")
|
||||
print("Response: ", end="", flush=True)
|
||||
|
||||
# Signal end of session right after the query. The server will
|
||||
# process the query first (it's already queued), then handle
|
||||
# video.done and reply with session.done.
|
||||
await ws.send(json.dumps({"type": "video.done"}))
|
||||
|
||||
# 5. Receive response until session.done
|
||||
recv_timeout = 120 # seconds — avoid infinite hang if server stalls
|
||||
while True:
|
||||
raw = await asyncio.wait_for(ws.recv(), timeout=recv_timeout)
|
||||
data = json.loads(raw)
|
||||
msg_type = data.get("type")
|
||||
|
||||
if msg_type == "response.text.delta":
|
||||
print(data.get("delta", ""), end="", flush=True)
|
||||
elif msg_type == "response.text.done":
|
||||
print() # newline
|
||||
elif msg_type == "response.evs_stats":
|
||||
retained = data.get("retained_count", 0)
|
||||
dropped = data.get("dropped_count", 0)
|
||||
rate = data.get("drop_rate", 0)
|
||||
print(f"\nEVS stats: retained={retained} dropped={dropped} drop_rate={rate:.1%}")
|
||||
elif msg_type == "session.done":
|
||||
print("Session complete.")
|
||||
break
|
||||
elif msg_type == "error":
|
||||
print(f"\nError: {data.get('message')}")
|
||||
break
|
||||
elif msg_type == "response.start":
|
||||
pass # expected
|
||||
else:
|
||||
print(f"\n[unknown message] {data}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Streaming video chat client")
|
||||
parser.add_argument("--host", default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--model", default="Qwen/Qwen3-Omni-MoE")
|
||||
parser.add_argument("--video", help="Path to video file (requires opencv-python)")
|
||||
parser.add_argument("--audio", help="Path to raw PCM 16kHz audio file (Phase 3)")
|
||||
parser.add_argument("--query", default="What do you see in this video?")
|
||||
parser.add_argument(
|
||||
"--synthetic-frames", type=int, default=10, help="Number of synthetic frames if --video is not set"
|
||||
)
|
||||
parser.add_argument("--max-frames", type=int, default=64)
|
||||
parser.add_argument("--num-sample-frames", type=int, default=16)
|
||||
parser.add_argument("--fps", type=int, default=2, help="Frame extraction rate from video")
|
||||
parser.add_argument(
|
||||
"--no-evs", dest="evs", action="store_false", help="Disable EVS frame filtering (enabled by default)"
|
||||
)
|
||||
parser.set_defaults(evs=True)
|
||||
parser.add_argument("--evs-threshold", type=float, default=0.95)
|
||||
args = parser.parse_args()
|
||||
asyncio.run(run(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in new issue
Block a user