Update src/pipeline.py

src/pipeline.py

@@ -1252,10 +1252,9 @@ def load_pipeline() -> StableDiffusionXLPipeline:
         torch_dtype=torch.float16,
         local_files_only=True,
     )
-    pipeline.scheduler = UniPCMultistepScheduler.from_config(
-    pipeline.scheduler.config,)
+    pipeline.scheduler = UniPCMultistepScheduler.from_config('./src',)
     pipeline.to("cuda")
-
+    pipeline.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.float16).to('cuda')
     config = CompilationConfig.Default()
     try:
         import xformers
@@ -1271,7 +1270,7 @@ def load_pipeline() -> StableDiffusionXLPipeline:
 
     pipeline = compile(pipeline, config)
     for _ in range(4):
-        pipeline(prompt="", num_inference_steps=10)
+        pipeline(prompt="kamala harris", num_inference_steps=20)
 
     return pipeline
 
@@ -1285,5 +1284,19 @@ def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> I
         width=request.width,
         height=request.height,
         generator=generator,
-        num_inference_steps=8,
+        num_inference_steps=2,
+    ).images[0]
+
+
+
+def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> Image:
+    generator = Generator(pipeline.device).manual_seed(request.seed) if request.seed else None
+
+    return pipeline(
+        prompt=request.prompt,
+        negative_prompt=request.negative_prompt,
+        width=request.width,
+        height=request.height,
+        generator=generator,
+        num_inference_steps=2,
     ).images[0]