Anonym26 commited on
Commit
a320656
·
verified ·
1 Parent(s): e81e176

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +29 -38
app.py CHANGED
@@ -21,42 +21,36 @@ MODELS = {
21
  "Leonardo AI": "goofyai/Leonardo_Ai_Style_Illustration",
22
  }
23
 
24
-
25
  # Асинхронная функция для отправки запроса к API
26
- async def generate_image(prompt, model_name, model_url):
27
- async with aiohttp.ClientSession() as session:
28
- try:
29
- async with session.post(
30
- f"https://api-inference.huggingface.co/models/{model_url}",
31
- headers=HEADERS,
32
- json={"inputs": prompt},
33
- ) as response:
34
- if response.status == 200:
35
- image_data = await response.read()
36
- return model_name, Image.open(BytesIO(image_data))
37
- elif response.status == 503:
38
- error_data = await response.json()
39
- print(f"Модель {model_name} перегружена: {error_data}")
40
- return model_name, None
41
- else:
42
- error_data = await response.text()
43
- print(f"Ошибка для модели {model_name}: {error_data}")
44
- return model_name, None
45
- except Exception as e:
46
- print(f"Ошибка соединения с моделью {model_name}: {e}")
47
- return model_name, None
48
-
49
 
50
- # Обработка запросов для всех моделей
51
  async def handle(prompt):
52
- tasks = [
53
- generate_image(prompt, model_name, model_url)
54
- for model_name, model_url in MODELS.items()
55
- ]
56
- results = await asyncio.gather(*tasks)
57
- outputs = {name: img for name, img in results}
58
- return outputs
59
-
 
60
 
61
  # Интерфейс Gradio
62
  with gr.Blocks() as demo:
@@ -74,11 +68,8 @@ with gr.Blocks() as demo:
74
 
75
  # Асинхронная обработка ввода
76
  async def on_submit(prompt):
77
- results = await handle(prompt)
78
- # Формируем вывод для каждой модели
79
- return [
80
- results.get(name, None) for name in MODELS.keys()
81
- ]
82
 
83
  generate_button.click(
84
  fn=on_submit,
 
21
  "Leonardo AI": "goofyai/Leonardo_Ai_Style_Illustration",
22
  }
23
 
 
24
  # Асинхронная функция для отправки запроса к API
25
+ async def query_model(session, prompt, model_name, model_url):
26
+ try:
27
+ async with session.post(
28
+ f"https://api-inference.huggingface.co/models/{model_url}",
29
+ headers=HEADERS,
30
+ json={"inputs": prompt},
31
+ ) as response:
32
+ if response.status == 200:
33
+ image_data = await response.read()
34
+ return model_name, Image.open(BytesIO(image_data))
35
+ else:
36
+ error_message = await response.text()
37
+ print(f"Ошибка для модели {model_name}: {error_message}")
38
+ return model_name, None
39
+ except Exception as e:
40
+ print(f"Ошибка соединения с моделью {model_name}: {e}")
41
+ return model_name, None
 
 
 
 
 
 
42
 
43
+ # Асинхронная обработка запросов для всех моделей
44
  async def handle(prompt):
45
+ async with aiohttp.ClientSession() as session:
46
+ # Создаём асинхронный поток результатов
47
+ results = {}
48
+ async for outputs in async_zip_stream(
49
+ *(query_model(session, prompt, model_name, model_url) for model_name, model_url in MODELS.items())
50
+ ):
51
+ for model_name, image in outputs:
52
+ results[model_name] = image
53
+ yield list(results.values())
54
 
55
  # Интерфейс Gradio
56
  with gr.Blocks() as demo:
 
68
 
69
  # Асинхронная обработка ввода
70
  async def on_submit(prompt):
71
+ async for result in handle(prompt):
72
+ return [result.get(name, None) for name in MODELS.keys()]
 
 
 
73
 
74
  generate_button.click(
75
  fn=on_submit,