Kirill Kirilenko
1c58359e44
Вся логика ИИ перенесена в модуль ai. Логика инструментов выделена в отдельные подмодули. Исправлены все проблемы, обнаруженные PyCharm.
85 lines
3.3 KiB
Python
85 lines
3.3 KiB
Python
from openrouter.components import ChatMessageContentItemTypedDict
|
||
from replicate import Client as ReplicateClient
|
||
from typing import Any, Dict, List
|
||
|
||
from ai.tool import Tool
|
||
from ai.utils import *
|
||
|
||
REPLICATE_MODEL = "ultracoderru/nova-anime-xl-17:8f702486aa2852a08564ede8c83a7f58e52c83f6698e7be0e061d79c113dc88b"
|
||
|
||
|
||
class GenerateImageAnimeTool(Tool):
|
||
def __init__(self, replicate_token: str):
|
||
self._client = ReplicateClient(api_token=replicate_token)
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "generate_image_anime"
|
||
|
||
@property
|
||
def description(self) -> str:
|
||
return "Генерация изображения в стиле аниме по описанию"
|
||
|
||
@property
|
||
def parameters(self) -> Dict[str, Any]:
|
||
return {
|
||
"type": "object",
|
||
"properties": {
|
||
"prompt": {
|
||
"type": "string",
|
||
"description": "Положительный запрос"
|
||
},
|
||
"negative_prompt": {
|
||
"type": "string",
|
||
"description": "Отрицательный запрос"
|
||
},
|
||
"aspect_ratio": {
|
||
"type": "string",
|
||
"enum": ["1:1", "4:3", "3:4", "16:9", "9:16", "9:20"],
|
||
"description": "Соотношение сторон"
|
||
}
|
||
},
|
||
"required": ["prompt", "negative_prompt"]
|
||
}
|
||
|
||
async def execute(self, args: Dict[str, Any], artifacts: Dict[str, Any]) -> List[ChatMessageContentItemTypedDict]:
|
||
prompt = args.get("prompt", "")
|
||
negative_prompt = args.get("negative_prompt", "")
|
||
aspect_ratio = args.get("aspect_ratio", None)
|
||
|
||
aspect_ratio_resolution_map = {
|
||
"1:1": (1280, 1280),
|
||
"4:3": (1280, 1024),
|
||
"3:4": (1024, 1280),
|
||
"16:9": (1280, 720),
|
||
"9:16": (720, 1280),
|
||
"9:20": (720, 1600)
|
||
}
|
||
width, height = aspect_ratio_resolution_map.get(aspect_ratio, (1280, 1024))
|
||
print(f"Генерация аниме-изображения {width}x{height}:\n+ {prompt}\n- {negative_prompt}")
|
||
|
||
arguments = {
|
||
"prompt": prompt,
|
||
"negative_prompt": negative_prompt,
|
||
"add_recommended_tags": False,
|
||
"width": width,
|
||
"height": height,
|
||
"guidance_scale": 4.5,
|
||
"num_inference_steps": 20,
|
||
"hires_enable": True,
|
||
"hires_num_inference_steps": 30,
|
||
"disable_safety_checker": True
|
||
}
|
||
|
||
try:
|
||
outputs = 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}")
|