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