Add handler_feed.py

handler_feed.py

@@ -0,0 +1,113 @@
+"""
+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})