feat: use replied photo as i2i source

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, пути к файлам.

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:
+        input_path = await self._download_attachment(
-            limit_mb = self._settings.max_photo_bytes / (1024 * 1024)
+            message, photo, user.id, suffix=".jpg"
-            await message.reply_text(f"Фото слишком большое. Максимум {limit_mb:.0f} МБ.")
+        )
-            return
+        if input_path is None:
-
-        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("Не удалось скачать фото, попробуй ещё раз.")
             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,