Rooni commited on
Commit
eff941a
·
verified ·
1 Parent(s): d3a9262

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +98 -0
app.py ADDED
@@ -0,0 +1,98 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from flask import Flask, request, jsonify
2
+ from gradio_client import Client, handle_file
3
+ import base64
4
+ import io
5
+ import requests
6
+ from PIL import Image
7
+
8
+ app = Flask(__name__)
9
+
10
+ # client = Client("multimodalart/flux-style-shaping")
11
+
12
+ def get_random_api_key():
13
+ keys = os.getenv("KEYS", "").split(",")
14
+ if keys and keys[0]: # Check if KEYS is set and not empty
15
+ return random.choice(keys).strip()
16
+ else:
17
+ raise ValueError("API keys not found. Please set the KEYS environment variable.")
18
+
19
+ def translate_to_english(prompt):
20
+ language = detect(prompt)
21
+ if language != 'en':
22
+ prompt = GoogleTranslator(source=language, target='en').translate(prompt)
23
+ return prompt
24
+
25
+ def handle_image_input(image_data):
26
+ """Функция для обработки разных форматов входных изображений."""
27
+ if image_data.startswith("http://") or image_data.startswith("https://"):
28
+ try:
29
+ response = requests.get(image_data, stream=True)
30
+ response.raise_for_status() # Проверяем, что запрос успешен
31
+ image = Image.open(io.BytesIO(response.content))
32
+ image_bytes = io.BytesIO()
33
+ image.save(image_bytes, format="PNG") # Сохраняем в PNG для единообразия
34
+ image_bytes = image_bytes.getvalue()
35
+ return handle_file(image_bytes)
36
+ except requests.exceptions.RequestException as e:
37
+ print(f"Ошибка при загрузке изображения по URL: {e}")
38
+ return None
39
+ elif image_data.startswith("data:image"): # Base64
40
+ try:
41
+ header, encoded = image_data.split(',', 1)
42
+ image_bytes = base64.b64decode(encoded)
43
+ return handle_file(image_bytes)
44
+ except Exception as e:
45
+ print(f"Ошибка при декодировании base64: {e}")
46
+ return None
47
+ else: # Предполгаем, что это просто путь к файлу
48
+ try:
49
+ with open(image_data, "rb") as f:
50
+ image_bytes = f.read()
51
+ return handle_file(image_bytes)
52
+ except Exception as e:
53
+ print(f"Ошибка при открытии файла: {e}")
54
+ return None
55
+
56
+
57
+ @app.route('/generate_image', methods=['GET'])
58
+ def generate_image():
59
+ prompt = request.args.get('prompt')
60
+ image1_data = request.args.get('image1')
61
+ image2_data = request.args.get('image2')
62
+ depth_strength = request.args.get('depth_strength', default=15, type=float)
63
+ style_strength = request.args.get('style_strength', default=0.5, type=float)
64
+
65
+ if not prompt or not image1_data or not image2_data:
66
+ return jsonify({"error": "Missing required parameters: prompt, image1, and image2."}), 400
67
+
68
+ structure_image = handle_image_input(image1_data)
69
+ style_image = handle_image_input(image2_data)
70
+
71
+ if not structure_image or not style_image:
72
+ return jsonify({"error": "Failed to process one or both of the input images."}), 400
73
+
74
+ prompt = translate_to_english(prompt) if prompt else ""
75
+
76
+ try:
77
+ client = Client("multimodalart/flux-style-shaping", hf_token=get_random_api_key())
78
+ result = client.predict(
79
+ prompt=prompt,
80
+ structure_image=structure_image,
81
+ style_image=style_image,
82
+ depth_strength=depth_strength,
83
+ style_strength=style_strength,
84
+ api_name="/generate_image"
85
+ )
86
+
87
+ # Преобразуем результат в base64
88
+ with open(result["path"], "rb") as image_file:
89
+ encoded_result = base64.b64encode(image_file.read()).decode('utf-8')
90
+
91
+ return jsonify({"generated_image": encoded_result})
92
+ except Exception as e:
93
+ return jsonify({"error": f"Error during image generation: {str(e)}"}), 500
94
+
95
+
96
+ if __name__ == '__main__':
97
+ app.run(host='0.0.0.0', port=7860, debug=True)
98
+