From 35a1d83aa7a1c6d2173f6ab6a26067c0f664137b Mon Sep 17 00:00:00 2001 From: MpakobecEx Date: Sun, 11 Oct 2026 21:47:00 +0500 Subject: [PATCH] =?UTF-8?q?=D0=A3=D0=BB=D1=83=D1=87=D1=88=D0=B5=D0=BD?= =?UTF-8?q?=D0=B8=D0=B5=20=D0=BF=D1=80=D0=BE=D0=BC=D0=BF=D1=82=D0=BE=D0=B2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Перевожу и переписываю запрос через qwen3vl в ComfyUI. Если модель не ответила, генерирую по исходному промпту. --- .env.example | 5 ++ README.md | 16 ++++ bot.py | 9 +++ checklist.md | 15 +++- comfy_client.py | 14 ++++ config.py | 17 +++++ db.py | 11 +++ prompt_enhancer.py | 135 ++++++++++++++++++++++++++++++++++ queue_manager.py | 27 ++++++- tests/test_comfy.py | 6 ++ tests/test_db.py | 4 + tests/test_prompt_enhancer.py | 77 +++++++++++++++++++ tests/test_queue.py | 44 ++++++++++- tests/test_web.py | 13 +++- web/static/style.css | 16 ++++ web/templates/detail.html | 6 +- web/templates/index.html | 9 ++- workflows/prompt_enhance.json | 43 +++++++++++ 18 files changed, 455 insertions(+), 12 deletions(-) create mode 100644 prompt_enhancer.py create mode 100644 tests/test_prompt_enhancer.py create mode 100644 workflows/prompt_enhance.json diff --git a/.env.example b/.env.example index 8bd62c5..777e235 100644 --- a/.env.example +++ b/.env.example @@ -10,6 +10,11 @@ COMFY_BASE_URL=http://192.168.31.113:8188 WORKFLOW_T2I=workflows/txt2img.json WORKFLOW_I2I=workflows/img2img.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 diff --git a/README.md b/README.md index d962568..9409514 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,7 @@ Telegram-бот для генерации и редактирования изо - **Image-to-image**: фото + подпись → редактирование (Krea2 edit). - **Reply-to-image**: ответ текстом на готовую картинку бота → редактирование этой картинки. - **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений. +- **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI. - Глобальная очередь с одной активной задачей на пользователя. - Позиция в очереди и прогресс генерации через WebSocket ComfyUI. - 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_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` доступна страница: diff --git a/bot.py b/bot.py index 1105d4b..63781e1 100644 --- a/bot.py +++ b/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 prompt_enhancer import PromptEnhancer from queue_manager import QueueEvent, QueueFull, QueueManager from workflows import WorkflowBuilder @@ -84,12 +85,18 @@ class MpakoBot: self._settings.comfy_ws_url, ) await self._db.connect() + enhancer = ( + PromptEnhancer(self._comfy, self._settings) + if self._settings.prompt_enhance_enabled + else None + ) self._queue = QueueManager( db=self._db, comfy=self._comfy, workflows=self._workflows, settings=self._settings, on_event=self._on_queue_event, + enhancer=enhancer, ) await self._queue.start() logger.info("Бот запущен, ComfyUI: %s", self._settings.comfy_base_url) @@ -359,6 +366,8 @@ class MpakoBot: async def _on_queue_event(self, event: QueueEvent) -> None: if event.kind == "started": await self._edit_status(event.request.id, "Генерация…", percent=0) + elif event.kind == "enhancing": + await self._edit_status(event.request.id, "Улучшаю промпт…") elif event.kind == "progress": percent = event.percent if event.percent is not None else 0 await self._edit_status( diff --git a/checklist.md b/checklist.md index fda4503..da0659c 100644 --- a/checklist.md +++ b/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. -- 15 тестов проходят: БД, workflows, очередь, вспомогательные функции ComfyUI. +- 25 тестов проходят: БД, workflows, очередь, веб-история, enhancer, вспомогательные функции ComfyUI. - Не проверено: сборка Docker (демон не запущен), очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, реальная отправка в Telegram (нет токена). ## 0. Допущения и решения по умолчанию @@ -307,6 +307,17 @@ ON requests(user_id) WHERE status IN ('queued', 'running'); **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. diff --git a/comfy_client.py b/comfy_client.py index 069f5bf..2d2bda2 100644 --- a/comfy_client.py +++ b/comfy_client.py @@ -272,6 +272,20 @@ class ComfyClient: return image 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: image = self.find_output_image(entry) if image is None: diff --git a/config.py b/config.py index 8b2f93f..89a3c64 100644 --- a/config.py +++ b/config.py @@ -38,6 +38,13 @@ def _float(name: str, default: float) -> float: 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]: raw = os.getenv(name, "") values: set[int] = set() @@ -60,6 +67,9 @@ class Settings: workflow_t2i: Path = PROJECT_ROOT / "workflows" / "txt2img.json" workflow_i2i: Path = PROJECT_ROOT / "workflows" / "img2img.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" db_path: Path = PROJECT_ROOT / "data" / "bot.db" worker_concurrency: int = 1 @@ -107,6 +117,13 @@ class Settings: workflow_t2i=_path("WORKFLOW_T2I", "workflows/txt2img.json"), workflow_i2i=_path("WORKFLOW_I2I", "workflows/img2img.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"), db_path=_path("DB_PATH", "data/bot.db"), worker_concurrency=max(1, _int("WORKER_CONCURRENCY", 1)), diff --git a/db.py b/db.py index 4402d16..b45eac0 100644 --- a/db.py +++ b/db.py @@ -32,6 +32,7 @@ class Request: username: str | None timestamp: str prompt_text: str + enhanced_prompt: str | None input_file_path: str | None output_file_path: str | None status: str @@ -53,6 +54,7 @@ class Request: username=row["username"], timestamp=row["timestamp"], prompt_text=row["prompt_text"], + enhanced_prompt=row["enhanced_prompt"], input_file_path=row["input_file_path"], output_file_path=row["output_file_path"], status=row["status"], @@ -83,6 +85,7 @@ CREATE TABLE IF NOT EXISTS requests ( username TEXT, timestamp TEXT NOT NULL, prompt_text TEXT NOT NULL, + enhanced_prompt TEXT, input_file_path TEXT, output_file_path TEXT, status TEXT NOT NULL DEFAULT 'queued' @@ -132,6 +135,8 @@ class Database: columns = {row["name"] async for row in cursor} if "username" not in columns: 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: if self._conn is not None: @@ -342,6 +347,12 @@ class Database: (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: await self._update( """ diff --git a/prompt_enhancer.py b/prompt_enhancer.py new file mode 100644 index 0000000..987a668 --- /dev/null +++ b/prompt_enhancer.py @@ -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 diff --git a/queue_manager.py b/queue_manager.py index 52f9d2e..2c977b3 100644 --- a/queue_manager.py +++ b/queue_manager.py @@ -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 prompt_enhancer import PromptEnhancer from workflows import WorkflowBuilder logger = logging.getLogger(__name__) @@ -22,7 +23,7 @@ EventCallback = Callable[["QueueEvent"], Awaitable[None] | None] @dataclass(frozen=True) class QueueEvent: - kind: str # started | progress | done | failed + kind: str # started | enhancing | progress | done | failed request: Request position: int | None = None percent: int | None = None @@ -50,12 +51,14 @@ class QueueManager: workflows: WorkflowBuilder, settings: Settings, on_event: EventCallback, + enhancer: PromptEnhancer | None = None, ) -> None: self._db = db self._comfy = comfy self._workflows = workflows self._settings = settings self._on_event = on_event + self._enhancer = enhancer self._queue: deque[int] = deque() self._active_users: set[int] = set() @@ -162,6 +165,23 @@ class QueueManager: client_id = uuid.uuid4().hex progress = None 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) progress = await self._comfy.open_progress(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: seed = request.seed if request.seed is not None else secrets.randbelow(2**32) negative = self._settings.default_negative + prompt = request.enhanced_prompt or request.prompt_text if request.mode == "t2i": return self._workflows.build_txt2img( - prompt=request.prompt_text, + prompt=prompt, negative=negative, seed=seed, ) @@ -235,7 +256,7 @@ class QueueManager: remote_name = f"{request.id}_input{suffix}" image_name = await self._comfy.upload_image(input_path, remote_name) return self._workflows.build_img2img( - prompt=request.prompt_text, + prompt=prompt, negative=negative, seed=seed, image_name=image_name, diff --git a/tests/test_comfy.py b/tests/test_comfy.py index 5cbfab7..5c70c77 100644 --- a/tests/test_comfy.py +++ b/tests/test_comfy.py @@ -38,3 +38,9 @@ def test_find_output_image() -> None: def test_find_output_image_empty() -> None: 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" diff --git a/tests/test_db.py b/tests/test_db.py index fd15241..1b8416d 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -202,5 +202,9 @@ async def test_migrates_legacy_schema_without_username(tmp_path) -> None: seed=1, ) 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: await database.close() diff --git a/tests/test_prompt_enhancer.py b/tests/test_prompt_enhancer.py new file mode 100644 index 0000000..50aa18f --- /dev/null +++ b/tests/test_prompt_enhancer.py @@ -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") == "" diff --git a/tests/test_queue.py b/tests/test_queue.py index f872015..4f1423a 100644 --- a/tests/test_queue.py +++ b/tests/test_queue.py @@ -44,6 +44,17 @@ class FakeComfy: 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 async def db(tmp_path): database = Database(tmp_path / "bot.db") @@ -65,7 +76,13 @@ def make_settings(tmp_path: Path, **overrides) -> Settings: 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: events.append(event) @@ -79,6 +96,7 @@ def make_queue(db: Database, settings: Settings, events: list[QueueEvent]) -> Qu ), settings=settings, 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" finally: 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() diff --git a/tests/test_web.py b/tests/test_web.py index ed5df56..5d664ac 100644 --- a/tests/test_web.py +++ b/tests/test_web.py @@ -29,17 +29,18 @@ def insert_request( mode: str = "t2i", output: str | None = None, username: str | None = None, + enhanced: str | None = None, ) -> int: connection = sqlite3.connect(path) try: cursor = connection.execute( """ INSERT INTO requests - (user_id, username, timestamp, prompt_text, status, mode, seed, - output_file_path) - VALUES (?, ?, '2026-10-18T12:00:00+00:00', ?, ?, ?, 1, ?) + (user_id, username, timestamp, prompt_text, enhanced_prompt, + status, mode, seed, output_file_path) + 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() return int(cursor.lastrowid) @@ -86,15 +87,19 @@ def test_detail_and_media(tmp_path: Path) -> None: settings.db_path, user_id=1, prompt="detail prompt", + enhanced="An enhanced English prompt", output=str(output), ) detail = client.get(f"/request/{request_id}") assert detail.status_code == 200 assert "detail prompt" in detail.text + assert "An enhanced English prompt" in detail.text index = client.get("/") 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") assert media.status_code == 200 diff --git a/web/static/style.css b/web/static/style.css index c792191..b28400e 100644 --- a/web/static/style.css +++ b/web/static/style.css @@ -120,6 +120,22 @@ tbody tr:last-child td { border-bottom: 0; } .empty { text-align: center; color: var(--muted); padding: 28px; } .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 { width: 64px; diff --git a/web/templates/detail.html b/web/templates/detail.html index 1605fed..0ef49eb 100644 --- a/web/templates/detail.html +++ b/web/templates/detail.html @@ -19,8 +19,12 @@
{{ item.seed if item.seed is not none else '—' }}
Comfy prompt_id
{{ item.comfy_prompt_id or '—' }}
-
Промпт
+
Промпт пользователя
{{ item.prompt_text }}
+ {% if item.enhanced_prompt and item.enhanced_prompt != item.prompt_text %} +
Промпт для модели
+
{{ item.enhanced_prompt }}
+ {% endif %} {% if item.error_text %}
Ошибка
{{ item.error_text }}
diff --git a/web/templates/index.html b/web/templates/index.html index 7508631..20632b9 100644 --- a/web/templates/index.html +++ b/web/templates/index.html @@ -79,7 +79,14 @@ {{ item.display_user }} {{ item.mode }} {{ item.status }} - {{ item.prompt_text|truncate(90, True) }} + + {% if item.enhanced_prompt %} +
{{ item.enhanced_prompt|truncate(80, True) }} EN
+
Исходный: {{ item.prompt_text|truncate(80, True) }}
+ {% else %} + {{ item.prompt_text|truncate(90, True) }} + {% endif %} + {{ item.seed if item.seed is not none else '—' }} Детали diff --git a/workflows/prompt_enhance.json b/workflows/prompt_enhance.json new file mode 100644 index 0000000..5e449e8 --- /dev/null +++ b/workflows/prompt_enhance.json @@ -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" + } + } +}