Update app.py

app.py

@@ -27,9 +27,9 @@ if not keys:
 def get_random_key():
     return random.choice(keys)
 
-def generate_cache_key(prompt, width, height, seed, model_name):
+def generate_cache_key(prompt, width, height, seed, model_name, api_key):
     # Создаем уникальный ключ на основе всех параметров
-    return hashlib.md5(f"{prompt}_{width}_{height}_{seed}_{model_name}".encode()).hexdigest()
+    return hashlib.md5(f"{prompt}_{width}_{height}_{seed}_{model_name}_{api_key}".encode()).hexdigest()
 
 def scale_dimensions(width, height, max_width, max_height):
     """Масштабирует размеры изображения, сохраняя соотношение сторон, и округляет до чисел, кратных 8."""
@@ -48,8 +48,7 @@ def scale_dimensions(width, height, max_width, max_height):
     return width, height
 
 @cached(cache)
-def generate_cached_image(cache_key, prompt, width, height, seed, model_name):
-    api_key = get_random_key()
+def generate_cached_image(cache_key, prompt, width, height, seed, model_name, api_key):
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
@@ -64,7 +63,12 @@ def generate_cached_image(cache_key, prompt, width, height, seed, model_name):
     }
 
     try:
-        response = requests.post(f"https://api-inference.huggingface.co/models/{model_name}", headers=headers, json=data)
+        response = requests.post(
+            f"https://api-inference.huggingface.co/models/{model_name}",
+            headers=headers,
+            json=data,
+            timeout=180  # Таймаут 3 минуты
+        )
         response.raise_for_status()
         image_data = response.content
         image = Image.open(io.BytesIO(image_data))
@@ -72,6 +76,9 @@ def generate_cached_image(cache_key, prompt, width, height, seed, model_name):
     except requests.exceptions.HTTPError as http_err:
         app.logger.error(f"HTTP error occurred: {http_err} - Response: {response.text}")
         return None
+    except requests.exceptions.Timeout as timeout_err:
+        app.logger.error(f"Timeout error occurred: {timeout_err}")
+        return None
     except requests.exceptions.RequestException as req_err:
         app.logger.error(f"Request error occurred: {req_err}")
         return None
@@ -82,6 +89,7 @@ def get_image(prompt):
     height = request.args.get('height', type=int, default=512)
     seed = request.args.get('seed', type=int, default=22)
     model_name = request.args.get('model', default="black-forest-labs_FLUX.1-dev").replace('_', '/')
+    api_key = request.args.get('key', default=None)
 
     # Декодируем URL-кодированный prompt
     prompt = urllib.parse.unquote(prompt)
@@ -105,11 +113,15 @@ def get_image(prompt):
     # Масштабируем размеры изображения, если они превышают максимальные значения, и округляем до чисел, кратных 8
     width, height = scale_dimensions(width, height, MAX_WIDTH, MAX_HEIGHT)
 
+    # Используем указанный ключ, если он предоставлен, иначе выбираем случайный ключ
+    if api_key is None:
+        api_key = get_random_key()
+
     # Генерируем уникальный ключ для кэша
-    cache_key = generate_cache_key(prompt, width, height, seed, model_name)
+    cache_key = generate_cache_key(prompt, width, height, seed, model_name, api_key)
 
     try:
-        image = generate_cached_image(cache_key, prompt, width, height, seed, model_name)
+        image = generate_cached_image(cache_key, prompt, width, height, seed, model_name, api_key)
         if image is None:
             return send_error_image()
     except Exception as e:
@@ -147,4 +159,4 @@ def send_error_image():
         abort(500, description="Error fetching error image")
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860, debug=False)
+    app.run(host='0.0.0.0', port=7860, debug=False)