Upload app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
-import requests
+import asyncio
+import aiohttp
 from PIL import Image
 from io import BytesIO
 from dotenv import load_dotenv
@@ -9,49 +10,95 @@ import os
 load_dotenv()
 API_TOKEN = os.getenv("HF_API_TOKEN")
 
-def generate_image(prompt, model_name):
-    API_URL = f"https://api-inference.huggingface.co/models/{model_name}"
-    headers = {"Authorization": f"Bearer {API_TOKEN}"}
-
-    response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
-    if response.status_code == 200:
-        return Image.open(BytesIO(response.content))
-    else:
-        return f"Ошибка: {response.status_code}, {response.text}"
-
-models = {
+# Конфигурация API
+HEADERS = {"Authorization": f"Bearer {API_TOKEN}"}
+MODELS = {
     "Stable Diffusion v1.5": "Yntec/stable-diffusion-v1-5",
-    "DreamBooth v2": "Yntec/stable-diffusion-v1-5",
-    "DALL-E Mini": "Yntec/stable-diffusion-v1-5",
-    "DeepAI Generator": "Yntec/stable-diffusion-v1-5",
-    "Realistic Vision": "Yntec/stable-diffusion-v1-5",
-    "Artistic AI": "Yntec/stable-diffusion-v1-5"
+    "Stable Diffusion v2.1": "stabilityai/stable-diffusion-2-1",
+    "Stable Diffusion v3.5 Large": "stabilityai/stable-diffusion-3.5-large",
+    "Midjourney": "Jovie/Midjourney",
+    "FLUX.1 [dev]": "black-forest-labs/FLUX.1-dev",
+    "Leonardo AI": "goofyai/Leonardo_Ai_Style_Illustration",
 }
 
-def handle_input(prompt):
-    outputs = {}
-    for name, model in models.items():
-        outputs[name] = generate_image(prompt, model)
-    return list(outputs.values())
 
+# Асинхронная функция для отправки запроса к API
+async def generate_image(prompt, model_name, model_url):
+    async with aiohttp.ClientSession() as session:
+        try:
+            async with session.post(
+                    f"https://api-inference.huggingface.co/models/{model_url}",
+                    headers=HEADERS,
+                    json={"inputs": prompt},
+            ) as response:
+                if response.status == 200:
+                    image_data = await response.read()
+                    return model_name, Image.open(BytesIO(image_data))
+                elif response.status == 503:
+                    error_data = await response.json()
+                    print(f"Модель {model_name} перегружена: {error_data}")
+                    return model_name, None
+                else:
+                    error_data = await response.text()
+                    print(f"Ошибка для модели {model_name}: {error_data}")
+                    return model_name, None
+        except Exception as e:
+            print(f"Ошибка соединения с моделью {model_name}: {e}")
+            return model_name, None
+
+
+# Обработка запросов для всех моделей
+async def handle(prompt):
+    tasks = [
+        generate_image(prompt, model_name, model_url)
+        for model_name, model_url in MODELS.items()
+    ]
+    results = await asyncio.gather(*tasks)
+    outputs = {name: img for name, img in results}
+    return outputs
 
+
+# Интерфейс Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("## Демо-генерация изображений с помощью различных моделей нейросетей")
-    with gr.Row():
-        user_input = gr.Textbox(label="Введите описание изображения",
-                                placeholder="Например, 'Астронавт верхом на лошади'")
+    gr.Markdown("## Генерация изображений с помощью различных моделей нейросетей")
+
+    # Поле ввода
+    user_input = gr.Textbox(label="Введите описание изображения", placeholder="Например, 'Астронавт верхом на лошади'")
+
+    # Вывод изображений
     with gr.Row():
-        outputs = [
-            gr.Image(label=name, interactive=False)
-            for name in models.keys()
-        ]
+        outputs = {name: gr.Image(label=name) for name in MODELS.keys()}
+
+    # Кнопка генерации
     generate_button = gr.Button("Сгенерировать")
 
-    # Обработка клика
+    # Асинхронная обработка ввода
+    async def on_submit(prompt):
+        results = await handle(prompt)
+        # Формируем вывод для каждой модели
+        return [
+            results.get(name, None) for name in MODELS.keys()
+        ]
+
     generate_button.click(
-        fn=handle_input,
+        fn=on_submit,
         inputs=[user_input],
-        outputs=outputs
+        outputs=list(outputs.values()),
     )
+    user_input.submit(
+        fn=on_submit,
+        inputs=[user_input],
+        outputs=list(outputs.values()),
+    )
+
+    # Ссылки ��а соцсети
+    with gr.Row():
+        gr.Markdown(
+            """
+            ### Поддержка проекта
+            - [Telegram](https://t.me/mlphys)
+            - [GitHub](https://github.com/freQuensy23-coder)
+            """
+        )
 
 demo.launch()