Compare commits

..

2 Commits

Author SHA1 Message Date
MpakobecEx 2161bbe015 Оплата генерации
Build and push Docker image / docker (push) Failing after 8m38s Details
2026-10-18 20:09:46 +05:00
MpakobecEx 7fcca6e39f Фильтрация промптов 2026-10-18 20:04:30 +05:00
12 changed files with 1054 additions and 13 deletions

View File

@ -34,6 +34,16 @@ DEFAULT_NEGATIVE=low quality, blurry, bad anatomy, extra digits
# Доступ: пусто — всем; иначе список id через запятую
ALLOWED_USER_IDS=
# Админы: их промпты не фильтруются
ADMIN_USER_IDS=
MODERATION_ENABLED=true
# Дополнительные имена публичных людей через запятую
BLOCKED_PUBLIC_FIGURES=
# Оплата Telegram Stars
PAYMENTS_ENABLED=true
STARS_PRICE=5
# Веб-история: только локальная сеть, без авторизации
WEB_BIND_IP=192.168.31.158
WEB_PAGE_SIZE=50

View File

@ -9,6 +9,8 @@ Telegram-бот для генерации и редактирования изо
- **Reply-to-image**: ответ текстом на готовую картинку бота → редактирование этой картинки.
- **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений.
- **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI.
- **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов.
- **Оплата**: 5 Telegram Stars за генерацию, админам бесплатно, возврат при ошибке.
- Глобальная очередь с одной активной задачей на пользователя.
- Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
- SQLite: запросы, статусы, seed, пути к файлам.
@ -94,6 +96,48 @@ Workflow: `.gitea/workflows/docker-build.yml`. Runner использует label
Отдельную модель скачивать не нужно, используется существующий CLIP-файл ComfyUI.
## Фильтрация промптов
Фильтр работает для всех, кроме пользователей из `ADMIN_USER_IDS`:
```dotenv
MODERATION_ENABLED=true
ADMIN_USER_IDS=7467236616
BLOCKED_PUBLIC_FIGURES=
```
`BLOCKED_PUBLIC_FIGURES` добавляет имена к встроенному списку (через запятую).
Что блокируется:
- CSAM: child/kid/minor/teen/young/loli/shota/school/underage + NSFW;
- зоофилия: sex/nude + animal/dog/cat/horse/beast;
- публичные люди + nude/naked/sex/undress/deepfake;
- undress/nude/naked + photo/real/person/woman/man;
- насилие + NSFW;
- экстремизм + инструкции/символика;
- обходные написания (`n_u_d_e`, `l0li`, `ch1ld` и т.п.).
Для i2i OpenCV проверяет наличие лица/человека на входном фото: NSFW-промпт в этом случае блокируется.
При блокировке пользователь получает «Запрос нарушает правила. Звёзды не списаны.», а запись с `user_id`, временем, промптом и причиной сохраняется в таблицу `blocked_requests`.
## Оплата
Генерация стоит 5 Telegram Stars:
```dotenv
PAYMENTS_ENABLED=true
STARS_PRICE=5
```
- Админы из `ADMIN_USER_IDS` генерируют бесплатно.
- Не-админу выставляется счёт через Telegram Stars (`XTR`).
- Фильтрация промпта выполняется до выставления счёта.
- Если поставить задачу в очередь не удалось, звёзды возвращаются.
- При ошибке генерации звёзды также возвращаются.
- Оплата и `telegram_payment_charge_id` сохраняются в `pending_payments` и `requests`.
## Веб-история
После `docker compose up -d` доступна страница:

230
bot.py
View File

