Upload app.py

app.py

@@ -1,11 +1,11 @@
 import gradio as gr
-import asyncio
 import aiohttp
+import asyncio
 from PIL import Image
 from io import BytesIO
+from asyncio import Semaphore
 from dotenv import load_dotenv
 import os
-
 # Загрузка токена из .env файла
 load_dotenv()
 API_TOKEN = os.getenv("HF_API_TOKEN")
@@ -16,48 +16,54 @@ MODELS = {
     "Stable Diffusion v1.5": "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",
 }
 
+# Настройки
+MAX_CONCURRENT_REQUESTS = 3
+
+
 # Асинхронная функция для отправки запроса к API
-async def query_model(session, prompt, model_name, model_url):
-    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))
-            else:
-                error_message = await response.text()
-                print(f"Ошибка для модели {model_name}: {error_message}")
-                return model_name, None
-    except Exception as e:
-        print(f"Ошибка соединения с моделью {model_name}: {e}")
-        return model_name, None
+async def query_model(prompt, model_name, model_url, semaphore):
+    async with semaphore:  # Ограничиваем количество одновременно выполняемых задач
+        try:
+            async with aiohttp.ClientSession() as session:
+                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))
+                    else:
+                        error_message = await response.json()
+                        warnings = error_message.get("warnings", [])
+                        print(f"Ошибка для модели {model_name}: {error_message.get('error', 'unknown error')}")
+                        if warnings:
+                            print(f"Предупреждения для модели {model_name}: {warnings}")
+                        return model_name, None
+        except Exception as e:
+            print(f"Ошибка соединения с моделью {model_name}: {e}")
+            return model_name, None
+
 
-# Асинхронная обработка запросов для всех моделей
+# Асинхронная обработка запросов
 async def handle(prompt):
-    async with aiohttp.ClientSession() as session:
-        # Создаём асинхронный поток результатов
-        results = {}
-        async for outputs in async_zip_stream(
-            *(query_model(session, prompt, model_name, model_url) for model_name, model_url in MODELS.items())
-        ):
-            for model_name, image in outputs:
-                results[model_name] = image
-            yield list(results.values())
+    semaphore = Semaphore(MAX_CONCURRENT_REQUESTS)  # Создаём локальный семафор
+    tasks = [
+        query_model(prompt, model_name, model_url, semaphore)
+        for model_name, model_url in MODELS.items()
+    ]
+    results = await asyncio.gather(*tasks)
+    return {model_name: image for model_name, image in results if image}
+
 
 # Интерфейс Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("## Генерация изображений с помощью различных моделей нейросетей")
+    gr.Markdown("## Генерация изображений с использованием моделей Hugging Face")
 
     # Поле ввода
-    user_input = gr.Textbox(label="Введите описание изображения", placeholder="Например, 'Астронавт верхом на лошади'")
+    user_input = gr.Textbox(label="Введите описание изображения", placeholder="Например, 'Красный автомобиль в лесу'")
 
     # Вывод изображений
     with gr.Row():
@@ -68,8 +74,8 @@ with gr.Blocks() as demo:
 
     # Асинхронная обработка ввода
     async def on_submit(prompt):
-        async for result in handle(prompt):
-            return [result.get(name, None) for name in MODELS.keys()]
+        results = await handle(prompt)
+        return [results.get(name, None) for name in MODELS.keys()]
 
     generate_button.click(
         fn=on_submit,