fixed image-generation code

image_generator.py

@@ -1,8 +1,27 @@
+# from diffusers import StableDiffusionPipeline
+
+# def generate_image(prompt):
+#     model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+#     model.to("cuda")  # Use GPU for faster generation
+#     image = model(prompt).images[0]
+#     image.save("output.png")
+#     return "output.png"
+
+# if __name__ == "__main__":
+#     prompt = "A friendly person saying 'How are you?'"
+#     print("Generated Image Path:", generate_image(prompt))
+
+
 from diffusers import StableDiffusionPipeline
+import torch
 
 def generate_image(prompt):
     model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
-    model.to("cuda")  # Use GPU for faster generation
+    
+    # Use GPU if available, otherwise fallback to CPU
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.to(device)
+    
     image = model(prompt).images[0]
     image.save("output.png")
     return "output.png"