FLUX-1-dev-serverless

Running

App Files Files Community

soiz commited on Nov 23, 2024

Commit

92b8b0d

verified ·

1 Parent(s): 2cfd327

Update app.py

Browse files

Files changed (1) hide show

app.py +62 -37

app.py CHANGED Viewed

@@ -1,31 +1,43 @@
-from flask import Flask, request, jsonify, send_file, render_template_string, make_response
-import requests
-import io
 import os
 import random
 from PIL import Image
 from deep_translator import GoogleTranslator
-app = Flask(__name__)
 API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
 API_TOKEN = os.getenv("HF_READ_TOKEN")
 headers = {"Authorization": f"Bearer {API_TOKEN}"}
 timeout = 50000  # タイムアウトを300秒に設定
-# Function to query the API and return the generated image
-def query(prompt, negative_prompt="", steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024, num_inference_steps=30, guidance_scale=7.5, top_k=50, top_p=0.9, eta=0.1):
     if not prompt:
         return None, "Prompt is required"
     key = random.randint(0, 999)
     # Translate the prompt from Russian to English if necessary
     prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
     print(f'Generation {key} translation: {prompt}')
-    # Add some extra flair to the prompt
     prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
     print(f'Generation {key}: {prompt}')
@@ -48,28 +60,24 @@ def query(prompt, negative_prompt="", steps=35, cfg_scale=7, sampler="DPM++ 2M K
         }
     }
-    for attempt in range(3):  # 最大3回の再試行
         try:
-            response = requests.post(API_URL, json=payload, headers=headers, timeout=timeout)
-            if response.status_code != 200:
-                return None, f"Error: Failed to get image. Status code: {response.status_code}, Details: {response.text}"
-            image_bytes = response.content
-            image = Image.open(io.BytesIO(image_bytes))
-            return image, None
-        except requests.exceptions.Timeout:
-            if attempt < 2:  # 最後の試行でない場合は再試行
-                print("Timeout occurred, retrying...")
-                continue
             return None, "Error: The request timed out. Please try again."
-        except requests.exceptions.RequestException as e:
-            return None, f"Request Exception: {str(e)}"
         except Exception as e:
-            return None, f"Error when trying to open the image: {e}"
 # Content-Security-Policyヘッダーを設定するための関数
 @app.after_request
-def add_security_headers(response):
     response.headers['Content-Security-Policy'] = (
         "default-src 'self'; "
         "connect-src 'self' ^https?:\/\/[\w.-]+\.[\w.-]+(\/[\w.-]*)*(\?[^\s]*)?$"
@@ -79,22 +87,30 @@ def add_security_headers(response):
     )
     return response
-# HTML template for the index page
 index_html = """
 """
 @app.route('/')
-def index():
-    return render_template_string(index_html)
 @app.route('/generate', methods=['GET'])
-def generate_image():
-    if request.headers.getlist("X-Forwarded-For"):
-        client_ip = request.headers.getlist("X-Forwarded-For")[0]
-    else:
-        client_ip = request.remote_addr
-    print(f"Client IP: {client_ip}")
     prompt = request.args.get("prompt", "")
     negative_prompt = request.args.get("negative_prompt", "")
     steps = int(request.args.get("steps", 35))
@@ -110,7 +126,12 @@ def generate_image():
     top_p = float(request.args.get("top_p", 0.9))
     eta = float(request.args.get("eta", 0.1))
-    image, error = query(prompt, negative_prompt, steps, cfg_scale, sampler, seed, strength, width, height, num_inference_steps, guidance_scale, top_k, top_p, eta)
     if error:
         return jsonify({"error": error}), 400
@@ -118,7 +139,11 @@ def generate_image():
     img_bytes = io.BytesIO()
     image.save(img_bytes, format='PNG')
     img_bytes.seek(0)
-    return send_file(img_bytes, mimetype='image/png')
 if __name__ == "__main__":
     app.run(host='0.0.0.0', port=7860)

 import os
 import random
+import io
 from PIL import Image
 from deep_translator import GoogleTranslator
+import aiohttp
+from quart import Quart, request, jsonify, send_file, render_template_string
+from werkzeug.contrib.cache import SimpleCache
+# アプリケーションの設定
+app = Quart(__name__)
+cache = SimpleCache()
 API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
 API_TOKEN = os.getenv("HF_READ_TOKEN")
 headers = {"Authorization": f"Bearer {API_TOKEN}"}
 timeout = 50000  # タイムアウトを300秒に設定
+# キャッシュの設定
+def get_cached_image(prompt):
+    cached = cache.get(prompt)
+    if cached is not None:
+        return cached
+    return None
+def set_cache_image(prompt, image):
+    cache.set(prompt, image, timeout=60*60*24)  # キャッシュを24時間保存
+# 非同期APIリクエストの実行関数
+async def query_async(prompt, negative_prompt="", steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024, num_inference_steps=30, guidance_scale=7.5, top_k=50, top_p=0.9, eta=0.1):
     if not prompt:
         return None, "Prompt is required"
     key = random.randint(0, 999)
     # Translate the prompt from Russian to English if necessary
     prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
     print(f'Generation {key} translation: {prompt}')
     prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
     print(f'Generation {key}: {prompt}')
         }
     }
+    async with aiohttp.ClientSession() as session:
         try:
+            async with session.post(API_URL, json=payload, headers=headers, timeout=timeout) as response:
+                if response.status != 200:
+                    return None, f"Error: Failed to get image. Status code: {response.status}, Details: {await response.text()}"
+                image_bytes = await response.read()
+                image = Image.open(io.BytesIO(image_bytes))
+                return image, None
+        except asyncio.TimeoutError:
             return None, "Error: The request timed out. Please try again."
         except Exception as e:
+            return None, f"Request Exception: {str(e)}"
 # Content-Security-Policyヘッダーを設定するための関数
 @app.after_request
+async def add_security_headers(response):
     response.headers['Content-Security-Policy'] = (
         "default-src 'self'; "
         "connect-src 'self' ^https?:\/\/[\w.-]+\.[\w.-]+(\/[\w.-]*)*(\?[^\s]*)?$"
     )
     return response
+# HTMLテンプレート
 index_html = """
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <title>Image Generator</title>
+</head>
+<body>
+    <h1>Welcome to the Image Generator</h1>
+    <form action="/generate" method="get">
+        <input type="text" name="prompt" placeholder="Enter prompt" required>
+        <input type="submit" value="Generate">
+    </form>
+</body>
+</html>
 """
 @app.route('/')
+async def index():
+    return await render_template_string(index_html)
 @app.route('/generate', methods=['GET'])
+async def generate_image():
     prompt = request.args.get("prompt", "")
     negative_prompt = request.args.get("negative_prompt", "")
     steps = int(request.args.get("steps", 35))
     top_p = float(request.args.get("top_p", 0.9))
     eta = float(request.args.get("eta", 0.1))
+    # キャッシュを確認
+    cached_image = get_cached_image(prompt)
+    if cached_image:
+        return await send_file(cached_image, mimetype='image/png')
+    image, error = await query_async(prompt, negative_prompt, steps, cfg_scale, sampler, seed, strength, width, height, num_inference_steps, guidance_scale, top_k, top_p, eta)
     if error:
         return jsonify({"error": error}), 400
     img_bytes = io.BytesIO()
     image.save(img_bytes, format='PNG')
     img_bytes.seek(0)
+    # 画像をキャッシュに保存
+    set_cache_image(prompt, img_bytes)
+    return await send_file(img_bytes, mimetype='image/png')
 if __name__ == "__main__":
     app.run(host='0.0.0.0', port=7860)