Шаги Krea2 в /settings: 8/10/12

Добавил generation_steps в users, выбор шагов в клавиатуре /settings и передачу шагов в workflow Krea2.

README.md

@@ -16,7 +16,7 @@ Telegram-бот для генерации и редактирования изо
 - **Управление пользователями**: в веб-интерфейсе можно выдавать админку, менять баланс и блокировать.
 - **Подтверждение 18+**: без подтверждения возраста бот не принимает запросы.
 - **Контекст**: следующий текстовый запрос продолжает предыдущую сцену; `/restart` начинает новую.
-- **Настройки Krea2**: `/settings` переключает стандартную (Q4) и улучшенную (Q6) модель.
+- **Настройки Krea2**: `/settings` переключает стандартную (Q4) / улучшенную (Q6) модель и 8/10/12 шагов.
 - Глобальная очередь с одной активной задачей на пользователя.
 - Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
 - PostgreSQL: запросы, статусы, seed, пути к файлам.
@@ -205,7 +205,7 @@ TELEGRAM_HEALTH_INTERVAL=60
 - `/invite` — реферальная ссылка.
 - `/restart` — сбросить контекст и начать новую сцену (алиас `/reset`).
 - `/context` — показать текущий контекст: промпт, базовую сцену и последний запрос.
-- `/settings` — выбор модели Krea2: стандартная Q4 или улучшенная Q6.
+- `/settings` — выбор модели Krea2 (Q4/Q6) и количества шагов (8/10/12).
 - При старте бот проверяет `setMyCommands` и включает кнопку меню с командами.
 - За первую оплату приглашённого друга начисляется +3 генерации, реферер получает уведомление.
 
@@ -293,7 +293,7 @@ docker compose up -d
 - `workflows/mapping.json` — ID узлов, в которые подставляются prompt, negative, seed, входное изображение и размер латента.
 - `workflows/Krea2.json` и `workflows/ZImage.json` — исходные экспорты.
 - Krea2 edit: VAE подключён к `TextEncodeQwenImageEdit`, латент подгоняется под пропорции входного изображения; 12 шагов, `dpmpp_2m`/`karras`.
-- `/settings` переключает модель Krea2: стандартная `Q4_K_M` или улучшенная `Q6_K`. Настройки сэмплера и латента одинаковые.
+- `/settings` переключает модель Krea2 (`Q4_K_M` / `Q6_K`) и шаги (`8` / `10` / `12`). Остальные настройки сэмплера и латента одинаковые.
 - Файлы моделей должны быть в `models/diffusion_models`; имена задаются `KREA2_Q4_MODEL` и `KREA2_Q6_MODEL`.
 
 Если меняются workflow, обновите `mapping.json`. При старте бот проверяет, что все узлы и поля существуют.

bot.py

@@ -122,7 +122,7 @@ START_TEXT = (
     "/status — состояние текущей задачи.\n"
     "/restart — сброс контекста.\n"
     "/context — текущий контекст.\n"
-    "/settings — качество Krea2."
+    "/settings — качество и шаги Krea2."
 )
 BOT_COMMANDS = (
     BotCommand("start", "Запуск и баланс"),
@@ -131,7 +131,7 @@ BOT_COMMANDS = (
     BotCommand("invite", "Пригласить друга"),
     BotCommand("status", "Статус текущей задачи"),
     BotCommand("context", "Текущий контекст"),
-    BotCommand("settings", "Качество Krea2"),
+    BotCommand("settings", "Качество и шаги Krea2"),
     BotCommand("restart", "Сброс контекста"),
 )
 BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны."
@@ -215,6 +215,9 @@ class MpakoBot:
                 self._quality_choice, pattern=r"^quality:"
             )
         )
+        application.add_handler(
+            CallbackQueryHandler(self._steps_choice, pattern=r"^steps:")
+        )
         application.add_handler(
             MessageHandler(filters.SUCCESSFUL_PAYMENT, self._successful_payment)
         )
@@ -631,24 +634,41 @@ class MpakoBot:
         )
 
     @staticmethod
