Grigo
126 lines
4.7 KiB
Python
126 lines
4.7 KiB
Python
import json
|
|
import os
|
|
import re
|
|
from services.llm import send_message
|
|
from services.personas import get_persona
|
|
|
|
PROMPT_BUILDER_SYSTEM = """You are a Stable Diffusion prompt engineer for anime illustration models.
|
|
Given a roleplay chat excerpt and character appearance hints, output ONLY valid JSON (no markdown):
|
|
{
|
|
"should_generate": true,
|
|
"shot_type": "first_person_pov" | "landscape" | "third_person",
|
|
"appearance_tags": "booru-style tags for character appearance extracted from hints, e.g. 'white hair, wolf ears, wolf tail, yellow eyes'",
|
|
"action_tags": "booru-style tags for pose/action, e.g. 'sitting, smiling, looking at viewer'",
|
|
"environment_tags": "booru-style tags for location/lighting, e.g. 'indoors, kitchen, sunlight'"
|
|
}
|
|
Rules:
|
|
- ONLY use real danbooru/e621 tags. Multi-word concepts MUST be written as single tags: 'white hair' not 'white, hair'. 'wolf ears' not 'wolf, ears'.
|
|
- Do NOT include quality tags, model names, style words, 'pov', or category/metadata words.
|
|
- Do NOT invent tags. If unsure — omit.
|
|
- Keep each field to 3-6 tags."""
|
|
|
|
|
|
def extract_image_prompt_tag(text: str) -> str | None:
|
|
if "[IMAGE_PROMPT:" not in text:
|
|
return None
|
|
try:
|
|
start = text.index("[IMAGE_PROMPT:") + len("[IMAGE_PROMPT:")
|
|
end = text.index("]", start)
|
|
return text[start:end].strip()
|
|
except ValueError:
|
|
return None
|
|
|
|
|
|
def strip_image_prompt_tag(text: str) -> str:
|
|
return re.sub(r"\[IMAGE_PROMPT:.*?\]", "", text, flags=re.DOTALL).strip()
|
|
|
|
|
|
PONY_CHECKPOINTS = {"ponyDiffusionV6XL_v6StartWithThisOne.safetensors"}
|
|
SD_CHECKPOINT = os.getenv("SD_CHECKPOINT", "")
|
|
PONY_NEGATIVE = "score_1, score_2, score_3, score_4, worst quality, low quality, blurry, bad anatomy, watermark, text, censored"
|
|
|
|
def build_positive_prompt(scene: dict, persona: dict | None) -> str:
|
|
is_pony = SD_CHECKPOINT in PONY_CHECKPOINTS
|
|
quality = "score_9, score_8_up, score_7_up, source_anime, highres" if is_pony else "masterpiece, best quality, highres"
|
|
parts = [quality]
|
|
|
|
# prefer LLM-extracted appearance over raw persona tags
|
|
appearance = scene.get("appearance_tags") or (persona or {}).get("appearance_tags", "")
|
|
if appearance:
|
|
parts.append(appearance)
|
|
|
|
if scene.get("shot_type") == "landscape":
|
|
parts.append(scene.get("environment_tags", ""))
|
|
else:
|
|
if scene.get("shot_type") == "first_person_pov":
|
|
parts.append("pov, first-person view, looking at viewer")
|
|
parts.append(scene.get("action_tags", ""))
|
|
parts.append(scene.get("environment_tags", ""))
|
|
|
|
lora = (persona or {}).get("lora_name", "")
|
|
weight = (persona or {}).get("lora_weight", 0.8)
|
|
if lora:
|
|
parts.append(f"<lora:{lora}:{weight}>")
|
|
|
|
positive = ", ".join(p.strip() for p in parts if p and p.strip())
|
|
seen, deduped = set(), []
|
|
for tag in positive.split(", "):
|
|
t = tag.strip()
|
|
if t and t not in seen:
|
|
seen.add(t)
|
|
deduped.append(t)
|
|
return ", ".join(deduped)
|
|
|
|
|
|
async def generate_sd_prompt(
|
|
messages: list,
|
|
persona_id: str,
|
|
) -> tuple[str | None, str | None]:
|
|
persona = await get_persona(persona_id)
|
|
if not persona or not persona.get("sd_enabled"):
|
|
return None, None
|
|
|
|
recent = [m for m in messages if m["role"] in ("user", "assistant")][-6:]
|
|
if not recent:
|
|
return None, None
|
|
|
|
excerpt = "\n".join(f"{m['role']}: {strip_image_prompt_tag(m['content'])}" for m in recent)
|
|
|
|
appearance = persona.get("appearance_tags", "")
|
|
# For card personas, also include description for better visual context
|
|
if persona_id.startswith("card_"):
|
|
from services.character_card import get_character
|
|
card = await get_character(persona_id[5:])
|
|
if card and card.get("description"):
|
|
appearance = f"{appearance}\nCharacter description: {card['description'][:400]}"
|
|
|
|
builder_messages = [
|
|
{"role": "system", "content": PROMPT_BUILDER_SYSTEM},
|
|
{
|
|
"role": "user",
|
|
"content": f"Persona appearance hints: {appearance}\n\nChat:\n{excerpt}",
|
|
},
|
|
]
|
|
|
|
try:
|
|
raw = await send_message(builder_messages)
|
|
raw = raw.strip()
|
|
if raw.startswith("```"):
|
|
raw = re.sub(r"^```\w*\n?", "", raw)
|
|
raw = re.sub(r"\n?```$", "", raw)
|
|
scene = json.loads(raw)
|
|
except (json.JSONDecodeError, Exception):
|
|
return None, None
|
|
|
|
|
|
positive = build_positive_prompt(scene, persona)
|
|
is_pony = SD_CHECKPOINT in PONY_CHECKPOINTS
|
|
negative = PONY_NEGATIVE if is_pony else "low quality, blurry, bad anatomy, watermark, text"
|
|
if scene.get("shot_type") == "first_person_pov":
|
|
negative += ", third person, over the shoulder"
|
|
|
|
full = positive
|
|
if negative:
|
|
full += f"\n\nNegative prompt: {negative}"
|
|
return full, negative
|