diffusion_lens.py

from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler

pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)

def get_images(prompt, skip_layers=0):
    print('inside get images')
    pipeline_output = pipeline(prompt, skip_layers=skip_layers, num_images_per_prompt=1)
    print('after pipeline')
    images = pipeline_output.images
    print('got images')
    return images