Update app.py

app.py

@@ -1,29 +1,33 @@
+from flask import Flask, request, jsonify, send_file
 from diffusers import StableDiffusionPipeline
 import torch
-from flask import Flask, request, jsonify
+from io import BytesIO
+from PIL import Image
 
 app = Flask(__name__)
 
 # Load the model
 model_id = "kothariyashhh/GenAi-Texttoimage"
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Use float32 if running on CPU
-torch_dtype = torch.float16 if device == "cuda" else torch.float32
-
-pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
-pipeline.to(device)
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pipe = pipe.to("cuda")  # Ensure you're using a GPU for inference
 
 @app.route("/generate", methods=["POST"])
 def generate_image():
-    data = request.get_json()
-    prompt = data.get("prompt", "A scenic landscape")
-    
-    image = pipeline(prompt).images[0]
-    image_path = "generated_image.png"
-    image.save(image_path)
-    
-    return jsonify({"image_url": image_path})
+    data = request.json
+    prompt = data.get("prompt")
+    if not prompt:
+        return jsonify({"error": "Prompt is required"}), 400
+
+    # Generate image
+    with torch.autocast("cuda"):
+        image = pipe(prompt).images[0]
+
+    # Save image to a BytesIO object
+    img_io = BytesIO()
+    image.save(img_io, format="PNG")
+    img_io.seek(0)
+
+    return send_file(img_io, mimetype="image/png")
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
+    app.run(host="0.0.0.0", port=5000)