File size: 1,004 Bytes
f18e889 928edd4 f18e889 |
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 |
from io import BytesIO
from socket import socket, AF_UNIX, SOCK_STREAM
from sys import byteorder
from pipelines.models import TextToImageRequest
from pipeline import load_pipeline, infer
SOCKET = "/api/inferences.sock"
def main():
pipeline = load_pipeline()
with socket(AF_UNIX, SOCK_STREAM) as inference_socket:
inference_socket.bind(SOCKET)
inference_socket.listen(1)
connection, _ = inference_socket.accept()
with connection:
while True:
size = int.from_bytes(connection.recv(2), byteorder)
request = TextToImageRequest.model_validate_json(connection.recv(size).decode("utf-8"))
image = infer(request, pipeline)
data = BytesIO()
image.save(data, format=image.format)
packet = data.getvalue()
connection.send(len(packet).to_bytes(4, byteorder))
connection.send(packet)
if __name__ == '__main__':
main()
|