Rooni commited on
Commit
cf44c0c
·
verified ·
1 Parent(s): de302d7

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +75 -0
app.py ADDED
@@ -0,0 +1,75 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from flask import Flask, request, send_file, abort
2
+ import requests
3
+ import io
4
+ from PIL import Image
5
+ from cachetools import TTLCache, cached
6
+ import random
7
+ import os
8
+ import urllib.parse
9
+
10
+ app = Flask(__name__)
11
+
12
+ # Кэш на 10 минут
13
+ cache = TTLCache(maxsize=100, ttl=600)
14
+
15
+ # Получаем ключи из переменной окружения
16
+ keys = os.getenv("keys", "").split(',')
17
+ if not keys:
18
+ raise ValueError("Environment variable 'keys' must be set with a comma-separated list of API keys.")
19
+
20
+ def get_random_key():
21
+ return random.choice(keys)
22
+
23
+ @cached(cache)
24
+ def generate_cached_image(prompt, width, height, seed, model_name):
25
+ api_key = get_random_key()
26
+ headers = {
27
+ "Authorization": f"Bearer {api_key}",
28
+ "Content-Type": "application/json"
29
+ }
30
+ data = {
31
+ "prompt": prompt,
32
+ "width": width,
33
+ "height": height,
34
+ "seed": seed,
35
+ "model": model_name
36
+ }
37
+
38
+ try:
39
+ response = requests.post("https://api.your-provider.com/generate", headers=headers, json=data)
40
+ response.raise_for_status()
41
+ image_data = response.content
42
+ image = Image.open(io.BytesIO(image_data))
43
+ return image
44
+ except requests.exceptions.RequestException as e:
45
+ app.logger.error(f"Error generating image: {e}")
46
+ abort(500, description="Error generating image")
47
+
48
+ @app.route('/prompt/<path:prompt>')
49
+ def get_image(prompt):
50
+ width = request.args.get('width', type=int, default=512)
51
+ height = request.args.get('height', type=int, default=512)
52
+ seed = request.args.get('seed', type=int, required=True)
53
+ model_name = request.args.get('model', required=True)
54
+
55
+ # Декодируем URL-кодированный prompt
56
+ prompt = urllib.parse.unquote(prompt)
57
+
58
+ try:
59
+ image = generate_cached_image(prompt, width, height, seed, model_name)
60
+ except Exception as e:
61
+ app.logger.error(f"Error generating image: {e}")
62
+ abort(500, description="Error generating image")
63
+
64
+ img_byte_arr = io.BytesIO()
65
+ image.save(img_byte_arr, format='PNG')
66
+ img_byte_arr = img_byte_arr.getvalue()
67
+
68
+ return send_file(
69
+ io.BytesIO(img_byte_arr),
70
+ mimetype='image/png'
71
+ )
72
+
73
+ if __name__ == '__main__':
74
+ app.run(host='0.0.0.0', port=5000)
75
+