Улучшение промптов
Build and push Docker image / docker (push) Successful in 11s
Details
Build and push Docker image / docker (push) Successful in 11s
Details
This commit is contained in:
parent
067f9f31c7
commit
cf5a3597c8
|
|
@ -10,6 +10,11 @@ COMFY_BASE_URL=http://192.168.31.113:8188
|
||||||
WORKFLOW_T2I=workflows/txt2img.json
|
WORKFLOW_T2I=workflows/txt2img.json
|
||||||
WORKFLOW_I2I=workflows/img2img.json
|
WORKFLOW_I2I=workflows/img2img.json
|
||||||
WORKFLOW_MAPPING=workflows/mapping.json
|
WORKFLOW_MAPPING=workflows/mapping.json
|
||||||
|
WORKFLOW_PROMPT=workflows/prompt_enhance.json
|
||||||
|
|
||||||
|
# Улучшение промптов через Qwen3-VL в ComfyUI
|
||||||
|
PROMPT_ENHANCE_ENABLED=true
|
||||||
|
PROMPT_ENHANCE_TIMEOUT=120
|
||||||
|
|
||||||
# Хранилище
|
# Хранилище
|
||||||
STORAGE_DIR=storage
|
STORAGE_DIR=storage
|
||||||
|
|
|
||||||
16
README.md
16
README.md
|
|
@ -8,6 +8,7 @@ Telegram-бот для генерации и редактирования изо
|
||||||
- **Image-to-image**: фото + подпись → редактирование (Krea2 edit).
|
- **Image-to-image**: фото + подпись → редактирование (Krea2 edit).
|
||||||
- **Reply-to-image**: ответ текстом на готовую картинку бота → редактирование этой картинки.
|
- **Reply-to-image**: ответ текстом на готовую картинку бота → редактирование этой картинки.
|
||||||
- **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений.
|
- **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений.
|
||||||
|
- **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI.
|
||||||
- Глобальная очередь с одной активной задачей на пользователя.
|
- Глобальная очередь с одной активной задачей на пользователя.
|
||||||
- Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
|
- Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
|
||||||
- SQLite: запросы, статусы, seed, пути к файлам.
|
- SQLite: запросы, статусы, seed, пути к файлам.
|
||||||
|
|
@ -78,6 +79,21 @@ docker compose up -d --build
|
||||||
|
|
||||||
Workflow: `.gitea/workflows/docker-build.yml`. Runner использует label `ubuntu-latest` и смонтированный Docker socket. Для входа в registry нужны секреты `REGISTRY_USER` и `REGISTRY_TOKEN`.
|
Workflow: `.gitea/workflows/docker-build.yml`. Runner использует label `ubuntu-latest` и смонтированный Docker socket. Для входа в registry нужны секреты `REGISTRY_USER` и `REGISTRY_TOKEN`.
|
||||||
|
|
||||||
|
- `WORKFLOW_PROMPT` — workflow улучшения промптов (по умолчанию `workflows/prompt_enhance.json`).
|
||||||
|
- `PROMPT_ENHANCE_ENABLED` — включить улучшение промптов.
|
||||||
|
- `PROMPT_ENHANCE_TIMEOUT` — таймаут генерации промпта в секундах.
|
||||||
|
|
||||||
|
## Улучшение промптов
|
||||||
|
|
||||||
|
Перед генерацией исходный текст проходит через ComfyUI `TextGenerate` с уже установленной моделью `qwen3vl_4b_fp8_scaled.safetensors`:
|
||||||
|
|
||||||
|
- T2I: перевод и переписывание в английский промпт для генерации.
|
||||||
|
- I2I: перевод и переписывание инструкции для редактирования.
|
||||||
|
- Исходный и улучшенный промпты сохраняются в БД и видны в веб-истории.
|
||||||
|
- При ошибке или таймауте используется исходный промпт.
|
||||||
|
|
||||||
|
Отдельную модель скачивать не нужно, используется существующий CLIP-файл ComfyUI.
|
||||||
|
|
||||||
## Веб-история
|
## Веб-история
|
||||||
|
|
||||||
После `docker compose up -d` доступна страница:
|
После `docker compose up -d` доступна страница:
|
||||||
|
|
|
||||||
9
bot.py
9
bot.py
|
|
@ -20,6 +20,7 @@ from telegram.ext import (
|
||||||
from comfy_client import ComfyClient
|
from comfy_client import ComfyClient
|
||||||
from config import Settings, setup_logging
|
from config import Settings, setup_logging
|
||||||
from db import Database, Request, UserAlreadyQueued
|
from db import Database, Request, UserAlreadyQueued
|
||||||
|
from prompt_enhancer import PromptEnhancer
|
||||||
from queue_manager import QueueEvent, QueueFull, QueueManager
|
from queue_manager import QueueEvent, QueueFull, QueueManager
|
||||||
from workflows import WorkflowBuilder
|
from workflows import WorkflowBuilder
|
||||||
|
|
||||||
|
|
@ -84,12 +85,18 @@ class MpakoBot:
|
||||||
self._settings.comfy_ws_url,
|
self._settings.comfy_ws_url,
|
||||||
)
|
)
|
||||||
await self._db.connect()
|
await self._db.connect()
|
||||||
|
enhancer = (
|
||||||
|
PromptEnhancer(self._comfy, self._settings)
|
||||||
|
if self._settings.prompt_enhance_enabled
|
||||||
|
else None
|
||||||
|
)
|
||||||
self._queue = QueueManager(
|
self._queue = QueueManager(
|
||||||
db=self._db,
|
db=self._db,
|
||||||
comfy=self._comfy,
|
comfy=self._comfy,
|
||||||
workflows=self._workflows,
|
workflows=self._workflows,
|
||||||
settings=self._settings,
|
settings=self._settings,
|
||||||
on_event=self._on_queue_event,
|
on_event=self._on_queue_event,
|
||||||
|
enhancer=enhancer,
|
||||||
)
|
)
|
||||||
await self._queue.start()
|
await self._queue.start()
|
||||||
logger.info("Бот запущен, ComfyUI: %s", self._settings.comfy_base_url)
|
logger.info("Бот запущен, ComfyUI: %s", self._settings.comfy_base_url)
|
||||||
|
|
@ -359,6 +366,8 @@ class MpakoBot:
|
||||||
async def _on_queue_event(self, event: QueueEvent) -> None:
|
async def _on_queue_event(self, event: QueueEvent) -> None:
|
||||||
if event.kind == "started":
|
if event.kind == "started":
|
||||||
await self._edit_status(event.request.id, "Генерация…", percent=0)
|
await self._edit_status(event.request.id, "Генерация…", percent=0)
|
||||||
|
elif event.kind == "enhancing":
|
||||||
|
await self._edit_status(event.request.id, "Улучшаю промпт…")
|
||||||
elif event.kind == "progress":
|
elif event.kind == "progress":
|
||||||
percent = event.percent if event.percent is not None else 0
|
percent = event.percent if event.percent is not None else 0
|
||||||
await self._edit_status(
|
await self._edit_status(
|
||||||
|
|
|
||||||
15
checklist.md
15
checklist.md
|
|
@ -6,9 +6,9 @@
|
||||||
|
|
||||||
## Текущий статус
|
## Текущий статус
|
||||||
|
|
||||||
- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12) и веб-история (13).
|
- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13) и улучшение промптов (14).
|
||||||
- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, переключение T2I → I2I.
|
- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, переключение T2I → I2I.
|
||||||
- 15 тестов проходят: БД, workflows, очередь, вспомогательные функции ComfyUI.
|
- 25 тестов проходят: БД, workflows, очередь, веб-история, enhancer, вспомогательные функции ComfyUI.
|
||||||
- Не проверено: сборка Docker (демон не запущен), очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, реальная отправка в Telegram (нет токена).
|
- Не проверено: сборка Docker (демон не запущен), очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, реальная отправка в Telegram (нет токена).
|
||||||
|
|
||||||
## 0. Допущения и решения по умолчанию
|
## 0. Допущения и решения по умолчанию
|
||||||
|
|
@ -307,6 +307,17 @@ ON requests(user_id) WHERE status IN ('queued', 'running');
|
||||||
|
|
||||||
**DoD:** история открывается по `http://IP:8099`, фильтр по пользователю работает, изображения доступны.
|
**DoD:** история открывается по `http://IP:8099`, фильтр по пользователю работает, изображения доступны.
|
||||||
|
|
||||||
|
## 14. Улучшение промптов
|
||||||
|
|
||||||
|
- [x] `workflows/prompt_enhance.json`: `CLIPLoader` + `TextGenerate` + `PreviewAny` на существующей `qwen3vl_4b`.
|
||||||
|
- [x] `prompt_enhancer.py`: ChatML-промпт для T2I/I2I, очистка ответа.
|
||||||
|
- [x] Улучшение до генерации, с fallback на исходный промпт.
|
||||||
|
- [x] Исходный и улучшенный промпты хранятся в БД, отображаются в веб-истории.
|
||||||
|
- [x] Событие «Улучшаю промпт…» в статусе бота.
|
||||||
|
- [x] Тесты enhancer, очереди и миграции БД.
|
||||||
|
|
||||||
|
**DoD:** русский запрос превращается в английский промпт, изображение генерируется по нему; при сбое LLM используется оригинал.
|
||||||
|
|
||||||
## Порядок работ
|
## Порядок работ
|
||||||
|
|
||||||
1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow.
|
1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow.
|
||||||
|
|
|
||||||
|
|
@ -272,6 +272,20 @@ class ComfyClient:
|
||||||
return image
|
return image
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
def find_output_text(self, entry: dict[str, Any]) -> str | None:
|
||||||
|
outputs = entry.get("outputs") or {}
|
||||||
|
if not isinstance(outputs, dict):
|
||||||
|
return None
|
||||||
|
for node_output in outputs.values():
|
||||||
|
if not isinstance(node_output, dict):
|
||||||
|
continue
|
||||||
|
value = node_output.get("text")
|
||||||
|
if isinstance(value, list):
|
||||||
|
value = value[0] if value else None
|
||||||
|
if isinstance(value, str) and value.strip():
|
||||||
|
return value.strip()
|
||||||
|
return None
|
||||||
|
|
||||||
async def download_output(self, entry: dict[str, Any], destination: Path) -> Path:
|
async def download_output(self, entry: dict[str, Any], destination: Path) -> Path:
|
||||||
image = self.find_output_image(entry)
|
image = self.find_output_image(entry)
|
||||||
if image is None:
|
if image is None:
|
||||||
|
|
|
||||||
17
config.py
17
config.py
|
|
@ -38,6 +38,13 @@ def _float(name: str, default: float) -> float:
|
||||||
raise ValueError(f"{name} должен быть числом") from exc
|
raise ValueError(f"{name} должен быть числом") from exc
|
||||||
|
|
||||||
|
|
||||||
|
def _bool(name: str, default: bool) -> bool:
|
||||||
|
raw = os.getenv(name)
|
||||||
|
if raw is None or not raw.strip():
|
||||||
|
return default
|
||||||
|
return raw.strip().lower() in {"1", "true", "yes", "on"}
|
||||||
|
|
||||||
|
|
||||||
def _user_ids(name: str) -> frozenset[int]:
|
def _user_ids(name: str) -> frozenset[int]:
|
||||||
raw = os.getenv(name, "")
|
raw = os.getenv(name, "")
|
||||||
values: set[int] = set()
|
values: set[int] = set()
|
||||||
|
|
@ -60,6 +67,9 @@ class Settings:
|
||||||
workflow_t2i: Path = PROJECT_ROOT / "workflows" / "txt2img.json"
|
workflow_t2i: Path = PROJECT_ROOT / "workflows" / "txt2img.json"
|
||||||
workflow_i2i: Path = PROJECT_ROOT / "workflows" / "img2img.json"
|
workflow_i2i: Path = PROJECT_ROOT / "workflows" / "img2img.json"
|
||||||
workflow_mapping: Path = PROJECT_ROOT / "workflows" / "mapping.json"
|
workflow_mapping: Path = PROJECT_ROOT / "workflows" / "mapping.json"
|
||||||
|
workflow_prompt: Path = PROJECT_ROOT / "workflows" / "prompt_enhance.json"
|
||||||
|
prompt_enhance_enabled: bool = False
|
||||||
|
prompt_enhance_timeout: float = 120.0
|
||||||
storage_dir: Path = PROJECT_ROOT / "storage"
|
storage_dir: Path = PROJECT_ROOT / "storage"
|
||||||
db_path: Path = PROJECT_ROOT / "data" / "bot.db"
|
db_path: Path = PROJECT_ROOT / "data" / "bot.db"
|
||||||
worker_concurrency: int = 1
|
worker_concurrency: int = 1
|
||||||
|
|
@ -107,6 +117,13 @@ class Settings:
|
||||||
workflow_t2i=_path("WORKFLOW_T2I", "workflows/txt2img.json"),
|
workflow_t2i=_path("WORKFLOW_T2I", "workflows/txt2img.json"),
|
||||||
workflow_i2i=_path("WORKFLOW_I2I", "workflows/img2img.json"),
|
workflow_i2i=_path("WORKFLOW_I2I", "workflows/img2img.json"),
|
||||||
workflow_mapping=_path("WORKFLOW_MAPPING", "workflows/mapping.json"),
|
workflow_mapping=_path("WORKFLOW_MAPPING", "workflows/mapping.json"),
|
||||||
|
workflow_prompt=_path(
|
||||||
|
"WORKFLOW_PROMPT", "workflows/prompt_enhance.json"
|
||||||
|
),
|
||||||
|
prompt_enhance_enabled=_bool("PROMPT_ENHANCE_ENABLED", False),
|
||||||
|
prompt_enhance_timeout=max(
|
||||||
|
10.0, _float("PROMPT_ENHANCE_TIMEOUT", 120.0)
|
||||||
|
),
|
||||||
storage_dir=_path("STORAGE_DIR", "storage"),
|
storage_dir=_path("STORAGE_DIR", "storage"),
|
||||||
db_path=_path("DB_PATH", "data/bot.db"),
|
db_path=_path("DB_PATH", "data/bot.db"),
|
||||||
worker_concurrency=max(1, _int("WORKER_CONCURRENCY", 1)),
|
worker_concurrency=max(1, _int("WORKER_CONCURRENCY", 1)),
|
||||||
|
|
|
||||||
11
db.py
11
db.py
|
|
@ -32,6 +32,7 @@ class Request:
|
||||||
username: str | None
|
username: str | None
|
||||||
timestamp: str
|
timestamp: str
|
||||||
prompt_text: str
|
prompt_text: str
|
||||||
|
enhanced_prompt: str | None
|
||||||
input_file_path: str | None
|
input_file_path: str | None
|
||||||
output_file_path: str | None
|
output_file_path: str | None
|
||||||
status: str
|
status: str
|
||||||
|
|
@ -53,6 +54,7 @@ class Request:
|
||||||
username=row["username"],
|
username=row["username"],
|
||||||
timestamp=row["timestamp"],
|
timestamp=row["timestamp"],
|
||||||
prompt_text=row["prompt_text"],
|
prompt_text=row["prompt_text"],
|
||||||
|
enhanced_prompt=row["enhanced_prompt"],
|
||||||
input_file_path=row["input_file_path"],
|
input_file_path=row["input_file_path"],
|
||||||
output_file_path=row["output_file_path"],
|
output_file_path=row["output_file_path"],
|
||||||
status=row["status"],
|
status=row["status"],
|
||||||
|
|
@ -83,6 +85,7 @@ CREATE TABLE IF NOT EXISTS requests (
|
||||||
username TEXT,
|
username TEXT,
|
||||||
timestamp TEXT NOT NULL,
|
timestamp TEXT NOT NULL,
|
||||||
prompt_text TEXT NOT NULL,
|
prompt_text TEXT NOT NULL,
|
||||||
|
enhanced_prompt TEXT,
|
||||||
input_file_path TEXT,
|
input_file_path TEXT,
|
||||||
output_file_path TEXT,
|
output_file_path TEXT,
|
||||||
status TEXT NOT NULL DEFAULT 'queued'
|
status TEXT NOT NULL DEFAULT 'queued'
|
||||||
|
|
@ -132,6 +135,8 @@ class Database:
|
||||||
columns = {row["name"] async for row in cursor}
|
columns = {row["name"] async for row in cursor}
|
||||||
if "username" not in columns:
|
if "username" not in columns:
|
||||||
await conn.execute("ALTER TABLE requests ADD COLUMN username TEXT")
|
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")
|
||||||
|
|
||||||
async def close(self) -> None:
|
async def close(self) -> None:
|
||||||
if self._conn is not None:
|
if self._conn is not None:
|
||||||
|
|
@ -342,6 +347,12 @@ class Database:
|
||||||
(comfy_prompt_id, request_id),
|
(comfy_prompt_id, request_id),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def set_enhanced_prompt(self, request_id: int, prompt: str) -> None:
|
||||||
|
await self._update(
|
||||||
|
"UPDATE requests SET enhanced_prompt = ? WHERE id = ?",
|
||||||
|
(prompt, request_id),
|
||||||
|
)
|
||||||
|
|
||||||
async def set_done(self, request_id: int, output_file_path: str) -> None:
|
async def set_done(self, request_id: int, output_file_path: str) -> None:
|
||||||
await self._update(
|
await self._update(
|
||||||
"""
|
"""
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,135 @@
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import copy
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import re
|
||||||
|
import uuid
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from comfy_client import ComfyClient
|
||||||
|
from config import Settings
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
SYSTEM_PROMPTS = {
|
||||||
|
"t2i": (
|
||||||
|
"You are a prompt engineer for an image generation model. "
|
||||||
|
"Translate the user's request into English and rewrite it as a concise, "
|
||||||
|
"descriptive image prompt. Preserve every requested subject, action, style, "
|
||||||
|
"color, composition and constraint. Do not invent new objects, camera angles "
|
||||||
|
"or details. Return only the final English prompt, without quotes or "
|
||||||
|
"explanations."
|
||||||
|
),
|
||||||
|
"i2i": (
|
||||||
|
"You are an image-editing prompt translator. Translate the user's editing "
|
||||||
|
"instruction into English and rewrite it clearly. Preserve the requested "
|
||||||
|
"changes, composition, pose, lighting and all constraints. Do not add new "
|
||||||
|
"objects, camera angles or details. Return only the final English "
|
||||||
|
"instruction, without quotes or explanations."
|
||||||
|
),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class PromptEnhanceError(RuntimeError):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class PromptEnhancer:
|
||||||
|
def __init__(self, comfy: ComfyClient, settings: Settings) -> None:
|
||||||
|
self._comfy = comfy
|
||||||
|
self._settings = settings
|
||||||
|
self._template = self._load_template(settings.workflow_prompt)
|
||||||
|
self._text_node_id = self._find_text_node(self._template)
|
||||||
|
if not any(
|
||||||
|
node.get("class_type") == "PreviewAny"
|
||||||
|
for node in self._template.values()
|
||||||
|
if isinstance(node, dict)
|
||||||
|
):
|
||||||
|
raise PromptEnhanceError(
|
||||||
|
"В workflow улучшения промпта нет узла PreviewAny"
|
||||||
|
)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def enabled(self) -> bool:
|
||||||
|
return self._settings.prompt_enhance_enabled
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _load_template(path: Path) -> dict[str, Any]:
|
||||||
|
try:
|
||||||
|
data = json.loads(path.read_text(encoding="utf-8"))
|
||||||
|
except FileNotFoundError as exc:
|
||||||
|
raise PromptEnhanceError(f"Не найден workflow: {path}") from exc
|
||||||
|
except json.JSONDecodeError as exc:
|
||||||
|
raise PromptEnhanceError(f"Некорректный JSON в {path}: {exc}") from exc
|
||||||
|
if not isinstance(data, dict) or not data:
|
||||||
|
raise PromptEnhanceError("Workflow улучшения промпта пустой")
|
||||||
|
return data
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _find_text_node(workflow: dict[str, Any]) -> str:
|
||||||
|
node_ids = [
|
||||||
|
node_id
|
||||||
|
for node_id, node in workflow.items()
|
||||||
|
if isinstance(node, dict) and node.get("class_type") == "TextGenerate"
|
||||||
|
]
|
||||||
|
if len(node_ids) != 1:
|
||||||
|
raise PromptEnhanceError(
|
||||||
|
"В workflow улучшения промпта должен быть один TextGenerate"
|
||||||
|
)
|
||||||
|
return str(node_ids[0])
|
||||||
|
|
||||||
|
async def enhance(self, prompt: str, mode: str) -> str:
|
||||||
|
if not self.enabled:
|
||||||
|
return prompt
|
||||||
|
|
||||||
|
user_prompt = prompt.strip()
|
||||||
|
for token in ("<|im_start|>", "<|im_end|>"):
|
||||||
|
user_prompt = user_prompt.replace(token, "")
|
||||||
|
if not user_prompt:
|
||||||
|
return prompt
|
||||||
|
|
||||||
|
workflow = copy.deepcopy(self._template)
|
||||||
|
workflow[self._text_node_id]["inputs"]["prompt"] = self._chat_prompt(
|
||||||
|
user_prompt, mode
|
||||||
|
)
|
||||||
|
prompt_id = await self._comfy.submit(workflow, uuid.uuid4().hex)
|
||||||
|
entry = await self._comfy.wait_for_history(
|
||||||
|
prompt_id,
|
||||||
|
timeout=self._settings.prompt_enhance_timeout,
|
||||||
|
)
|
||||||
|
text = self._comfy.find_output_text(entry)
|
||||||
|
if not text:
|
||||||
|
raise PromptEnhanceError("ComfyUI вернул пустой промпт")
|
||||||
|
|
||||||
|
cleaned = self._clean(text)
|
||||||
|
if not cleaned:
|
||||||
|
raise PromptEnhanceError("Не удалось очистить сгенерированный промпт")
|
||||||
|
return cleaned[: self._settings.max_prompt_length]
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _chat_prompt(prompt: str, mode: str) -> str:
|
||||||
|
system = SYSTEM_PROMPTS.get(mode, SYSTEM_PROMPTS["t2i"])
|
||||||
|
return (
|
||||||
|
f"<|im_start|>system\n{system}<|im_end|>\n"
|
||||||
|
f"<|im_start|>user\n{prompt}<|im_end|>\n"
|
||||||
|
"<|im_start|>assistant\n"
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _clean(text: str) -> str:
|
||||||
|
text = " ".join(text.split()).strip()
|
||||||
|
for token in ("<|im_end|>", "<|im_start|>"):
|
||||||
|
text = text.split(token)[0].strip()
|
||||||
|
text = text.strip("`").strip()
|
||||||
|
for prefix in ("English prompt:", "Prompt:", "Assistant:"):
|
||||||
|
if text.lower().startswith(prefix.lower()):
|
||||||
|
text = text[len(prefix):].strip()
|
||||||
|
text = text.strip(" \"'")
|
||||||
|
text = text.lstrip(" ,.;:-")
|
||||||
|
if text.lower() in {"user", "assistant", "system"}:
|
||||||
|
return ""
|
||||||
|
if not re.search(r"[A-Za-z]", text):
|
||||||
|
return ""
|
||||||
|
return text
|
||||||
|
|
@ -13,6 +13,7 @@ from pathlib import Path
|
||||||
from comfy_client import ComfyClient, ComfyError
|
from comfy_client import ComfyClient, ComfyError
|
||||||
from config import Settings
|
from config import Settings
|
||||||
from db import Database, Request, UserAlreadyQueued
|
from db import Database, Request, UserAlreadyQueued
|
||||||
|
from prompt_enhancer import PromptEnhancer
|
||||||
from workflows import WorkflowBuilder
|
from workflows import WorkflowBuilder
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
@ -22,7 +23,7 @@ EventCallback = Callable[["QueueEvent"], Awaitable[None] | None]
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
class QueueEvent:
|
class QueueEvent:
|
||||||
kind: str # started | progress | done | failed
|
kind: str # started | enhancing | progress | done | failed
|
||||||
request: Request
|
request: Request
|
||||||
position: int | None = None
|
position: int | None = None
|
||||||
percent: int | None = None
|
percent: int | None = None
|
||||||
|
|
@ -50,12 +51,14 @@ class QueueManager:
|
||||||
workflows: WorkflowBuilder,
|
workflows: WorkflowBuilder,
|
||||||
settings: Settings,
|
settings: Settings,
|
||||||
on_event: EventCallback,
|
on_event: EventCallback,
|
||||||
|
enhancer: PromptEnhancer | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
self._db = db
|
self._db = db
|
||||||
self._comfy = comfy
|
self._comfy = comfy
|
||||||
self._workflows = workflows
|
self._workflows = workflows
|
||||||
self._settings = settings
|
self._settings = settings
|
||||||
self._on_event = on_event
|
self._on_event = on_event
|
||||||
|
self._enhancer = enhancer
|
||||||
|
|
||||||
self._queue: deque[int] = deque()
|
self._queue: deque[int] = deque()
|
||||||
self._active_users: set[int] = set()
|
self._active_users: set[int] = set()
|
||||||
|
|
@ -162,6 +165,23 @@ class QueueManager:
|
||||||
client_id = uuid.uuid4().hex
|
client_id = uuid.uuid4().hex
|
||||||
progress = None
|
progress = None
|
||||||
try:
|
try:
|
||||||
|
if self._enhancer is not None and self._enhancer.enabled:
|
||||||
|
await self._emit(
|
||||||
|
QueueEvent(kind="enhancing", request=request)
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
enhanced = await self._enhancer.enhance(
|
||||||
|
request.prompt_text, request.mode
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
logger.exception(
|
||||||
|
"Не удалось улучшить промпт запроса %s", request_id
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
if enhanced and enhanced != request.prompt_text:
|
||||||
|
await self._db.set_enhanced_prompt(request_id, enhanced)
|
||||||
|
request = await self._refresh(request_id)
|
||||||
|
|
||||||
workflow = await self._build_workflow(request)
|
workflow = await self._build_workflow(request)
|
||||||
progress = await self._comfy.open_progress(client_id)
|
progress = await self._comfy.open_progress(client_id)
|
||||||
prompt_id = await self._comfy.submit(workflow, client_id)
|
prompt_id = await self._comfy.submit(workflow, client_id)
|
||||||
|
|
@ -221,9 +241,10 @@ class QueueManager:
|
||||||
async def _build_workflow(self, request: Request) -> dict:
|
async def _build_workflow(self, request: Request) -> dict:
|
||||||
seed = request.seed if request.seed is not None else secrets.randbelow(2**32)
|
seed = request.seed if request.seed is not None else secrets.randbelow(2**32)
|
||||||
negative = self._settings.default_negative
|
negative = self._settings.default_negative
|
||||||
|
prompt = request.enhanced_prompt or request.prompt_text
|
||||||
if request.mode == "t2i":
|
if request.mode == "t2i":
|
||||||
return self._workflows.build_txt2img(
|
return self._workflows.build_txt2img(
|
||||||
prompt=request.prompt_text,
|
prompt=prompt,
|
||||||
negative=negative,
|
negative=negative,
|
||||||
seed=seed,
|
seed=seed,
|
||||||
)
|
)
|
||||||
|
|
@ -235,7 +256,7 @@ class QueueManager:
|
||||||
remote_name = f"{request.id}_input{suffix}"
|
remote_name = f"{request.id}_input{suffix}"
|
||||||
image_name = await self._comfy.upload_image(input_path, remote_name)
|
image_name = await self._comfy.upload_image(input_path, remote_name)
|
||||||
return self._workflows.build_img2img(
|
return self._workflows.build_img2img(
|
||||||
prompt=request.prompt_text,
|
prompt=prompt,
|
||||||
negative=negative,
|
negative=negative,
|
||||||
seed=seed,
|
seed=seed,
|
||||||
image_name=image_name,
|
image_name=image_name,
|
||||||
|
|
|
||||||
|
|
@ -38,3 +38,9 @@ def test_find_output_image() -> None:
|
||||||
|
|
||||||
def test_find_output_image_empty() -> None:
|
def test_find_output_image_empty() -> None:
|
||||||
assert ComfyClient.find_output_image(None, {}) is None # type: ignore[arg-type]
|
assert ComfyClient.find_output_image(None, {}) is None # type: ignore[arg-type]
|
||||||
|
|
||||||
|
|
||||||
|
def test_find_output_text() -> None:
|
||||||
|
entry = {"outputs": {"3": {"text": ["A red apple"]}}}
|
||||||
|
text = ComfyClient.find_output_text(None, entry) # type: ignore[arg-type]
|
||||||
|
assert text == "A red apple"
|
||||||
|
|
|
||||||
|
|
@ -202,5 +202,9 @@ async def test_migrates_legacy_schema_without_username(tmp_path) -> None:
|
||||||
seed=1,
|
seed=1,
|
||||||
)
|
)
|
||||||
assert request.username == "alice"
|
assert request.username == "alice"
|
||||||
|
await database.set_enhanced_prompt(request.id, "a red apple")
|
||||||
|
loaded = await database.get_request(request.id)
|
||||||
|
assert loaded is not None
|
||||||
|
assert loaded.enhanced_prompt == "a red apple"
|
||||||
finally:
|
finally:
|
||||||
await database.close()
|
await database.close()
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,77 @@
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from config import Settings
|
||||||
|
from prompt_enhancer import PromptEnhancer
|
||||||
|
|
||||||
|
BASE = Path(__file__).resolve().parents[1] / "workflows"
|
||||||
|
|
||||||
|
|
||||||
|
class FakeEnhanceComfy:
|
||||||
|
def __init__(self, text: str = "A red apple on a wooden table, soft light") -> None:
|
||||||
|
self.text = text
|
||||||
|
self.workflow: dict | None = None
|
||||||
|
self.submits = 0
|
||||||
|
|
||||||
|
async def submit(self, workflow: dict, client_id: str) -> str:
|
||||||
|
self.workflow = workflow
|
||||||
|
self.submits += 1
|
||||||
|
return "prompt-1"
|
||||||
|
|
||||||
|
async def wait_for_history(self, prompt_id: str, *, timeout: float) -> dict:
|
||||||
|
return {"outputs": {"3": {"text": [self.text]}}}
|
||||||
|
|
||||||
|
def find_output_text(self, entry: dict) -> str | None:
|
||||||
|
return entry["outputs"]["3"]["text"][0]
|
||||||
|
|
||||||
|
|
||||||
|
def make_settings(enabled: bool = True) -> Settings:
|
||||||
|
return Settings(
|
||||||
|
bot_token="",
|
||||||
|
prompt_enhance_enabled=enabled,
|
||||||
|
workflow_prompt=BASE / "prompt_enhance.json",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def test_enhance_t2i_prompt() -> None:
|
||||||
|
comfy = FakeEnhanceComfy()
|
||||||
|
enhancer = PromptEnhancer(comfy, make_settings())
|
||||||
|
|
||||||
|
result = await enhancer.enhance("красное яблоко на столе, мягкий свет", "t2i")
|
||||||
|
|
||||||
|
assert result == "A red apple on a wooden table, soft light"
|
||||||
|
assert comfy.submits == 1
|
||||||
|
assert comfy.workflow is not None
|
||||||
|
prompt = comfy.workflow["2"]["inputs"]["prompt"]
|
||||||
|
assert "<|im_start|>system" in prompt
|
||||||
|
assert "image generation model" in prompt
|
||||||
|
assert "красное яблоко" in prompt
|
||||||
|
assert prompt.endswith("<|im_start|>assistant\n")
|
||||||
|
|
||||||
|
|
||||||
|
async def test_enhance_i2i_prompt() -> None:
|
||||||
|
comfy = FakeEnhanceComfy("Make the apple blue, keep the composition")
|
||||||
|
enhancer = PromptEnhancer(comfy, make_settings())
|
||||||
|
|
||||||
|
result = await enhancer.enhance("Сделай яблоко синим", "i2i")
|
||||||
|
|
||||||
|
assert result == "Make the apple blue, keep the composition"
|
||||||
|
assert comfy.workflow is not None
|
||||||
|
prompt = comfy.workflow["2"]["inputs"]["prompt"]
|
||||||
|
assert "image-editing prompt translator" in prompt
|
||||||
|
|
||||||
|
|
||||||
|
async def test_disabled_returns_original_without_comfy_call() -> None:
|
||||||
|
comfy = FakeEnhanceComfy()
|
||||||
|
enhancer = PromptEnhancer(comfy, make_settings(enabled=False))
|
||||||
|
|
||||||
|
result = await enhancer.enhance("красное яблоко", "t2i")
|
||||||
|
|
||||||
|
assert result == "красное яблоко"
|
||||||
|
assert comfy.submits == 0
|
||||||
|
|
||||||
|
|
||||||
|
async def test_clean_strips_prefix_and_empty_output() -> None:
|
||||||
|
assert PromptEnhancer._clean('Prompt: "A red apple"') == "A red apple"
|
||||||
|
assert PromptEnhancer._clean("user") == ""
|
||||||
|
|
@ -44,6 +44,17 @@ class FakeComfy:
|
||||||
return remote_name
|
return remote_name
|
||||||
|
|
||||||
|
|
||||||
|
class FakeEnhancer:
|
||||||
|
enabled = True
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.calls: list[tuple[str, str]] = []
|
||||||
|
|
||||||
|
async def enhance(self, prompt: str, mode: str) -> str:
|
||||||
|
self.calls.append((prompt, mode))
|
||||||
|
return f"enhanced: {prompt}"
|
||||||
|
|
||||||
|
|
||||||
@pytest_asyncio.fixture
|
@pytest_asyncio.fixture
|
||||||
async def db(tmp_path):
|
async def db(tmp_path):
|
||||||
database = Database(tmp_path / "bot.db")
|
database = Database(tmp_path / "bot.db")
|
||||||
|
|
@ -65,7 +76,13 @@ def make_settings(tmp_path: Path, **overrides) -> Settings:
|
||||||
return Settings(**values)
|
return Settings(**values)
|
||||||
|
|
||||||
|
|
||||||
def make_queue(db: Database, settings: Settings, events: list[QueueEvent]) -> QueueManager:
|
def make_queue(
|
||||||
|
db: Database,
|
||||||
|
settings: Settings,
|
||||||
|
events: list[QueueEvent],
|
||||||
|
*,
|
||||||
|
enhancer=None,
|
||||||
|
) -> QueueManager:
|
||||||
async def on_event(event: QueueEvent) -> None:
|
async def on_event(event: QueueEvent) -> None:
|
||||||
events.append(event)
|
events.append(event)
|
||||||
|
|
||||||
|
|
@ -79,6 +96,7 @@ def make_queue(db: Database, settings: Settings, events: list[QueueEvent]) -> Qu
|
||||||
),
|
),
|
||||||
settings=settings,
|
settings=settings,
|
||||||
on_event=on_event,
|
on_event=on_event,
|
||||||
|
enhancer=enhancer,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -176,3 +194,27 @@ async def test_image_to_image_uploads_input(db: Database, tmp_path: Path) -> Non
|
||||||
assert source == f"{request.id}_input.png"
|
assert source == f"{request.id}_input.png"
|
||||||
finally:
|
finally:
|
||||||
await queue.stop()
|
await queue.stop()
|
||||||
|
|
||||||
|
|
||||||
|
async def test_prompt_enhancer_is_used(db: Database, tmp_path: Path) -> None:
|
||||||
|
events: list[QueueEvent] = []
|
||||||
|
settings = make_settings(tmp_path)
|
||||||
|
enhancer = FakeEnhancer()
|
||||||
|
queue = make_queue(db, settings, events, enhancer=enhancer)
|
||||||
|
comfy = queue._comfy # type: ignore[attr-defined]
|
||||||
|
|
||||||
|
await queue.start()
|
||||||
|
try:
|
||||||
|
request, _ = await queue.enqueue(
|
||||||
|
user_id=1, mode="t2i", prompt="apple"
|
||||||
|
)
|
||||||
|
await wait_for_event(events, "done")
|
||||||
|
|
||||||
|
assert enhancer.calls == [("apple", "t2i")]
|
||||||
|
assert comfy.workflows[0]["67"]["inputs"]["text"] == "enhanced: apple"
|
||||||
|
stored = await db.get_request(request.id)
|
||||||
|
assert stored is not None
|
||||||
|
assert stored.enhanced_prompt == "enhanced: apple"
|
||||||
|
assert any(event.kind == "enhancing" for event in events)
|
||||||
|
finally:
|
||||||
|
await queue.stop()
|
||||||
|
|
|
||||||
|
|
@ -29,17 +29,18 @@ def insert_request(
|
||||||
mode: str = "t2i",
|
mode: str = "t2i",
|
||||||
output: str | None = None,
|
output: str | None = None,
|
||||||
username: str | None = None,
|
username: str | None = None,
|
||||||
|
enhanced: str | None = None,
|
||||||
) -> int:
|
) -> int:
|
||||||
connection = sqlite3.connect(path)
|
connection = sqlite3.connect(path)
|
||||||
try:
|
try:
|
||||||
cursor = connection.execute(
|
cursor = connection.execute(
|
||||||
"""
|
"""
|
||||||
INSERT INTO requests
|
INSERT INTO requests
|
||||||
(user_id, username, timestamp, prompt_text, status, mode, seed,
|
(user_id, username, timestamp, prompt_text, enhanced_prompt,
|
||||||
output_file_path)
|
status, mode, seed, output_file_path)
|
||||||
VALUES (?, ?, '2026-10-18T12:00:00+00:00', ?, ?, ?, 1, ?)
|
VALUES (?, ?, '2026-10-18T12:00:00+00:00', ?, ?, ?, ?, 1, ?)
|
||||||
""",
|
""",
|
||||||
(user_id, username, prompt, status, mode, output),
|
(user_id, username, prompt, enhanced, status, mode, output),
|
||||||
)
|
)
|
||||||
connection.commit()
|
connection.commit()
|
||||||
return int(cursor.lastrowid)
|
return int(cursor.lastrowid)
|
||||||
|
|
@ -86,15 +87,19 @@ def test_detail_and_media(tmp_path: Path) -> None:
|
||||||
settings.db_path,
|
settings.db_path,
|
||||||
user_id=1,
|
user_id=1,
|
||||||
prompt="detail prompt",
|
prompt="detail prompt",
|
||||||
|
enhanced="An enhanced English prompt",
|
||||||
output=str(output),
|
output=str(output),
|
||||||
)
|
)
|
||||||
|
|
||||||
detail = client.get(f"/request/{request_id}")
|
detail = client.get(f"/request/{request_id}")
|
||||||
assert detail.status_code == 200
|
assert detail.status_code == 200
|
||||||
assert "detail prompt" in detail.text
|
assert "detail prompt" in detail.text
|
||||||
|
assert "An enhanced English prompt" in detail.text
|
||||||
|
|
||||||
index = client.get("/")
|
index = client.get("/")
|
||||||
assert f"/thumb/{request_id}" in index.text
|
assert f"/thumb/{request_id}" in index.text
|
||||||
|
assert "detail prompt" in index.text
|
||||||
|
assert "An enhanced English prompt" in index.text
|
||||||
|
|
||||||
media = client.get(f"/media/{request_id}/output")
|
media = client.get(f"/media/{request_id}/output")
|
||||||
assert media.status_code == 200
|
assert media.status_code == 200
|
||||||
|
|
|
||||||
|
|
@ -120,6 +120,22 @@ tbody tr:last-child td { border-bottom: 0; }
|
||||||
.empty { text-align: center; color: var(--muted); padding: 28px; }
|
.empty { text-align: center; color: var(--muted); padding: 28px; }
|
||||||
.details { white-space: nowrap; }
|
.details { white-space: nowrap; }
|
||||||
|
|
||||||
|
.mark {
|
||||||
|
display: inline-block;
|
||||||
|
padding: 0 5px;
|
||||||
|
border-radius: 4px;
|
||||||
|
background: #26324a;
|
||||||
|
color: #9db8e8;
|
||||||
|
font-size: 10px;
|
||||||
|
vertical-align: middle;
|
||||||
|
}
|
||||||
|
.enhanced { color: #b9c8ff; }
|
||||||
|
.prompt-original {
|
||||||
|
margin-top: 4px;
|
||||||
|
color: var(--muted);
|
||||||
|
font-size: 12px;
|
||||||
|
}
|
||||||
|
|
||||||
.thumb-cell { width: 78px; }
|
.thumb-cell { width: 78px; }
|
||||||
.thumb {
|
.thumb {
|
||||||
width: 64px;
|
width: 64px;
|
||||||
|
|
|
||||||
|
|
@ -19,8 +19,12 @@
|
||||||
<dd class="mono">{{ item.seed if item.seed is not none else '—' }}</dd>
|
<dd class="mono">{{ item.seed if item.seed is not none else '—' }}</dd>
|
||||||
<dt>Comfy prompt_id</dt>
|
<dt>Comfy prompt_id</dt>
|
||||||
<dd class="mono">{{ item.comfy_prompt_id or '—' }}</dd>
|
<dd class="mono">{{ item.comfy_prompt_id or '—' }}</dd>
|
||||||
<dt>Промпт</dt>
|
<dt>Промпт пользователя</dt>
|
||||||
<dd class="prompt-block">{{ item.prompt_text }}</dd>
|
<dd class="prompt-block">{{ item.prompt_text }}</dd>
|
||||||
|
{% if item.enhanced_prompt and item.enhanced_prompt != item.prompt_text %}
|
||||||
|
<dt>Промпт для модели</dt>
|
||||||
|
<dd class="prompt-block enhanced">{{ item.enhanced_prompt }}</dd>
|
||||||
|
{% endif %}
|
||||||
{% if item.error_text %}
|
{% if item.error_text %}
|
||||||
<dt>Ошибка</dt>
|
<dt>Ошибка</dt>
|
||||||
<dd class="error-text">{{ item.error_text }}</dd>
|
<dd class="error-text">{{ item.error_text }}</dd>
|
||||||
|
|
|
||||||
|
|
@ -79,7 +79,14 @@
|
||||||
<td><a href="/?user_id={{ item.user_id }}" title="ID: {{ item.user_id }}">{{ item.display_user }}</a></td>
|
<td><a href="/?user_id={{ item.user_id }}" title="ID: {{ item.user_id }}">{{ item.display_user }}</a></td>
|
||||||
<td class="mono">{{ item.mode }}</td>
|
<td class="mono">{{ item.mode }}</td>
|
||||||
<td><span class="badge {{ item.status }}">{{ item.status }}</span></td>
|
<td><span class="badge {{ item.status }}">{{ item.status }}</span></td>
|
||||||
<td class="prompt" title="{{ item.prompt_text }}">{{ item.prompt_text|truncate(90, True) }}</td>
|
<td class="prompt" title="{{ item.prompt_text }}">
|
||||||
|
{% if item.enhanced_prompt %}
|
||||||
|
<div>{{ item.enhanced_prompt|truncate(80, True) }} <span class="mark">EN</span></div>
|
||||||
|
<div class="prompt-original">Исходный: {{ item.prompt_text|truncate(80, True) }}</div>
|
||||||
|
{% else %}
|
||||||
|
{{ item.prompt_text|truncate(90, True) }}
|
||||||
|
{% endif %}
|
||||||
|
</td>
|
||||||
<td class="mono">{{ item.seed if item.seed is not none else '—' }}</td>
|
<td class="mono">{{ item.seed if item.seed is not none else '—' }}</td>
|
||||||
<td><a class="details" href="/request/{{ item.id }}">Детали</a></td>
|
<td><a class="details" href="/request/{{ item.id }}">Детали</a></td>
|
||||||
</tr>
|
</tr>
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,43 @@
|
||||||
|
{
|
||||||
|
"1": {
|
||||||
|
"class_type": "CLIPLoader",
|
||||||
|
"inputs": {
|
||||||
|
"clip_name": "qwen3vl_4b_fp8_scaled.safetensors",
|
||||||
|
"type": "krea2",
|
||||||
|
"device": "default"
|
||||||
|
},
|
||||||
|
"_meta": {
|
||||||
|
"title": "Load Qwen3-VL CLIP"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"2": {
|
||||||
|
"class_type": "TextGenerate",
|
||||||
|
"inputs": {
|
||||||
|
"clip": [
|
||||||
|
"1",
|
||||||
|
0
|
||||||
|
],
|
||||||
|
"prompt": "__PROMPT__",
|
||||||
|
"max_length": 300,
|
||||||
|
"sampling_mode": "off",
|
||||||
|
"thinking": false,
|
||||||
|
"use_default_template": false,
|
||||||
|
"mtp": "auto"
|
||||||
|
},
|
||||||
|
"_meta": {
|
||||||
|
"title": "Generate Prompt"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"3": {
|
||||||
|
"class_type": "PreviewAny",
|
||||||
|
"inputs": {
|
||||||
|
"source": [
|
||||||
|
"2",
|
||||||
|
0
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"_meta": {
|
||||||
|
"title": "Prompt Text"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue