Update app.py

app.py

@@ -1,10 +1,12 @@
 from flask import Flask, request, jsonify
-from transformers import pipeline
+from diffusers import StableDiffusionPipeline
+import torch
 
 app = Flask(__name__)
 
-# Load Stable Diffusion model (You can use OpenAI DALL·E API instead if preferred)
-generator = pipeline('text-to-image', model='CompVis/stable-diffusion-v-1-4-original')
+# Load Stable Diffusion model from Hugging Face
+model_id = "CompVis/stable-diffusion-v-1-4-original"
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
 
 @app.route('/generate-image', methods=['POST'])
 def generate_image():
@@ -12,19 +14,20 @@ def generate_image():
     if not description:
         return jsonify({'error': 'No description provided'}), 400
     
-    # Generate image from description
-    image = generator(description)
+    # Generate image from description using Stable Diffusion
+    image = pipe(description).images[0]
     
-    # Assuming image is returned as a URL or base64-encoded string
-    image_url = save_image(image)  # Implement a function to save or serve the image
+    # Save or serve the image
+    image_path = save_image(image)  # Implement a function to save or serve the image
     
-    return jsonify({'image_url': image_url})
+    return jsonify({'image_url': image_path})
 
 def save_image(image):
-    # Save image and return the URL or path for it
-    image_path = '/path/to/save/image.jpg'
-    image[0].save(image_path)
+    # Save the generated image to a path
+    image_path = '/path/to/save/generated_image.png'
+    image.save(image_path)
     return image_path
 
 if __name__ == '__main__':
     app.run(debug=True)
+