diffusers/benchmark-pt2.1
0
1import gradio as gr2import torch3from diffusers import (4 AutoPipelineForText2Image,5 StableDiffusionXLControlNetPipeline,6 DiffusionPipeline,7 StableDiffusionImg2ImgPipeline,8 StableDiffusionInpaintPipeline,9 StableDiffusionAdapterPipeline,10 StableDiffusionControlNetPipeline,11 StableDiffusionXLAdapterPipeline,12 StableDiffusionXLImg2ImgPipeline,13 StableDiffusionXLInpaintPipeline,14 ControlNetModel,15 T2IAdapter,16)17import time18import utils19 20 21dtype = torch.float1622device = torch.device("cuda")23 24# pipeline_to_benchmark, batch_size, use_channels_last, do_torch_compile25# examples = [["SD T2I", 4, True, True]]26 27pipeline_mapping = {28 "SD T2I": (DiffusionPipeline, "runwayml/stable-diffusion-v1-5"),29 "SD I2I": (StableDiffusionImg2ImgPipeline, "runwayml/stable-diffusion-v1-5"),30 "SD Inpainting": (31 StableDiffusionInpaintPipeline,32 "runwayml/stable-diffusion-inpainting",33 ),34 "SD ControlNet": (35 StableDiffusionControlNetPipeline,36 "runwayml/stable-diffusion-v1-5",37 "lllyasviel/sd-controlnet-canny",38 ),39 "SD T2I Adapters": (40 StableDiffusionAdapterPipeline,41 "CompVis/stable-diffusion-v1-4",42 "TencentARC/t2iadapter_canny_sd14v1",43 ),44 "SDXL T2I": (DiffusionPipeline, "stabilityai/stable-diffusion-xl-base-1.0"),45 "SDXL I2I": (46 StableDiffusionXLImg2ImgPipeline,47 "stabilityai/stable-diffusion-xl-base-1.0",48 ),49 "SDXL Inpainting": (50 StableDiffusionXLInpaintPipeline,51 "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",52 ),53 "SDXL ControlNet": (54 StableDiffusionXLControlNetPipeline,55 "stabilityai/stable-diffusion-xl-base-1.0",56 "diffusers/controlnet-canny-sdxl-1.0",57 ),58 "SDXL T2I Adapters": (59 StableDiffusionXLAdapterPipeline,60 "stabilityai/stable-diffusion-xl-base-1.0",61 "TencentARC/t2i-adapter-canny-sdxl-1.0",62 ),63 "Kandinsky 2.2 (T2I)": (64 AutoPipelineForText2Image,65 "kandinsky-community/kandinsky-2-2-decoder",66 ),67 "Würstchen (T2I)": (AutoPipelineForText2Image, "warp-ai/wuerstchen"),68}69 70 71def load_pipeline(72 pipeline_to_benchmark: str,73 use_channels_last: bool = False,74 do_torch_compile: bool = False,75):76 # Get pipeline details.77 print(f"Loading pipeline: {pipeline_to_benchmark}")78 pipeline_details = pipeline_mapping[pipeline_to_benchmark]79 pipeline_cls = pipeline_details[0]80 pipeline_ckpt = pipeline_details[1]81 82 # Load adapter if needed.83 if "ControlNet" in pipeline_to_benchmark:84 controlnet_ckpt = pipeline_details[2]85 controlnet = ControlNetModel.from_pretrained(86 controlnet_ckpt, torch_dtype=dtype87 ).to(device)88 elif "Adapters" in pipeline_to_benchmark:89 adapter_clpt = pipeline_details[2]90 adapter = T2IAdapter.from_pretrained(adapter_clpt, torch_dtype=dtype).to(device)91 92 # Load pipeline.93 if (94 "ControlNet" not in pipeline_to_benchmark95 and "Adapters" not in pipeline_to_benchmark96 ):97 pipeline = pipeline_cls.from_pretrained(pipeline_ckpt, torch_dtype=dtype)98 99 elif "ControlNet" in pipeline_to_benchmark:100 pipeline = pipeline_cls.from_pretrained(101 pipeline_ckpt, controlnet=controlnet, torch_dtype=dtype102 )103 elif "Adapters" in pipeline_to_benchmark:104 pipeline = pipeline_cls.from_pretrained(105 pipeline_ckpt, adapter=adapter, torch_dtype=dtype106 )107 108 pipeline.to(device)109 110 # Optionally set memory layout.111 if use_channels_last:112 print("Setting memory layout.")113 if pipeline_to_benchmark != "Würstchen (T2I)":114 pipeline.unet.to(memory_format=torch.channels_last)115 elif pipeline_to_benchmark == "Würstchen (T2I)":116 pipeline.prior_prior.to(memory_format=torch.channels_last)117 pipeline.decoder.to(memory_format=torch.channels_last)118 119 if hasattr(pipeline, "controlnet"):120 pipeline.controlnet.to(memory_format=torch.channels_last)121 elif hasattr(pipeline, "adapter"):122 pipeline.adapter.to(memory_format=torch.channels_last)123 124 # Optional torch compilation.125 if do_torch_compile:126 print("Compiling pipeline.")127 if pipeline_to_benchmark != "Würstchen (T2I)":128 pipeline.unet = torch.compile(129 pipeline.unet, mode="reduce-overhead", fullgraph=True130 )131 elif pipeline_to_benchmark == "Würstchen (T2I)":132 pipeline.prior_prior = torch.compile(133 pipeline.prior_prior, mode="reduce-overhead", fullgraph=True134 )135 pipeline.decoder = torch.compile(136 pipeline.decoder, mode="reduce-overhead", fullgraph=True137 )138 139 if hasattr(pipeline, "controlnet"):140 pipeline.controlnet = torch.compile(141 pipeline.controlnet, mode="reduce-overhead", fullgraph=True142 )143 elif hasattr(pipeline, "adapter"):144 pipeline.adapter = torch.compile(145 pipeline.adapter, mode="reduce-overhead", fullgraph=True146 )147 148 print("Pipeline loaded.")149 pipeline.set_progress_bar_config(disable=True)150 return pipeline151 152 153def generate(154 pipeline_to_benchmark: str,155 num_images_per_prompt: int = 1,156 use_channels_last: bool = False,157 do_torch_compile: bool = False,158):159 if isinstance(pipeline_to_benchmark, list):160 # It can only happen when we don't select a pipeline to benchmark.161 raise ValueError(162 "pipeline_to_benchmark cannot be None. Please select a pipeline to benchmark."163 )164 print("Start...")165 print("Torch version", torch.__version__)166 print("Torch CUDA version", torch.version.cuda)167 168 pipeline = load_pipeline(169 pipeline_to_benchmark=pipeline_to_benchmark,170 use_channels_last=use_channels_last,171 do_torch_compile=do_torch_compile,172 )173 for _ in range(3):174 prompt = 77 * "a"175 num_inference_steps = 20176 call_args = dict(177 prompt=prompt,178 num_images_per_prompt=num_images_per_prompt,179 num_inference_steps=num_inference_steps,180 )181 182 if pipeline_to_benchmark in ["SD I2I", "SDXL I2I"]:183 image = utils.get_image_for_img_to_img(pipeline_to_benchmark)184 call_args.update({"image": image})185 elif "Inpainting" in pipeline_to_benchmark:186 image, mask_image = utils.get_image_for_inpainting(pipeline_to_benchmark)187 call_args.update({"image": image, "mask_image": mask_image})188 elif "ControlNet" in pipeline_to_benchmark:189 image = utils.get_image_for_controlnet(pipeline_to_benchmark)190 call_args.update({"image": image})191 elif "Adapters" in pipeline_to_benchmark:192 image = utils.get_image_for_adapters(pipeline_to_benchmark)193 call_args.update({"image": image})194 195 start_time = time.time()196 _ = pipeline(**call_args).images197 end_time = time.time()198 199 print(f"For {num_inference_steps} steps", end_time - start_time)200 print("Avg per step", (end_time - start_time) / num_inference_steps)201 202 return (203 f"Avg per step: {((end_time - start_time) / num_inference_steps):.4f} seconds."204 )205 206 207with gr.Blocks(css="style.css") as demo:208 do_torch_compile = gr.Checkbox(label="Enable torch.compile()?")209 use_channels_last = gr.Checkbox(label="Use `channels_last` memory layout?")210 pipeline_to_benchmark = gr.Dropdown(211 list(pipeline_mapping.keys()),212 value=None,213 multiselect=False,214 label="Pipeline to benchmark",215 )216 batch_size = gr.Slider(217 label="Number of images per prompt",218 minimum=1,219 maximum=16,220 step=1,221 value=1,222 )223 224 btn = gr.Button("Benchmark!").style(225 margin=False,226 rounded=(False, True, True, False),227 full_width=False,228 )229 result = gr.Text(label="Result")230 231 # gr.Examples(232 # examples=examples,233 # inputs=[pipeline_to_benchmark, batch_size, use_channels_last, do_torch_compile],234 # outputs=result,235 # fn=generate,236 # cache_examples=True,237 # )238 239 btn.click(240 fn=generate,241 inputs=[pipeline_to_benchmark, batch_size, use_channels_last, do_torch_compile],242 outputs=result,243 )244 245demo.launch(show_error=True)246