DiffusionLens / diffusion_lens.py
tokeron's picture
Update diffusion_lens.py
ea8200b verified
raw
history blame
710 Bytes
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
import torch
pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
# Check if CUDA is available and set the device accordingly
device = "cuda" if torch.cuda.is_available() else "cpu"
# Move the pipeline to the device
pipeline.to(device)
def get_images(prompt, skip_layers=0):
print('inside get images')
pipeline_output = pipeline(prompt, skip_layers=skip_layers, num_images_per_prompt=1, return_tensors=False)
print('after pipeline')
images = pipeline_output.images
print('got images')
return images