src/main.py

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()