136 lines
4.9 KiB
Python
136 lines
4.9 KiB
Python
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
|