超越AITemplate，打平TensorRT，SD全系列模型加速框架stable-fast隆重登场

| Framework                              | SD 1.5    | SD 2.1    | SD XL (1024x1024) |
| -------------------------------------- | --------- | ----------| ----------------- |
| Vanilla PyTorch (2.1.0+cu118)          | 29.5 it/s | 32.4 it/s | 4.6 it/s          |
| torch.compile (2.1.0+cu118, NHWC UNet) | 40.0 it/s | 44.0 it/s | 6.1 it/s          |
| AITemplate                             | 44.2 it/s | untested  | untested          |
| OneFlow                                | 50.3 it/s | untested  | untested          |
| AUTO1111 WebUI                         | 17.2 it/s | 15.2 it/s | 3.6 it/s          |
| AUTO1111 WebUI (with SDPA)             | 24.5 it/s | 26.1 it/s | 4.3 it/s          |
| TensorRT (AUTO1111 WebUI)              | 40.8 it/s | untested  | untested          |
| TensorRT Official Demo                 | 52.6 it/s | untested  | untested          |
| Stable Fast (with xformers & Triton)   | 50.5 it/s | 53.3 it/s | 8.3 it/s          |

# Make sure you have CUDNN/CUBLAS installed.# https://developer.nvidia.com/cudnn# https://developer.nvidia.com/cublas# Install PyTorch with CUDA and other packages at firstpip3 install 'torch>=1.12.0' 'diffusers>=0.19.3' 'xformers>=0.0.20' 'triton>=2.1.0'# Windows user: Triton might be not available, you could skip it.# (Optional) Makes the build much fasterpip3 install ninja# Set TORCH_CUDA_ARCH_LIST if running and building on different GPU types# You can also install the latest stable release from PyPI# pip3 install -v -U stable-fastpip3 install -v -U git+https://github.com/chengzeyi/stable-fast.git@main#egg=stable-fast# (this can take dozens of minutes)

import timeimport torchfrom diffusers import (StableDiffusionPipeline,
                       EulerAncestralDiscreteScheduler)from sfast.compilers.stable_diffusion_pipeline_compiler import (
    compile, CompilationConfig)def load_model():
    model = StableDiffusionPipeline.from_pretrained(
        'runwayml/stable-diffusion-v1-5',
        torch_dtype=torch.float16)

    model.scheduler = EulerAncestralDiscreteScheduler.from_config(
        model.scheduler.config)
    model.safety_checker = None
    model.to(torch.device('cuda'))
    return modelmodel = load_model()config = CompilationConfig.Default()# xformers and Triton are suggested for achieving best performance.try:
    import xformers
    config.enable_xformers = Trueexcept ImportError:
    print('xformers not installed, skip')try:
    import triton
    config.enable_triton = Trueexcept ImportError:
    print('Triton not installed, skip')# CUDA Graph is suggested for small batch sizes and small resolutions to reduce CPU overhead.config.enable_cuda_graph = Truemodel = compile(model, config)kwarg_inputs = dict(
    prompt=
    '(masterpiece:1,2), best quality, masterpiece, best detail face, a beautiful girl',
    height=512,
    width=512,
    num_inference_steps=30,
    num_images_per_prompt=1,)# NOTE: Warm it up.# The initial calls will trigger compilation and might be very slow.# After that, it should be very fast.for _ in range(3):
    output_image = model(**kwarg_inputs).images[0]# Let's see it!# Note: Progress bar might work incorrectly due to the async nature of CUDA.begin = time.time()output_image = model(**kwarg_inputs).images[0]print(f'Inference time: {time.time() - begin:.3f}s')# Let's view it in terminal!from sfast.utils.term_image import print_imageprint_image(output_image, max_width=80)