diffusers/tools

1#!/usr/bin/env python32import torch3from diffusers import DiffusionPipeline4from diffusers import EulerAncestralDiscreteScheduler5 6import cProfile7import pstats8import io9from pstats import SortKey10 11path = 'stabilityai/stable-diffusion-2-1-base'12prompt = "Women standing on a mountain top"13 14torch.set_grad_enabled(False)15torch.backends.cudnn.benchmark = True16 17with torch.inference_mode():18    pipe = DiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False)19    pipe.to('cuda')20    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)21    pipe.unet.to(device='cuda', dtype=torch.float16, memory_format=torch.channels_last)22 23    for bi in range(7):24        if bi == 2:      # Start profiler on 3rd image25            ob = cProfile.Profile()26            ob.enable()27        images = pipe(prompt=prompt, width=512, height=512, num_inference_steps=20, num_images_per_prompt=1).images28    ob.disable()29    sec = io.StringIO()30    sortby = SortKey.TIME31    ps = pstats.Stats(ob, stream=sec).sort_stats(sortby)32    ps.print_stats()33    print(sec.getvalue()[0:1000])34