from time import perf_counterimport numpy as np  def measure_latency(pipe, prompt):    latencies = []    # warm up    pipe.set_progress_bar_config(disable=True)    for _ in range(2):        _ =  pipe(prompt)    # Timed run    for _ in range(10):        start_time = perf_counter()        _ = pipe(prompt)        latency = perf_counter() - start_time        latencies.append(latency)    # Compute run statistics    time_avg_s = np.mean(latencies)    time_std_s = np.std(latencies)    time_p95_s = np.percentile(latencies,95)    return f"P95 latency (seconds) - {time_p95_s:.2f}; Average latency (seconds) - {time_avg_s:.2f} +\- {time_std_s:.2f};", time_p95_s

prompt = "a photo of an astronaut riding a horse on mars" vanilla_results = measure_latency(pipeline,prompt) print(f"Vanilla pipeline: {vanilla_results[0]}")#     Vanilla pipeline: P95 latency (seconds) - 4.57; Average latency (seconds) - 4.56 +\- 0.00;

import torchimport deepspeedfrom diffusers import DiffusionPipeline  # Model Repository on huggingface.coHF_MODEL_ID="CompVis/stable-diffusion-v1-4"HF_TOKEN="" # your hf token: https://huggingface.co/settings/tokensassert len(HF_TOKEN) > 0, "Please set HF_TOKEN to your huggingface token. You can find it here: https://huggingface.co/settings/tokens" # load vanilla pipelineds_pipeline = DiffusionPipeline.from_pretrained(HF_MODEL_ID, torch_dtype=torch.float16, revision="fp16",use_auth_token=HF_TOKEN) # init deepspeed inference enginedeepspeed.init_inference(    model=getattr(ds_pipeline,"model", ds_pipeline),      # Transformers models    mp_size=1,        # Number of GPU    dtype=torch.float16, # dtype of the weights (fp16)    replace_method="auto", # Lets DS autmatically identify the layer to replace    replace_with_kernel_inject=False, # replace the model with the kernel injector)print("DeepSpeed Inference Engine initialized") 

DSUNet(  (unet): UNet2DConditionModel(    (conv_in): Conv2d(4, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))    (time_proj): Timesteps()    (time_embedding): TimestepEmbedding(      (linear_1): Linear(in_features=320, out_features=1280, bias=True)      (act): SiLU()      (linear_2): Linear(in_features=1280, out_features=1280, bias=True)    )    (down_blocks): ModuleList(      (0): CrossAttnDownBlock2D(        (attentions): ModuleList(          (0): SpatialTransformer(            (norm): GroupNorm(32, 320, eps=1e-06, affine=True)            (proj_in): Conv2d(320, 320, kernel_size=(1, 1), stride=(1, 1))            (transformer_blocks): ModuleList(              (0): BasicTransformerBlock(                (attn1): DeepSpeedAttention(                  (triton_flash_attn_kernel): triton_flash_attn()                )

from deepspeed.ops.transformer.inference.triton_ops import triton_flash_attnfrom deepspeed.ops.transformer.inference import DeepSpeedAttention assert isinstance(ds_pipeline.unet.unet.down_blocks[0].attentions[0].transformer_blocks[0].attn1, DeepSpeedAttention) == True, "Model not sucessfully initalized"assert isinstance(ds_pipeline.unet.unet.down_blocks[0].attentions[0].transformer_blocks[0].attn1.triton_flash_attn_kernel, triton_flash_attn) == True, "Model not sucessfully initalized"

prompt = "a photo of an astronaut riding a horse on mars" ds_results = measure_latency(ds_pipeline,prompt) print(f"DeepSpeed model: {ds_results[0]}")#     DeepSpeed model: P95 latency (seconds) - 2.68; Average latency (seconds) - 2.67 +\- 0.00;

ds_pipeline("a photo of an astronaut riding a horse on mars").images[0]
We managed to accelerate the CompVis/stable-diffusion-v1-4 pipeline latency from 4.57s to 2.68s for generating a 512x512
Conclusion
We successfully optimized our Stable Diffusion with DeepSpeed-inference and managed to decrease our model latency from 4.57s to 2.68s or 1.7x.Those are good results results thinking of that we only needed to add 1 additional line of code, but applying the optimization was as easy as adding one additional call to deepspeed.init_inference.But I have to say that this isn't a plug-and-play process you can transfer to any Transformers model, task, or dataset. Also, make sure to check if your model is compatible with DeepSpeed-Inference.
If you want to learn more about Stable Diffusion you should check out:
Thanks for reading! If you have any questions, feel free to contact me on Twitter or LinkedIn.