Update src/pipeline.py

src/pipeline.py

@@ -1,15 +1,11 @@
 import torch
 from PIL.Image import Image
-from diffusers import StableDiffusionXLPipeline
-from sfast.compilers.diffusion_pipeline_compiler import (compile,
-                                                         CompilationConfig)
 from pipelines.models import TextToImageRequest
 from torch import Generator
 
 import inspect
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
-import torch
 from transformers import (
     CLIPImageProcessor,
     CLIPTextModel,
@@ -809,7 +805,7 @@ class StableDiffusionXLPipeline(
         return_dict: bool = True,
         cross_attention_kwargs: Optional[Dict[str, Any]] = None,
         guidance_rescale: float = 0.0,
-        end_cfg: float = 0.4,
+        end_cfg: float = 0.73,
         original_size: Optional[Tuple[int, int]] = None,
         crops_coords_top_left: Tuple[int, int] = (0, 0),
         target_size: Optional[Tuple[int, int]] = None,
@@ -1245,31 +1241,26 @@ class StableDiffusionXLPipeline(
 
         return StableDiffusionXLPipelineOutput(images=image)
 
+from onediff.schedulers import EulerDiscreteScheduler
+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,
+        )
+    pipeline.scheduler = EulerDiscreteScheduler.from_config(
+    pipeline.scheduler.config,)
 
-def load_pipeline() -> StableDiffusionXLPipeline:
-    pipeline = StableDiffusionXLPipeline.from_pretrained(
-        "./models/newdream-sdxl-20",
-        torch_dtype=torch.float16,
-        local_files_only=True,
-    )
-#    pipeline.scheduler = UniPCMultistepScheduler.from_config('./src',)
     pipeline.to("cuda")
-    config = CompilationConfig.Default()
-    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(4):
-        pipeline(prompt="kamala harris", num_inference_steps=20)
+        pipeline(prompt="kamala harris flying to the moon", num_inference_steps=20, end_cfg=0.73)
+
+
 
     return pipeline
 
@@ -1286,5 +1277,6 @@ def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> I
         width=request.width,
         height=request.height,
         generator=generator,
-        num_inference_steps=20,
+        end_cfg=0.73,
+        num_inference_steps=17,
     ).images[0]