Rooni commited on
Commit
d8b1d45
·
verified ·
1 Parent(s): 0402965

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +20 -8
app.py CHANGED
@@ -27,9 +27,9 @@ if not keys:
27
  def get_random_key():
28
  return random.choice(keys)
29
 
30
- def generate_cache_key(prompt, width, height, seed, model_name):
31
  # Создаем уникальный ключ на основе всех параметров
32
- return hashlib.md5(f"{prompt}_{width}_{height}_{seed}_{model_name}".encode()).hexdigest()
33
 
34
  def scale_dimensions(width, height, max_width, max_height):
35
  """Масштабирует размеры изображения, сохраняя соотношение сторон, и округляет до чисел, кратных 8."""
@@ -48,8 +48,7 @@ def scale_dimensions(width, height, max_width, max_height):
48
  return width, height
49
 
50
  @cached(cache)
51
- def generate_cached_image(cache_key, prompt, width, height, seed, model_name):
52
- api_key = get_random_key()
53
  headers = {
54
  "Authorization": f"Bearer {api_key}",
55
  "Content-Type": "application/json"
@@ -64,7 +63,12 @@ def generate_cached_image(cache_key, prompt, width, height, seed, model_name):
64
  }
65
 
66
  try:
67
- response = requests.post(f"https://api-inference.huggingface.co/models/{model_name}", headers=headers, json=data)
 
 
 
 
 
68
  response.raise_for_status()
69
  image_data = response.content
70
  image = Image.open(io.BytesIO(image_data))
@@ -72,6 +76,9 @@ def generate_cached_image(cache_key, prompt, width, height, seed, model_name):
72
  except requests.exceptions.HTTPError as http_err:
73
  app.logger.error(f"HTTP error occurred: {http_err} - Response: {response.text}")
74
  return None
 
 
 
75
  except requests.exceptions.RequestException as req_err:
76
  app.logger.error(f"Request error occurred: {req_err}")
77
  return None
@@ -82,6 +89,7 @@ def get_image(prompt):
82
  height = request.args.get('height', type=int, default=512)
83
  seed = request.args.get('seed', type=int, default=22)
84
  model_name = request.args.get('model', default="black-forest-labs_FLUX.1-dev").replace('_', '/')
 
85
 
86
  # Декодируем URL-кодированный prompt
87
  prompt = urllib.parse.unquote(prompt)
@@ -105,11 +113,15 @@ def get_image(prompt):
105
  # Масштабируем размеры изображения, если они превышают максимальные значения, и округляем до чисел, кратных 8
106
  width, height = scale_dimensions(width, height, MAX_WIDTH, MAX_HEIGHT)
107
 
 
 
 
 
108
  # Генерируем уникальный ключ для кэша
109
- cache_key = generate_cache_key(prompt, width, height, seed, model_name)
110
 
111
  try:
112
- image = generate_cached_image(cache_key, prompt, width, height, seed, model_name)
113
  if image is None:
114
  return send_error_image()
115
  except Exception as e:
@@ -147,4 +159,4 @@ def send_error_image():
147
  abort(500, description="Error fetching error image")
148
 
149
  if __name__ == '__main__':
150
- app.run(host='0.0.0.0', port=7860, debug=False)
 
27
  def get_random_key():
28
  return random.choice(keys)
29
 
30
+ def generate_cache_key(prompt, width, height, seed, model_name, api_key):
31
  # Создаем уникальный ключ на основе всех параметров
32
+ return hashlib.md5(f"{prompt}_{width}_{height}_{seed}_{model_name}_{api_key}".encode()).hexdigest()
33
 
34
  def scale_dimensions(width, height, max_width, max_height):
35
  """Масштабирует размеры изображения, сохраняя соотношение сторон, и округляет до чисел, кратных 8."""
 
48
  return width, height
49
 
50
  @cached(cache)
51
+ def generate_cached_image(cache_key, prompt, width, height, seed, model_name, api_key):
 
52
  headers = {
53
  "Authorization": f"Bearer {api_key}",
54
  "Content-Type": "application/json"
 
63
  }
64
 
65
  try:
66
+ response = requests.post(
67
+ f"https://api-inference.huggingface.co/models/{model_name}",
68
+ headers=headers,
69
+ json=data,
70
+ timeout=180 # Таймаут 3 минуты
71
+ )
72
  response.raise_for_status()
73
  image_data = response.content
74
  image = Image.open(io.BytesIO(image_data))
 
76
  except requests.exceptions.HTTPError as http_err:
77
  app.logger.error(f"HTTP error occurred: {http_err} - Response: {response.text}")
78
  return None
79
+ except requests.exceptions.Timeout as timeout_err:
80
+ app.logger.error(f"Timeout error occurred: {timeout_err}")
81
+ return None
82
  except requests.exceptions.RequestException as req_err:
83
  app.logger.error(f"Request error occurred: {req_err}")
84
  return None
 
89
  height = request.args.get('height', type=int, default=512)
90
  seed = request.args.get('seed', type=int, default=22)
91
  model_name = request.args.get('model', default="black-forest-labs_FLUX.1-dev").replace('_', '/')
92
+ api_key = request.args.get('key', default=None)
93
 
94
  # Декодируем URL-кодированный prompt
95
  prompt = urllib.parse.unquote(prompt)
 
113
  # Масштабируем размеры изображения, если они превышают максимальные значения, и округляем до чисел, кратных 8
114
  width, height = scale_dimensions(width, height, MAX_WIDTH, MAX_HEIGHT)
115
 
116
+ # Используем указанный ключ, если он предоставлен, иначе выбираем случайный ключ
117
+ if api_key is None:
118
+ api_key = get_random_key()
119
+
120
  # Генерируем уникальный ключ для кэша
121
+ cache_key = generate_cache_key(prompt, width, height, seed, model_name, api_key)
122
 
123
  try:
124
+ image = generate_cached_image(cache_key, prompt, width, height, seed, model_name, api_key)
125
  if image is None:
126
  return send_error_image()
127
  except Exception as e:
 
159
  abort(500, description="Error fetching error image")
160
 
161
  if __name__ == '__main__':
162
+ app.run(host='0.0.0.0', port=7860, debug=False)