| import argparse |
| import asyncio |
| import json |
| import struct |
| import time |
| import urllib.request |
|
|
| import numpy as np |
|
|
| FRAME_DIMENSION = 133 |
| BLOCK_HEADER = struct.Struct("<III") |
|
|
|
|
| class MotionCollector: |
| def __init__(self) -> None: |
| self.started_at = time.perf_counter() |
| self.first_block_seconds = None |
| self.total_frames_received = 0 |
| self.blocks_received = 0 |
| self.last_block = None |
|
|
| def elapsed(self) -> float: |
| return time.perf_counter() - self.started_at |
|
|
| def consume(self, message) -> bool: |
| if isinstance(message, (bytes, bytearray)): |
| request_index, frame_offset, frame_count = BLOCK_HEADER.unpack_from(message, 0) |
| expected_length = BLOCK_HEADER.size + frame_count * FRAME_DIMENSION * 4 |
| if len(message) != expected_length: |
| print("Bad block length", len(message), "expected", expected_length) |
| return True |
|
|
| frames = np.frombuffer(message, dtype="<f4", offset=BLOCK_HEADER.size) |
| frames = frames.reshape(frame_count, FRAME_DIMENSION) |
|
|
| if self.first_block_seconds is None: |
| self.first_block_seconds = self.elapsed() |
|
|
| self.total_frames_received += frame_count |
| self.blocks_received += 1 |
| self.last_block = frames |
| return False |
|
|
| data = json.loads(message) |
| message_type = data.get("type") |
|
|
| if message_type == "hello": |
| print("hello:", "ready =", data.get("ready"), "fps =", data.get("fps"), "transport =", data.get("transport")) |
| return False |
|
|
| if message_type == "sentence": |
| print( |
| "sentence", data["sentence_index"] + 1, "of", data["sentence_count"], |
| "frames =", data["frame_count"], |
| "generation_ms =", data["generation_ms"], |
| "cache_hit =", data["cache_hit"], |
| "text =", data["text"], |
| ) |
| return False |
|
|
| if message_type == "done": |
| print("done: total_frames =", data["total_frames"], "server_elapsed_ms =", data["elapsed_ms"]) |
| return True |
|
|
| if message_type == "error": |
| print("error:", data) |
| return True |
|
|
| print("message:", data) |
| return False |
|
|
| def report(self) -> None: |
| print("blocks received:", self.blocks_received) |
| print("frames received:", self.total_frames_received) |
| if self.first_block_seconds is not None: |
| print("seconds until first motion block:", round(self.first_block_seconds, 3)) |
| print("seconds until finished:", round(self.elapsed(), 3)) |
| if self.last_block is not None: |
| print("last block shape:", self.last_block.shape) |
| print("last block value range:", float(self.last_block.min()), float(self.last_block.max())) |
|
|
|
|
| def build_generate_message(text: str) -> str: |
| return json.dumps({"type": "generate", "id": "test-1", "text": text}) |
|
|
|
|
| def http_request(url: str, token: str, payload=None) -> dict: |
| headers = {"Content-Type": "application/json"} |
| if token: |
| headers["Authorization"] = "Bearer " + token |
|
|
| request_body = None |
| if payload is not None: |
| request_body = json.dumps(payload).encode("utf-8") |
|
|
| request = urllib.request.Request(url, data=request_body, headers=headers) |
| with urllib.request.urlopen(request, timeout=60) as response: |
| return json.loads(response.read().decode("utf-8")) |
|
|
|
|
| async def run_websocket_test(base_url: str, text: str, token: str) -> None: |
| import websockets |
|
|
| websocket_url = base_url.replace("http", "ws", 1) + "/ws" |
| if token: |
| websocket_url += "?token=" + token |
|
|
| print("WEBSOCKET TEST", websocket_url) |
| collector = MotionCollector() |
|
|
| async with websockets.connect(websocket_url, max_size=None) as socket: |
| await socket.send(build_generate_message(text)) |
| collector.started_at = time.perf_counter() |
|
|
| while True: |
| message = await asyncio.wait_for(socket.recv(), timeout=180) |
| if collector.consume(message): |
| break |
|
|
| collector.report() |
|
|
|
|
| async def run_webrtc_test(base_url: str, text: str, token: str) -> None: |
| from aiortc import RTCConfiguration, RTCIceServer, RTCPeerConnection, RTCSessionDescription |
|
|
| print("WEBRTC TEST", base_url) |
|
|
| ice_response = await asyncio.to_thread(http_request, base_url + "/ice", token) |
| ice_servers = [] |
| for entry in ice_response["iceServers"]: |
| ice_servers.append( |
| RTCIceServer( |
| urls=entry["urls"], |
| username=entry.get("username"), |
| credential=entry.get("credential"), |
| ) |
| ) |
|
|
| peer = RTCPeerConnection(configuration=RTCConfiguration(iceServers=ice_servers)) |
| channel = peer.createDataChannel("motion", ordered=True) |
| collector = MotionCollector() |
| finished = asyncio.get_event_loop().create_future() |
|
|
| @channel.on("open") |
| def on_open() -> None: |
| print("data channel open") |
| collector.started_at = time.perf_counter() |
| channel.send(build_generate_message(text)) |
|
|
| @channel.on("message") |
| def on_message(message) -> None: |
| if collector.consume(message) and not finished.done(): |
| finished.set_result(True) |
|
|
| offer = await peer.createOffer() |
| await peer.setLocalDescription(offer) |
|
|
| answer = await asyncio.to_thread( |
| http_request, |
| base_url + "/webrtc/offer", |
| token, |
| {"sdp": peer.localDescription.sdp, "type": peer.localDescription.type}, |
| ) |
|
|
| await peer.setRemoteDescription(RTCSessionDescription(sdp=answer["sdp"], type=answer["type"])) |
|
|
| try: |
| await asyncio.wait_for(finished, timeout=180) |
| finally: |
| await peer.close() |
|
|
| collector.report() |
|
|
|
|
| async def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--base", default="http://127.0.0.1:7860") |
| parser.add_argument("--text", default="Hello. How are you today? The doctor will help you now.") |
| parser.add_argument("--mode", choices=["websocket", "webrtc", "both"], default="both") |
| parser.add_argument("--token", default="") |
| arguments = parser.parse_args() |
|
|
| if arguments.mode in ("websocket", "both"): |
| await run_websocket_test(arguments.base, arguments.text, arguments.token) |
|
|
| if arguments.mode in ("webrtc", "both"): |
| await run_webrtc_test(arguments.base, arguments.text, arguments.token) |
|
|
|
|
| if __name__ == "__main__": |
| asyncio.run(main()) |
|
|
