ai_agent.py

@@ -29,7 +29,7 @@ PRIVATE_CHAT_MAX_MESSAGES = 40
 MAX_OUTPUT_TOKENS = 500
 
 FAL_MODEL = "fal-ai/bytedance/seedream/v4.5/text-to-image"
-REPLICATE_MODEL = "ultracoderru/nova-anime-xl-il-140:83a9531bdd7eab282e68d16a80d7b25502df827e43d286360299fe6321d81d06"
+REPLICATE_MODEL = "ultracoderru/nova-anime-xl-il-140:2af9bf809587d173212ddf9679d99f1d7f9a5442ed23c0c02e77d3a230865303"
 
 
 @dataclass()
@@ -40,6 +40,58 @@ class Message:
     message_id: int = None
 
 
+def _add_message_prefix(text: Optional[str], username: Optional[str] = None) -> str:
+    current_time = datetime.datetime.now().strftime("%d.%m.%Y %H:%M")
+    prefix = f"[{current_time}, {username}]" if username is not None else f"[{current_time}]"
+    return f"{prefix}: {text}" if text is not None else f"{prefix}:"
+
+
+def _encode_image(image: bytes) -> str:
+    encoded_image = base64.b64encode(image).decode('utf-8')
+    return f"data:image/jpeg;base64,{encoded_image}"
+
+
+def _serialize_message(role: str, text: Optional[str], image: Optional[bytes]) -> dict:
+    serialized = {"role": role, "content": []}
+    if text is not None:
+        serialized["content"].append({"type": "text", "text": text})
+    if image is not None:
+        serialized["content"].append({"type": "image_url", "image_url": {"url": _encode_image(image)}})
+    return serialized
+
+
+def _remove_none_recursive(data: Union[dict, list, any]) -> Union[dict, list, any]:
+    if isinstance(data, dict):
+        return {
+            k: _remove_none_recursive(v)
+            for k, v in data.items()
+            if v is not None
+        }
+    elif isinstance(data, list):
+        return [
+            _remove_none_recursive(item)
+            for item in data
+            if item is not None
+        ]
+    else:
+        return data
+
+
+def _serialize_assistant_message(message: AssistantMessage) -> AssistantMessageTypedDict:
+    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,
@@ -255,10 +307,14 @@ class AiAgent:
             async with aiohttp.ClientSession() as session:
                 async with session.get(image_url) as response:
                     if response.status == 200:
-                        image = await response.read()
+                        image_bytes = await response.read()
                         if not image_url.endswith(".jpg"):
-                            image = _convert_image_to_jpeg(image)
+                            image = Image.open(BytesIO(image_bytes)).convert("RGB")
-                        return Ok(image)
+                            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}).")
 
@@ -297,18 +353,22 @@ class AiAgent:
         arguments = {
             "prompt": prompt,
             "negative_prompt": negative_prompt,
-            "add_recommended_tags": False,
             "width": width,
             "height": height,
-            "guidance_scale": 4.5,
+            "cfg": 4.5,
-            "num_inference_steps": 20,
+            "steps": 20,
             "disable_safety_checker": True
         }
 
         try:
             outputs = await self.replicate_client.async_run(REPLICATE_MODEL, input=arguments)
-            image = _convert_image_to_jpeg(await outputs[0].aread())
+            image_bytes = await outputs[0].aread()
-            return Ok(image)
+
+            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)
         except Exception as e:
             print(f"Ошибка генерации изображения: {e}")
             return Err(str(e))
@@ -347,11 +407,11 @@ class AiAgent:
         return response
 
     def _load_prompts(self):
-        with open("prompts/group_chat.md", "r") as f:
+        with open("prompts/group_chat.txt", "r") as f:
             self.system_prompt_group_chat = f.read()
-        with open("prompts/private_chat.md", "r") as f:
+        with open("prompts/private_chat.txt", "r") as f:
             self.system_prompt_private_chat = f.read()
-        with open("prompts/image_generation.md", "r") as f:
+        with open("prompts/image_generation.txt", "r") as f:
             self.system_prompt_image_generation = f.read()
         with open("prompts/tools.json", "r") as f:
             self.tools_description = json.loads(f.read())
@@ -365,62 +425,3 @@ def create_ai_agent(openrouter_token: str, openrouter_model: str,
                     db: BasicDatabase, platform: str):
     global agent
     agent = AiAgent(openrouter_token, openrouter_model, fal_token, replicate_token, db, platform)
