65 lines
2.5 KiB
Python
65 lines
2.5 KiB
Python
from typing import Any, Dict, List
|
||
|
||
from openrouter.components import ChatToolMessageContentTypedDict
|
||
from replicate import Client as ReplicateClient
|
||
|
||
from ai.tool import Tool
|
||
from ai.utils import *
|
||
|
||
REPLICATE_MODEL = "bytedance/seedream-4.5"
|
||
|
||
|
||
class GenerateImageTool(Tool):
|
||
def __init__(self, replicate_token: str):
|
||
self._client = ReplicateClient(api_token=replicate_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")
|
||
print(f"Генерация изображения {aspect_ratio}: {prompt}")
|
||
|
||
arguments = {
|
||
"prompt": prompt,
|
||
"aspect_ratio": aspect_ratio,
|
||
"disable_safety_checker": True
|
||
}
|
||
|
||
try:
|
||
outputs: Any = await self._client.async_run(REPLICATE_MODEL, input=arguments)
|
||
artifacts["generated_image_hires"] = await outputs[0].aread()
|
||
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}")
|