From 64d51f06ffbcfd9c408bb3e7daa3a2373ab1230d Mon Sep 17 00:00:00 2001 From: MpakobecEx Date: Tue, 20 Oct 2026 01:35:11 +0500 Subject: [PATCH] =?UTF-8?q?=D0=A8=D0=B0=D0=B3=D0=B8=20Krea2=20=D0=B2=20/se?= =?UTF-8?q?ttings:=208/10/12?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Добавил generation_steps в users, выбор шагов в клавиатуре /settings и передачу шагов в workflow Krea2. --- README.md | 6 ++-- bot.py | 74 +++++++++++++++++++++++++++++++++++------ checklist.md | 5 +-- db.py | 24 +++++++++++++ queue_manager.py | 2 ++ tests/test_db.py | 5 +++ tests/test_queue.py | 2 ++ tests/test_settings.py | 27 +++++++++++---- tests/test_workflows.py | 2 ++ workflows.py | 2 ++ 10 files changed, 128 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index e8768c2..206b5d1 100644 --- a/README.md +++ b/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`. При старте бот проверяет, что все узлы и поля существуют. diff --git a/bot.py b/bot.py index 7a9ef3e..1a6075e 100644 --- a/bot.py +++ b/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 label(name: str, value: str) -> str: + def _settings_keyboard( + 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", - reply_markup=self._quality_keyboard(row.generation_quality), + "Улучшенная — Q6\n" + "Шаги: 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: diff --git a/checklist.md b/checklist.md index ba90b28..f4fbd24 100644 --- a/checklist.md +++ b/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. diff --git a/db.py b/db.py index 5c75ed4..1cc9eee 100644 --- a/db.py +++ b/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: diff --git a/queue_manager.py b/queue_manager.py index 2a09699..bf04b92 100644 --- a/queue_manager.py +++ b/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, ) diff --git a/tests/test_db.py b/tests/test_db.py index d480e53..c82c39c 100644 --- a/tests/test_db.py +++ b/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: diff --git a/tests/test_queue.py b/tests/test_queue.py index c8102fb..ddf923c 100644 --- a/tests/test_queue.py +++ b/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 diff --git a/tests/test_settings.py b/tests/test_settings.py index 4b6bc7a..0875ab2 100644 --- a/tests/test_settings.py +++ b/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: - keyboard = MpakoBot._quality_keyboard("improved") +def test_quality_keyboard_marks_improved_and_steps() -> None: + 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 шагов", + ] diff --git a/tests/test_workflows.py b/tests/test_workflows.py index 83b40c0..338de86 100644 --- a/tests/test_workflows.py +++ b/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" diff --git a/workflows.py b/workflows.py index 5168d9b..a6156be 100644 --- a/workflows.py +++ b/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