From 2e67865b1d6a984ae080d1192ca395e4547d8bba Mon Sep 17 00:00:00 2001 From: Kirill Kirilenko Date: Wed, 8 Apr 2026 21:04:58 +0300 Subject: [PATCH] =?UTF-8?q?=D0=93=D0=B5=D0=BD=D0=B5=D1=80=D0=B0=D1=86?= =?UTF-8?q?=D0=B8=D1=8F=20=D0=BE=D0=B1=D1=8B=D1=87=D0=BD=D1=8B=D1=85=20?= =?UTF-8?q?=D0=B8=D0=B7=D0=BE=D0=B1=D1=80=D0=B0=D0=B6=D0=B5=D0=BD=D0=B8?= =?UTF-8?q?=D0=B9=20=D1=87=D0=B5=D1=80=D0=B5=D0=B7=20Replicate.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ai/__init__.py | 4 +-- ai/agent.py | 4 +-- ai/tools/image_generation/__init__.py | 4 +-- ai/tools/image_generation/generate_image.py | 37 ++++++--------------- requirements.txt | 1 - tg/__main__.py | 4 +-- vk/__main__.py | 4 +-- 7 files changed, 21 insertions(+), 37 deletions(-) diff --git a/ai/__init__.py b/ai/__init__.py index baf5ec4..828a267 100644 --- a/ai/__init__.py +++ b/ai/__init__.py @@ -9,11 +9,11 @@ agent_instance: ai.agent.AiAgent def create_ai_agent(openrouter_token: str, openrouter_model: str, - fal_token: str, replicate_token: str, tavily_token: str, + replicate_token: str, tavily_token: str, db: BasicDatabase, platform: str): global agent_instance agent_instance = ai.agent.AiAgent(openrouter_token, openrouter_model, - fal_token, replicate_token, tavily_token, + replicate_token, tavily_token, db, platform) diff --git a/ai/agent.py b/ai/agent.py index 64c6410..45535ea 100644 --- a/ai/agent.py +++ b/ai/agent.py @@ -37,7 +37,7 @@ class Message: class AiAgent: def __init__(self, openrouter_token: str, openrouter_model: str, - fal_token: str, replicate_token: str, tavily_token: str, + replicate_token: str, tavily_token: str, db: BasicDatabase, platform: str): retry_config = RetryConfig(strategy="backoff", @@ -58,7 +58,7 @@ class AiAgent: # Создание наборов инструментов self.toolsets: list[ai.tool.ToolSet] = [] self.toolsets.append( - ImageGenerationToolSet(fal_token=fal_token, replicate_token=replicate_token) + ImageGenerationToolSet(replicate_token=replicate_token) ) self.toolsets.append(TavilySearchToolSet(tavily_token=tavily_token)) diff --git a/ai/tools/image_generation/__init__.py b/ai/tools/image_generation/__init__.py index 42683be..db2ca8e 100644 --- a/ai/tools/image_generation/__init__.py +++ b/ai/tools/image_generation/__init__.py @@ -5,9 +5,9 @@ from .generate_image_anime import GenerateImageAnimeTool class ImageGenerationToolSet(ToolSet): - def __init__(self, fal_token: str, replicate_token: str): + def __init__(self, replicate_token: str): functions = [ - GenerateImageTool(fal_token), + GenerateImageTool(replicate_token), GenerateImageAnimeTool(replicate_token) ] with open("ai/tools/image_generation/prompt.md", "r") as f: diff --git a/ai/tools/image_generation/generate_image.py b/ai/tools/image_generation/generate_image.py index c9d3531..9802119 100644 --- a/ai/tools/image_generation/generate_image.py +++ b/ai/tools/image_generation/generate_image.py @@ -1,17 +1,17 @@ from typing import Any, Dict, List from openrouter.components import ChatToolMessageContentTypedDict -from fal_client import AsyncClient as FalClient +from replicate import Client as ReplicateClient from ai.tool import Tool from ai.utils import * -FAL_MODEL = "fal-ai/bytedance/seedream/v4.5/text-to-image" +REPLICATE_MODEL = "bytedance/seedream-4.5" class GenerateImageTool(Tool): - def __init__(self, fal_token: str): - self._client = FalClient(key=fal_token) + def __init__(self, replicate_token: str): + self._client = ReplicateClient(api_token=replicate_token) @property def name(self) -> str: @@ -43,34 +43,19 @@ class GenerateImageTool(Tool): 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}") + print(f"Генерация изображения {aspect_ratio}: {prompt}") arguments = { "prompt": prompt, - "image_size": image_size, - "enable_safety_checker": False + "aspect_ratio": aspect_ratio, + "disable_safety_checker": True } - + 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) + 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 diff --git a/requirements.txt b/requirements.txt index 7c48962..b003dac 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,7 +5,6 @@ vkbottle-types~=5.199.99.18 pyodbc~=5.3.0 openrouter==0.8.1 replicate~=1.0.7 -fal_client~=0.13.2 tavily~=1.1.0 pillow~=12.2.0 pymorphy3~=2.0.6 diff --git a/tg/__main__.py b/tg/__main__.py index 0b58ad2..c5bda8c 100644 --- a/tg/__main__.py +++ b/tg/__main__.py @@ -25,8 +25,8 @@ async def main() -> None: database.create_database(config['db_connection_string']) create_ai_agent(config['openrouter_token'], config['openrouter_model'], - config['fal_token'], config['replicate_token'], - config['tavily_token'], database.DB, 'tg') + config['replicate_token'], config['tavily_token'], + database.DB, 'tg') bots: list[Bot] = [] for item in database.DB.get_bots(): diff --git a/vk/__main__.py b/vk/__main__.py index b304e72..49cb904 100644 --- a/vk/__main__.py +++ b/vk/__main__.py @@ -25,8 +25,8 @@ if __name__ == '__main__': database.create_database(config['db_connection_string']) create_ai_agent(config['openrouter_token'], config['openrouter_model'], - config['fal_token'], config['replicate_token'], - config['tavily_token'], database.DB, 'vk') + config['replicate_token'], config['tavily_token'], + database.DB, 'vk') bot = Bot(labeler=handlers.labeler)