Add files using upload-large-folder tool

src/pipeline.py

@@ -0,0 +1,54 @@
+import torch
+from pathlib import Path
+from PIL.Image import Image
+from diffusers import StableDiffusionXLPipeline, DDIMScheduler
+from pipelines.models import TextToImageRequest
+from torch import Generator
+from cache_diffusion import cachify
+from pipe.deploy import compile
+from loss import SchedulerWrapper
+
+generator = Generator(torch.device("cuda")).manual_seed(6969)
+prompt = "Make submissions great again"
+SDXL_DEFAULT_CONFIG = [
+            {
+                "wildcard_or_filter_func": lambda name: "down_blocks.2" not in name and"down_blocks.3" not in name and "up_blocks.2" not in name,
+                "select_cache_step_func": lambda step: (step % 2 != 0) and (step >= 8),
+                }]
+def load_pipeline() -> StableDiffusionXLPipeline:
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        "models/newdream-sdxl-20",torch_dtype=torch.float16, use_safetensors=True, local_files_only=True
+    ).to("cuda") 
+    compile(
+        pipe,
+        onnx_path=Path("./onnx"),
+        engine_path=Path("./engine"),
+        batch_size=1,
+    )
+    cachify.prepare(pipe, SDXL_DEFAULT_CONFIG)
+    cachify.enable(pipe)
+    pipe.scheduler = SchedulerWrapper(DDIMScheduler.from_config(pipe.scheduler.config))
+    with cachify.infer(pipe) as cached_pipe:
+        for _ in range(4):
+            pipe(prompt=prompt, num_inference_steps=20)
+        pipe.scheduler.prepare_loss()
+    cachify.disable(pipe)
+    return pipe
+
+def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> Image:
+
+    if request.seed is None:
+        generator = None
+    else:
+        generator = Generator(pipeline.device).manual_seed(request.seed)
+    cachify.enable(pipeline)
+    with cachify.infer(pipeline) as cached_pipe:
+        image = cached_pipe(
+            prompt=request.prompt,
+            negative_prompt=request.negative_prompt,
+            width=request.width,
+            height=request.height,
+            generator=generator,
+            num_inference_steps=13,
+            ).images[0]
+    return image