Update src/pipeline.py

src/pipeline.py

@@ -1314,35 +1314,24 @@ class StableDiffusionXLPipeline(
 
         return StableDiffusionXLPipelineOutput(images=image)
 
-
-def load_pipeline() -> StableDiffusionXLPipeline:
-    pipeline = StableDiffusionXLPipeline.from_pretrained(
-        "./models/newdream-sdxl-20",
-        torch_dtype=torch.float16,
-        local_files_only=True,
-    ).to("cuda")
+from onediffx import compile_pipe
+
+def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
+    if not pipeline:
+        pipeline = StableDiffusionXLPipeline.from_pretrained(
+            "./models/newdream-sdxl-20",
+            torch_dtype=torch.float16,
+            local_files_only=True,
+        ).to("cuda")
     pipeline.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.float16).to('cuda')
-    config = CompilationConfig.Default()
-    # xformers and Triton are suggested for achieving best performance.
-    try:
-        import xformers
-        config.enable_xformers = True
-    except ImportError:
-        print('xformers not installed, skip')
-    try:
-        import triton
-        config.enable_triton = True
-    except ImportError:
-        print('Triton not installed, skip')
-    config.enable_cuda_graph = True
-
-    pipeline = compile(pipeline, config)
+    pipeline = compile_pipe(pipeline)
     for _ in range(2):
         pipeline(prompt="", num_inference_steps=10)
 
     return pipeline
 
 
+
 def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> Image:
     generator = Generator(pipeline.device).manual_seed(request.seed) if request.seed else None
 
@@ -1352,5 +1341,5 @@ def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> I
         width=request.width,
         height=request.height,
         generator=generator,
-        num_inference_steps=9,
+        num_inference_steps=10,
     ).images[0]