from diffusers import DiffusionPipeline
import torch

# Load the pipeline
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float16)

# If running on CUDA, ensure the model is moved to GPU for faster inference
pipe.to("cuda")

# Get the prompt from the user
prompt = input("What do you want to see? ")

# Generate the image
image = pipe(prompt).images[0]

# Display the image
image.show()

# Print a message
print("How cool is this!")