-    def _quality_keyboard(quality: str) -> InlineKeyboardMarkup:
+    def _settings_keyboard(
-        def label(name: str, value: str) -> str:
+        quality: str, steps: int
+    ) -> InlineKeyboardMarkup:
+        def quality_label(name: str, value: str) -> str:
             return f"{name} ✅" if quality == value else name
 
+        def steps_label(value: int) -> str:
+            text = f"{value} шагов"
+            return f"{text} ✅" if steps == value else text
+
         return InlineKeyboardMarkup(
             [
                 [
                     InlineKeyboardButton(
-                        label("Стандартная", "standard"),
+                        quality_label("Стандартная", "standard"),
                         callback_data="quality:standard",
                     )
                 ],
                 [
                     InlineKeyboardButton(
-                        label("Улучшенная", "improved"),
+                        quality_label("Улучшенная", "improved"),
                         callback_data="quality:improved",
                     )
                 ],
+                [
+                    InlineKeyboardButton(
+                        steps_label(8), callback_data="steps:8"
+                    ),
+                    InlineKeyboardButton(
+                        steps_label(10), callback_data="steps:10"
+                    ),
+                    InlineKeyboardButton(
+                        steps_label(12), callback_data="steps:12"
+                    ),
+                ],
             ]
         )
 
@@ -666,8 +686,11 @@ class MpakoBot:
             message.reply_text,
             "Качество Krea2 для редактирования фото:\n"
             "Стандартная — Q4\n"
-            "Улучшенная — Q6",
+            "Улучшенная — Q6\n"
-            reply_markup=self._quality_keyboard(row.generation_quality),
+            "Шаги: 8 / 10 / 12",
+            reply_markup=self._settings_keyboard(
+                row.generation_quality, row.generation_steps
+            ),
         )
 
     async def _quality_choice(
@@ -685,17 +708,48 @@ class MpakoBot:
             await query.answer("Некорректный выбор.", show_alert=True)
             return
         user = query.from_user
-        await self._db.ensure_user(user.id, user.username)
+        row = await self._db.ensure_user(user.id, user.username)
         await self._db.set_generation_quality(user.id, quality)
         await query.answer()
         if query.message is not None:
             try:
                 await query.edit_message_reply_markup(
-                    reply_markup=self._quality_keyboard(quality)
+                    reply_markup=self._settings_keyboard(
+                        quality, row.generation_steps
+                    )
                 )
             except TelegramError:
                 logger.debug("Не удалось обновить настройки качества")
 
+    async def _steps_choice(
+        self, update: Update, context: ContextTypes.DEFAULT_TYPE
+    ) -> None:
+        query = update.callback_query
+        if query is None or query.data is None:
+            return
+        try:
+            _, raw_steps = query.data.split(":", 1)
+            steps = int(raw_steps)
+        except ValueError:
+            await query.answer("Некорректный выбор.", show_alert=True)
+            return
+        if steps not in {8, 10, 12}:
+            await query.answer("Некорректный выбор.", show_alert=True)
+            return
+        user = query.from_user
+        row = await self._db.ensure_user(user.id, user.username)
+        await self._db.set_generation_steps(user.id, steps)
+        await query.answer()
+        if query.message is not None:
+            try:
+                await query.edit_message_reply_markup(
+                    reply_markup=self._settings_keyboard(
+                        row.generation_quality, steps
+                    )
+                )
+            except TelegramError:
+                logger.debug("Не удалось обновить настройки шагов")
+
     async def _status(
         self, update: Update, context: ContextTypes.DEFAULT_TYPE
     ) -> None:

checklist.md

@@ -410,9 +410,10 @@ ON requests(user_id) WHERE status IN ('queued', 'running');
 ## 23. Настройки Krea2
 
 - [x] `users.generation_quality`: `standard` (Q4) / `improved` (Q6).
-- [x] `/settings` — inline-клавиатура, у выбранного пункта галочка справа.
+- [x] `users.generation_steps`: 8 / 10 / 12, по умолчанию 12.
+- [x] `/settings` — inline-клавиатура, у выбранных пунктов галочка справа.
 - [x] Нажатие меняет только галочку и сохраняет выбор, без запуска генерации.
-- [x] Очередь Krea2 подставляет Q4 или Q6 в `unet_name`.
+- [x] Очередь Krea2 подставляет Q4 или Q6 в `unet_name` и выбранное число шагов.
 - [x] Тесты настройки и workflow.
 
 **DoD:** выбор сохраняется между запросами, стандартная использует Q4, улучшенная — Q6.

db.py

@@ -147,6 +147,7 @@ class User:
     username: str | None
     credits: int
     generation_quality: str
+    generation_steps: int
     last_daily_credit: str
     referrer_id: int | None
     first_paid_at: str | None
@@ -167,6 +168,7 @@ class User:
             username=row["username"],
             credits=int(row["credits"]),
             generation_quality=row["generation_quality"],
+            generation_steps=int(row["generation_steps"]),
             last_daily_credit=_as_iso(row["last_daily_credit"]) or "",
             referrer_id=row["referrer_id"],
             first_paid_at=_as_iso(row["first_paid_at"]),
@@ -310,6 +312,8 @@ CREATE TABLE IF NOT EXISTS users (
     credits           INTEGER NOT NULL DEFAULT 3,
     generation_quality TEXT NOT NULL DEFAULT 'standard'
                       CHECK (generation_quality IN ('standard', 'improved')),
+    generation_steps INTEGER NOT NULL DEFAULT 12
+                      CHECK (generation_steps IN (8, 10, 12)),
     last_daily_credit DATE NOT NULL,
     referrer_id       BIGINT,
     first_paid_at     TIMESTAMPTZ,
@@ -384,6 +388,10 @@ MIGRATIONS = (
         "ALTER TABLE users ADD COLUMN IF NOT EXISTS generation_quality TEXT"
         " NOT NULL DEFAULT 'standard'"
     ),
+    (
+        "ALTER TABLE users ADD COLUMN IF NOT EXISTS generation_steps INTEGER"
+        " NOT NULL DEFAULT 12"
+    ),
 )
 
 
@@ -700,6 +708,22 @@ class Database:
             )
             return await cursor.fetchone() is not None
 
