Upload app.py

app.py

@@ -21,42 +21,36 @@ MODELS = {
     "Leonardo AI": "goofyai/Leonardo_Ai_Style_Illustration",
 }
 
-
 # Асинхронная функция для отправки запроса к 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 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 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
-
+    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())
 
 # Интерфейс Gradio
 with gr.Blocks() as demo:
@@ -74,11 +68,8 @@ with gr.Blocks() as demo:
 
     # Асинхронная обработка ввода
     async def on_submit(prompt):
-        results = await handle(prompt)
-        # Формируем вывод для каждой модели
-        return [
-            results.get(name, None) for name in MODELS.keys()
-        ]
+        async for result in handle(prompt):
+            return [result.get(name, None) for name in MODELS.keys()]
 
     generate_button.click(
         fn=on_submit,