64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
import json
|
|
from typing import Any, Optional
|
|
|
|
from .settings import MemoirImageSettings
|
|
|
|
|
|
class MemoirImagePromptService:
|
|
CATEGORY_STYLE_MAP = {
|
|
"childhood": "watercolor",
|
|
"family": "watercolor",
|
|
"career_early": "realistic",
|
|
"career_achievement": "realistic",
|
|
"career_challenge": "realistic",
|
|
"beliefs": "editorial illustration",
|
|
"summary": "editorial illustration",
|
|
}
|
|
|
|
def __init__(self, llm: Optional[Any], settings: MemoirImageSettings):
|
|
self.llm = llm
|
|
self.settings = settings
|
|
|
|
def build_prompt(
|
|
self,
|
|
chapter_title: str,
|
|
chapter_category: str,
|
|
description: str,
|
|
context_excerpt: str,
|
|
) -> dict[str, str]:
|
|
style = self.CATEGORY_STYLE_MAP.get(chapter_category, self.settings.default_style)
|
|
prompt_context = f"{chapter_category}: {chapter_title}"
|
|
|
|
llm_input = {
|
|
"chapter_title": chapter_title,
|
|
"chapter_category": chapter_category,
|
|
"description": description,
|
|
"context_excerpt": context_excerpt,
|
|
"default_style": style,
|
|
"default_size": self.settings.default_size,
|
|
}
|
|
|
|
if self.llm:
|
|
try:
|
|
response = self.llm.invoke(
|
|
"Return JSON only with keys prompt, style, size. "
|
|
"Convert the memoir scene into an image-generation prompt.\n"
|
|
+ json.dumps(llm_input, ensure_ascii=False)
|
|
)
|
|
parsed = json.loads(response.content)
|
|
return {
|
|
"prompt": parsed["prompt"],
|
|
"style": parsed.get("style", style),
|
|
"size": parsed.get("size", self.settings.default_size),
|
|
"prompt_context": prompt_context,
|
|
}
|
|
except Exception:
|
|
pass
|
|
|
|
return {
|
|
"prompt": f"{description}\n\nScene context: {context_excerpt}",
|
|
"style": style,
|
|
"size": self.settings.default_size,
|
|
"prompt_context": prompt_context,
|
|
}
|