Фильтрация промптов
This commit is contained in:
parent
cf5a3597c8
commit
7fcca6e39f
|
|
@ -34,6 +34,12 @@ DEFAULT_NEGATIVE=low quality, blurry, bad anatomy, extra digits
|
|||
# Доступ: пусто — всем; иначе список id через запятую
|
||||
ALLOWED_USER_IDS=
|
||||
|
||||
# Админы: их промпты не фильтруются
|
||||
ADMIN_USER_IDS=
|
||||
MODERATION_ENABLED=true
|
||||
# Дополнительные имена публичных людей через запятую
|
||||
BLOCKED_PUBLIC_FIGURES=
|
||||
|
||||
# Веб-история: только локальная сеть, без авторизации
|
||||
WEB_BIND_IP=192.168.31.158
|
||||
WEB_PAGE_SIZE=50
|
||||
|
|
|
|||
27
README.md
27
README.md
|
|
@ -9,6 +9,7 @@ Telegram-бот для генерации и редактирования изо
|
|||
- **Reply-to-image**: ответ текстом на готовую картинку бота → редактирование этой картинки.
|
||||
- **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений.
|
||||
- **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI.
|
||||
- **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов.
|
||||
- Глобальная очередь с одной активной задачей на пользователя.
|
||||
- Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
|
||||
- SQLite: запросы, статусы, seed, пути к файлам.
|
||||
|
|
@ -94,6 +95,32 @@ 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`.
|
||||
|
||||
## Веб-история
|
||||
|
||||
После `docker compose up -d` доступна страница:
|
||||
|
|
|
|||
17
bot.py
17
bot.py
|
|
@ -20,6 +20,7 @@ from telegram.ext import (
|
|||
from comfy_client import ComfyClient
|
||||
from config import Settings, setup_logging
|
||||
from db import Database, Request, UserAlreadyQueued
|
||||
from moderation import PromptBlocked, detect_person
|
||||
from prompt_enhancer import PromptEnhancer
|
||||
from queue_manager import QueueEvent, QueueFull, QueueManager
|
||||
from workflows import WorkflowBuilder
|
||||
|
|
@ -32,6 +33,7 @@ START_TEXT = (
|
|||
"Или ответь текстом на готовую картинку — возьму её как исходник.\n\n"
|
||||
"/status — узнать состояние текущей задачи."
|
||||
)
|
||||
BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны."
|
||||
|
||||
|
||||
class MpakoBot:
|
||||
|
|
@ -294,6 +296,16 @@ class MpakoBot:
|
|||
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)
|
||||
|
||||
try:
|
||||
request, position = await self._queue.enqueue(
|
||||
user_id=user.id,
|
||||
|
|
@ -301,7 +313,12 @@ class MpakoBot:
|
|||
mode=mode,
|
||||
prompt=prompt,
|
||||
input_file_path=input_path,
|
||||
has_person=has_person,
|
||||
is_admin=is_admin,
|
||||
)
|
||||
except PromptBlocked:
|
||||
await message.reply_text(BLOCKED_TEXT)
|
||||
return False
|
||||
except UserAlreadyQueued:
|
||||
active = await self._db.get_active_by_user(user.id)
|
||||
text = "Ты уже в очереди."
|
||||
|
|
|
|||
19
checklist.md
19
checklist.md
|
|
@ -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).
|
||||
- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история.
|
||||
- 45 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, вспомогательные функции ComfyUI.
|
||||
- Не проверено: очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото.
|
||||
|
||||
## 0. Допущения и решения по умолчанию
|
||||
|
||||
|
|
@ -318,6 +318,17 @@ 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:** небезопасный промпт не попадает в очередь, причина пишется в БД, админ может отправлять без фильтра.
|
||||
|
||||
## Порядок работ
|
||||
|
||||
1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow.
|
||||
|
|
|
|||
16
config.py
16
config.py
|
|
@ -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,15 @@ 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)
|
||||
web_page_size: int = 50
|
||||
web_bind_ip: str = "127.0.0.1"
|
||||
|
||||
|
|
@ -129,6 +142,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 +150,8 @@ 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"),
|
||||
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",
|
||||
|
|
|
|||
66
db.py
66
db.py
|
|
@ -78,6 +78,29 @@ 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"]),
|
||||
)
|
||||
|
||||
|
||||
SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS requests (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
|
|
@ -105,6 +128,19 @@ 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);
|
||||
"""
|
||||
|
||||
|
||||
|
|
@ -378,6 +414,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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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,28 @@ class QueueManager:
|
|||
prompt: str,
|
||||
input_file_path: str | None = None,
|
||||
username: str | None = None,
|
||||
has_person: bool = False,
|
||||
is_admin: bool = False,
|
||||
) -> 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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in New Issue