Montey/python-edge

1import torch2from PIL.Image import Image3from onediffx.deep_cache import StableDiffusionXLPipeline4from pipelines.models import TextToImageRequest5from torch import Generator6import oneflow as flow7from onediff.infer_compiler import oneflow_compile8from onediffx import compile_pipe, save_pipe, load_pipe9from diffusers import DDIMScheduler10from loss import SchedulerWrapper11 12def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:13    if not pipeline:14        pipeline = StableDiffusionXLPipeline.from_pretrained(15            "./models/newdream-sdxl-20",16            torch_dtype=torch.float16,17            local_files_only=True,18        )19    pipeline.to("cuda")20    pipeline.scheduler = SchedulerWrapper(DDIMScheduler.from_config(pipeline.scheduler.config))21    pipeline = compile_pipe(pipeline)22    pipeline.unet = oneflow_compile(pipeline.unet)23    24    load_pipe(pipeline,dir="cached_pipe")25    for _ in range(5):26        deepcache_output = pipeline(prompt="make submissions great again", cache_interval=1, cache_layer_id=0, cache_block_id=0, num_inference_steps=20)27    pipeline.scheduler.prepare_loss()28    return pipeline29 30def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> Image:31    if request.seed is None:32        generator = None33    else:34        generator = Generator(pipeline.device).manual_seed(request.seed)35 36    return pipeline(37        prompt=request.prompt,38        negative_prompt=request.negative_prompt,39        width=request.width,40        height=request.height,41        generator=generator,42        num_inference_steps=10,43        cache_interval=1,44        cache_layer_id=0,45        cache_block_id=0,46    ).images[0]47 48