Kirill Kirilenko
4b265b5405
Добавлена возможность пересылки сообщений в личных чатах. Обновлены зависимости. Добавлен requirements.txt. Исправлены предупреждения PyCharm 2026.1.
80 lines
3.1 KiB
Python
80 lines
3.1 KiB
Python
from typing import Any, Dict, List
|
||
|
||
from openrouter.components import ChatToolMessageContentTypedDict
|
||
from fal_client import AsyncClient as FalClient
|
||
|
||
from ai.tool import Tool
|
||
from ai.utils import *
|
||
|
||
FAL_MODEL = "fal-ai/bytedance/seedream/v4.5/text-to-image"
|
||
|
||
|
||
class GenerateImageTool(Tool):
|
||
def __init__(self, fal_token: str):
|
||
self._client = FalClient(key=fal_token)
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "generate_image"
|
||
|
||
@property
|
||
def description(self) -> str:
|
||
return "Генерация изображения по описанию"
|
||
|
||
@property
|
||
def parameters(self) -> Dict[str, Any]:
|
||
return {
|
||
"type": "object",
|
||
"properties": {
|
||
"prompt": {
|
||
"type": "string",
|
||
"description": "Подробное описание сцены на английском языке БЕЗ технических параметров "
|
||
"(соотношение сторон, разрешение)"
|
||
},
|
||
"aspect_ratio": {
|
||
"type": "string",
|
||
"enum": ["1:1", "4:3", "3:4", "16:9", "9:16"],
|
||
"description": "Соотношение сторон"
|
||
}
|
||
},
|
||
"required": ["prompt"]
|
||
}
|
||
|
||
async def execute(self, args: Dict[str, Any], artifacts: Dict[str, Any]) -> List[ChatToolMessageContentTypedDict]:
|
||
prompt = args.get("prompt", "")
|
||
aspect_ratio = args.get("aspect_ratio", "4:3")
|
||
|
||
aspect_ratio_size_map = {
|
||
"1:1": "square",
|
||
"4:3": "landscape_4_3",
|
||
"3:4": "portrait_4_3",
|
||
"16:9": "landscape_16_9",
|
||
"9:16": "portrait_16_9",
|
||
"9:20": "portrait_16_9"
|
||
}
|
||
image_size = aspect_ratio_size_map.get(aspect_ratio, "landscape_4_3")
|
||
print(f"Генерация изображения {image_size}: {prompt}")
|
||
|
||
arguments = {
|
||
"prompt": prompt,
|
||
"image_size": image_size,
|
||
"enable_safety_checker": False
|
||
}
|
||
|
||
try:
|
||
result = await self._client.run(FAL_MODEL, arguments=arguments)
|
||
if "images" not in result:
|
||
raise RuntimeError("Неожиданный ответ от сервера.")
|
||
image_url = result["images"][0]["url"]
|
||
|
||
from utils import download_file
|
||
artifacts["generated_image_hires"] = await download_file(image_url)
|
||
artifacts["generated_image"] = compress_image(artifacts["generated_image_hires"], 1280)
|
||
|
||
return serialize_message_content(
|
||
text="Изображение сгенерировано и будет показано пользователю.",
|
||
image=None
|
||
)
|
||
except Exception as e:
|
||
print(f"Ошибка генерации изображения: {e}")
|
||
return serialize_message_content(text=f"Не удалось сгенерировать изображение: {e}")
|