Update app.py

app.py

@@ -1,7 +1,9 @@
-from flask import Flask, request, send_file, render_template, abort
+from flask import Flask, request, send_file, abort
 import requests
 import io
 from PIL import Image
+from cachetools import TTLCache, cached
+import random
 import os
 import urllib.parse
 import hashlib
@@ -14,6 +16,9 @@ app = Flask(__name__)
 MAX_WIDTH = 1384
 MAX_HEIGHT = 1384
 
+# Кэш на 10 минут
+cache = TTLCache(maxsize=100, ttl=600)
+
 # Получаем ключи из переменной окружения
 keys = os.getenv("keys", "").split(',')
 if not keys:
@@ -22,6 +27,10 @@ if not keys:
 def get_random_key():
     return random.choice(keys)
 
+def generate_cache_key(prompt, width, height, seed, model_name):
+    # Создаем уникальный ключ на основе всех параметров, кроме ключа API
+    return hashlib.md5(f"{prompt}_{width}_{height}_{seed}_{model_name}".encode()).hexdigest()
+
 def scale_dimensions(width, height, max_width, max_height):
     """Масштабирует размеры изображения, сохраняя соотношение сторон, и округляет до чисел, кратных 8."""
     aspect_ratio = width / height
@@ -38,7 +47,8 @@ def scale_dimensions(width, height, max_width, max_height):
     height = (height + 3) // 8 * 8
     return width, height
 
-def generate_image(prompt, width, height, seed, model_name, api_key):
+@cached(cache)
+def generate_cached_image(cache_key, prompt, width, height, seed, model_name, api_key):
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
@@ -127,56 +137,6 @@ def get_image(prompt):
         mimetype='image/png'
     )
 
-@app.route('/play', methods=['GET', 'POST'])
-def play():
-    if request.method == 'POST':
-        prompt = request.form.get('prompt', default="")
-        width = request.form.get('width', type=int, default=512)
-        height = request.form.get('height', type=int, default=512)
-        seed = request.form.get('seed', type=int, default=22)
-        model_name = request.form.get('model', default="black-forest-labs_FLUX.1-dev").replace('_', '/')
-        api_key = request.form.get('key', default=None)
-
-        # Определяем язык промпта
-        try:
-            language = detect(prompt)
-        except Exception as e:
-            app.logger.error(f"Error detecting language: {e}")
-            return send_error_image()
-
-        # Переводим промпт, если он не на английском языке
-        if language != 'en':
-            try:
-                translator = GoogleTranslator(source=language, target='en')
-                prompt = translator.translate(prompt)
-            except Exception as e:
-                app.logger.error(f"Error translating prompt: {e}")
-                return send_error_image()
-
-        # Масштабируем размеры изображения, если они превышают максимальные значения, и округляем до чисел, кратных 8
-        width, height = scale_dimensions(width, height, MAX_WIDTH, MAX_HEIGHT)
-
-        # Используем указанный ключ, если он предоставлен, иначе выбираем случайный ключ
-        if api_key is None:
-            api_key = get_random_key()
-
-        # Генерируем изображение без кэширования
-        image = generate_image(prompt, width, height, seed, model_name, api_key)
-        if image is None:
-            return send_error_image()
-
-        img_byte_arr = io.BytesIO()
-        image.save(img_byte_arr, format='PNG')
-        img_byte_arr = img_byte_arr.getvalue()
-
-        return send_file(
-            io.BytesIO(img_byte_arr),
-            mimetype='image/png'
-        )
-
-    # GET запрос - отображаем форму
-    return render_template('play.html')
-
 @app.route('/')
 def health_check():
     return "OK", 200