Update src/main.py

src/main.py

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