Files
AI_A4000/video_worker/scripts/ws_smoke_test.py
2026-04-07 01:00:56 +08:00

38 lines
1018 B
Python

import asyncio
import json
import sys
import websockets
WS_URL = sys.argv[1] if len(sys.argv) > 1 else "ws://127.0.0.1:8010/ws/generate"
async def main() -> None:
payload = {
"action": "generate",
"payload": {
"prompt": "a lonely man walking in a rainy neon street, cinematic, handheld camera",
"negative_prompt": "blurry, deformed face, extra limbs, flicker",
"quality_mode": "preview",
"duration_sec": 1,
"width": 320,
"height": 240,
"fps": 8,
"steps": 8,
"seed": 123456,
},
}
async with websockets.connect(WS_URL, max_size=2**22) as ws:
await ws.send(json.dumps(payload, ensure_ascii=False))
while True:
msg = await ws.recv()
data = json.loads(msg)
print(json.dumps(data, ensure_ascii=False))
if data.get("event") == "result":
break
if __name__ == "__main__":
asyncio.run(main())