Spaces:
Sleeping
Sleeping
File size: 3,968 Bytes
eff941a f5f4805 eff941a |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
from flask import Flask, request, jsonify
from gradio_client import Client, handle_file
import base64
import io
import requests
from PIL import Image
app = Flask(__name__)
# client = Client("multimodalart/flux-style-shaping")
def get_random_api_key():
keys = os.getenv("KEYS", "").split(",")
if keys and keys[0]: # Check if KEYS is set and not empty
return random.choice(keys).strip()
else:
raise ValueError("API keys not found. Please set the KEYS environment variable.")
def translate_to_english(prompt):
language = detect(prompt)
if language != 'en':
prompt = GoogleTranslator(source=language, target='en').translate(prompt)
return prompt
def handle_image_input(image_data):
"""Функция для обработки разных форматов входных изображений."""
if image_data.startswith("http://") or image_data.startswith("https://"):
try:
response = requests.get(image_data, stream=True)
response.raise_for_status() # Проверяем, что запрос успешен
image = Image.open(io.BytesIO(response.content))
image_bytes = io.BytesIO()
image.save(image_bytes, format="PNG") # Сохраняем в PNG для единообразия
image_bytes = image_bytes.getvalue()
return handle_file(image_bytes)
except requests.exceptions.RequestException as e:
print(f"Ошибка при загрузке изображения по URL: {e}")
return None
elif image_data.startswith("data:image"): # Base64
try:
header, encoded = image_data.split(',', 1)
image_bytes = base64.b64decode(encoded)
return handle_file(image_bytes)
except Exception as e:
print(f"Ошибка при декодировании base64: {e}")
return None
else: # Предполгаем, что это просто путь к файлу
try:
with open(image_data, "rb") as f:
image_bytes = f.read()
return handle_file(image_bytes)
except Exception as e:
print(f"Ошибка при открытии файла: {e}")
return None
@app.route('/generate_image/', methods=['GET'])
def generate_image():
prompt = request.args.get('prompt')
image1_data = request.args.get('image1')
image2_data = request.args.get('image2')
depth_strength = request.args.get('depth_strength', default=15, type=float)
style_strength = request.args.get('style_strength', default=0.5, type=float)
if not prompt or not image1_data or not image2_data:
return jsonify({"error": "Missing required parameters: prompt, image1, and image2."}), 400
structure_image = handle_image_input(image1_data)
style_image = handle_image_input(image2_data)
if not structure_image or not style_image:
return jsonify({"error": "Failed to process one or both of the input images."}), 400
prompt = translate_to_english(prompt) if prompt else ""
try:
client = Client("multimodalart/flux-style-shaping", hf_token=get_random_api_key())
result = client.predict(
prompt=prompt,
structure_image=structure_image,
style_image=style_image,
depth_strength=depth_strength,
style_strength=style_strength,
api_name="/generate_image"
)
# Преобразуем результат в base64
with open(result["path"], "rb") as image_file:
encoded_result = base64.b64encode(image_file.read()).decode('utf-8')
return jsonify({"generated_image": encoded_result})
except Exception as e:
return jsonify({"error": f"Error during image generation: {str(e)}"}), 500
if __name__ == '__main__':
app.run(host='0.0.0.0', port=7860, debug=True)
|