Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
85 lines
4.5 KiB
Python
85 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Exercise real FlashHead video artifacts through offline HTTP and streaming SSE."""
|
|
import argparse
|
|
import base64
|
|
import json
|
|
from pathlib import Path
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
import wave
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'server'))
|
|
from server_test_utils import Server
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('--server', type=Path, required=True)
|
|
parser.add_argument('--model', type=Path, required=True)
|
|
parser.add_argument('--portrait', type=Path, required=True)
|
|
parser.add_argument('--audio', type=Path, required=True)
|
|
parser.add_argument('--output', type=Path, required=True)
|
|
args = parser.parse_args()
|
|
out = args.output.resolve()
|
|
models = [{'id': mode, 'family': 'soulx_flashhead', 'path': str(args.model.resolve()),
|
|
'task': 'gen', 'mode': mode} for mode in ('offline', 'streaming')]
|
|
request = {'audio': str(args.audio.resolve()), 'options': {'reference_image_file': str(args.portrait.resolve()),
|
|
'height': 256, 'width': 256, 'seed': 42}}
|
|
with wave.open(str(args.audio), 'rb') as audio:
|
|
assert audio.getframerate() == 16000 and audio.getnchannels() == 1
|
|
expected_chunks = (audio.getnframes() + 15359) // 15360
|
|
streamed = []
|
|
events = []
|
|
with Server(args.server.resolve(), models, 'cuda', 0, out,
|
|
extra_args=('--log', '--log-file', str(out / 'timing.log'))) as server:
|
|
started = time.perf_counter()
|
|
body = {'model': 'streaming', 'stream_format': 'sse', 'include_artifacts': True, 'request': request}
|
|
http = urllib.request.Request(server.url + '/v1/tasks/stream', data=json.dumps(body).encode(),
|
|
headers={'Content-Type': 'application/json'})
|
|
completed = False
|
|
done_result = None
|
|
with urllib.request.urlopen(http, timeout=180) as response:
|
|
assert response.status == 200
|
|
assert response.headers['Content-Type'].startswith('text/event-stream')
|
|
for line in response:
|
|
if not line.startswith(b'data: '):
|
|
continue
|
|
payload = line[6:].strip()
|
|
if payload == b'[DONE]':
|
|
completed = True
|
|
break
|
|
event = json.loads(payload)
|
|
assert event.get('type') != 'error', event
|
|
events.append({'type': event['type'], 'http_ms': (time.perf_counter()-started)*1000})
|
|
if event['type'] == 'task.stream.event':
|
|
assert done_result is None, 'video event arrived after completion'
|
|
streamed.extend(event['event'].get('artifacts', []))
|
|
elif event['type'] == 'task.stream.done':
|
|
done_result = event['result']
|
|
assert completed and done_result is not None
|
|
assert done_result['timing']['ttft_ms'] is not None, 'artifact-only output has no TTFA'
|
|
assert not done_result.get('artifacts'), 'video chunks duplicated in the final result'
|
|
assert len(streamed) == expected_chunks, (len(streamed), expected_chunks)
|
|
assert events[0]['type'] == 'task.stream.event' and events[-1]['type'] == 'task.stream.done'
|
|
status, offline = server.api('/v1/tasks/run', {'model': 'offline', 'request': request})
|
|
assert status == 200, offline
|
|
assert len(offline['artifacts']) == len(streamed)
|
|
artifacts_dir = out / 'artifacts'
|
|
artifacts_dir.mkdir(exist_ok=True)
|
|
for index, (actual, expected) in enumerate(zip(streamed, offline['artifacts'])):
|
|
data = base64.b64decode(actual['payload'], validate=True)
|
|
assert data == base64.b64decode(expected['payload'], validate=True), 'HTTP offline/streaming parity failed'
|
|
meta = actual['meta']
|
|
assert int(meta['chunk_index']) == index and int(meta['frames']) == 24
|
|
assert len(data) == int(meta['width'])*int(meta['height'])*int(meta['frames'])*3
|
|
artifact = {'id': actual['id'], 'kind': actual['kind'], 'meta': meta, 'payload_hex': data.hex()}
|
|
(artifacts_dir / f'{actual["id"]}.json').write_text(json.dumps(artifact))
|
|
report = {'events': events, 'chunks': len(streamed), 'ttfa_ms': done_result['timing']['ttft_ms'],
|
|
'offline_streaming_byte_identical': True}
|
|
(out / 'results.json').write_text(json.dumps(report, indent=2)+'\n')
|
|
print(json.dumps(report, indent=2))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|