diffusers_stable_diffusion_xl_text_to_image_generation.py

1import torch
2from diffusers import DiffusionPipeline
3
4# Load the pipeline
5pipe = DiffusionPipeline.from_pretrained(
6    "stabilityai/stable-diffusion-xl-base-1.0", 
7    torch_dtype=torch.float16, 
8    use_safetensors=True, 
9    variant="fp16"
10)
11
12# Move to GPU
13pipe.to("cuda")
14
15# Run the pipeline
16prompt = "An astronaut riding a green horse"
17image = pipe(prompt=prompt).images[0]
18
19# Save the image
20image.save("astronaut_rides_horse.png")