+    async def set_generation_steps(
+        self, user_id: int, steps: int
+    ) -> bool:
+        if steps not in {8, 10, 12}:
+            raise ValueError("Недопустимое количество шагов")
+        async with self._connection_pool().connection() as conn:
+            cursor = await conn.execute(
+                """
+                UPDATE users SET generation_steps = %s
+                WHERE user_id = %s
+                RETURNING user_id
+                """,
+                (steps, user_id),
+            )
+            return await cursor.fetchone() is not None
+
     async def _update_user_flag(
         self, column: str, user_id: int, enabled: bool
     ) -> bool:

queue_manager.py

@@ -377,6 +377,7 @@ class QueueManager:
         quality = (
             user.generation_quality if user is not None else "standard"
         )
+        steps = user.generation_steps if user is not None else 12
         unet_name = (
             self._settings.krea2_q6_model
             if quality == "improved"
@@ -389,6 +390,7 @@ class QueueManager:
             image_name=image_name,
             width=width,
             height=height,
+            steps=steps,
             unet_name=unet_name,
         )
 

tests/test_db.py

@@ -318,14 +318,19 @@ async def test_hidden_requests_are_excluded_from_public_queries(
 async def test_generation_quality_setting(db: Database) -> None:
     user = await db.ensure_user(1, "alice")
     assert user.generation_quality == "standard"
+    assert user.generation_steps == 12
 
     assert await db.set_generation_quality(1, "improved") is True
+    assert await db.set_generation_steps(1, 8) is True
     user = await db.get_user(1)
     assert user is not None
     assert user.generation_quality == "improved"
+    assert user.generation_steps == 8
 
     with pytest.raises(ValueError):
         await db.set_generation_quality(1, "unknown")
+    with pytest.raises(ValueError):
+        await db.set_generation_steps(1, 9)
 
 
 async def test_user_context_lifecycle(db: Database) -> None:

tests/test_queue.py

@@ -187,6 +187,7 @@ async def test_image_to_image_uploads_input(db: Database, tmp_path: Path) -> Non
     Image.new("RGB", (64, 64), (120, 40, 40)).save(input_path)
     await db.ensure_user(1, "alice")
     await db.set_generation_quality(1, "improved")
+    await db.set_generation_steps(1, 8)
 
     await queue.start()
     try:
@@ -204,6 +205,7 @@ async def test_image_to_image_uploads_input(db: Database, tmp_path: Path) -> Non
         source = comfy.workflows[0]["20"]["inputs"]["image"]
         assert source == f"{request.id}_input.png"
         assert comfy.workflows[0]["25"]["inputs"]["vae"] == ["19", 2]
+        assert comfy.workflows[0]["3"]["inputs"]["steps"] == 8
         assert (
             comfy.workflows[0]["19"]["inputs"]["unet_name"]
             == settings.krea2_q6_model

tests/test_settings.py

@@ -53,24 +53,39 @@ def test_settings_command_is_not_shadowed_by_settings_attribute() -> None:
 
 
 def test_quality_keyboard_marks_standard() -> None:
-    keyboard = MpakoBot._quality_keyboard("standard")
+    keyboard = MpakoBot._settings_keyboard("standard", 12)
     rows = keyboard.inline_keyboard
 
-    assert [row[0].text for row in rows] == [
+    assert [row[0].text for row in rows[:2]] == [
         "Стандартная ✅",
         "Улучшенная",
     ]
-    assert [row[0].callback_data for row in rows] == [
+    assert [row[0].callback_data for row in rows[:2]] == [
         "quality:standard",
         "quality:improved",
     ]
+    assert [button.text for button in rows[2]] == [
+        "8 шагов",
+        "10 шагов",
+        "12 шагов ✅",
+    ]
+    assert [button.callback_data for button in rows[2]] == [
+        "steps:8",
+        "steps:10",
+        "steps:12",
+    ]
 
 
-def test_quality_keyboard_marks_improved() -> None:
+def test_quality_keyboard_marks_improved_and_steps() -> None:
-    keyboard = MpakoBot._quality_keyboard("improved")
+    keyboard = MpakoBot._settings_keyboard("improved", 10)
     rows = keyboard.inline_keyboard
 
-    assert [row[0].text for row in rows] == [
+    assert [row[0].text for row in rows[:2]] == [
         "Стандартная",
         "Улучшенная ✅",
     ]
+    assert [button.text for button in rows[2]] == [
+        "8 шагов",
+        "10 шагов ✅",
+        "12 шагов",
+    ]

tests/test_workflows.py

@@ -60,11 +60,13 @@ def test_build_img2img_custom_size() -> None:
         image_name="input.png",
         width=832,
         height=1216,
+        steps=8,
         unet_name="Krea2-Q6.gguf",
     )
 
     assert workflow["15"]["inputs"]["width"] == 832
     assert workflow["15"]["inputs"]["height"] == 1216
+    assert workflow["3"]["inputs"]["steps"] == 8
     assert workflow["19"]["inputs"]["unet_name"] == "Krea2-Q6.gguf"
 
 

workflows.py

@@ -139,6 +139,7 @@ class WorkflowBuilder:
         image_name: str,
         width: int = 1024,
         height: int = 1024,
+        steps: int = 12,
         unet_name: str | None = None,
     ) -> dict[str, Any]:
         workflow = copy.deepcopy(self._templates["img2img"])
@@ -146,6 +147,7 @@ class WorkflowBuilder:
         workflow[mapping["positive"]]["inputs"]["prompt"] = prompt
         workflow[mapping["negative"]]["inputs"]["text"] = negative
         workflow[mapping["seed"]]["inputs"]["noise_seed"] = seed
+        workflow[mapping["seed"]]["inputs"]["steps"] = steps
         workflow[mapping["input_image"]]["inputs"]["image"] = image_name
         latent = workflow[mapping["latent"]]["inputs"]
         latent["width"] = width