soft987/Alpha-ff50ca

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