Montey/python-edge

1import atexit2from io import BytesIO3from multiprocessing.connection import Listener4from os import chmod, remove5from os.path import abspath, exists6from pathlib import Path7 8import torch9 10from PIL.JpegImagePlugin import JpegImageFile11from pipelines.models import TextToImageRequest12 13from pipeline import load_pipeline, infer14 15SOCKET = abspath(Path(__file__).parent.parent / "inferences.sock")16 17 18def at_exit():19    torch.cuda.empty_cache()20 21 22def main():23    atexit.register(at_exit)24 25    print(f"Loading pipeline")26    pipeline = load_pipeline()27 28    print(f"Pipeline loaded, creating socket at '{SOCKET}'")29 30    if exists(SOCKET):31        remove(SOCKET)32 33    with Listener(SOCKET) as listener:34        chmod(SOCKET, 0o777)35 36        print(f"Awaiting connections")37        with listener.accept() as connection:38            print(f"Connected")39 40            while True:41                try:42                    request = TextToImageRequest.model_validate_json(connection.recv_bytes().decode("utf-8"))43                except EOFError:44                    print(f"Inference socket exiting")45 46                    return47 48                image = infer(request, pipeline)49 50                data = BytesIO()51                image.save(data, format=JpegImageFile.format)52 53                packet = data.getvalue()54 55                connection.send_bytes(packet)56 57 58if __name__ == '__main__':59    main()60