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