| from diffusers import AutoencoderTiny, StableDiffusionXLPipeline |
| from .hacked_sdxl_pipeline import HackedSDXLPipeline |
| import torch |
|
|
| def fast_diffusion_pipeline(model_id = "stabilityai/sdxl-turbo", vae_id = "madebyollin/taesdxl", compile = False): |
| """ |
| :param compile: If true, does a bunch of stuff to make calls fast, but the first call will be very slow as a consequence |
| - If you use this, don't vary the batch size (probably) |
| """ |
|
|
| pipe = HackedSDXLPipeline.from_pretrained(model_id, torch_dtype = torch.float16) |
| pipe.set_progress_bar_config(disable=True) |
| pipe.cached_encode = None |
| pipe.vae = AutoencoderTiny.from_pretrained(vae_id, torch_dtype=torch.float16) |
|
|
| pipe.to('cuda') |
|
|
| if compile: |
| pipe.unet = torch.compile(pipe.unet) |
| pipe.vae.decode = torch.compile(pipe.vae.decode) |
| """ |
| from sfast.compilers.stable_diffusion_pipeline_compiler import (compile, CompilationConfig) |
| |
| config = CompilationConfig() |
| config.enable_jit = True |
| config.enable_jit_freeze = True |
| config.trace_scheduler = True |
| config.enable_cnn_optimization = True |
| config.preserve_parameters = False |
| config.prefer_lowp_gemm = True |
| |
| pipe = compile(pipe, config) |
| """ |
| return pipe |
|
|
