Update src/pipeline.py

src/pipeline.py

@@ -6,6 +6,8 @@ from onediffx.deep_cache import StableDiffusionXLPipeline
 #from diffusers import StableDiffusionXLPipeline
 from pipelines.models import TextToImageRequest
 from torch import Generator
+from diffusers import DDIMScheduler
+from loss import SchedulerWrapper
 
 #from onediff.schedulers import EulerDiscreteScheduler
 from onediffx import compile_pipe
@@ -16,14 +18,11 @@ def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
         pipeline = StableDiffusionXLPipeline.from_pretrained(
             "stablediffusionapi/newdream-sdxl-20",
             torch_dtype=torch.float16,
-            local_files_only=True,
         )
- #   pipeline.scheduler = EulerDiscreteScheduler.from_config(
- #   pipeline.scheduler.config,)
-
     pipeline.to("cuda")
-    
+    pipeline.scheduler = SchedulerWrapper(DDIMScheduler.from_config(pipeline.scheduler.config))    
     pipeline = compile_pipe(pipeline)
+    pipeline.scheduler.prepare_loss()
     for _ in range(4):
         deepcache_output = pipeline(prompt="kamala harris defends my submission", output_type="pil", cache_interval=1, cache_layer_id=1, cache_block_id=0)
     return pipeline
@@ -41,7 +40,7 @@ def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> I
         height=request.height,
         generator=generator,
         end_cfg=0.5,
-        num_inference_steps=20,
+        num_inference_steps=14,
         cache_interval=1,
         cache_layer_id=1,
         cache_block_id=0,