diff --git a/.env.example b/.env.example index 2a350e0..c58a4dc 100644 --- a/.env.example +++ b/.env.example @@ -25,6 +25,10 @@ PROMPT_ENHANCE_TIMEOUT=120 # Хранилище STORAGE_DIR=storage +# Модели Krea2 для /settings +KREA2_Q4_MODEL=Krea2_turbo_uncensored_edit_v1.1-Q4_K_M.gguf +KREA2_Q6_MODEL=Krea2_turbo_uncensored_edit-v1.1-Q6_K.gguf + # PostgreSQL # В Docker хост базы — db. При локальном запуске замени db на localhost. POSTGRES_DB=mpakoimageaibot diff --git a/README.md b/README.md index 7c46376..25d3b76 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,7 @@ Telegram-бот для генерации и редактирования изо - **Управление пользователями**: в веб-интерфейсе можно выдавать админку, менять баланс и блокировать. - **Подтверждение 18+**: без подтверждения возраста бот не принимает запросы. - **Контекст**: следующий текстовый запрос продолжает предыдущую сцену; `/restart` начинает новую. +- **Настройки Krea2**: `/settings` переключает стандартную (Q4) и улучшенную (Q6) модель. - Глобальная очередь с одной активной задачей на пользователя. - Позиция в очереди и прогресс генерации через WebSocket ComfyUI. - PostgreSQL: запросы, статусы, seed, пути к файлам. @@ -203,6 +204,7 @@ TELEGRAM_HEALTH_INTERVAL=60 - `/balance` — текущий баланс. - `/invite` — реферальная ссылка. - `/restart` — сбросить контекст и начать новую сцену (алиас `/reset`). +- `/settings` — выбор модели Krea2: стандартная Q4 или улучшенная Q6. - За первую оплату приглашённого друга начисляется +3 генерации, реферер получает уведомление. Баланс и рефералы хранятся в таблице `users`. Пропущенные дни начисляются при следующем запуске/полуночи. @@ -287,7 +289,9 @@ docker compose up -d - `workflows/img2img.json` — Krea2 edit: `LoadImage → TextEncodeQwenImageEdit → KSamplerAdvanced → VAEDecode → SaveImage`. - `workflows/mapping.json` — ID узлов, в которые подставляются prompt, negative, seed, входное изображение и размер латента. - `workflows/Krea2.json` и `workflows/ZImage.json` — исходные экспорты. -- Krea2 edit: VAE подключён к `TextEncodeQwenImageEdit`, латент подгоняется под пропорции входного изображения; 10 шагов, `dpmpp_2m`/`karras`. +- Krea2 edit: VAE подключён к `TextEncodeQwenImageEdit`, латент подгоняется под пропорции входного изображения; 12 шагов, `dpmpp_2m`/`karras`. +- `/settings` переключает модель Krea2: стандартная `Q4_K_M` или улучшенная `Q6_K`. Настройки сэмплера и латента одинаковые. +- Файлы моделей должны быть в `models/diffusion_models`; имена задаются `KREA2_Q4_MODEL` и `KREA2_Q6_MODEL`. Если меняются workflow, обновите `mapping.json`. При старте бот проверяет, что все узлы и поля существуют. diff --git a/bot.py b/bot.py index eeaec2b..231f94e 100644 --- a/bot.py +++ b/bot.py @@ -77,7 +77,8 @@ START_TEXT = ( "/balance — баланс генераций.\n" "/invite — пригласить друга.\n" "/status — состояние текущей задачи.\n" - "/restart — сброс контекста." + "/restart — сброс контекста.\n" + "/settings — качество Krea2." ) BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны." AGE_TEXT = ( @@ -138,6 +139,7 @@ class MpakoBot: application.add_handler(CommandHandler("status", self._status)) application.add_handler(CommandHandler("balance", self._balance)) application.add_handler(CommandHandler("invite", self._invite)) + application.add_handler(CommandHandler("settings", self._settings)) application.add_handler(CommandHandler("restart", self._reset)) application.add_handler(CommandHandler("reset", self._reset)) application.add_handler( @@ -149,6 +151,11 @@ class MpakoBot: application.add_handler( CallbackQueryHandler(self._age_choice, pattern=r"^age:") ) + application.add_handler( + CallbackQueryHandler( + self._quality_choice, pattern=r"^quality:" + ) + ) application.add_handler( MessageHandler(filters.SUCCESSFUL_PAYMENT, self._successful_payment) ) @@ -524,6 +531,72 @@ class MpakoBot: ) await self._retry(message.reply_text, text) + @staticmethod + def _quality_keyboard(quality: str) -> InlineKeyboardMarkup: + def label(name: str, value: str) -> str: + return f"{name} ✅" if quality == value else name + + return InlineKeyboardMarkup( + [ + [ + InlineKeyboardButton( + label("Стандартная", "standard"), + callback_data="quality:standard", + ) + ], + [ + InlineKeyboardButton( + label("Улучшенная", "improved"), + callback_data="quality:improved", + ) + ], + ] + ) + + async def _settings( + self, update: Update, context: ContextTypes.DEFAULT_TYPE + ) -> None: + if not await self._check_access(update): + return + message = update.effective_message + user = update.effective_user + if message is None or user is None: + return + row = await self._db.ensure_user(user.id, user.username) + await self._retry( + message.reply_text, + "Качество Krea2 для редактирования фото:\n" + "Стандартная — Q4\n" + "Улучшенная — Q6", + reply_markup=self._quality_keyboard(row.generation_quality), + ) + + async def _quality_choice( + self, update: Update, context: ContextTypes.DEFAULT_TYPE + ) -> None: + query = update.callback_query + if query is None or query.data is None: + return + try: + _, quality = query.data.split(":", 1) + except ValueError: + await query.answer("Некорректный выбор.", show_alert=True) + return + if quality not in {"standard", "improved"}: + await query.answer("Некорректный выбор.", show_alert=True) + return + user = query.from_user + 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) + ) + except TelegramError: + logger.debug("Не удалось обновить настройки качества") + async def _status( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: diff --git a/checklist.md b/checklist.md index 3b14041..314fae0 100644 --- a/checklist.md +++ b/checklist.md @@ -8,7 +8,7 @@ - Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15), оплата Stars (16), баланс/рефералы (17), управление пользователями (18), возрастной gate (19) и надёжность Telegram (20). - Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история. -- 69 тестов: 42 проходят локально без БД, 27 тестов БД выполняются в CI на PostgreSQL через `TEST_DATABASE_URL`. Покрыты workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, баланс/рефералы, управление пользователями, retry, миграция, watchdog, контекст, админ-панель, скрытие истории. +- 75 тестов: 47 проходят локально без БД, 28 тестов БД выполняются в CI на PostgreSQL через `TEST_DATABASE_URL`. Покрыты workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, баланс/рефералы, управление пользователями, retry, миграция, watchdog, контекст, админ-панель, скрытие истории, настройки Krea2. - Не проверено: реальная оплата Stars и реферальная выплата в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. PostgreSQL-миграция и тесты БД ещё не запускались вживую — это сделает CI и первый деплой. ## 0. Допущения и решения по умолчанию @@ -406,6 +406,16 @@ ON requests(user_id) WHERE status IN ('queued', 'running'); **DoD:** публичная история не показывает скрытые запросы, админ-панель требует пароль, история ComfyUI чистится после генерации. +## 23. Настройки Krea2 + +- [x] `users.generation_quality`: `standard` (Q4) / `improved` (Q6). +- [x] `/settings` — inline-клавиатура, у выбранного пункта галочка справа. +- [x] Нажатие меняет только галочку и сохраняет выбор, без запуска генерации. +- [x] Очередь Krea2 подставляет Q4 или Q6 в `unet_name`. +- [x] Тесты настройки и workflow. + +**DoD:** выбор сохраняется между запросами, стандартная использует Q4, улучшенная — Q6. + ## Порядок работ 1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow. diff --git a/config.py b/config.py index a0c0066..219f8c4 100644 --- a/config.py +++ b/config.py @@ -106,6 +106,12 @@ class Settings: web_bind_ip: str = "127.0.0.1" web_admin_username: str = "admin" web_admin_password: str = "" + krea2_q4_model: str = ( + "Krea2_turbo_uncensored_edit_v1.1-Q4_K_M.gguf" + ) + krea2_q6_model: str = ( + "Krea2_turbo_uncensored_edit-v1.1-Q6_K.gguf" + ) @property def comfy_ws_url(self) -> str: @@ -184,6 +190,16 @@ class Settings: ).strip() or "admin", web_admin_password=os.getenv("WEB_ADMIN_PASSWORD", "").strip(), + krea2_q4_model=os.getenv( + "KREA2_Q4_MODEL", + cls.krea2_q4_model, + ).strip() + or cls.krea2_q4_model, + krea2_q6_model=os.getenv( + "KREA2_Q6_MODEL", + cls.krea2_q6_model, + ).strip() + or cls.krea2_q6_model, ) return settings diff --git a/db.py b/db.py index 865a18a..5c75ed4 100644 --- a/db.py +++ b/db.py @@ -146,6 +146,7 @@ class User: user_id: int username: str | None credits: int + generation_quality: str last_daily_credit: str referrer_id: int | None first_paid_at: str | None @@ -165,6 +166,7 @@ class User: user_id=int(row["user_id"]), username=row["username"], credits=int(row["credits"]), + generation_quality=row["generation_quality"], last_daily_credit=_as_iso(row["last_daily_credit"]) or "", referrer_id=row["referrer_id"], first_paid_at=_as_iso(row["first_paid_at"]), @@ -306,6 +308,8 @@ CREATE TABLE IF NOT EXISTS users ( user_id BIGINT PRIMARY KEY, username TEXT, credits INTEGER NOT NULL DEFAULT 3, + generation_quality TEXT NOT NULL DEFAULT 'standard' + CHECK (generation_quality IN ('standard', 'improved')), last_daily_credit DATE NOT NULL, referrer_id BIGINT, first_paid_at TIMESTAMPTZ, @@ -376,6 +380,10 @@ MIGRATIONS = ( "ALTER TABLE users ADD COLUMN IF NOT EXISTS age_declined BOOLEAN" " NOT NULL DEFAULT FALSE" ), + ( + "ALTER TABLE users ADD COLUMN IF NOT EXISTS generation_quality TEXT" + " NOT NULL DEFAULT 'standard'" + ), ) @@ -676,6 +684,22 @@ class Database: ) return await cursor.fetchone() is not None + async def set_generation_quality( + self, user_id: int, quality: str + ) -> bool: + if quality not in {"standard", "improved"}: + raise ValueError("Недопустимое качество генерации") + async with self._connection_pool().connection() as conn: + cursor = await conn.execute( + """ + UPDATE users SET generation_quality = %s + WHERE user_id = %s + RETURNING user_id + """, + (quality, 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 3811c72..2a09699 100644 --- a/queue_manager.py +++ b/queue_manager.py @@ -373,6 +373,15 @@ class QueueManager: width, height = await asyncio.to_thread( _image_size, request.input_file_path ) + user = await self._db.get_user(request.user_id) + quality = ( + user.generation_quality if user is not None else "standard" + ) + unet_name = ( + self._settings.krea2_q6_model + if quality == "improved" + else self._settings.krea2_q4_model + ) return self._workflows.build_img2img( prompt=prompt, negative=negative, @@ -380,6 +389,7 @@ class QueueManager: image_name=image_name, width=width, height=height, + unet_name=unet_name, ) async def _refresh(self, request_id: int) -> Request: diff --git a/tests/test_db.py b/tests/test_db.py index 357b069..d480e53 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -315,6 +315,19 @@ async def test_hidden_requests_are_excluded_from_public_queries( assert loaded.hidden is True +async def test_generation_quality_setting(db: Database) -> None: + user = await db.ensure_user(1, "alice") + assert user.generation_quality == "standard" + + assert await db.set_generation_quality(1, "improved") is True + user = await db.get_user(1) + assert user is not None + assert user.generation_quality == "improved" + + with pytest.raises(ValueError): + await db.set_generation_quality(1, "unknown") + + async def test_user_context_lifecycle(db: Database) -> None: assert await db.get_user_context(1) is None diff --git a/tests/test_queue.py b/tests/test_queue.py index c46c6ad..c8102fb 100644 --- a/tests/test_queue.py +++ b/tests/test_queue.py @@ -185,6 +185,8 @@ async def test_image_to_image_uploads_input(db: Database, tmp_path: Path) -> Non input_path = tmp_path / "input.png" 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 queue.start() try: @@ -201,6 +203,13 @@ async def test_image_to_image_uploads_input(db: Database, tmp_path: Path) -> Non assert comfy.uploads[0][1] == f"{request.id}_input.png" 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]["19"]["inputs"]["unet_name"] + == settings.krea2_q6_model + ) + assert comfy.workflows[0]["15"]["inputs"]["width"] == 1024 + assert comfy.workflows[0]["15"]["inputs"]["height"] == 1024 finally: await queue.stop() diff --git a/tests/test_settings.py b/tests/test_settings.py new file mode 100644 index 0000000..10bda1a --- /dev/null +++ b/tests/test_settings.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from bot import MpakoBot + + +def test_quality_keyboard_marks_standard() -> None: + keyboard = MpakoBot._quality_keyboard("standard") + rows = keyboard.inline_keyboard + + assert [row[0].text for row in rows] == [ + "Стандартная ✅", + "Улучшенная", + ] + assert [row[0].callback_data for row in rows] == [ + "quality:standard", + "quality:improved", + ] + + +def test_quality_keyboard_marks_improved() -> None: + keyboard = MpakoBot._quality_keyboard("improved") + rows = keyboard.inline_keyboard + + assert [row[0].text for row in rows] == [ + "Стандартная", + "Улучшенная ✅", + ] diff --git a/tests/test_workflows.py b/tests/test_workflows.py index b4163df..83b40c0 100644 --- a/tests/test_workflows.py +++ b/tests/test_workflows.py @@ -42,7 +42,7 @@ def test_build_img2img() -> None: assert workflow["25"]["inputs"]["vae"] == ["19", 2] assert workflow["10"]["inputs"]["text"] == "bad" assert workflow["3"]["inputs"]["noise_seed"] == 7 - assert workflow["3"]["inputs"]["steps"] == 10 + assert workflow["3"]["inputs"]["steps"] == 12 assert workflow["3"]["inputs"]["sampler_name"] == "dpmpp_2m" assert workflow["3"]["inputs"]["scheduler"] == "karras" assert workflow["20"]["class_type"] == "LoadImage" @@ -60,10 +60,12 @@ def test_build_img2img_custom_size() -> None: image_name="input.png", width=832, height=1216, + unet_name="Krea2-Q6.gguf", ) assert workflow["15"]["inputs"]["width"] == 832 assert workflow["15"]["inputs"]["height"] == 1216 + assert workflow["19"]["inputs"]["unet_name"] == "Krea2-Q6.gguf" def test_fit_latent_size_preserves_aspect_ratio() -> None: diff --git a/workflows.py b/workflows.py index 562f398..5168d9b 100644 --- a/workflows.py +++ b/workflows.py @@ -22,6 +22,7 @@ ROLE_FIELDS: dict[str, dict[str, RoleField]] = { "negative": ("text", lambda value: isinstance(value, str)), "seed": ("noise_seed", lambda value: isinstance(value, int)), "input_image": ("image", lambda value: isinstance(value, str)), + "model": ("unet_name", lambda value: isinstance(value, str)), "latent": (None, None), "save": (None, None), }, @@ -138,6 +139,7 @@ class WorkflowBuilder: image_name: str, width: int = 1024, height: int = 1024, + unet_name: str | None = None, ) -> dict[str, Any]: workflow = copy.deepcopy(self._templates["img2img"]) mapping = self._mapping["img2img"] @@ -148,4 +150,6 @@ class WorkflowBuilder: latent = workflow[mapping["latent"]]["inputs"] latent["width"] = width latent["height"] = height + if unet_name is not None: + workflow[mapping["model"]]["inputs"]["unet_name"] = unet_name return workflow diff --git a/workflows/Krea2.json b/workflows/Krea2.json index 3d178ed..a4bca55 100644 --- a/workflows/Krea2.json +++ b/workflows/Krea2.json @@ -3,7 +3,7 @@ "inputs": { "add_noise": "enable", "noise_seed": 782408632091183, - "steps": 10, + "steps": 12, "cfg": 1, "sampler_name": "dpmpp_2m", "scheduler": "karras", diff --git a/workflows/img2img.json b/workflows/img2img.json index 1ab4a5f..805f2f4 100644 --- a/workflows/img2img.json +++ b/workflows/img2img.json @@ -3,7 +3,7 @@ "inputs": { "add_noise": "enable", "noise_seed": 0, - "steps": 10, + "steps": 12, "cfg": 1, "sampler_name": "dpmpp_2m", "scheduler": "karras", diff --git a/workflows/mapping.json b/workflows/mapping.json index 5528f71..89a3008 100644 --- a/workflows/mapping.json +++ b/workflows/mapping.json @@ -11,6 +11,7 @@ "negative": "10", "seed": "3", "input_image": "20", + "model": "19", "latent": "15", "save": "24" }