revamped code

image_generator.py

@@ -31,20 +31,34 @@
 #     print("Generated Image Path:", generate_image(prompt))
 
 
+# from diffusers import StableDiffusionPipeline
+# import torch
+
+# # Preload the model globally
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+# model = StableDiffusionPipeline.from_pretrained(
+#     "stabilityai/stable-diffusion-2-1-base",
+#     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+# )
+# model.to(device)
+
+# def generate_image(prompt):
+#     """Generate an image from a text prompt."""
+#     image = model(prompt).images[0]
+#     output_path = "output.png"
+#     image.save(output_path)
+#     return output_path
+
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Preload the model globally
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = StableDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-2-1-base",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-)
-model.to(device)
+# Load the Stable Diffusion pipeline on the CPU
+pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32)
+pipe.to("cpu")  # Force the pipeline to run on the CPU
 
 def generate_image(prompt):
     """Generate an image from a text prompt."""
-    image = model(prompt).images[0]
-    output_path = "output.png"
-    image.save(output_path)
-    return output_path
+    image = pipe(prompt).images[0]
+    image_path = "output_image.png"
+    image.save(image_path)
+    return image_path

requirements.txt

@@ -6,4 +6,5 @@ sentencepiece
 accelerate
 sacremoses
 --extra-index-url https://download.pytorch.org/whl/cu113
-torch
+torch
+pillow