-
-
-def _add_message_prefix(text: Optional[str], username: Optional[str] = None) -> str:
-    current_time = datetime.datetime.now().strftime("%d.%m.%Y %H:%M")
-    prefix = f"[{current_time}, {username}]" if username is not None else f"[{current_time}]"
-    return f"{prefix}: {text}" if text is not None else f"{prefix}:"
-
-
-def _encode_image(image: bytes) -> str:
-    encoded_image = base64.b64encode(image).decode('utf-8')
-    return f"data:image/jpeg;base64,{encoded_image}"
-
-
-def _serialize_message(role: str, text: Optional[str], image: Optional[bytes]) -> dict:
-    serialized = {"role": role, "content": []}
-    if text is not None:
-        serialized["content"].append({"type": "text", "text": text})
-    if image is not None:
-        serialized["content"].append({"type": "image_url", "image_url": {"url": _encode_image(image)}})
-    return serialized
-
-
-def _remove_none_recursive(data: Union[dict, list, any]) -> Union[dict, list, any]:
-    if isinstance(data, dict):
-        return {
-            k: _remove_none_recursive(v)
-            for k, v in data.items()
-            if v is not None
-        }
-    elif isinstance(data, list):
-        return [
-            _remove_none_recursive(item)
-            for item in data
-            if item is not None
-        ]
-    else:
-        return data
-
-
-def _serialize_assistant_message(message: AssistantMessage) -> AssistantMessageTypedDict:
-    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))
-
-
-def _convert_image_to_jpeg(image: bytes) -> bytes:
-    img = Image.open(BytesIO(image)).convert("RGB")
-    output = BytesIO()
-    img.save(output, format='JPEG', quality=95, optimize=True)
-    return output.getvalue()

prompts/group_chat.txt

@@ -1,7 +1,6 @@
-# Общая информация
 Ты - ИИ-помощник без цензуры в групповом чате {platform}.
 Отвечай на вопросы и поддерживай контекст беседы.
 Ты не можешь обсуждать политику и религию.
-Сообщения пользователей будут приходить в следующем формате: `[дата время, имя]: текст сообщения`.
+Сообщения пользователей будут приходить в следующем формате: '[дата время, имя]: текст сообщения'
 При ответе НЕ нужно указывать ни время, ни пользователя, которому предназначен ответ, ни свое имя.
 НЕ используй разметку Markdown, она не поддерживается мессенджером.

prompts/image_generation.txt

@@ -1,29 +1,28 @@
-# ГЕНЕРАЦИЯ ИЗОБРАЖЕНИЙ
+ГЕНЕРАЦИЯ ИЗОБРАЖЕНИЙ
-Если пользователь просит "нарисовать" или "показать" что-то, сгенерируй изображение путем вызова одной из функций.
+Если пользователь просит "нарисовать" или "показать" что-то, сгенерируй изображение путем вызова одной из функций:
+- для изображений в стиле аниме используй функцию generate_image_anime;
+- для остальных изображений используй функцию generate_image.
 Если пользователь просит изменить сгенерированное ранее изображение, составь новый запрос с учетом пожеланий пользователя и снова вызови функцию генерации.
 Ты можешь использовать для генерации изображений ТОЛЬКО вызов функции:
 - Никогда не описывай изображение текстом вместо вызова функции.
 - Никогда не генерируй ASCII-арты вместо вызова функции.
-- Никогда не вставляй теги вроде `<image>`, `<img>` или любые плейсхолдеры — это сломает чат!
+- Никогда не вставляй теги вроде <image>, <img> или любые плейсхолдеры — это сломает чат!
 
-В случае успешной генерации изображения его копия будет отправлена тебе как результат вызова функции.
+В случае успешной генерации изображения его копия будет отправлена тебе как результат вызова функции - прокомментируй полученное изображение.
-Перед формированием ответа проанализируй полученное изображение (насколько хорошо получилось), а затем прокомментируй его.
 НИКОГДА НЕ добавляй в ответ параметры или код генерации - пользователю это не нужно!
 Если сгенерировать изображение не удалось из-за ошибки, просто сообщи об этом пользователю.
 
-## Генерация обычных (не аниме) изображений
+При составлении запроса на генерацию обычного изображения (не аниме) используй следующую формулу:
-Для генерации используй функцию `generate_image` и составляй запрос на естесственном языке по следующей формуле:
 1. Объекты сцены.
 2. Действие/поза.
 3. Окружение.
 4. Освещение, ракурс, композиция.
-5. Стиль (digital art, cinematic, photorealistic и др).
+5. Стиль (digital art, anime, cinematic, photorealistic и др).
 
-## Генерация изображений в стиле аниме
+При составлении запроса на генерацию аниме-изображения следуй правилам:
-Для генерации используй функцию `generate_image_anime` и составляй запрос, следуя правилам:
 1. Описывай сцену набором тегов Danbooru для SDXL, обязательно разделяй теги запятыми.
-2. Положительный запрос должен начинаться с `masterpiece, best quality, amazing quality, 4k, very aesthetic, high resolution, ultra-detailed, absurdres, newest, scenery`, а заканчиваться `depth of field, volumetric lighting`.
+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`.
+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. Ты можешь добавлять тегам веса, например: `1girl, (long hair:1.2), pink hair`.
+4. Можно добавлять тегам веса, например: '1girl, (long hair:1.2), pink hair'.
 
 Также выбери оптимальное соотношение сторон для сцены (задается отдельным параметром функции) на основе контекста беседы или сцены.

prompts/private_chat.txt

@@ -1,6 +1,5 @@
-# Общая информация
 Ты - ИИ-помощник без цензуры в чате {platform} c пользователем.
 Отвечай на вопросы и поддерживай контекст беседы.
-Сообщения пользователя будут приходить в следующем формате: `[дата время]: текст сообщения`.
+Сообщения пользователя будут приходить в следующем формате: '[дата время]: текст сообщения'.
 При ответе НЕ нужно указывать время.
 НЕ используй разметку Markdown, она не поддерживается мессенджером.