File size: 1,928 Bytes
14bd491 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
from pprint import pprint
from PIL import Image
import os
import time
import requests
import base64
import numpy as np
from io import BytesIO
import concurrent.futures
NUM_SEGMENTS = 10
tic = time.time()
video_dir = '/home/ubuntu/shared_storage/images'
frames = [(os.path.splitext(item)[0], os.path.join(video_dir, item)) for item in os.listdir(video_dir)]
frames = [item[1] for item in sorted(frames, key=lambda x: x[0])]
indices = np.linspace(start=0, stop=len(frames)-1, num=NUM_SEGMENTS).astype(int)
image_paths = [frames[ind] for ind in indices]
request = {}
byte_images = []
def image_path_handler(image_path):
img = Image.open(image_path)
byte_io = BytesIO()
img.save(byte_io, format='PNG')
encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
return encoded_image
with concurrent.futures.ThreadPoolExecutor(16) as executor:
byte_images = list(executor.map(image_path_handler, image_paths))
payload = {
"images": byte_images * 16,
"parameters": {
"max_new_tokens": 90,
"top_k": 4,
"top_p": None,
"temperature": 0.01,
"no_repeat_ngram_size": None,
}
}
IP = '127.0.0.1'
headers = {'Content-Type': 'application/json'}
tic1 = time.time()
def loop_request(port):
url = f'http://{IP}:{port}/generate'
response = requests.post(url, json=payload, headers=headers)
return response.text
num = 6
ports = []
for port in range(num):
ports.append(str(8000 + port))
# with concurrent.futures.ProcessPoolExecutor(num) as executor:
# results = list(executor.map(loop_request, ports))
# toc = time.time()
# pprint(results)
with concurrent.futures.ThreadPoolExecutor(num) as executor:
results = [executor.submit(loop_request, port) for port in ports]
toc = time.time()
for result in results:
print(result.result())
print('stage 1', tic1 - tic)
print('stage 2', toc - tic1)
|