diff --git a/README.md b/README.md index 9b969b4..4aaaafc 100644 --- a/README.md +++ b/README.md @@ -6,6 +6,7 @@ Telegram-бот для генерации и редактирования изо - **Text-to-image**: текст → изображение (Z-Image Turbo, 1024×1024). - **Image-to-image**: фото + подпись → редактирование (Krea2 edit). +- **Reply-to-image**: ответ текстом на готовую картинку бота → редактирование этой картинки. - Глобальная очередь с одной активной задачей на пользователя. - Позиция в очереди и прогресс генерации через WebSocket ComfyUI. - SQLite: запросы, статусы, seed, пути к файлам. diff --git a/bot.py b/bot.py index fdd8fd6..d16a41f 100644 --- a/bot.py +++ b/bot.py @@ -27,7 +27,8 @@ logger = logging.getLogger(__name__) START_TEXT = ( "Пришли текст — сгенерирую картинку.\n" - "Пришли фото с подписью — отредактирую его по описанию.\n\n" + "Пришли фото с подписью — отредактирую его по описанию.\n" + "Или ответь текстом на готовую картинку — возьму её как исходник.\n\n" "/status — узнать состояние текущей задачи." ) @@ -160,7 +161,37 @@ class MpakoBot: ) return - await self._submit(update, mode="t2i", prompt=prompt, input_path=None) + reply = message.reply_to_message + attachment = None + suffix = ".jpg" + if reply is not None: + if reply.photo: + attachment = reply.photo[-1] + elif reply.document is not None and ( + reply.document.mime_type or "" + ).startswith("image/"): + attachment = reply.document + suffix = Path(reply.document.file_name or "").suffix.lower() or ".png" + + if attachment is None: + await self._submit( + update, mode="t2i", prompt=prompt, input_path=None + ) + return + + input_path = await self._download_attachment( + message, attachment, user.id, suffix=suffix + ) + if input_path is None: + return + accepted = await self._submit( + update, + mode="i2i", + prompt=prompt, + input_path=str(input_path), + ) + if not accepted: + input_path.unlink(missing_ok=True) async def _photo( self, update: Update, context: ContextTypes.DEFAULT_TYPE @@ -185,22 +216,10 @@ class MpakoBot: return photo = message.photo[-1] - if photo.file_size and photo.file_size > self._settings.max_photo_bytes: - limit_mb = self._settings.max_photo_bytes / (1024 * 1024) - await message.reply_text(f"Фото слишком большое. Максимум {limit_mb:.0f} МБ.") - return - - user_dir = self._settings.storage_dir / str(user.id) - user_dir.mkdir(parents=True, exist_ok=True) - input_path = user_dir / f"{uuid.uuid4().hex}_input.jpg" - try: - await message.chat.send_action(ChatAction.TYPING) - file = await photo.get_file() - await file.download_to_drive(custom_path=input_path) - except TelegramError: - logger.exception("Не удалось скачать фото пользователя %s", user.id) - input_path.unlink(missing_ok=True) - await message.reply_text("Не удалось скачать фото, попробуй ещё раз.") + input_path = await self._download_attachment( + message, photo, user.id, suffix=".jpg" + ) + if input_path is None: return accepted = await self._submit( @@ -212,6 +231,49 @@ class MpakoBot: if not accepted: input_path.unlink(missing_ok=True) + async def _download_attachment( + self, + message, + attachment, + user_id: int, + *, + suffix: str = ".jpg", + ) -> Path | None: + if attachment.file_size and attachment.file_size > self._settings.max_photo_bytes: + limit_mb = self._settings.max_photo_bytes / (1024 * 1024) + try: + await self._retry( + message.reply_text, + f"Файл слишком большой. Максимум {limit_mb:.0f} МБ.", + ) + except TelegramError: + logger.debug("Не удалось сообщить о размере файла") + return None + + suffix = suffix.lower() + if not suffix.startswith(".") or len(suffix) > 8: + suffix = ".png" + + user_dir = self._settings.storage_dir / str(user_id) + input_path = user_dir / f"{uuid.uuid4().hex}_input{suffix}" + try: + user_dir.mkdir(parents=True, exist_ok=True) + await message.chat.send_action(ChatAction.TYPING) + file = await attachment.get_file() + await file.download_to_drive(custom_path=input_path) + except (TelegramError, OSError): + logger.exception("Не удалось скачать вложение пользователя %s", user_id) + input_path.unlink(missing_ok=True) + try: + await self._retry( + message.reply_text, + "Не удалось скачать файл, попробуй ещё раз.", + ) + except TelegramError: + pass + return None + return input_path + async def _submit( self, update: Update,