Update app.py

app.py

@@ -1,9 +1,10 @@
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, send_from_directory
 import requests
 import io
 import random
 import os
 from PIL import Image
+from datetime import datetime
 
 app = Flask(__name__)
 
@@ -11,6 +12,8 @@ API_URL = "https://api-inference.huggingface.co/models/openskyml/dalle-3-xl"
 API_TOKEN = os.getenv("HF_READ_TOKEN")  # it is free
 headers = {"Authorization": f"Bearer {API_TOKEN}"}
 
+TEMP_DIR = "temp"
+
 
 def query(prompt, is_negative=False, steps=1, cfg_scale=6, seed=None):
     payload = {
@@ -26,6 +29,14 @@ def query(prompt, is_negative=False, steps=1, cfg_scale=6, seed=None):
     return image
 
 
+def save_image(image, prompt):
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    filename = f"{prompt}-{random.randint(1, 100000)}-{timestamp}.png"
+    filepath = os.path.join(TEMP_DIR, filename)
+    image.save(filepath, format='PNG')
+    return filename
+
+
 @app.route("/generate", methods=["POST"])
 def generate():
     try:
@@ -35,15 +46,11 @@ def generate():
         is_negative = True if negative_prompt else False
 
         image = query(prompt, is_negative=is_negative)
-
-        # Convert PIL Image to bytes
-        img_byte_array = io.BytesIO()
-        image.save(img_byte_array, format='PNG')
-        img_byte_array = img_byte_array.getvalue()
+        filename = save_image(image, prompt)
 
         response = {
             "success": True,
-            "image": img_byte_array.decode('latin-1')
+            "image_filename": filename
         }
     except Exception as e:
         response = {
@@ -54,5 +61,13 @@ def generate():
     return jsonify(response)
 
 
+@app.route("/temp/<filename>")
+def show_image(filename):
+    return send_from_directory(TEMP_DIR, filename)
+
+
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
+    if not os.path.exists(TEMP_DIR):
+        os.makedirs(TEMP_DIR)
+
+    app.run(debug=True, host="0.0.0.0", port=7860)