Update main.py

main.py

@@ -1,4 +1,4 @@
-from flask import request
+from flask import url_for
 from diffusers import StableDiffusionPipeline
 import torch
 from fastapi import FastAPI, Response  
@@ -26,13 +26,18 @@ pipe.safety_checker = dummy
 
 @app.route('/')
 def hello():
-    return "Hello Art"
+    return {"Hello": Artist}
+
 
 @app.route('/gen')
 def generate_image():
     prompt = request.args.get('prompt')
     image = pipe(prompt).images[0]
+    # Save the image
+    image.save('static/image.png')
+
     # do something with the generated image
     image_data = image.tobytes().hex()
+    image_url = url_for('static', filename='image.png')
 
-    return {'image_data': image_data}
+    return {'image_data': image_data, 'image_url': image_url}