@ -6,7 +6,7 @@ import time
import uuid
from pathlib import Path
from telegram import InputFile, Update
from telegram import InputFile, LabeledPrice, Update
from telegram.constants import ChatAction, ChatType
from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut
from telegram.ext import (
@ -14,12 +14,14 @@ from telegram.ext import (
CommandHandler,
ContextTypes,
MessageHandler,
PreCheckoutQueryHandler,
filters,
)
from comfy_client import ComfyClient
from config import Settings, setup_logging
from db import Database, Request, UserAlreadyQueued
from moderation import PromptBlocked, detect_person, is_prompt_allowed
from prompt_enhancer import PromptEnhancer
from queue_manager import QueueEvent, QueueFull, QueueManager
from workflows import WorkflowBuilder
@ -32,6 +34,7 @@ START_TEXT = (
"Или ответь текстом на готовую картинку — возьму её как исходник.\n\n"
"/status — узнать состояние текущей задачи."
)
BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны."
class MpakoBot:
@ -65,6 +68,12 @@ class MpakoBot:
application.add_handler(CommandHandler("start", self._start))
application.add_handler(CommandHandler("help", self._help))
application.add_handler(CommandHandler("status", self._status))
application.add_handler(
PreCheckoutQueryHandler(self._pre_checkout)
)
application.add_handler(
MessageHandler(filters.SUCCESSFUL_PAYMENT, self._successful_payment)
)
application.add_handler(MessageHandler(filters.PHOTO, self._photo))
application.add_handler(
MessageHandler(filters.TEXT & ~filters.COMMAND, self._text)
@ -109,6 +118,15 @@ class MpakoBot:
await self._db.close()
logger.info("Бот остановлен")
def _start_text(self) -> str:
text = START_TEXT
if self._settings.payments_enabled:
text += (
f"\n\nСтоимость генерации: {self._settings.stars_price} звёзд. "
"Админам бесплатно."
)
return text
async def _start(
self, update: Update, context: ContextTypes.DEFAULT_TYPE
) -> None:
@ -116,7 +134,7 @@ class MpakoBot:
return
message = update.effective_message
if message is not None:
await message.reply_text(START_TEXT)
await message.reply_text(self._start_text())
async def _help(
self, update: Update, context: ContextTypes.DEFAULT_TYPE
@ -125,7 +143,7 @@ class MpakoBot:
return
message = update.effective_message
if message is not None:
await message.reply_text(START_TEXT)
await message.reply_text(self._start_text())
async def _status(
self, update: Update, context: ContextTypes.DEFAULT_TYPE
@ -148,6 +166,76 @@ class MpakoBot:
f"Задача в очереди. Перед тобой: {position} {self._people(position)}."
)
async def _pre_checkout(
self, update: Update, context: ContextTypes.DEFAULT_TYPE
) -> None:
query = update.pre_checkout_query
if query is None:
return
pending = await self._db.get_pending_payment(query.invoice_payload)
allowed = (
pending is not None
and pending.user_id == query.from_user.id
and pending.status == "pending"
and query.currency == "XTR"
and query.total_amount == pending.amount
)
try:
await query.answer(
ok=allowed,
error_message=None
if allowed
else "Счёт устарел. Отправь запрос заново.",
)
except TelegramError:
logger.exception("Не удалось ответить на pre_checkout_query")
async def _successful_payment(
self, update: Update, context: ContextTypes.DEFAULT_TYPE
) -> None:
message = update.effective_message
user = update.effective_user
payment = message.successful_payment if message is not None else None
if message is None or user is None or payment is None:
return
payload = payment.invoice_payload
pending = await self._db.get_pending_payment(payload)
if pending is None or pending.user_id != user.id:
logger.error(
"Платёж %s не найден или принадлежит другому пользователю",
payload,
)
return
if pending.status == "paid":
await message.reply_text("Этот платёж уже обработан.")
return
if pending.status != "pending":
return
charge_id = payment.telegram_payment_charge_id
if payment.currency != "XTR" or payment.total_amount != pending.amount:
refunded = await self._refund_stars(user.id, charge_id)
if refunded:
await self._db.mark_pending_refunded(payload)
note = " Звёзды возвращены." if refunded else ""
await message.reply_text("Платёж не принят." + note)
return
if not await self._db.mark_pending_paid(payload, charge_id):
await message.reply_text("Этот платёж уже обработан.")
return
await self._submit(
update,
mode=pending.mode,
prompt=pending.prompt,
input_path=pending.input_file_path,
skip_payment=True,
payment_charge_id=charge_id,
payment_payload=payload,
)
async def _text(
self, update: Update, context: ContextTypes.DEFAULT_TYPE
) -> None:
@ -288,12 +376,78 @@ class MpakoBot:
mode: str,
prompt: str,
input_path: str | None,
skip_payment: bool = False,
payment_charge_id: str | None = None,
payment_payload: str | None = None,
) -> bool:
message = update.effective_message
user = update.effective_user
if message is None or user is None or self._queue is None:
return False
is_admin = user.id in self._settings.admin_user_ids
has_person = False
if (
mode == "i2i"
and input_path
and not is_admin
and self._settings.moderation_enabled
):
has_person = await asyncio.to_thread(detect_person, input_path)
if (
not skip_payment
and not is_admin
and self._settings.payments_enabled
):
allowed, reason = is_prompt_allowed(
prompt,
mode=mode,
has_person=has_person,
extra_public_figures=self._settings.blocked_public_figures,
)
if not allowed:
await self._db.log_blocked(
user_id=user.id,
prompt=prompt,
reason=reason or "unknown",
mode=mode,
has_person=has_person,
)
await message.reply_text(BLOCKED_TEXT)
return False
payload = uuid.uuid4().hex
amount = self._settings.stars_price
await self._db.create_pending_payment(
payload=payload,
user_id=user.id,
mode=mode,
prompt=prompt,
username=user.username,
input_file_path=input_path,
amount=amount,
)
try:
await self._retry(
self._bot.send_invoice,
chat_id=user.id,
title="Генерация изображения",
description="Одна генерация или редактирование изображения.",
payload=payload,
provider_token="",
currency="XTR",
prices=[LabeledPrice(label="Генерация", amount=amount)],
)
except TelegramError:
logger.exception("Не удалось выставить счёт пользователю %s", user.id)
await self._db.mark_pending_cancelled(payload)
await message.reply_text(
"Не удалось выставить счёт. Попробуй позже."
)
return False
return True
try:
request, position = await self._queue.enqueue(
user_id=user.id,
@ -301,8 +455,20 @@ class MpakoBot:
mode=mode,
prompt=prompt,
input_file_path=input_path,
has_person=has_person,
is_admin=is_admin,
payment_charge_id=payment_charge_id,
)
except PromptBlocked:
note = await self._refund_if_paid(
user.id, payment_charge_id, payment_payload
)
await message.reply_text(BLOCKED_TEXT + note)
return False
except UserAlreadyQueued:
note = await self._refund_if_paid(
user.id, payment_charge_id, payment_payload
)
active = await self._db.get_active_by_user(user.id)
text = "Ты уже в очереди."
if active is not None:
@ -311,16 +477,27 @@ class MpakoBot:
else:
ahead = await self._queue.position(active.id)
text += f" Перед тобой: {ahead} {self._people(ahead)}."
await message.reply_text(text)
await message.reply_text(text + note)
return False
except QueueFull as exc:
await message.reply_text(exc.user_message)
note = await self._refund_if_paid(
user.id, payment_charge_id, payment_payload
)
await message.reply_text(exc.user_message + note)
return False
except Exception:
logger.exception("Не удалось поставить задачу пользователю %s", user.id)
await message.reply_text("Не удалось поставить задачу. Попробуй позже.")
note = await self._refund_if_paid(
user.id, payment_charge_id, payment_payload
)
await message.reply_text(
"Не удалось поставить задачу. Попробуй позже." + note
)
return False
if payment_payload:
await self._db.attach_pending_request(payment_payload, request.id)
try:
status_message = await self._retry(
message.reply_text, self._ack_text(position)
@ -350,6 +527,35 @@ class MpakoBot:
)
return True
async def _refund_if_paid(
self,
user_id: int,
charge_id: str | None,
payload: str | None,
) -> str:
if not charge_id:
return ""
refunded = await self._refund_stars(user_id, charge_id)
if refunded and payload:
await self._db.mark_pending_refunded(payload)
if refunded:
return " Звёзды возвращены."
return " Не удалось вернуть звёзды, напиши админу."
async def _refund_stars(self, user_id: int, charge_id: str) -> bool:
if self._bot is None:
return False
try:
return await self._retry(
self._bot.refund_star_payment,
user_id=user_id,
telegram_payment_charge_id=charge_id,
attempts=2,
)
except TelegramError:
logger.exception("Не удалось вернуть звёзды пользователю %s", user_id)
return False
def _ack_text(self, position: int) -> str:
if position <= 0:
return "Принято. Ты первый, начинаю…"
@ -496,7 +702,17 @@ class MpakoBot:
async def _finish_failed(self, request: Request, error: str) -> None:
target = self._status_messages.get(request.id)
chat_id = target[0] if target is not None else request.user_id
text = f"Не получилось: {error}"
refund_note = ""
if request.payment_charge_id and not request.payment_refunded:
refunded = await self._refund_stars(
request.user_id, request.payment_charge_id
)
if refunded:
await self._db.set_refunded(request.id)
refund_note = " Звёзды возвращены."
else:
refund_note = " Не удалось вернуть звёзды, напиши админу."
text = f"Не получилось: {error}{refund_note}"
if target is not None and self._bot is not None:
try:
await self._retry(

View File

@ -6,10 +6,10 @@
## Текущий статус
- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13) и улучшение промптов (14).
- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, переключение T2I → I2I.
- 25 тестов проходят: БД, workflows, очередь, веб-история, enhancer, вспомогательные функции ComfyUI.
- Не проверено: сборка Docker (демон не запущен), очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, реальная отправка в Telegram (нет токена).
- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15) и оплата Stars (16).
- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история.
- 46 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, вспомогательные функции ComfyUI.
- Не проверено: реальная оплата Stars в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото.
## 0. Допущения и решения по умолчанию
@ -318,6 +318,28 @@ ON requests(user_id) WHERE status IN ('queued', 'running');
**DoD:** русский запрос превращается в английский промпт, изображение генерируется по нему; при сбое LLM используется оригинал.
## 15. Фильтрация промптов
- [x] `is_prompt_allowed`: нормализация, hard/context-блоки, обходные паттерны.
- [x] Админы из `ADMIN_USER_IDS` обходят фильтр.
- [x] Логирование блокировок в `blocked_requests`.
- [x] Ответ «Запрос нарушает правила. Звёзды не списаны.».
- [x] i2i: детекция лица/человека через OpenCV.
- [x] Тесты правил и очереди.
**DoD:** небезопасный промпт не попадает в очередь, причина пишется в БД, админ может отправлять без фильтра.
## 16. Оплата Stars
- [x] `PAYMENTS_ENABLED` / `STARS_PRICE`, админы бесплатно.
- [x] Invoice `XTR` + `pre_checkout_query` + `successful_payment`.
- [x] `pending_payments`: payload, сумма, статус, charge_id, request_id.
- [x] Возврат звёзд при ошибке постановки в очередь и генерации.
- [x] `payment_charge_id` и `payment_refunded` в `requests`.
- [x] Тесты жизненного цикла платежа.
**DoD:** не-админ получает счёт на 5 звёзд, после оплаты задача ставится в очередь; при сбое звёзды возвращаются.
## Порядок работ
1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow.

View File

@ -59,6 +59,16 @@ def _user_ids(name: str) -> frozenset[int]:
return frozenset(values)
def _csv(name: str) -> frozenset[str]:
raw = os.getenv(name, "")
values = {
part.strip().lower()
for part in raw.replace(";", ",").split(",")
if part.strip()
}
return frozenset(values)
@dataclass(frozen=True)
class Settings:
bot_token: str
@ -75,12 +85,17 @@ class Settings:
worker_concurrency: int = 1
max_queue_size: int = 20
allowed_user_ids: frozenset[int] = field(default_factory=frozenset)
admin_user_ids: frozenset[int] = field(default_factory=frozenset)
progress_interval: float = 3.0
generation_timeout: float = 600.0
log_level: str = "INFO"
default_negative: str = "low quality, blurry, bad anatomy, extra digits"
max_prompt_length: int = 1000
max_photo_bytes: int = 15 * 1024 * 1024
moderation_enabled: bool = True
blocked_public_figures: frozenset[str] = field(default_factory=frozenset)
payments_enabled: bool = True
stars_price: int = 5
web_page_size: int = 50
web_bind_ip: str = "127.0.0.1"
@ -129,6 +144,7 @@ class Settings:
worker_concurrency=max(1, _int("WORKER_CONCURRENCY", 1)),
max_queue_size=max(1, _int("MAX_QUEUE_SIZE", 20)),
allowed_user_ids=_user_ids("ALLOWED_USER_IDS"),
admin_user_ids=_user_ids("ADMIN_USER_IDS"),
progress_interval=max(1.0, _float("PROGRESS_INTERVAL", 3.0)),
generation_timeout=max(30.0, _float("GENERATION_TIMEOUT", 600.0)),
log_level=os.getenv("LOG_LEVEL", "INFO").strip().upper() or "INFO",
@ -136,6 +152,10 @@ class Settings:
or cls.default_negative,
max_prompt_length=max(1, _int("MAX_PROMPT_LENGTH", 1000)),
max_photo_bytes=max(1, _int("MAX_PHOTO_BYTES", 15 * 1024 * 1024)),
moderation_enabled=_bool("MODERATION_ENABLED", True),
blocked_public_figures=_csv("BLOCKED_PUBLIC_FIGURES"),
payments_enabled=_bool("PAYMENTS_ENABLED", True),
stars_price=max(1, _int("STARS_PRICE", 5)),
web_page_size=max(10, min(200, _int("WEB_PAGE_SIZE", 50))),
web_bind_ip=os.getenv("WEB_BIND_IP", "127.0.0.1").strip()
or "127.0.0.1",

219
db.py
View File

@ -37,6 +37,8 @@ class Request:
output_file_path: str | None
status: str
comfy_prompt_id: str | None
payment_charge_id: str | None
payment_refunded: bool
mode: str
seed: int | None
error_text: str | None
@ -59,6 +61,8 @@ class Request:
output_file_path=row["output_file_path"],
status=row["status"],
comfy_prompt_id=row["comfy_prompt_id"],
payment_charge_id=row["payment_charge_id"],
payment_refunded=bool(row["payment_refunded"]),
mode=row["mode"],
seed=row["seed"],
error_text=row["error_text"],
@ -78,6 +82,60 @@ class UserStat:
return f"@{self.username}" if self.username else str(self.user_id)
@dataclass(frozen=True)
class BlockedRequest:
id: int
user_id: int
timestamp: str
prompt: str
reason: str
mode: str | None
has_person: bool
@classmethod
def from_row(cls, row: aiosqlite.Row) -> BlockedRequest:
return cls(
id=row["id"],
user_id=row["user_id"],
timestamp=row["timestamp"],
prompt=row["prompt"],
reason=row["reason"],
mode=row["mode"],
has_person=bool(row["has_person"]),
)
@dataclass(frozen=True)
class PendingPayment:
payload: str
user_id: int
timestamp: str
mode: str
prompt: str
username: str | None
input_file_path: str | None
amount: int
status: str
telegram_payment_charge_id: str | None
request_id: int | None
@classmethod
def from_row(cls, row: aiosqlite.Row) -> PendingPayment:
return cls(
payload=row["payload"],
user_id=row["user_id"],
timestamp=row["timestamp"],
mode=row["mode"],
prompt=row["prompt"],
username=row["username"],
input_file_path=row["input_file_path"],
amount=int(row["amount"]),
status=row["status"],
telegram_payment_charge_id=row["telegram_payment_charge_id"],
request_id=row["request_id"],
)
SCHEMA = """
CREATE TABLE IF NOT EXISTS requests (
id INTEGER PRIMARY KEY AUTOINCREMENT,
@ -91,6 +149,8 @@ CREATE TABLE IF NOT EXISTS requests (
status TEXT NOT NULL DEFAULT 'queued'
CHECK (status IN ('queued', 'running', 'done', 'failed')),
comfy_prompt_id TEXT,
payment_charge_id TEXT,
payment_refunded INTEGER NOT NULL DEFAULT 0,
mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')),
seed INTEGER,
error_text TEXT,
@ -105,6 +165,37 @@ CREATE INDEX IF NOT EXISTS idx_requests_status_id
CREATE UNIQUE INDEX IF NOT EXISTS uq_requests_one_active
ON requests(user_id) WHERE status IN ('queued', 'running');
CREATE TABLE IF NOT EXISTS blocked_requests (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER NOT NULL,
timestamp TEXT NOT NULL,
prompt TEXT NOT NULL,
reason TEXT NOT NULL,
mode TEXT,
has_person INTEGER NOT NULL DEFAULT 0
);
CREATE INDEX IF NOT EXISTS idx_blocked_user_time
ON blocked_requests(user_id, id);
CREATE TABLE IF NOT EXISTS pending_payments (
payload TEXT PRIMARY KEY,
user_id INTEGER NOT NULL,
timestamp TEXT NOT NULL,
mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')),
prompt TEXT NOT NULL,
username TEXT,
input_file_path TEXT,
amount INTEGER NOT NULL,
status TEXT NOT NULL DEFAULT 'pending'
CHECK (status IN ('pending', 'paid', 'cancelled', 'refunded')),
telegram_payment_charge_id TEXT,
request_id INTEGER
);
CREATE INDEX IF NOT EXISTS idx_pending_user_time
ON pending_payments(user_id, timestamp);
"""
@ -137,6 +228,12 @@ class Database:
await conn.execute("ALTER TABLE requests ADD COLUMN username TEXT")
if "enhanced_prompt" not in columns:
await conn.execute("ALTER TABLE requests ADD COLUMN enhanced_prompt TEXT")
if "payment_charge_id" not in columns:
await conn.execute("ALTER TABLE requests ADD COLUMN payment_charge_id TEXT")
if "payment_refunded" not in columns:
await conn.execute(
"ALTER TABLE requests ADD COLUMN payment_refunded INTEGER NOT NULL DEFAULT 0"
)
async def close(self) -> None:
if self._conn is not None:
@ -157,6 +254,7 @@ class Database:
prompt_text: str,
input_file_path: str | None,
seed: int,
payment_charge_id: str | None = None,
) -> Request:
conn = self._connection()
async with self._write_lock:
@ -165,8 +263,8 @@ class Database:
"""
INSERT INTO requests
(user_id, username, timestamp, prompt_text,
input_file_path, status, mode, seed)
VALUES (?, ?, ?, ?, ?, 'queued', ?, ?)
input_file_path, status, mode, seed, payment_charge_id)
VALUES (?, ?, ?, ?, ?, 'queued', ?, ?, ?)
""",
(
user_id,
@ -176,6 +274,7 @@ class Database:
input_file_path,
mode,
seed,
payment_charge_id,
),
)
request_id = cursor.lastrowid
@ -353,6 +452,92 @@ class Database:
(prompt, request_id),
)
async def set_refunded(self, request_id: int) -> None:
await self._update(
"UPDATE requests SET payment_refunded = 1 WHERE id = ?",
(request_id,),
)
async def create_pending_payment(
self,
*,
payload: str,
user_id: int,
mode: str,
prompt: str,
username: str | None,
input_file_path: str | None,
amount: int,
) -> PendingPayment:
conn = self._connection()
async with self._write_lock:
await conn.execute(
"""
INSERT INTO pending_payments
(payload, user_id, timestamp, mode, prompt, username,
input_file_path, amount, status)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending')
""",
(
payload,
user_id,
utc_now(),
mode,
prompt,
username,
input_file_path,
amount,
),
)
await conn.commit()
pending = await self.get_pending_payment(payload)
if pending is None:
raise RuntimeError("Не удалось прочитать счёт")
return pending
async def get_pending_payment(self, payload: str) -> PendingPayment | None:
conn = self._connection()
async with conn.execute(
"SELECT * FROM pending_payments WHERE payload = ?", (payload,)
) as cursor:
row = await cursor.fetchone()
return PendingPayment.from_row(row) if row is not None else None
async def mark_pending_paid(self, payload: str, charge_id: str) -> bool:
conn = self._connection()
async with self._write_lock:
cursor = await conn.execute(
"""
UPDATE pending_payments
SET status = 'paid', telegram_payment_charge_id = ?
WHERE payload = ? AND status = 'pending'
""",
(charge_id, payload),
)
await conn.commit()
return cursor.rowcount > 0
async def mark_pending_refunded(self, payload: str) -> None:
await self._update(
"UPDATE pending_payments SET status = 'refunded' WHERE payload = ?",
(payload,),
)
async def mark_pending_cancelled(self, payload: str) -> None:
await self._update(
"""
UPDATE pending_payments SET status = 'cancelled'
WHERE payload = ? AND status = 'pending'
""",
(payload,),
)
async def attach_pending_request(self, payload: str, request_id: int) -> None:
await self._update(
"UPDATE pending_payments SET request_id = ? WHERE payload = ?",
(request_id, payload),
)
async def set_done(self, request_id: int, output_file_path: str) -> None:
await self._update(
"""
@ -378,6 +563,36 @@ class Database:
(error_text[:1000], utc_now(), request_id),
)
async def log_blocked(
self,
*,
user_id: int,
prompt: str,
reason: str,
mode: str | None = None,
has_person: bool = False,
) -> int:
conn = self._connection()
async with self._write_lock:
cursor = await conn.execute(
"""
INSERT INTO blocked_requests
(user_id, timestamp, prompt, reason, mode, has_person)
VALUES (?, ?, ?, ?, ?, ?)
""",
(user_id, utc_now(), prompt, reason, mode, int(has_person)),
)
await conn.commit()
return int(cursor.lastrowid)
async def list_blocked(self, limit: int = 100) -> list[BlockedRequest]:
conn = self._connection()
async with conn.execute(
"SELECT * FROM blocked_requests ORDER BY id DESC LIMIT ?", (limit,)
) as cursor:
rows = await cursor.fetchall()
return [BlockedRequest.from_row(row) for row in rows]
async def recover(self, reason: str = "прервано перезапуском") -> list[Request]:
"""Помечает running как failed и возвращает queued для повторной постановки."""
conn = self._connection()

344
moderation.py Normal file
View File

@ -0,0 +1,344 @@
from __future__ import annotations
import logging
import re
from pathlib import Path
logger = logging.getLogger(__name__)
DIGIT_MAP = str.maketrans(
{
"0": "o",
"1": "i",
"3": "e",
"4": "a",
"5": "s",
"7": "t",
}
)
CSAM_TERMS = {
"child",
"kid",
"minor",
"teen",
"young",
"loli",
"shota",
"school",
"underage",
"ребен",
"ребён",
"дети",
"детск",
"школь",
"малолет",
"подрост",
}
NSFW_TERMS = {
"nude",
"naked",
"sex",
"erotic",
"explicit",
"porn",
"nsfw",
"нюд",
"голый",
"голая",
"голое",
"голые",
"голую",
"голой",
"голых",
"секс",
"эрот",
"порн",
"обнаж",
"раздет",
}
ANIMAL_TERMS = {
"animal",
"dog",
"cat",
"horse",
"beast",
"животн",
"собак",
"кошк",
"лошад",
"звер",
}
PUBLIC_FIGURES = {
"trump",
"biden",
"putin",
"zelensky",
"obama",
"clinton",
"harris",
"macron",
"merkel",
"modi",
"netanyahu",
"kimjongun",
"elonmusk",
"taylorswift",
"beyonce",
"kanyewest",
"ronaldo",
"messi",
"madonna",
"bieber",
"kardashian",
"spears",
"michaeljackson",
"queenelizabeth",
"kingcharles",
"princewilliam",
"katemiddleton",
"трамп",
"байден",
"путин",
"зеленский",
"обама",
"клинтон",
"маск",
}
REAL_PERSON_TERMS = {
"nude",
"naked",
"sex",
"undress",
"deepfake",
"нюд",
"голый",
"голая",
"голое",
"голые",
"голую",
"голой",
"голых",
"секс",
"раздет",
"дипфейк",
}
NONCONSENSUAL_TERMS = {
"undress",
"nude",
"naked",
"раздет",
"нюд",
"голый",
"голая",
"голое",
"голые",
"голую",
"голой",
"голых",
}
PERSON_TERMS = {
"photo",
"real",
"person",
"woman",
"man",
"фото",
"реальн",
"человек",
"женщин",
"мужчин",
"девуш",
"парн",
}
VIOLENCE_TERMS = {
"gore",
"snuff",
"torture",
"dismember",
"blood",
"кров",
"пытк",
"расчлен",
}
EXTREMISM_TERMS = {
"terrorist",
"nazi",
"isis",
"jihad",
"bomb",
"террорист",
"нацист",
"игил",
"джихад",
"бомб",
}
INSTRUCTION_TERMS = {
"how",
"instruction",
"instructions",
"guide",
"tutorial",
"symbol",
"symbols",
"flag",
"uniform",
"manifesto",
"recruit",
"attack",
"make",
"build",
"инструкц",
"символ",
"флаг",
"как",
"сделать",
"создать",
"собрать",
}
OBFUSCATED_PATTERN = re.compile(
r"l[0o1]l[i1]|sh[0o]t[a4]|ch[i1]ld|n[uv]de|undr[e3]ss",
re.IGNORECASE,
)
class PromptBlocked(RuntimeError):
def __init__(self, reason: str) -> None:
super().__init__(f"Промпт заблокирован: {reason}")
self.reason = reason
def _normalize(prompt: str) -> tuple[str, str, set[str]]:
lowered = prompt.lower()
normalized = lowered.translate(DIGIT_MAP)
compact = re.sub(r"[\s_\-]+", "", normalized)
tokens = re.split(r"[\s_\-]+", normalized)
merged: list[str] = []
buffer = ""
for token in tokens:
if not token:
continue
if len(token) == 1:
buffer += token
continue
if buffer:
merged.append(buffer)
buffer = ""
merged.append(token)
if buffer:
merged.append(buffer)
return lowered, compact, set(merged)
def _has(term: str, compact: str, words: set[str]) -> bool:
return term in compact or term in words
def is_prompt_allowed(
prompt: str,
*,
mode: str = "t2i",
has_person: bool = False,
is_admin: bool = False,
extra_public_figures: set[str] | frozenset[str] | None = None,
) -> tuple[bool, str | None]:
"""Проверка промпта. Возвращает (allowed, reason_code)."""
if is_admin:
return True, None
lowered, compact, words = _normalize(prompt)
figures = PUBLIC_FIGURES
if extra_public_figures:
extra = {
re.sub(r"[\s_\-]+", "", name.lower().translate(DIGIT_MAP))
for name in extra_public_figures
if name.strip()
}
figures = PUBLIC_FIGURES | extra
nsfw = any(_has(term, compact, words) for term in NSFW_TERMS)
csam = any(_has(term, compact, words) for term in CSAM_TERMS)
if csam and nsfw:
return False, "csam"
if _has("loli", compact, words) or _has("shota", compact, words):
return False, "csam"
if (
any(_has(term, compact, words) for term in ("sex", "nude", "naked"))
and any(_has(term, compact, words) for term in ANIMAL_TERMS)
):
return False, "bestiality"
if any(name in compact for name in figures) and any(
_has(term, compact, words) for term in REAL_PERSON_TERMS
):
return False, "real_person"
if any(_has(term, compact, words) for term in NONCONSENSUAL_TERMS) and any(
_has(term, compact, words) for term in PERSON_TERMS
):
return False, "nonconsensual_sexualization"
if mode == "i2i" and has_person and any(
_has(term, compact, words) for term in ("nude", "naked", "undress", "sex")
):
return False, "i2i_person_nsfw"
if any(_has(term, compact, words) for term in VIOLENCE_TERMS) and nsfw:
return False, "violence_sexual"
if any(_has(term, compact, words) for term in EXTREMISM_TERMS) and any(
_has(term, compact, words) for term in INSTRUCTION_TERMS
):
return False, "extremism"
if re.search(r"[0-9_\-]", lowered) and OBFUSCATED_PATTERN.search(compact):
return False, "obfuscated_nsfw"
return True, None
def detect_person(path: str | Path) -> bool:
"""Ищет лицо или человека через OpenCV. При ошибке считает, что человека нет."""
try:
import cv2 # type: ignore[import-not-found]
except ImportError:
logger.warning("OpenCV не установлен, детекция людей отключена")
return False
try:
image = cv2.imread(str(path))
if image is None:
return False
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
cascade = cv2.CascadeClassifier(
str(Path(cv2.data.haarcascades) / "haarcascade_frontalface_default.xml")
)
faces = cascade.detectMultiScale(
gray,
scaleFactor=1.1,
minNeighbors=5,
minSize=(24, 24),
)
if len(faces) > 0:
return True
height, width = image.shape[:2]
scale = min(1.0, 640.0 / max(height, width))
if scale < 1.0:
image = cv2.resize(image, (int(width * scale), int(height * scale)))
hog = cv2.HOGDescriptor()
hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())
_, weights = hog.detectMultiScale(
image,
winStride=(8, 8),
padding=(8, 8),
scale=1.05,
)
return any(float(weight) > 0.5 for weight in weights)
except Exception:
logger.warning("Ошибка детекции человека в %s", path, exc_info=True)
return False

View File

@ -13,6 +13,7 @@ from pathlib import Path
from comfy_client import ComfyClient, ComfyError
from config import Settings
from db import Database, Request, UserAlreadyQueued
from moderation import PromptBlocked, is_prompt_allowed
from prompt_enhancer import PromptEnhancer
from workflows import WorkflowBuilder
@ -104,10 +105,29 @@ class QueueManager:
prompt: str,
input_file_path: str | None = None,
username: str | None = None,
has_person: bool = False,
is_admin: bool = False,
payment_charge_id: str | None = None,
) -> tuple[Request, int]:
async with self._lock:
if self._stopping:
raise QueueStopped()
if self._settings.moderation_enabled and not is_admin:
allowed, reason = is_prompt_allowed(
prompt,
mode=mode,
has_person=has_person,
extra_public_figures=self._settings.blocked_public_figures,
)
if not allowed:
await self._db.log_blocked(
user_id=user_id,
prompt=prompt,
reason=reason or "unknown",
mode=mode,
has_person=has_person,
)
raise PromptBlocked(reason or "unknown")
active = await self._db.get_active_by_user(user_id)
if active is not None:
raise UserAlreadyQueued(user_id)
@ -122,6 +142,7 @@ class QueueManager:
prompt_text=prompt,
input_file_path=input_file_path,
seed=seed,
payment_charge_id=payment_charge_id,
)
self._queue.append(request.id)
self._active_users.add(user_id)

View File

@ -7,3 +7,4 @@ fastapi>=0.115,<1
uvicorn>=0.30,<1
jinja2>=3.1,<4
pillow>=10,<13
opencv-python-headless>=4.10,<5

View File

@ -165,6 +165,51 @@ async def test_history_filters_and_stats(db: Database) -> None:
assert alice.display_user == "@alice"
async def test_pending_payment_lifecycle(db: Database) -> None:
pending = await db.create_pending_payment(
payload="payload-1",
user_id=1,
mode="t2i",
prompt="apple",
username="alice",
input_file_path=None,
amount=5,
)
assert pending.status == "pending"
assert pending.amount == 5
assert await db.mark_pending_paid("payload-1", "charge-1") is True
loaded = await db.get_pending_payment("payload-1")
assert loaded is not None
assert loaded.status == "paid"
assert loaded.telegram_payment_charge_id == "charge-1"
await db.attach_pending_request("payload-1", 42)
loaded = await db.get_pending_payment("payload-1")
assert loaded is not None
assert loaded.request_id == 42
await db.mark_pending_refunded("payload-1")
loaded = await db.get_pending_payment("payload-1")
assert loaded is not None
assert loaded.status == "refunded"
request = await db.create_request(
user_id=2,
mode="t2i",
prompt_text="paid apple",
input_file_path=None,
seed=1,
payment_charge_id="charge-2",
)
assert request.payment_charge_id == "charge-2"
assert request.payment_refunded is False
await db.set_refunded(request.id)
loaded_request = await db.get_request(request.id)
assert loaded_request is not None
assert loaded_request.payment_refunded is True
async def test_migrates_legacy_schema_without_username(tmp_path) -> None:
path = tmp_path / "legacy.db"
connection = sqlite3.connect(path)

70
tests/test_moderation.py Normal file
View File

@ -0,0 +1,70 @@
from __future__ import annotations
import pytest
from PIL import Image
from moderation import detect_person, is_prompt_allowed
@pytest.mark.parametrize(
("prompt", "mode", "has_person", "reason"),
[
("child nude", "t2i", False, "csam"),
("l0li", "t2i", False, "csam"),
("sex with dog", "t2i", False, "bestiality"),
("Trump nude", "t2i", False, "real_person"),
("nude woman", "t2i", False, "nonconsensual_sexualization"),
("голая женщина", "t2i", False, "nonconsensual_sexualization"),
("blood sex", "t2i", False, "violence_sexual"),
("how to make a bomb", "t2i", False, "extremism"),
("nazi symbol", "t2i", False, "extremism"),
("ch1ld", "t2i", False, "obfuscated_nsfw"),
("nude", "i2i", True, "i2i_person_nsfw"),
],
)
def test_blocked_prompts(
prompt: str, mode: str, has_person: bool, reason: str
) -> None:
allowed, actual_reason = is_prompt_allowed(
prompt, mode=mode, has_person=has_person
)
assert not allowed
assert actual_reason == reason
@pytest.mark.parametrize(
"prompt",
[
"red apple on a wooden table, soft light",
"nude beach landscape",
"cat sitting on a chair",
"голубое небо, женщина на пляже",
],
)
def test_allowed_prompts(prompt: str) -> None:
allowed, reason = is_prompt_allowed(prompt)
assert allowed
assert reason is None
def test_admin_bypasses_filters() -> None:
allowed, reason = is_prompt_allowed("child nude", is_admin=True)
assert allowed
assert reason is None
def test_extra_public_figures() -> None:
allowed, reason = is_prompt_allowed(
"customperson nude",
extra_public_figures={"customperson"},
)
assert not allowed
assert reason == "real_person"
def test_detect_person_blank_image(tmp_path) -> None:
path = tmp_path / "blank.png"
with Image.new("RGB", (320, 240), (255, 255, 255)) as image:
image.save(path)
assert detect_person(path) is False
assert detect_person(tmp_path / "missing.png") is False

View File

@ -9,6 +9,7 @@ import pytest_asyncio
from config import Settings
from db import Database, UserAlreadyQueued
from moderation import PromptBlocked
from queue_manager import QueueEvent, QueueFull, QueueManager
from workflows import WorkflowBuilder
@ -218,3 +219,35 @@ async def test_prompt_enhancer_is_used(db: Database, tmp_path: Path) -> None:
assert any(event.kind == "enhancing" for event in events)
finally:
await queue.stop()
async def test_blocked_prompt_is_logged(db: Database, tmp_path: Path) -> None:
events: list[QueueEvent] = []
queue = make_queue(db, make_settings(tmp_path), events)
with pytest.raises(PromptBlocked):
await queue.enqueue(user_id=1, mode="t2i", prompt="child nude")
blocked = await db.list_blocked()
assert len(blocked) == 1
assert blocked[0].user_id == 1
assert blocked[0].reason == "csam"
assert blocked[0].prompt == "child nude"
await queue.stop()
async def test_admin_prompt_is_not_filtered(db: Database, tmp_path: Path) -> None:
events: list[QueueEvent] = []
queue = make_queue(db, make_settings(tmp_path), events)
request, _ = await queue.enqueue(
user_id=1,
mode="t2i",
prompt="child nude",
is_admin=True,
)
assert request.id > 0
assert await db.list_blocked() == []
await queue.stop()