Update app.py

app.py

@@ -10,7 +10,12 @@ from typing import Tuple
 from diffusers import DiffusionPipeline
 import torch
 from flask import Flask, request, jsonify
-app = Flask(__name__)
+from fastapi import FastAPI, Request, HTTPException
+
+
+#app = Flask(__name__)
+app=FastApi()
+
 
 style_list = [
     {
@@ -92,9 +97,9 @@ def infer(prompt, negative="low_quality", style_name=None, guidance_scale=None):
     return image_b64
 
 
-@app.route('/', methods=['POST'])
-def generate_image():
-    data = request.get_json()
+@app.post("/")
+async def generate_image(request: Request):
+    data = await request.json()
     if 'prompt' in data and 'style_name' in data and 'guidance_scale' in data:
         prompt = data['prompt']
         style_name = data['style_name']
@@ -105,10 +110,11 @@ def generate_image():
         # Convert the first generated image to base64
         image_b64 = image_urls[0].split(",")[1]
         
-        return jsonify({'image_base64': image_b64})
+        return {"image_base64": image_b64}
     else:
-        return jsonify({'error': 'Missing required parameters'}), 400
+        raise HTTPException(status_code=400, detail="Missing required parameters")
   
 
-if __name__ == '__main__':
-    app.run(host='0.0.0.0',port=7860,debug=True)
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)