Use multiprocessing.connection

src/main.py

@@ -1,6 +1,5 @@
 from io import BytesIO
-from socket import socket, AF_UNIX, SOCK_STREAM
-from sys import byteorder
+from multiprocessing.connection import Listener
 from os import chmod
 
 from PIL.JpegImagePlugin import JpegImageFile
@@ -17,34 +16,15 @@ def main():
 
     print(f"Pipeline loaded")
 
-    with socket(AF_UNIX, SOCK_STREAM) as inference_socket:
-        inference_socket.bind(SOCKET)
-
+    with Listener(SOCKET) as listener:
         chmod(SOCKET, 0o777)
 
-        inference_socket.listen(1)
-
         print(f"Awaiting connections")
-
-        connection, _ = inference_socket.accept()
-
-        print(f"Connected")
-
-        with connection:
-            connection.setblocking(True)
+        with listener.accept() as connection:
+            print(f"Connected")
 
             while True:
-                size = int.from_bytes(connection.recv(2), byteorder)
-
-                print(f"Awaiting message of size {size}")
-
-                message = connection.recv(size).decode("utf-8")
-
-                if not message:
-                    print(f"Empty message received")
-                    continue
-
-                request = TextToImageRequest.model_validate_json(message)
+                request = TextToImageRequest.model_validate_json(connection.recv_bytes().decode("utf-8"))
 
                 image = infer(request, pipeline)
 
@@ -53,8 +33,7 @@ def main():
 
                 packet = data.getvalue()
 
-                connection.send(len(packet).to_bytes(4, byteorder))
-                connection.send(packet)
+                connection.send_bytes(packet)
 
 
 if __name__ == '__main__':