from diffusers import StableDiffusionPipeline
import matplotlib.pyplot as plt

# Find models in https://huggingface.co/models?pipeline_tag=text-to-image&library=diffusers&sort=trending
model_id = "stabilityai/stable-diffusion-2-1"

pipe = StableDiffusionPipeline.from_pretrained(model_id)
pipe = pipe.to("cpu")

prompt = "beautiful horse"

image = pipe(prompt).images[0]

# Convert the torch Tensor to a NumPy array and move to CPU
image_np = image.cpu().numpy()

print("[PROMPT]: ", prompt)
plt.axis('off')
plt.imsave('generated_image.png', image_np.transpose(1, 2, 0))