File size: 1,292 Bytes
c70fcb3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
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