Update app.py

app.py

@@ -1,5 +1,6 @@
 import json
 import time
+
 import requests
 import base64
 from io import BytesIO
@@ -22,6 +23,8 @@ styles = {
     "Кандинский": "KANDINSKY"
 }
 
+timeout = 125
+
 api_key = os.getenv("api_key")
 secret_key = os.getenv("secret_key")
 
@@ -39,10 +42,10 @@ class Text2ImageAPI:
         data = response.json()
         return data[0]['id']
 
-    def generate(self, prompt, negative, style, width, height, model, num_images):
+    def generate(self, prompt, negative, style, width, height, model):
         params = {
             "type": "GENERATE",
-            "numImages": num_images,
+            "numImages": 1,
             "style": f"{style}",
             "width": width,
             "height": height,
@@ -57,7 +60,7 @@ class Text2ImageAPI:
             'model_id': (None, model),
             'params': (None, json.dumps(params), 'application/json')
         }
-        response = requests.post(self.URL + 'key/api/v1/text2image/run', headers=self.AUTH_HEADERS, files=data)
+        response = requests.post(self.URL + 'key/api/v1/text2image/run', headers=self.AUTH_HEADERS, files=data, timeout=timeout)
         data = response.json()
         return data['uuid']
 
@@ -72,20 +75,23 @@ class Text2ImageAPI:
             time.sleep(delay)
 
 
-def api_gradio(prompt, negative, style, width, height, num_images):
+def api_gradio(prompt, negative, style, width, height):
     api = Text2ImageAPI('https://api-key.fusionbrain.ai/', api_key, secret_key)
     model_id = api.get_model()
-    uuid = api.generate(prompt, negative, styles[style], width, height, model_id, num_images)
+    uuid = api.generate(prompt, negative, styles[style], width, height, model_id)
     images = api.check_generation(uuid)
 
-    decoded_images = [Image.open(BytesIO(base64.b64decode(image))) for image in images]
-    return decoded_images
+    decoded_data = base64.b64decode(images[0])
+    image = Image.open(BytesIO(decoded_data))
+
+    return image
+
 
 
 with gr.Blocks(css=css) as demo:
     with gr.Column():
         with gr.Row():
-            prompt = gr.Textbox(show_label=False, placeholder="Описание изображения", lines=3, interactive=True, scale=20)
+            prompt = gr.Textbox(show_label=False, placeholder="Описание изображения", max_lines=3, lines=1, interactive=True, scale=20)
         with gr.Row():
             style = gr.Radio(show_label=False, value="Свой стиль", choices=list(styles.keys()))
         with gr.Row():
@@ -93,14 +99,14 @@ with gr.Blocks(css=css) as demo:
     with gr.Accordion("Дополнительные настройки", open=False):
         with gr.Row():
             negative = gr.Textbox(label="Отрицательная подсказка", placeholder="Исключения, чего не должно быть на фото", max_lines=3, lines=1, interactive=True, scale=20)
-        with gr.Row():
-            width = gr.Slider(label="Ширина", minimum=128, maximum=1024, step=1, value=1024, interactive=True)
-            height = gr.Slider(label="Высота", minimum=128, maximum=1024, step=1, value=1024, interactive=True)
-        with gr.Row():
-            num_images = gr.Slider(label="Количество изображений", minimum=1, maximum=4, step=1, value=1, interactive=True)
+            width = gr.Slider(label="Ширина", minimum=512, maximum=1024, step=1, value=1024, interactive=True)
+            height = gr.Slider(label="Высота", minimum=512, maximum=1024, step=1, value=1024, interactive=True)
+            
+      #  with gr.Row():
+         #   images = gr.Slider(label="Количество изображений", minimum=1, maximum=4, step=1, value=1, interactive=True)
     with gr.Row():
-        gallery = gr.Gallery(show_label=False, columns=2)
+        gallery = gr.Image(show_label=False)
 
-    button.click(api_gradio, inputs=[prompt, negative, style, width, height, num_images], outputs=gallery)
+    button.click(api_gradio, inputs=[prompt, negative, style, width, height], outputs=gallery)
 
-demo.queue(max_size=100).launch(show_api=False)
+demo.queue(max_size=100).launch(show_api=False)