diff --git a/ai_agent.py b/ai_agent.py index b586e9f..39b910b 100644 --- a/ai_agent.py +++ b/ai_agent.py @@ -77,6 +77,17 @@ def _serialize_assistant_message(message: AssistantMessage) -> AssistantMessageT return _remove_none_recursive(message.model_dump(by_alias=True)) +def _get_resolution_for_aspect_ratio(aspect_ratio: str) -> Tuple[int, int]: + aspect_ratio_resolution_map = { + "1:1": (1280, 1280), + "4:3": (1280, 1024), + "3:4": (1024, 1280), + "16:9": (1280, 720), + "9:16": (720, 1280) + } + return aspect_ratio_resolution_map.get(aspect_ratio, (1280, 1024)) + + class AiAgent: def __init__(self, openrouter_token: str, openrouter_model: str, @@ -236,7 +247,8 @@ class AiAgent: functions_map: Dict[str, Callable[[int, int, Dict, AiAgent._ToolsArtifacts], Awaitable[List[ChatMessageContentItemTypedDict]]]] = { - "generate_image": self._process_tool_generate_image + "generate_image": self._process_tool_generate_image, + "generate_image_anime": self._process_tool_generate_image_anime } for tool_call in tool_calls: @@ -273,14 +285,7 @@ class AiAgent: return content async def _generate_image(self, prompt: str, aspect_ratio: Optional[str]) -> Result[bytes, str]: - aspect_ratio_resolution_map = { - "1:1": (1280, 1280), - "4:3": (1280, 1024), - "3:4": (1024, 1280), - "16:9": (1280, 720), - "9:16": (720, 1280) - } - width, height = aspect_ratio_resolution_map.get(aspect_ratio, (1280, 1024)) + width, height = _get_resolution_for_aspect_ratio(aspect_ratio) print(f"Генерация изображения {width}x{height}: {prompt}") arguments = { @@ -313,6 +318,61 @@ class AiAgent: print(f"Ошибка генерации изображения: {e}") return Err(str(e)) + async def _process_tool_generate_image_anime(self, _bot_id: int, _chat_id: int, + args: dict, artifacts: _ToolsArtifacts) \ + -> List[ChatMessageContentItemTypedDict]: + prompt = args.get("prompt", "") + negative_prompt = args.get("negative_prompt", "") + aspect_ratio = args.get("aspect_ratio", None) + result = await self._generate_image_anime(prompt=prompt, negative_prompt=negative_prompt, + aspect_ratio=aspect_ratio) + + content = [] + if result.is_ok(): + content.append( + {"type": "text", + "text": "Изображение сгенерировано и будет показано пользователю."}) + content.append( + {"type": "image_url", "image_url": {"url": _encode_image(result.ok_value)}}) + artifacts.generated_image = result.ok_value + else: + content.append( + {"type": "text", + "text": f"Не удалось сгенерировать изображение: {result.err_value}"}) + return content + + @staticmethod + async def _generate_image_anime(prompt: str, negative_prompt: str, aspect_ratio: Optional[str]) \ + -> Result[bytes, str]: + width, height = _get_resolution_for_aspect_ratio(aspect_ratio) + print(f"Генерация изображения {width}x{height}: positive='{prompt}', negative='{negative_prompt}'") + + arguments = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "width": width, + "height": height + } + + try: + async with aiohttp.ClientSession() as session: + async with session.post("http://192.168.64.2:8787/sdapi/v1/txt2img", + json=arguments, timeout=120) as response: + if response.status == 200: + data = await response.json() + image_base64 = data["images"][0] + image_bytes = base64.b64decode(image_base64) + image = Image.open(BytesIO(image_bytes)).convert("RGB") + output = BytesIO() + image.save(output, format="JPEG", quality=80, optimize=True) + image_bytes = output.getvalue() + return Ok(image_bytes) + else: + raise RuntimeError(f"Сервер вернул код {response.status}") + except Exception as e: + print(f"Ошибка генерации изображения: {e}") + return Err(str(e)) + async def _async_chat_completion_request(self, **kwargs): try: return await self.client_openrouter.chat.send_async(**kwargs) diff --git a/prompts/image_generation.txt b/prompts/image_generation.txt index 27d9f5d..d9998cb 100644 --- a/prompts/image_generation.txt +++ b/prompts/image_generation.txt @@ -1,5 +1,7 @@ ГЕНЕРАЦИЯ ИЗОБРАЖЕНИЙ -Если пользователь просит "нарисовать" или "показать" что-то, сгенерируй изображение путем вызова функции generate_image. +Если пользователь просит "нарисовать" или "показать" что-то, сгенерируй изображение путем вызова одной из функций: +- для изображений в стиле аниме используй функцию generate_image_anime; +- для остальных изображений используй функцию generate_image. Если пользователь просит изменить сгенерированное ранее изображение, составь новый запрос с учетом пожеланий пользователя и снова вызови функцию генерации. Ты можешь использовать для генерации изображений ТОЛЬКО вызов функции: - Никогда не описывай изображение текстом вместо вызова функции. @@ -7,10 +9,17 @@ - Никогда не вставляй теги вроде , или любые плейсхолдеры — это сломает чат! Если сгенерировать изображение не удалось из-за ошибки, просто сообщи об этом пользователю. -При составлении запроса на генерацию изображения используй следующую формулу: +При составлении запроса на генерацию обычного изображения (не аниме) используй следующую формулу: 1. Объекты сцены. 2. Действие/поза. 3. Окружение. 4. Освещение, ракурс, композиция. 5. Стиль (digital art, anime, cinematic, photorealistic и др). + +При составлении запроса на генерацию аниме-изображения следуй правилам: +1. Описывай сцену набором тегов Danbooru для SDXL, обязательно разделяй теги запятыми. +2. Положительный запрос должен начинаться с 'masterpiece, best quality, amazing quality, 4k, very aesthetic, high resolution, ultra-detailed, absurdres, newest, scenery', а заканчиваться 'depth of field, volumetric lighting'. +3. Отрицательный запрос должен включать 'modern, recent, old, oldest, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured, long body, lowres, bad anatomy, bad hands, missing fingers, extra digits, fewer digits, cropped, very displeasing, (worst quality, bad quality:1.2), bad anatomy, sketch, jpeg artifacts, signature, watermark, username, signature, simple background, conjoined, bad ai-generated'. +4. Количество токенов в каждом запросе не должно превышать 75. + Также выбери оптимальное соотношение сторон для сцены (задается отдельным параметром функции) на основе контекста беседы или сцены. diff --git a/prompts/tools.json b/prompts/tools.json index b421a22..0f4fff7 100644 --- a/prompts/tools.json +++ b/prompts/tools.json @@ -20,5 +20,31 @@ "required": ["prompt"] } } + }, + { + "type": "function", + "function": { + "name": "generate_image_anime", + "description": "Генерация изображения в стиле аниме по описанию", + "parameters": { + "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"], + "description": "Соотношение сторон" + } + }, + "required": ["prompt", "negative_prompt"] + } + } } ]