Update src/main.py

src/main.py

@@ -7,7 +7,7 @@ from pathlib import Path
 from PIL.JpegImagePlugin import JpegImageFile
 from pipelines.models import TextToImageRequest
 
-from pipeline import PipelineManager, TorchOptimizer
+from pipeline import load_pipeline, infer
 
 SOCKET = abspath(Path(__file__).parent.parent / "inferences.sock")
 
@@ -15,7 +15,7 @@ SOCKET = abspath(Path(__file__).parent.parent / "inferences.sock")
 def main():
     print(f"Loading pipeline")
 
-    pipeline_manager = PipelineManager()
+    pipeline = load_pipeline()
 
     print(f"Pipeline loaded! , creating socket at '{SOCKET}'")
 
@@ -37,7 +37,7 @@ def main():
 
                     return
 
-                result = pipeline_manager.run_inference(request)
+                result = infer(pipeline, request)
 
                 data = BytesIO()
                 image.save(data, format=JpegImageFile.format)