Did an oopsie

main.py

@@ -1,17 +1,19 @@
 from diffusers import StableDiffusionPipeline
 import matplotlib.pyplot as plt
-import torch
 
 # Find models in https://huggingface.co/models?pipeline_tag=text-to-image&library=diffusers&sort=trending
-model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+model_id = "stabilityai/stable-diffusion-2-1"
 
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipe = pipe.to("cuda")
+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.imshow(image)
-plt.axis('off')
+plt.axis('off')
+plt.imsave('generated_image.png', image_np.transpose(1, 2, 0))

requirements.txt

@@ -1,2 +1,3 @@
 diffusers
-transformers
+transformers
+matplotlib