File size: 3,195 Bytes
575c8eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """
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})
|