File size: 2,229 Bytes
3c436c6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
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 >= 12),
}]
def load_pipeline() -> StableDiffusionXLPipeline:
pipe = StableDiffusionXLPipeline.from_pretrained(
"stablediffusionapi/newdream-sdxl-20",torch_dtype=torch.float16, use_safetensors=True
).to("cuda")
compile(
pipe,
onnx_path=Path("/home/sandbox/.cache/huggingface/hub/models--RobertML--edge-onnx/snapshots/d56fa8ea1dc675b87de08eece735bc5ec80a247f"),
engine_path=Path("/home/sandbox/.cache/huggingface/hub/models--RobertML--edge-engine/snapshots/e0dd02be0c58057947801857c41839f76df2fc88"),
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
|