49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
"""Run ComfyUI generation from SdPromptBundle (single hybrid prompt for Anima)."""
|
|
|
|
import logging
|
|
import os
|
|
|
|
from services import sdbackend as sd_service
|
|
from services.memory import update_message_image, update_message_prompt, update_message_prompt_alt
|
|
from services.sd_prompt import SdPromptBundle
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
SD_AUTO_GENERATE = os.getenv("SD_AUTO_GENERATE", "false").lower() in ("1", "true", "yes")
|
|
|
|
|
|
async def run_sd_for_message(bundle: SdPromptBundle | None, msg_id: int | None) -> dict:
|
|
"""Generate image, persist prompts/paths on message. Returns fields for API/SSE."""
|
|
out = {
|
|
"image_prompt": None,
|
|
"image_prompt_alt": None,
|
|
"image_path": None,
|
|
"image_path_alt": None,
|
|
"image_error": None,
|
|
"image_error_alt": None,
|
|
}
|
|
if not bundle or not bundle.tag_full:
|
|
return out
|
|
|
|
out["image_prompt"] = bundle.tag_full
|
|
if bundle.desc_full and bundle.desc_full != bundle.tag_full:
|
|
out["image_prompt_alt"] = bundle.desc_full
|
|
|
|
if msg_id:
|
|
await update_message_prompt(msg_id, bundle.tag_full)
|
|
if out["image_prompt_alt"]:
|
|
await update_message_prompt_alt(msg_id, out["image_prompt_alt"])
|
|
|
|
if not SD_AUTO_GENERATE:
|
|
return out
|
|
|
|
rel, err = await sd_service.generate_from_full_prompt(bundle.tag_full)
|
|
if rel:
|
|
out["image_path"] = f"/static/{rel}"
|
|
if msg_id:
|
|
await update_message_image(msg_id, rel)
|
|
else:
|
|
out["image_error"] = err
|
|
|
|
return out
|