handler_feed.py

"""
RunPod handler for Mongle feed generation.

Input:
  {
    "input": {
      "appearance": {...},
      "quest_ko": "공원에서 30분 달리기를 완료했어요!",
      "quest_en": "optional English scene prompt",
      "name": "optional_job_name",
      "feed_seed": 123
    }
  }

Output:
  {
    "name": "...",
    "quest": {
      "quest_ko": "...",
      "quest_en": "..."
    },
    "feed_image": "<base64 png>"
  }
"""

import base64
import io
import json
import os
import traceback
from pathlib import Path

import runpod
from PIL import Image

from test_text2feed_pipeline import DEFAULT_QUEST_KO, translate_quest
from src.feed.feed_generator_1 import generate, load_pipeline, unload_pipeline


OUTPUT_ROOT = Path("outputs/handler_feed")


def _encode_png(image_or_path) -> str:
    if isinstance(image_or_path, (str, Path)):
        image = Image.open(image_or_path).convert("RGB")
    else:
        image = image_or_path.convert("RGB")
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def _load_appearance(payload: dict):
    appearance = payload.get("appearance")
    if appearance is not None:
        if isinstance(appearance, str):
            return json.loads(appearance)
        return appearance

    appearance_path = payload.get("appearance_path")
    if appearance_path:
        return json.loads(Path(appearance_path).read_text(encoding="utf-8"))

    raise ValueError("Missing input.appearance or input.appearance_path.")


def handler(event):
    payload = event.get("input", {}) if isinstance(event, dict) else {}
    name = payload.get("name") or "feed_job"
    out_dir = OUTPUT_ROOT / name
    out_dir.mkdir(parents=True, exist_ok=True)

    try:
        appearance = _load_appearance(payload)
        quest_ko = payload.get("quest_ko") or DEFAULT_QUEST_KO
        quest_en = payload.get("quest_en") or translate_quest(quest_ko)
        feed_seed = int(payload.get("feed_seed", payload.get("seed", 123)))

        pipe = load_pipeline()
        try:
            image = generate(appearance, quest_en, pipe, seed=feed_seed)
        finally:
            unload_pipeline(pipe)

        feed_path = out_dir / "feed.png"
        image.save(feed_path)

        quest = {"quest_ko": quest_ko, "quest_en": quest_en}
        result_data = {
            "name": name,
            "appearance": appearance,
            "quest": quest,
            "feed_image_path": str(feed_path).replace("\\", "/"),
        }
        (out_dir / "result.json").write_text(
            json.dumps(result_data, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )

        return {
            "name": name,
            "quest": quest,
            "feed_image": _encode_png(feed_path),
            "output_dir": str(out_dir).replace("\\", "/"),
        }
    except Exception:
        error_text = traceback.format_exc()
        (out_dir / "error.log").write_text(error_text, encoding="utf-8")
        return {"error": error_text, "name": name, "output_dir": str(out_dir).replace("\\", "/")}


if os.getenv("RUNPOD_SERVERLESS", "1") == "1":
    runpod.serverless.start({"handler": handler})