Настройки Krea2: Q4/Q6 и 12 шагов
Добавил /settings с галочкой выбора и хранением generation_quality в users. Очередь подставляет Q4 или Q6 в unet_name. Krea2: 12 шагов dpmpp_2m/karras.
This commit is contained in:
parent
34806cf820
commit
6338a47333
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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`. При старте бот проверяет, что все узлы и поля существуют.
|
||||
|
||||
|
|
|
|||
75
bot.py
75
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:
|
||||
|
|
|
|||
12
checklist.md
12
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.
|
||||
|
|
|
|||
16
config.py
16
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
|
||||
|
||||
|
|
|
|||
24
db.py
24
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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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] == [
|
||||
"Стандартная",
|
||||
"Улучшенная ✅",
|
||||
]
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
"inputs": {
|
||||
"add_noise": "enable",
|
||||
"noise_seed": 782408632091183,
|
||||
"steps": 10,
|
||||
"steps": 12,
|
||||
"cfg": 1,
|
||||
"sampler_name": "dpmpp_2m",
|
||||
"scheduler": "karras",
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
"inputs": {
|
||||
"add_noise": "enable",
|
||||
"noise_seed": 0,
|
||||
"steps": 10,
|
||||
"steps": 12,
|
||||
"cfg": 1,
|
||||
"sampler_name": "dpmpp_2m",
|
||||
"scheduler": "karras",
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@
|
|||
"negative": "10",
|
||||
"seed": "3",
|
||||
"input_image": "20",
|
||||
"model": "19",
|
||||
"latent": "15",
|
||||
"save": "24"
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue