diffusers_stable_diffusion_text_to_image_generation.py

1import torch
2from diffusers import DiffusionPipeline
3
4# Load the pipeline from a pre-trained model checkpoint
5pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
6
7# Move the pipeline to GPU (if available)
8pipe.to("cuda")
9
10# Define your prompt
11prompt = "An image of a squirrel in Picasso style"
12
13# Generate the image
14image = pipe(prompt).images[0]
15
16# Save the generated image
17image.save("image_of_squirrel_picasso.png")