diffusers/tools

1#!/usr/bin/env python32import matplotlib.pyplot as plt3import numpy as np4import seaborn as sns5 6sns.set(font_scale=1.1)7 8batch_sizes = {9    "fp16_4": {10       "A100": [4.75, 3.26, 3.24, 3.10],  # those values are made up11       "A10": [13.94, 9.81, 10.01, 9.35],12       "T4": [38.81, 30.09, 29.74, 27.55],13       "V100": [9.84, 8.16, 8.09, 7.65],14       "3090": [10.04, 7.82, 7.89, 7.47],15       "3090TI": [9.07, 7.14, 7.15, 6.81],16    },17    "fp16_16": {18       "A100": [18.95, 13.57, 13.67, 12.25],19       "A10": [0, 37.55, 38.31, 36.81],20       "T4": [0, 111.47, 113.26, 106.93],21       "V100": [0, 30.29, 29.84, 28.22],22       "3090": [0, 29.06, 29.06, 28.2],23       "3090TI": [0, 26.1, 26.28, 25.46],24    },25    "fp32_4": {26       "A100": [16.56, 12.42, 12.2, 11.84],27       "A10": [34.77, 27.63, 22.77, 22.07],28       "T4": [0, 85.72, 85.78, 84.48],29       "V100": [0, 25.73, 25.31, 24.7],30       "3090": [22.69, 21.45, 18.67, 18.09],31       "3090TI": [20.32, 19.31, 16.9, 16.37],32    },33    "fp32_16": {34       "A100": [0, 47.08, 46.27, 44.8],35       "A10": [0, 116.49, 88.56, 86.64],36       "T4": [0, 276.47, 280.26, 270.93],  # numbers are made up37       "V100": [0, 84.99, 84.73, 82.55],38       "3090": [0, 85.35, 72.37, 70.25],39       "3090TI": [0, 75.37, 65.25, 64.32],40    },41}42 43batch_size = 1644dtype = "fp32"45 46key = f"{dtype}_{batch_size}"47 48methods = {49    "Vanilla Attention": [x[0] for x in batch_sizes[key].values()],50    "xFormers": [x[1] for x in batch_sizes[key].values()],51    "PyTorch2.0 SDPA": [x[2] for x in batch_sizes[key].values()],52    "SDPA + torch.compile": [x[3] for x in batch_sizes[key].values()],53}54 55x = np.arange(len(batch_sizes[key]))  # the label locations56width = 0.1  # the width of the bars57multiplier = 058 59fig, ax = plt.subplots(constrained_layout=True)60 61for attribute, measurement in methods.items():62    offset = width * multiplier63    rects = ax.bar(x + offset, measurement, width, label=attribute)64    multiplier += 165 66# Add some text for labels, title and custom x-axis tick labels, etc.67ax.set_ylabel('Time (s)')68ax.set_title(f'Inference Speed at Batch Size={batch_size} for {dtype}')69ax.set_xticks(x + width, batch_sizes[key])70ax.legend(loc='upper left')71 72plt.show()73