196 lines
6.9 KiB
Python
196 lines
6.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."
|
||
),
|
||
}
|
||
|
||
REFINE_SYSTEM_PROMPTS = {
|
||
"t2i": (
|
||
"You are editing an existing image prompt. The previous prompt describes "
|
||
"the current image. Apply only the user's requested changes. Keep all "
|
||
"subjects, style, colors, composition and constraints from the previous "
|
||
"prompt unless the user explicitly asks to change them. Return only the "
|
||
"final English prompt, without quotes or explanations."
|
||
),
|
||
"i2i": (
|
||
"You are updating an image-editing instruction. The previous prompt "
|
||
"describes the current image. Apply only the user's requested changes. "
|
||
"Keep the existing composition, subjects, style, lighting and constraints "
|
||
"unless the user explicitly asks to change them. 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,
|
||
*,
|
||
previous_prompt: str | None = None,
|
||
base_prompt: str | None = None,
|
||
) -> str:
|
||
if not self.enabled:
|
||
return prompt
|
||
|
||
user_prompt = self._sanitize(prompt)
|
||
if not user_prompt:
|
||
return prompt
|
||
|
||
previous = self._sanitize(previous_prompt or "")
|
||
base = self._sanitize(base_prompt or "")
|
||
if previous:
|
||
chat_prompt = self._refine_chat_prompt(
|
||
user_prompt,
|
||
mode,
|
||
previous_prompt=previous,
|
||
base_prompt=base or previous,
|
||
)
|
||
else:
|
||
chat_prompt = self._chat_prompt(user_prompt, mode)
|
||
|
||
workflow = copy.deepcopy(self._template)
|
||
workflow[self._text_node_id]["inputs"]["prompt"] = chat_prompt
|
||
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 _sanitize(text: str) -> str:
|
||
value = text.strip()
|
||
for token in ("<|im_start|>", "<|im_end|>"):
|
||
value = value.replace(token, "")
|
||
return value
|
||
|
||
@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 _refine_chat_prompt(
|
||
prompt: str,
|
||
mode: str,
|
||
*,
|
||
previous_prompt: str,
|
||
base_prompt: str,
|
||
) -> str:
|
||
system = REFINE_SYSTEM_PROMPTS.get(
|
||
mode, REFINE_SYSTEM_PROMPTS["t2i"]
|
||
)
|
||
return (
|
||
f"<|im_start|>system\n{system}<|im_end|>\n"
|
||
"<|im_start|>user\n"
|
||
f"Base scene: {base_prompt}\n"
|
||
f"Current prompt: {previous_prompt}\n"
|
||
f"Requested change: {prompt}\n"
|
||
"<|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
|